interpolate.cpp 4.08 KB
Newer Older
wuyuefeng's avatar
Credit  
wuyuefeng committed
1
2
3
// Modified from
// https://github.com/sshaoshuai/Pointnet2.PyTorch/tree/master/pointnet2/src/interpolate.cpp

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

#include <vector>
wuyuefeng's avatar
wuyuefeng committed
14
15
16
17

extern THCState *state;

void three_nn_wrapper(int b, int n, int m, at::Tensor unknown_tensor,
zhangwenwei's avatar
zhangwenwei committed
18
19
                      at::Tensor known_tensor, at::Tensor dist2_tensor,
                      at::Tensor idx_tensor);
wuyuefeng's avatar
wuyuefeng committed
20
21

void three_nn_kernel_launcher(int b, int n, int m, const float *unknown,
zhangwenwei's avatar
zhangwenwei committed
22
23
                              const float *known, float *dist2, int *idx,
                              cudaStream_t stream);
wuyuefeng's avatar
wuyuefeng committed
24

zhangwenwei's avatar
zhangwenwei committed
25
26
27
void three_interpolate_wrapper(int b, int c, int m, int n,
                               at::Tensor points_tensor, at::Tensor idx_tensor,
                               at::Tensor weight_tensor, at::Tensor out_tensor);
wuyuefeng's avatar
wuyuefeng committed
28
29

void three_interpolate_kernel_launcher(int b, int c, int m, int n,
zhangwenwei's avatar
zhangwenwei committed
30
31
32
                                       const float *points, const int *idx,
                                       const float *weight, float *out,
                                       cudaStream_t stream);
wuyuefeng's avatar
wuyuefeng committed
33

zhangwenwei's avatar
zhangwenwei committed
34
35
36
37
38
void three_interpolate_grad_wrapper(int b, int c, int n, int m,
                                    at::Tensor grad_out_tensor,
                                    at::Tensor idx_tensor,
                                    at::Tensor weight_tensor,
                                    at::Tensor grad_points_tensor);
wuyuefeng's avatar
wuyuefeng committed
39

zhangwenwei's avatar
zhangwenwei committed
40
41
42
43
44
void three_interpolate_grad_kernel_launcher(int b, int c, int n, int m,
                                            const float *grad_out,
                                            const int *idx, const float *weight,
                                            float *grad_points,
                                            cudaStream_t stream);
wuyuefeng's avatar
wuyuefeng committed
45
46

void three_nn_wrapper(int b, int n, int m, at::Tensor unknown_tensor,
zhangwenwei's avatar
zhangwenwei committed
47
48
49
50
51
52
53
                      at::Tensor known_tensor, at::Tensor dist2_tensor,
                      at::Tensor idx_tensor) {
  const float *unknown = unknown_tensor.data_ptr<float>();
  const float *known = known_tensor.data_ptr<float>();
  float *dist2 = dist2_tensor.data_ptr<float>();
  int *idx = idx_tensor.data_ptr<int>();

54
  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
zhangwenwei's avatar
zhangwenwei committed
55
  three_nn_kernel_launcher(b, n, m, unknown, known, dist2, idx, stream);
wuyuefeng's avatar
wuyuefeng committed
56
57
58
}

void three_interpolate_wrapper(int b, int c, int m, int n,
zhangwenwei's avatar
zhangwenwei committed
59
60
61
62
63
64
65
66
                               at::Tensor points_tensor, at::Tensor idx_tensor,
                               at::Tensor weight_tensor,
                               at::Tensor out_tensor) {
  const float *points = points_tensor.data_ptr<float>();
  const float *weight = weight_tensor.data_ptr<float>();
  float *out = out_tensor.data_ptr<float>();
  const int *idx = idx_tensor.data_ptr<int>();

67
  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
zhangwenwei's avatar
zhangwenwei committed
68
69
  three_interpolate_kernel_launcher(b, c, m, n, points, idx, weight, out,
                                    stream);
wuyuefeng's avatar
wuyuefeng committed
70
71
72
}

void three_interpolate_grad_wrapper(int b, int c, int n, int m,
zhangwenwei's avatar
zhangwenwei committed
73
74
75
76
77
78
79
80
81
                                    at::Tensor grad_out_tensor,
                                    at::Tensor idx_tensor,
                                    at::Tensor weight_tensor,
                                    at::Tensor grad_points_tensor) {
  const float *grad_out = grad_out_tensor.data_ptr<float>();
  const float *weight = weight_tensor.data_ptr<float>();
  float *grad_points = grad_points_tensor.data_ptr<float>();
  const int *idx = idx_tensor.data_ptr<int>();

82
  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
zhangwenwei's avatar
zhangwenwei committed
83
84
  three_interpolate_grad_kernel_launcher(b, c, n, m, grad_out, idx, weight,
                                         grad_points, stream);
wuyuefeng's avatar
wuyuefeng committed
85
86
87
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
zhangwenwei's avatar
zhangwenwei committed
88
89
90
91
92
  m.def("three_nn_wrapper", &three_nn_wrapper, "three_nn_wrapper");
  m.def("three_interpolate_wrapper", &three_interpolate_wrapper,
        "three_interpolate_wrapper");
  m.def("three_interpolate_grad_wrapper", &three_interpolate_grad_wrapper,
        "three_interpolate_grad_wrapper");
wuyuefeng's avatar
wuyuefeng committed
93
}