[TOPI] Fix data race of batch multibox detection

python/tvm/topi/cuda/ssd/multibox.py

@@ -234,7 +234,7 @@ def transform_loc_pre(cls_prob, valid_count, temp_valid_count, temp_cls_id, temp
                     temp_valid_count[tid * num_anchors + k] += temp_valid_count[
                         tid * num_anchors + k - 1
                     ]
-            valid_count[i] = temp_valid_count[tid * num_anchors + num_anchors - 1]
+            valid_count[tid] = temp_valid_count[tid * num_anchors + num_anchors - 1]
 
     return ib.get()
 
@@ -342,7 +342,7 @@ def transform_loc(loc, loc_base_idx, anchor, anchor_base_idx, clip, vx, vy, vw,
         j = idxm(tid, num_anchors)
 
         with ib.if_scope(cls_id[tid] > 0):
-            with ib.if_scope(tid == 0):
+            with ib.if_scope(j == 0):
                 out_base_idx = i * num_anchors * 6
                 out_loc[out_base_idx] = cls_id[tid] - 1.0
                 out_loc[out_base_idx + 1] = score[tid]

python/tvm/topi/vision/ssd/multibox.py

@@ -137,17 +137,17 @@ def multibox_prior(data, sizes=(1,), ratios=(1,), steps=(-1, -1), offsets=(0.5,
 
 
 @hybrid.script
-def _hybridy_transform_loc(box, pred_loc, variance, clip):
+def _hybridy_transform_loc(anchor, pred_loc, variance, clip, batch_idx, anchor_idx):
     """Transform prior anchor box to output box through location predictions."""
-    al = box[0]
-    at = box[1]
-    ar = box[2]
-    ab = box[3]
+    al = anchor[0, anchor_idx, 0]
+    at = anchor[0, anchor_idx, 1]
+    ar = anchor[0, anchor_idx, 2]
+    ab = anchor[0, anchor_idx, 3]
 
-    px = pred_loc[0]
-    py = pred_loc[1]
-    pw = pred_loc[2]
-    ph = pred_loc[3]
+    px = pred_loc[batch_idx, 0]
+    py = pred_loc[batch_idx, 1]
+    pw = pred_loc[batch_idx, 2]
+    ph = pred_loc[batch_idx, 3]
 
     vx = variance[0]
     vy = variance[1]
@@ -206,8 +206,13 @@ def hybrid_multibox_transform_loc(cls_prob, loc_pred, anchor, clip, threshold, v
     batch_size = cls_prob.shape[0]
     num_classes = cls_prob.shape[1]
     num_anchors = cls_prob.shape[2]
-    box_coord = allocate((4,), loc_pred.dtype)
-    pred_coord = allocate((4,), loc_pred.dtype)
+    pred_coord = allocate(
+        (
+            batch_size,
+            4,
+        ),
+        loc_pred.dtype,
+    )
     out_loc = output_tensor((batch_size, num_anchors, 6), loc_pred.dtype)
     valid_count = output_tensor((batch_size,), "int32")
 
@@ -230,9 +235,8 @@ def hybrid_multibox_transform_loc(cls_prob, loc_pred, anchor, clip, threshold, v
                 out_loc[i, valid_count[i], 0] = cls_id - 1.0
                 out_loc[i, valid_count[i], 1] = score
                 for l in range(4):
-                    box_coord[l] = anchor[0, j, l]
-                    pred_coord[l] = loc_pred[i, j * 4 + l]
-                out_coord = _hybridy_transform_loc(box_coord, pred_coord, variances, clip)
+                    pred_coord[i, l] = loc_pred[i, j * 4 + l]
+                out_coord = _hybridy_transform_loc(anchor, pred_coord, variances, clip, i, j)
                 out_loc[i, valid_count[i], 2] = out_coord[0]
                 out_loc[i, valid_count[i], 3] = out_coord[1]
                 out_loc[i, valid_count[i], 4] = out_coord[2]

tests/python/topi/python/test_topi_vision.py

@@ -364,41 +364,57 @@ def test_multibox_prior(
         tvm.testing.assert_allclose(tvm_out.numpy(), np_out, rtol=1e-3)
 
 
-def test_multibox_detection(target, dev):
-    batch_size = 1
-    num_anchors = 3
-    num_classes = 3
-    cls_prob = te.placeholder((batch_size, num_anchors, num_classes), name="cls_prob")
-    loc_preds = te.placeholder((batch_size, num_anchors * 4), name="loc_preds")
-    anchors = te.placeholder((1, num_anchors, 4), name="anchors")
-
-    # Manually create test case
-    np_cls_prob = np.array([[[0.2, 0.5, 0.3], [0.25, 0.3, 0.45], [0.7, 0.1, 0.2]]])
-    np_loc_preds = np.array([[0.1, -0.2, 0.3, 0.2, 0.2, 0.4, 0.5, -0.3, 0.7, -0.2, -0.4, -0.8]])
-    np_anchors = np.array([[[-0.1, -0.1, 0.1, 0.1], [-0.2, -0.2, 0.2, 0.2], [1.2, 1.2, 1.5, 1.5]]])
-
-    expected_np_out = np.array(
-        [
+class TestMultiboxDetection:
+    (batch_size,) = tvm.testing.parameters((1,), (6,))
+
+    @tvm.testing.fixture(cache_return_value=True)
+    def ref_data(
+        self,
+        batch_size,
+    ):
+        # Manually create test case
+        np_cls_prob = np.array([[[0.2, 0.5, 0.3], [0.25, 0.3, 0.45], [0.7, 0.1, 0.2]]] * batch_size)
+        np_loc_preds = np.array(
+            [[0.1, -0.2, 0.3, 0.2, 0.2, 0.4, 0.5, -0.3, 0.7, -0.2, -0.4, -0.8]] * batch_size
+        )
+        np_anchors = np.array(
+            [[[-0.1, -0.1, 0.1, 0.1], [-0.2, -0.2, 0.2, 0.2], [1.2, 1.2, 1.5, 1.5]]] * batch_size
+        )
+        expected_np_out = np.array(
             [
-                [1, 0.69999999, 0, 0, 0.10818365, 0.10008108],
-                [0, 0.44999999, 1, 1, 1, 1],
-                [0, 0.30000001, 0, 0, 0.22903419, 0.20435292],
+                [
+                    [1, 0.69999999, 0, 0, 0.10818365, 0.10008108],
+                    [0, 0.44999999, 1, 1, 1, 1],
+                    [0, 0.30000001, 0, 0, 0.22903419, 0.20435292],
+                ]
             ]
-        ]
-    )
+            * batch_size
+        )
+        return np_cls_prob, np_loc_preds, np_anchors, expected_np_out
 
-    fcompute, fschedule = tvm.topi.testing.dispatch(target, _multibox_detection_implement)
-    with tvm.target.Target(target):
-        out = fcompute(cls_prob, loc_preds, anchors)
-        s = fschedule(out)
+    def test_multibox_detection(self, target, dev, ref_data):
+
+        np_cls_prob, np_loc_preds, np_anchors, expected_np_out = ref_data
+
+        batch_size = np_cls_prob.shape[0]
+        num_anchors = 3
+        num_classes = 3
+        cls_prob = te.placeholder((batch_size, num_anchors, num_classes), name="cls_prob")
+        loc_preds = te.placeholder((batch_size, num_anchors * 4), name="loc_preds")
+        anchors = te.placeholder((batch_size, num_anchors, 4), name="anchors")
+
+        fcompute, fschedule = tvm.topi.testing.dispatch(target, _multibox_detection_implement)
+        with tvm.target.Target(target):
+            out = fcompute(cls_prob, loc_preds, anchors)
+            s = fschedule(out)
 
-    tvm_cls_prob = tvm.nd.array(np_cls_prob.astype(cls_prob.dtype), dev)
-    tvm_loc_preds = tvm.nd.array(np_loc_preds.astype(loc_preds.dtype), dev)
-    tvm_anchors = tvm.nd.array(np_anchors.astype(anchors.dtype), dev)
-    tvm_out = tvm.nd.array(np.zeros((batch_size, num_anchors, 6)).astype(out.dtype), dev)
-    f = tvm.build(s, [cls_prob, loc_preds, anchors, out], target)
-    f(tvm_cls_prob, tvm_loc_preds, tvm_anchors, tvm_out)
-    tvm.testing.assert_allclose(tvm_out.numpy(), expected_np_out, rtol=1e-4)
+        tvm_cls_prob = tvm.nd.array(np_cls_prob.astype(cls_prob.dtype), dev)
+        tvm_loc_preds = tvm.nd.array(np_loc_preds.astype(loc_preds.dtype), dev)
+        tvm_anchors = tvm.nd.array(np_anchors.astype(anchors.dtype), dev)
+        tvm_out = tvm.nd.array(np.zeros((batch_size, num_anchors, 6)).astype(out.dtype), dev)
+        f = tvm.build(s, [cls_prob, loc_preds, anchors, out], target)
+        f(tvm_cls_prob, tvm_loc_preds, tvm_anchors, tvm_out)
+        tvm.testing.assert_allclose(tvm_out.numpy(), expected_np_out, rtol=1e-4)
 
 
 class TestRoiAlign: