kernel.cu 1.25 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
#include <THC.h>

#include "kernel.h"

rusty1s's avatar
rusty1s committed
5
6
7
#include "common.cuh"
#include "THCIndex.cuh"

rusty1s's avatar
bugfix  
rusty1s committed
8
#define cluster_(NAME) TH_CONCAT_4(cluster_, NAME, _kernel_, Real)
rusty1s's avatar
rusty1s committed
9
#define thc_(NAME) TH_CONCAT_4(thc_, NAME, _, Real)
rusty1s's avatar
rusty1s committed
10

rusty1s's avatar
rusty1s committed
11
#include "generic/common.cu"
rusty1s's avatar
rusty1s committed
12
#include "THCGenerateAllTypes.h"
rusty1s's avatar
rusty1s committed
13
14

template<typename Real, int Dims>
15
__global__ void gridKernel(int64_t *output, TensorInfo<Real> position, Real *size, Real *maxPosition, const int N) {
rusty1s's avatar
rusty1s committed
16
  KERNEL_LOOP(i, N) {
17
    int positionOffset = 0;
rusty1s's avatar
rusty1s committed
18
    IndexToOffset<Real, Dims>::compute(i, position, &positionOffset);
19
20
21
22
23
24

    int D = position.size[position.dims - 1];
    int weight = 1; int64_t cluster = 0;
    for (int d = D - 1; d >= 0; d--) {
      cluster += weight * (int64_t) (position.data[positionOffset + d] / size[d]);
      weight *= (int64_t) (maxPosition[d] / size[d]) + 1;
rusty1s's avatar
rusty1s committed
25
    }
26
    output[i] = cluster;
rusty1s's avatar
rusty1s committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
  }
}

#include "generic/kernel.cu"
#include "THCGenerateFloatType.h"
#include "generic/kernel.cu"
#include "THCGenerateDoubleType.h"
#include "generic/kernel.cu"
#include "THCGenerateByteType.h"
#include "generic/kernel.cu"
#include "THCGenerateCharType.h"
#include "generic/kernel.cu"
#include "THCGenerateShortType.h"
#include "generic/kernel.cu"
#include "THCGenerateIntType.h"
#include "generic/kernel.cu"
#include "THCGenerateLongType.h"