Support thread number for both modes

Author: nicolas-chaulet

cuda/include/cuda_utils.h

@@ -10,18 +10,19 @@
 
 #include <vector>
 
-#define TOTAL_THREADS 512
+#define TOTAL_THREADS_DENSE 512
+#define TOTAL_THREADS_SPARSE 1024
 
 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, TOTAL_THREADS), 1);
+  return max(min(1 << pow_2, TOTAL_THREADS_DENSE), 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);
+      max(min(opt_n_threads(y), TOTAL_THREADS_DENSE / x_threads), 1);
   dim3 block_config(x_threads, y_threads, 1);
 
   return block_config;

cuda/src/ball_query_gpu.cu

@@ -66,7 +66,7 @@ __global__ void query_ball_point_kernel_partial_dense(int size_x,
 	const ptrdiff_t end_idx_y = batch_y[batch_idx + 1];
 	float radius2 = radius * radius;
 
-	for (ptrdiff_t n_x = start_idx_x + idx; n_x < end_idx_x; n_x += TOTAL_THREADS) {
+	for (ptrdiff_t n_x = start_idx_x + idx; n_x < end_idx_x; n_x += TOTAL_THREADS_SPARSE) {
 		int64_t count = 0;
 		for (ptrdiff_t n_y = start_idx_y; n_y < end_idx_y; n_y++) {
 			float dist = 0;
@@ -108,7 +108,7 @@ void query_ball_point_kernel_partial_wrapper(long batch_size,
 						int64_t *idx_out,
 						float *dist_out) {
 
-	query_ball_point_kernel_partial_dense<<<batch_size, TOTAL_THREADS>>>(
+	query_ball_point_kernel_partial_dense<<<batch_size, TOTAL_THREADS_SPARSE>>>(
 		size_x, size_y, radius, nsample, x, y,
 		batch_x, batch_y, idx_out, dist_out);