tkonolige commented on a change in pull request #6580:
URL: https://github.com/apache/incubator-tvm/pull/6580#discussion_r498365231
##########
File path: tests/python/topi/python/test_topi_sparse.py
##########
@@ -342,38 +343,29 @@ def verify_sparse_dense_bsr(M, N, K, BS_R, BS_C, density,
use_relu):
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))
- def check_device(device):
- ctx = tvm.context(device, 0)
- if not tvm.testing.device_enabled(device):
- print("Skip because %s is not enabled" % device)
- return
- print("Running on target: %s" % device)
- fcompute, fschedule = tvm.topi.testing.dispatch(device,
_sparse_dense_implement)
- with tvm.target.Target(device):
- Y = fcompute(X, W_data, W_indices, W_indptr)
- 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),
ctx=ctx)
- func(
- tvm.nd.array(X_np, ctx=ctx),
- tvm.nd.array(W_sp_np.data, ctx=ctx),
- tvm.nd.array(W_sp_np.indices, ctx=ctx),
- tvm.nd.array(W_sp_np.indptr, ctx=ctx),
- Y_tvm,
- )
- tvm.testing.assert_allclose(Y_tvm.asnumpy(), Y_np, atol=1e-4,
rtol=1e-4)
-
- for device in ["llvm", "cuda"]:
- check_device(device)
+ fcompute, fschedule = tvm.topi.testing.dispatch(target,
_sparse_dense_implement)
+ with tvm.target.Target(target):
+ Y = fcompute(X, W_data, W_indices, W_indptr)
+ 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), ctx=ctx)
+ func(
+ tvm.nd.array(X_np, ctx=ctx),
+ tvm.nd.array(W_sp_np.data, ctx=ctx),
+ tvm.nd.array(W_sp_np.indices, ctx=ctx),
+ tvm.nd.array(W_sp_np.indptr, ctx=ctx),
+ Y_tvm,
+ )
+ tvm.testing.assert_allclose(Y_tvm.asnumpy(), Y_np, atol=1e-4,
rtol=1e-4)
[email protected]_gpu
-def test_sparse_dense_bsr():
[email protected]_targets("llvm", "cuda")
+def test_sparse_dense_bsr_relu(ctx, target):
Review comment:
This test is already enabled.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
