"vscode:/vscode.git/clone" did not exist on "a466b3c2fe22c866523afea8cffbee242f4489a5"
group_points.cpp 2.5 KB
Newer Older
wuyuefeng's avatar
Credit  
wuyuefeng committed
1
2
3
// Modified from
// https://github.com/sshaoshuai/Pointnet2.PyTorch/tree/master/pointnet2/src/group_points.cpp

zhangwenwei's avatar
zhangwenwei committed
4
#include <THC/THC.h>
wuyuefeng's avatar
wuyuefeng committed
5
6
7
#include <cuda.h>
#include <cuda_runtime_api.h>
#include <torch/extension.h>
zhangwenwei's avatar
zhangwenwei committed
8
9
10
#include <torch/serialize/tensor.h>

#include <vector>
wuyuefeng's avatar
wuyuefeng committed
11
12
13
14

extern THCState *state;

int group_points_wrapper(int b, int c, int n, int npoints, int nsample,
zhangwenwei's avatar
zhangwenwei committed
15
16
                         at::Tensor points_tensor, at::Tensor idx_tensor,
                         at::Tensor out_tensor);
wuyuefeng's avatar
wuyuefeng committed
17
18

void group_points_kernel_launcher(int b, int c, int n, int npoints, int nsample,
zhangwenwei's avatar
zhangwenwei committed
19
20
                                  const float *points, const int *idx,
                                  float *out, cudaStream_t stream);
wuyuefeng's avatar
wuyuefeng committed
21
22

int group_points_grad_wrapper(int b, int c, int n, int npoints, int nsample,
zhangwenwei's avatar
zhangwenwei committed
23
24
                              at::Tensor grad_out_tensor, at::Tensor idx_tensor,
                              at::Tensor grad_points_tensor);
wuyuefeng's avatar
wuyuefeng committed
25

zhangwenwei's avatar
zhangwenwei committed
26
27
28
29
void group_points_grad_kernel_launcher(int b, int c, int n, int npoints,
                                       int nsample, const float *grad_out,
                                       const int *idx, float *grad_points,
                                       cudaStream_t stream);
wuyuefeng's avatar
wuyuefeng committed
30
31

int group_points_grad_wrapper(int b, int c, int n, int npoints, int nsample,
zhangwenwei's avatar
zhangwenwei committed
32
33
34
35
36
                              at::Tensor grad_out_tensor, at::Tensor idx_tensor,
                              at::Tensor grad_points_tensor) {
  float *grad_points = grad_points_tensor.data_ptr<float>();
  const int *idx = idx_tensor.data_ptr<int>();
  const float *grad_out = grad_out_tensor.data_ptr<float>();
wuyuefeng's avatar
wuyuefeng committed
37

38
  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
wuyuefeng's avatar
wuyuefeng committed
39

zhangwenwei's avatar
zhangwenwei committed
40
41
42
  group_points_grad_kernel_launcher(b, c, n, npoints, nsample, grad_out, idx,
                                    grad_points, stream);
  return 1;
wuyuefeng's avatar
wuyuefeng committed
43
44
45
}

int group_points_wrapper(int b, int c, int n, int npoints, int nsample,
zhangwenwei's avatar
zhangwenwei committed
46
47
48
49
50
                         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>();
wuyuefeng's avatar
wuyuefeng committed
51

52
  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
wuyuefeng's avatar
wuyuefeng committed
53

zhangwenwei's avatar
zhangwenwei committed
54
55
56
  group_points_kernel_launcher(b, c, n, npoints, nsample, points, idx, out,
                               stream);
  return 1;
wuyuefeng's avatar
wuyuefeng committed
57
58
59
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
zhangwenwei's avatar
zhangwenwei committed
60
61
  m.def("forward", &group_points_wrapper, "group_points_wrapper");
  m.def("backward", &group_points_grad_wrapper, "group_points_grad_wrapper");
wuyuefeng's avatar
wuyuefeng committed
62
}