diff --git a/pcdet/ops/pointnet2/pointnet2_stack/src/vector_pool_gpu.h b/pcdet/ops/pointnet2/pointnet2_stack/src/vector_pool_gpu.h new file mode 100644 index 0000000..febfb85 --- /dev/null +++ b/pcdet/ops/pointnet2/pointnet2_stack/src/vector_pool_gpu.h @@ -0,0 +1,71 @@ +/* +Vector-pool aggregation based local feature aggregation for point cloud. +PV-RCNN++: Point-Voxel Feature Set Abstraction With Local Vector Representation for 3D Object Detection +https://arxiv.org/abs/2102.00463 + +Written by Shaoshuai Shi +All Rights Reserved 2020. +*/ + + +#ifndef _STACK_VECTOR_POOL_GPU_H +#define _STACK_VECTOR_POOL_GPU_H + +#include +#include +#include +#include + + +int query_stacked_local_neighbor_idxs_kernel_launcher_stack( + const float *support_xyz, const int *xyz_batch_cnt, const float *new_xyz, const int *new_xyz_batch_cnt, + int *stack_neighbor_idxs, int *start_len, int *cumsum, int avg_length_of_neighbor_idxs, + float max_neighbour_distance, int batch_size, int M, int nsample, int neighbor_type); + +int query_stacked_local_neighbor_idxs_wrapper_stack(at::Tensor support_xyz_tensor, at::Tensor xyz_batch_cnt_tensor, + at::Tensor new_xyz_tensor, at::Tensor new_xyz_batch_cnt_tensor, + at::Tensor stack_neighbor_idxs_tensor, at::Tensor start_len_tensor, at::Tensor cumsum_tensor, + int avg_length_of_neighbor_idxs, float max_neighbour_distance, int nsample, int neighbor_type); + + +int query_three_nn_by_stacked_local_idxs_kernel_launcher_stack( + const float *support_xyz, const float *new_xyz, const float *new_xyz_grid_centers, + int *new_xyz_grid_idxs, float *new_xyz_grid_dist2, + const int *stack_neighbor_idxs, const int *start_len, + int M, int num_total_grids); + +int query_three_nn_by_stacked_local_idxs_wrapper_stack(at::Tensor support_xyz_tensor, + at::Tensor new_xyz_tensor, at::Tensor new_xyz_grid_centers_tensor, + at::Tensor new_xyz_grid_idxs_tensor, at::Tensor new_xyz_grid_dist2_tensor, + at::Tensor stack_neighbor_idxs_tensor, at::Tensor start_len_tensor, + int M, int num_total_grids); + + +int vector_pool_wrapper_stack(at::Tensor support_xyz_tensor, at::Tensor xyz_batch_cnt_tensor, + at::Tensor support_features_tensor, at::Tensor new_xyz_tensor, at::Tensor new_xyz_batch_cnt_tensor, + at::Tensor new_features_tensor, at::Tensor new_local_xyz, + at::Tensor point_cnt_of_grid_tensor, at::Tensor grouped_idxs_tensor, + int num_grid_x, int num_grid_y, int num_grid_z, float max_neighbour_distance, int use_xyz, + int num_max_sum_points, int nsample, int neighbor_type, int pooling_type); + + +int vector_pool_kernel_launcher_stack( + const float *support_xyz, const float *support_features, const int *xyz_batch_cnt, + const float *new_xyz, float *new_features, float * new_local_xyz, const int *new_xyz_batch_cnt, + int *point_cnt_of_grid, int *grouped_idxs, + int num_grid_x, int num_grid_y, int num_grid_z, float max_neighbour_distance, + int batch_size, int N, int M, int num_c_in, int num_c_out, int num_total_grids, int use_xyz, + int num_max_sum_points, int nsample, int neighbor_type, int pooling_type); + + +int vector_pool_grad_wrapper_stack(at::Tensor grad_new_features_tensor, + at::Tensor point_cnt_of_grid_tensor, at::Tensor grouped_idxs_tensor, + at::Tensor grad_support_features_tensor); + + +void vector_pool_grad_kernel_launcher_stack( + const float *grad_new_features, const int *point_cnt_of_grid, const int *grouped_idxs, + float *grad_support_features, int N, int M, int num_c_out, int num_c_in, int num_total_grids, + int num_max_sum_points); + +#endif