gather_points.cpp 2.37 KB
Newer Older
wuyuefeng's avatar
wuyuefeng committed
1
2
3
#include <ATen/cuda/CUDAContext.h>
#include <THC/THC.h>
#include <torch/extension.h>
zhangwenwei's avatar
zhangwenwei committed
4
5
6
#include <torch/serialize/tensor.h>

#include <vector>
wuyuefeng's avatar
wuyuefeng committed
7
8
9
10

extern THCState *state;

int gather_points_wrapper(int b, int c, int n, int npoints,
zhangwenwei's avatar
zhangwenwei committed
11
12
                          at::Tensor points_tensor, at::Tensor idx_tensor,
                          at::Tensor out_tensor);
wuyuefeng's avatar
wuyuefeng committed
13
14

void gather_points_kernel_launcher(int b, int c, int n, int npoints,
zhangwenwei's avatar
zhangwenwei committed
15
16
                                   const float *points, const int *idx,
                                   float *out, cudaStream_t stream);
wuyuefeng's avatar
wuyuefeng committed
17
18

int gather_points_grad_wrapper(int b, int c, int n, int npoints,
zhangwenwei's avatar
zhangwenwei committed
19
20
21
                               at::Tensor grad_out_tensor,
                               at::Tensor idx_tensor,
                               at::Tensor grad_points_tensor);
wuyuefeng's avatar
wuyuefeng committed
22
23

void gather_points_grad_kernel_launcher(int b, int c, int n, int npoints,
zhangwenwei's avatar
zhangwenwei committed
24
25
26
                                        const float *grad_out, const int *idx,
                                        float *grad_points,
                                        cudaStream_t stream);
wuyuefeng's avatar
wuyuefeng committed
27
28

int gather_points_wrapper(int b, int c, int n, int npoints,
zhangwenwei's avatar
zhangwenwei committed
29
30
31
32
33
34
35
36
37
                          at::Tensor points_tensor, at::Tensor idx_tensor,
                          at::Tensor out_tensor) {
  const float *points = points_tensor.data_ptr<float>();
  const int *idx = idx_tensor.data_ptr<int>();
  float *out = out_tensor.data_ptr<float>();

  cudaStream_t stream = THCState_getCurrentStream(state);
  gather_points_kernel_launcher(b, c, n, npoints, points, idx, out, stream);
  return 1;
wuyuefeng's avatar
wuyuefeng committed
38
39
40
}

int gather_points_grad_wrapper(int b, int c, int n, int npoints,
zhangwenwei's avatar
zhangwenwei committed
41
42
43
44
45
46
47
48
49
50
51
                               at::Tensor grad_out_tensor,
                               at::Tensor idx_tensor,
                               at::Tensor grad_points_tensor) {
  const float *grad_out = grad_out_tensor.data_ptr<float>();
  const int *idx = idx_tensor.data_ptr<int>();
  float *grad_points = grad_points_tensor.data_ptr<float>();

  cudaStream_t stream = THCState_getCurrentStream(state);
  gather_points_grad_kernel_launcher(b, c, n, npoints, grad_out, idx,
                                     grad_points, stream);
  return 1;
wuyuefeng's avatar
wuyuefeng committed
52
53
54
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
zhangwenwei's avatar
zhangwenwei committed
55
56
57
58
  m.def("gather_points_wrapper", &gather_points_wrapper,
        "gather_points_wrapper");
  m.def("gather_points_grad_wrapper", &gather_points_grad_wrapper,
        "gather_points_grad_wrapper");
wuyuefeng's avatar
wuyuefeng committed
59
}