| #include <torch/serialize/tensor.h> |
| #include <vector> |
| |
| #include <cuda.h> |
| #include <cuda_runtime_api.h> |
| #include "ball_query_gpu.h" |
|
|
| |
|
|
| #include <ATen/cuda/CUDAContext.h> |
| #include <ATen/cuda/CUDAEvent.h> |
| |
|
|
| #define CHECK_CUDA(x) do { \ |
| if (!x.type().is_cuda()) { \ |
| fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \ |
| exit(-1); \ |
| } \ |
| } while (0) |
| #define CHECK_CONTIGUOUS(x) do { \ |
| if (!x.is_contiguous()) { \ |
| fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \ |
| exit(-1); \ |
| } \ |
| } while (0) |
| #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x) |
|
|
| int ball_query_wrapper_fast(int b, int n, int m, float radius, int nsample, |
| at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor) { |
| CHECK_INPUT(new_xyz_tensor); |
| CHECK_INPUT(xyz_tensor); |
| const float *new_xyz = new_xyz_tensor.data<float>(); |
| const float *xyz = xyz_tensor.data<float>(); |
| int *idx = idx_tensor.data<int>(); |
|
|
| ball_query_kernel_launcher_fast(b, n, m, radius, nsample, new_xyz, xyz, idx); |
| return 1; |
| } |
|
|
|
|
| int ball_center_query_wrapper_fast(int b, int n, int m, float radius, |
| at::Tensor point_tensor, at::Tensor key_point_tensor, at::Tensor idx_tensor) { |
| CHECK_INPUT(point_tensor); |
| CHECK_INPUT(key_point_tensor); |
| const float *point = point_tensor.data<float>(); |
| const float *key_point = key_point_tensor.data<float>(); |
| int *idx = idx_tensor.data<int>(); |
|
|
| ball_center_query_kernel_launcher_fast(b, n, m, radius, point, key_point, idx); |
| return 1; |
| } |
|
|
|
|
| int knn_query_wrapper_fast(int b, int n, int m, int nsample, |
| at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor dist2_tensor, at::Tensor idx_tensor) { |
| CHECK_INPUT(new_xyz_tensor); |
| CHECK_INPUT(xyz_tensor); |
| const float *new_xyz = new_xyz_tensor.data<float>(); |
| const float *xyz = xyz_tensor.data<float>(); |
| float *dist2 = dist2_tensor.data<float>(); |
| int *idx = idx_tensor.data<int>(); |
|
|
| knn_query_kernel_launcher_fast(b, n, m, nsample, new_xyz, xyz, dist2, idx); |
| return 1; |
| } |
|
|
|
|
| int ball_query_wrapper_stack(int B, int M, float radius, int nsample, |
| at::Tensor new_xyz_tensor, at::Tensor new_xyz_batch_cnt_tensor, |
| at::Tensor xyz_tensor, at::Tensor xyz_batch_cnt_tensor, at::Tensor idx_tensor) { |
| CHECK_INPUT(new_xyz_tensor); |
| CHECK_INPUT(xyz_tensor); |
| CHECK_INPUT(new_xyz_batch_cnt_tensor); |
| CHECK_INPUT(xyz_batch_cnt_tensor); |
|
|
| const float *new_xyz = new_xyz_tensor.data<float>(); |
| const float *xyz = xyz_tensor.data<float>(); |
| const int *new_xyz_batch_cnt = new_xyz_batch_cnt_tensor.data<int>(); |
| const int *xyz_batch_cnt = xyz_batch_cnt_tensor.data<int>(); |
| int *idx = idx_tensor.data<int>(); |
|
|
| ball_query_kernel_launcher_stack(B, M, radius, nsample, new_xyz, new_xyz_batch_cnt, xyz, xyz_batch_cnt, idx); |
| return 1; |
| } |
