File size: 1,308 Bytes
28c256d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
#include "pytorch_cpp_helper.hpp"
#include "pytorch_device_registry.hpp"

void gather_points_forward_impl(int b, int c, int n, int npoints,
                                const Tensor points, const Tensor idx,
                                Tensor out) {
  DISPATCH_DEVICE_IMPL(gather_points_forward_impl, b, c, n, npoints, points,
                       idx, out);
}

void gather_points_backward_impl(int b, int c, int n, int npoints,
                                 const Tensor grad_out, const Tensor idx,
                                 Tensor grad_points) {
  DISPATCH_DEVICE_IMPL(gather_points_backward_impl, b, c, n, npoints, grad_out,
                       idx, grad_points);
}

void gather_points_forward(Tensor points_tensor, Tensor idx_tensor,
                           Tensor out_tensor, int b, int c, int n,
                           int npoints) {
  gather_points_forward_impl(b, c, n, npoints, points_tensor, idx_tensor,
                             out_tensor);
}

void gather_points_backward(Tensor grad_out_tensor, Tensor idx_tensor,
                            Tensor grad_points_tensor, int b, int c, int n,
                            int npoints) {
  gather_points_backward_impl(b, c, n, npoints, grad_out_tensor, idx_tensor,
                              grad_points_tensor);
}