kernel.cu 1.51 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
#ifndef THC_GENERIC_FILE
#define THC_GENERIC_FILE "generic/kernel.cu"
#else

rusty1s's avatar
rusty1s committed
5
void check(THCState *state, THCTensor *output, THCudaLongTensor *index, THCTensor *input) {
rusty1s's avatar
max dim  
rusty1s committed
6
7
8
  THCAssertSameGPU(THCTensor_(checkGPU)(state, 1, output, input));
  THCAssertSameGPU(THCudaLongTensor_checkGPU(state, 2, index));
  THArgCheck(THCTensor_(nDimension)(state, output) <= MAX_DIMS, 1, "Tensor too large or too many dimensions");
rusty1s's avatar
rusty1s committed
9
}
rusty1s's avatar
max dim  
rusty1s committed
10

rusty1s's avatar
rusty1s committed
11
12
void scatter_(mul)(THCState *state, int dim, THCTensor *output, THCudaLongTensor *index, THCTensor *input) {
  check(state, output, index, input);
rusty1s's avatar
rusty1s committed
13
14
15
16
  printf("mul");
}

void scatter_(div)(THCState *state, int dim, THCTensor *output, THCudaLongTensor *index, THCTensor *input) {
rusty1s's avatar
rusty1s committed
17
  check(state, output, index, input);
rusty1s's avatar
rusty1s committed
18
19
20
21
  printf("div");
}

void scatter_(mean)(THCState *state, int dim, THCTensor *output, THCudaLongTensor *index, THCTensor *input, THCTensor *num_output) {
rusty1s's avatar
rusty1s committed
22
  check(state, output, index, input);
rusty1s's avatar
rusty1s committed
23
24
25
26
  printf("mean");
}

void scatter_(max)(THCState *state, int dim, THCTensor *output, THCudaLongTensor *index, THCTensor *input, THCudaLongTensor *arg_output) {
rusty1s's avatar
rusty1s committed
27
  check(state, output, index, input);
rusty1s's avatar
rusty1s committed
28
29
30
31
  printf("max");
}

void scatter_(min)(THCState *state, int dim, THCTensor *output, THCudaLongTensor *index, THCTensor *input, THCudaLongTensor *arg_output) {
rusty1s's avatar
rusty1s committed
32
  check(state, output, index, input);
rusty1s's avatar
rusty1s committed
33
34
35
36
  printf("min");
}

void index_backward(THCState *state, int dim, THCTensor *output, THCudaLongTensor *index, THCTensor *grad, THCudaLongTensor *arg_grad) {
rusty1s's avatar
rusty1s committed
37
  check(state, output, index, grad);
rusty1s's avatar
rusty1s committed
38
39
40
41
  printf("index_backward");
}

#endif