cpu.c 1.94 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
#ifndef TH_GENERIC_FILE
rusty1s's avatar
rusty1s committed
2
#define TH_GENERIC_FILE "generic/cpu.c"
rusty1s's avatar
rusty1s committed
3
4
5
#else

void scatter_(add)(THTensor *output, THLongTensor *index, THTensor *input, int dim) {
rusty1s's avatar
rusty1s committed
6
  int64_t idx;
rusty1s's avatar
rusty1s committed
7
8
  TH_TENSOR_DIM_APPLY3(real, output, real, input, int64_t, index, dim,
    TH_TENSOR_DIM_APPLY3_SIZE_EQ_EXCEPT_DIM,
rusty1s's avatar
rusty1s committed
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
    for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) {
      idx = *(index_data + i * index_stride);
      assertIndexInBoundaries(idx, output_size, TH_TENSOR_DIM_APPLY_counter);
      output_data[idx] += *(input_data + i * input_stride);
    })
}

void scatter_(sub)(THTensor *output, THLongTensor *index, THTensor *input, int dim) {
  int64_t idx;
  TH_TENSOR_DIM_APPLY3(real, output, real, input, int64_t, index, dim,
    TH_TENSOR_DIM_APPLY3_SIZE_EQ_EXCEPT_DIM,
    for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) {
      idx = *(index_data + i * index_stride);
      assertIndexInBoundaries(idx, output_size, TH_TENSOR_DIM_APPLY_counter);
      output_data[idx] += *(input_data + i * input_stride);
    })
}

void scatter_(mul)(THTensor *output, THLongTensor *index, THTensor *input, int dim) {
  int64_t idx;
  TH_TENSOR_DIM_APPLY3(real, output, real, input, int64_t, index, dim,
    TH_TENSOR_DIM_APPLY3_SIZE_EQ_EXCEPT_DIM,
    for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) {
      idx = *(index_data + i * index_stride);
      assertIndexInBoundaries(idx, output_size, TH_TENSOR_DIM_APPLY_counter);
      output_data[idx] += *(input_data + i * input_stride);
    })
}

void scatter_(div)(THTensor *output, THLongTensor *index, THTensor *input, int dim) {
  int64_t idx;
  TH_TENSOR_DIM_APPLY3(real, output, real, input, int64_t, index, dim,
    TH_TENSOR_DIM_APPLY3_SIZE_EQ_EXCEPT_DIM,
    for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) {
rusty1s's avatar
rusty1s committed
43
      idx = *(index_data + i * index_stride);
rusty1s's avatar
rusty1s committed
44
      assertIndexInBoundaries(idx, output_size, TH_TENSOR_DIM_APPLY_counter);
rusty1s's avatar
rusty1s committed
45
46
47
48
49
      output_data[idx] += *(input_data + i * input_stride);
    })
}

#endif