fix bug

tests/python/topi/python/test_topi_sparse.py

@@ -34,8 +34,6 @@
     "x86": (topi.nn.sparse_dense, topi.x86.schedule_sparse_dense),
 }
 
-_sparse_conv2d_implement = {"generic": (topi.nn.sparse_conv2d, topi.generic.schedule_sparse_conv2d)}
-
 
 def verify_dynamic_csrmv(batch, in_dim, out_dim, use_bias=True):
     nr, nc, n = te.var("nr"), te.var("nc"), te.var("n")
@@ -554,52 +552,47 @@ def test_sparse_add_csr():
             tvm.testing.assert_allclose(Z_tvm.asnumpy(), Z_np, atol=1e-4, rtol=1e-4)
 
 
-def verify_sparse_conv2d_bsr(M, H, W, N, K, BS_R, BS_C, density, use_relu, device, target, layout):
+def test_sparse_conv2d_bsr(M, H, W, N, K, BS_R, BS_C, density, layout):
     if layout == "NHWC":
         X_np = np.random.randn(M, H, W, K).astype("float32")
     elif layout == "NCHW":
         X_np = np.random.randn(M, K, H, W).astype("float32")
     W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
     W_np = W_sp_np.todense()
     if layout == "NHWC":
-        Y_np = tvm.topi.testing.conv2d_nhwc_python(X_np, W_np.reshape(1, 1, K, N), 1, 0)
+        Y_np = tvm.topi.testing.conv2d_nhwc_python(X_np, np.array(W_np).T.reshape(1, 1, K, N), 1, 0)
     elif layout == "NCHW":
-        Y_np = tvm.topi.testing.conv2d_nchw_python(X_np, W_np.reshape(N, K, 1, 1), 1, 0)
-    if use_relu:
-        Y_np = np.maximum(Y_np, 0.0)
+        Y_np = tvm.topi.testing.conv2d_nchw_python(X_np, np.array(W_np).reshape(N, K, 1, 1), 1, 0)
 
     if BS_C == 1:
         W_data = te.placeholder(shape=W_sp_np.data.shape[:-1], dtype=str(W_sp_np.data.dtype))
+        W_sp_np_data = W_sp_np.data.reshape(W_sp_np.data.shape[0], BS_R)
     else:
         W_data = te.placeholder(shape=W_sp_np.data.shape, dtype=str(W_sp_np.data.dtype))
+        W_sp_np_data = W_sp_np.data
     W_indices = te.placeholder(shape=W_sp_np.indices.shape, dtype=str(W_sp_np.indices.dtype))
     W_indptr = te.placeholder(shape=W_sp_np.indptr.shape, dtype=str(W_sp_np.indptr.dtype))
     X = te.placeholder(shape=X_np.shape, dtype=str(X_np.dtype))
 
-    fcompute, fschedule = tvm.topi.testing.dispatch(target, _sparse_conv2d_implement)
-    with tvm.target.Target(target):
-        Y = fcompute(X, W_data, W_indices, W_indptr, layout)
-        if use_relu:
-            Y = topi.nn.relu(Y)
-        s = fschedule([Y])
-        func = tvm.build(s, [X, W_data, W_indices, W_indptr, Y])
-        Y_tvm = tvm.nd.array(np.zeros(Y_np.shape, dtype=Y_np.dtype), device=device)
-        func(
-            tvm.nd.array(X_np, device=device),
-            tvm.nd.array(W_sp_np.data, device=device),
-            tvm.nd.array(W_sp_np.indices, device=device),
-            tvm.nd.array(W_sp_np.indptr, device=device),
-            Y_tvm,
-        )
-        tvm.testing.assert_allclose(Y_tvm.asnumpy(), Y_np, atol=1e-4, rtol=1e-4)
+    Y = topi.nn.sparse_conv2d(X, W_data, W_indices, W_indptr, layout)
+    s = te.create_schedule(Y.op)
+    func = tvm.build(s, [X, W_data, W_indices, W_indptr, Y])
+    Y_tvm = tvm.nd.array(np.zeros(Y_np.shape, dtype="float32"))
+    func(
+        tvm.nd.array(X_np),
+        tvm.nd.array(W_sp_np_data),
+        tvm.nd.array(W_sp_np.indices),
+        tvm.nd.array(W_sp_np.indptr),
+        Y_tvm,
+    )
+    tvm.testing.assert_allclose(Y_tvm.asnumpy(), Y_np.astype("float32"), atol=1e-4, rtol=1e-4)
 
 
-@tvm.testing.parametrize_targets("llvm", "cuda")
-def test_sparse_conv2d_bsr_relu(dev, target):
-    M, H, W, N, K, BS_R, BS_C, density = 1, 32, 32, 64, 128, 8, 16, 0.9
-    verify_sparse_conv2d_bsr(M, H, W, N, K, BS_R, BS_C, density, True, dev, target, "NHWC")
-    verify_sparse_conv2d_bsr(M, H, W, N, K, BS_R, BS_C, density, False, dev, target, "NCHW")
-    verify_sparse_conv2d_bsr(M, H, W, N, K, BS_R, 1, density, False, dev, target, "NHWC")
+def test_sparse_conv2d():
+    M, H, W, N, K, BS_R, BS_C, density = 1, 32, 32, 128, 64, 8, 16, 0.9
+    test_sparse_conv2d_bsr(M, H, W, N, K, BS_R, BS_C, density, "NHWC")
+    test_sparse_conv2d_bsr(M, H, W, N, K, BS_R, BS_C, density, "NCHW")
+    test_sparse_conv2d_bsr(M, H, W, N, K, BS_R, 1, density, "NHWC")
 
 
 if __name__ == "__main__":
@@ -614,4 +607,4 @@ def test_sparse_conv2d_bsr_relu(dev, target):
     test_sparse_dense_csr_reverse()
     test_sparse_dense_bsr_reverse()
     test_sparse_add_csr()
-    test_sparse_conv2d_bsr_relu()
+    test_sparse_conv2d()