Fix furthest_sample_point (#1405)

b484abac · q.yao · GitHub · 5c25ae1a · b484abac
Unverified Commit b484abac authored Oct 15, 2021 by q.yao Committed by GitHub Oct 15, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 8 additions and 6 deletions

mmcv/ops/csrc/common/cuda/furthest_point_sample_cuda_kernel.cuh ...ps/csrc/common/cuda/furthest_point_sample_cuda_kernel.cuh +8 -6

No files found.
--- a/mmcv/ops/csrc/common/cuda/furthest_point_sample_cuda_kernel.cuh
+++ b/mmcv/ops/csrc/common/cuda/furthest_point_sample_cuda_kernel.cuh
@@ -67,10 +67,11 @@ __global__ void furthest_point_sampling_forward_cuda_kernel(
    dists_i[tid] = besti;
    __syncthreads();

+#pragma unroll
    for (int block_size_thres = 1024; block_size_thres >= 2;
-         block_size_thres /= 2) {
-      int tid_thres = block_size_thres / 2;
-      if (block_size >= block_size_thres) {
+         block_size_thres >>= 1) {
+      const int tid_thres = block_size_thres / 2;
+      if (block_size >= block_size_thres && tid < tid_thres) {
        __update(dists, dists_i, tid, tid + tid_thres);
      }
      __syncthreads();
@@ -133,10 +134,11 @@ __global__ void furthest_point_sampling_with_dist_forward_cuda_kernel(
    dists_i[tid] = besti;
    __syncthreads();

+#pragma unroll
    for (int block_size_thres = 1024; block_size_thres >= 2;
-         block_size_thres /= 2) {
-      int tid_thres = block_size_thres / 2;
-      if (block_size >= block_size_thres) {
+         block_size_thres >>= 1) {
+      const int tid_thres = block_size_thres / 2;
+      if (block_size >= block_size_thres && tid < tid_thres) {
        __update(dists, dists_i, tid, tid + tid_thres);
      }
      __syncthreads();