temp disable write by only thread 0

neo-ai · Dec 28, 2020 · d75ee0a · d75ee0a
1 parent 20b5630
commit d75ee0a
Showing 1 changed file with 11 additions and 8 deletions.
diff --git a/python/tvm/topi/cuda/nms.py b/python/tvm/topi/cuda/nms.py
@@ -530,12 +530,15 @@ def calculate_overlap(out_tensor, box_a_idx, box_b_idx):
         def nms_inner_loop(ib, j):
             # the box j is valid, invalidate other boxes that overlap with j above iou_threshold
 
-            # When return_indices is False, no need to populate box_indices
-            if return_indices:
-                # Only one thread needs to this write
-                with ib.if_scope(tx == 0):
-                    orig_idx = sorted_index[i * num_anchors + j]
-                    box_indices[i, num_valid_boxes_local[0]] = indices[i, orig_idx]
+            # # When return_indices is False, no need to populate box_indices
+            # if return_indices:
+            #     # Only one thread needs to this write
+            #     with ib.if_scope(tx == 0):
+            #         orig_idx = sorted_index[i * num_anchors + j]
+            #         box_indices[i, num_valid_boxes_local[0]] = indices[i, orig_idx]
+
+            orig_idx = sorted_index[i * num_anchors + j]
+            box_indices[i, num_valid_boxes_local[0]] = indices[i, orig_idx]
 
             num_valid_boxes_local[0] += 1
 
@@ -593,8 +596,8 @@ def nms_inner_loop(ib, j):
                     with ib.else_scope():
                         nms_inner_loop(ib, j)
 
-            with ib.if_scope(tx == 0):
-                num_valid_boxes[i] = num_valid_boxes_local[0]
+            # with ib.if_scope(tx == 0):
+            num_valid_boxes[i] = num_valid_boxes_local[0]
 
         with ib.else_scope():
             num_valid_boxes[i] = 0