Some faster/better kernels. Tensors with points are now kept in (b, c, ...) format as this is easier for pytorch

utils/cinclude/cuda_utils.h

@@ -3,10 +3,21 @@
 
 #include <cmath>
 
+#define TOTAL_THREADS 512
+
 inline int opt_n_threads(int work_size) {
     const int pow_2 = std::log(static_cast<double>(work_size)) / std::log(2.0);
 
-    return max(min(1 << pow_2, 512), 32);
+    return max(min(1 << pow_2, TOTAL_THREADS), 1);
+}
+
+inline dim3 opt_block_config(int x, int y) {
+    const int x_threads = opt_n_threads(x);
+    const int y_threads =
+        max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1);
+    dim3 block_config(x_threads, y_threads, 1);
+
+    return block_config;
 }
 
 #endif

utils/cinclude/group_points_gpu.h

@@ -5,11 +5,11 @@
 extern "C" {
 #endif
 
-void group_points_kernel_wrapper(int b, int n, int c, int npoints, int nsample,
+void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample,
 				 const float *points, const int *idx,
 				 float *out, cudaStream_t stream);
 
-void group_points_grad_kernel_wrapper(int b, int n, int c, int npoints,
+void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
 				      int nsample, const float *grad_out,
 				      const int *idx, float *grad_points,
 				      cudaStream_t stream);

utils/cinclude/group_points_wrapper.h

@@ -1,8 +1,7 @@
-
-int group_points_wrapper(int b, int n, int c, int npoints, int nsample,
+int group_points_wrapper(int b, int c, int n, int npoints, int nsample,
 			 THCudaTensor *points_tensor,
 			 THCudaIntTensor *idx_tensor, THCudaTensor *out);
-int group_points_grad_wrapper(int b, int n, int c, int npoints, int nsample,
+int group_points_grad_wrapper(int b, int c, int n, int npoints, int nsample,
 			      THCudaTensor *grad_out_tensor,
 			      THCudaIntTensor *idx_tensor,
 			      THCudaTensor *grad_points_tensor);

utils/cinclude/interpolate_gpu.h

@@ -9,12 +9,12 @@ void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown,
 			     const float *known, float *dist2, int *idx,
 			     cudaStream_t stream);
 
-void three_interpolate_kernel_wrapper(int b, int m, int c, int n,
+void three_interpolate_kernel_wrapper(int b, int c, int m, int n,
 				      const float *points, const int *idx,
 				      const float *weight, float *out,
 				      cudaStream_t stream);
 
-void three_interpolate_grad_kernel_wrapper(int b, int n, int c, int m,
+void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m,
 					   const float *grad_out,
 					   const int *idx, const float *weight,
 					   float *grad_points,

utils/cinclude/interpolate_wrapper.h

@@ -3,13 +3,13 @@
 void three_nn_wrapper(int b, int n, int m, THCudaTensor *unknown_tensor,
 		      THCudaTensor *known_tensor, THCudaTensor *dist2_tensor,
 		      THCudaIntTensor *idx_tensor);
-void three_interpolate_wrapper(int b, int m, int c, int n,
+void three_interpolate_wrapper(int b, int c, int m, int n,
 			       THCudaTensor *points_tensor,
 			       THCudaIntTensor *idx_tensor,
 			       THCudaTensor *weight_tensor,
 			       THCudaTensor *out_tensor);
 
-void three_interpolate_grad_wrapper(int b, int n, int c, int m,
+void three_interpolate_grad_wrapper(int b, int c, int n, int m,
 				    THCudaTensor *grad_out_tensor,
 				    THCudaIntTensor *idx_tensor,
 				    THCudaTensor *weight_tensor,

utils/cinclude/sampling_gpu.h

@@ -5,10 +5,14 @@
 extern "C" {
 #endif
 
-void gather_points_kernel_wrapper(int b, int n, int c, int npoints,
+void gather_points_kernel_wrapper(int b, int c, int n, int npoints,
 				  const float *points, const int *idx,
 				  float *out, cudaStream_t stream);
 
+void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
+				       const float *grad_out, const int *idx,
+				       float *grad_points, cudaStream_t stream);
+
 void furthest_point_sampling_kernel_wrapper(int b, int n, int m,
 					    const float *dataset, float *temp,
 					    int *idxs, cudaStream_t stream);

utils/cinclude/sampling_wrapper.h

@@ -1,8 +1,12 @@
 
-int gather_points_wrapper(int b, int n, int c, int npoints,
+int gather_points_wrapper(int b, int c, int n, int npoints,
 			  THCudaTensor *points_tensor,
 			  THCudaIntTensor *idx_tensor,
 			  THCudaTensor *out_tensor);
+int gather_points_grad_wrapper(int b, int c, int n, int npoints,
+			       THCudaTensor *grad_out_tensor,
+			       THCudaIntTensor *idx_tensor,
+			       THCudaTensor *grad_points_tensor);
 
 int furthest_point_sampling_wrapper(int b, int n, int m,
 				    THCudaTensor *points_tensor,