[8] Review comments handled

Author: ANSHUMAN TRIPATHY

python/tvm/relay/op/nn/nn.py

@@ -2000,17 +2000,17 @@ def sparse_dense(dense_mat, sparse_mat, sparse_lhs=False):
     a dense matrix and `sparse_mat` is a sparse (either BSR or CSR) namedtuple with
     fields `data`, `indices`, and `indptr`.
 
-    .. math::
-
-       if sparse_lhs=True
+    \if sparse_lhs=False:
+        .. math::
 
-        \mbox{sparse_dense}(dense_mat, sparse_mat)[m, n]
-        = \mbox{matmul}(D, \mbox{as_dense}(S)^T)[m, n]
+            \mbox{sparse_dense}(dense_mat, sparse_mat)[m, n]
+            = \mbox{matmul}(D, \mbox{as_dense}(S)^T)[m, n]
 
-       if sparse_lhs=False
+    \if sparse_lhs=True:
+        .. math::
 
-        \mbox{sparse_dense}(dense_mat, sparse_mat)[m, n]
-        = \mbox{matmul}(\mbox{as_dense}(S), (D)^T)[m, n]
+            \mbox{sparse_dense}(dense_mat, sparse_mat)[m, n]
+            = \mbox{matmul}(\mbox{as_dense}(S), (D)^T)[m, n]
 
     where `as_dense` returns dense equivalent of the given S(sparse matrix)
     while performing matmul with given D(dense matrix).

python/tvm/topi/cuda/sparse.py

@@ -363,7 +363,7 @@ def _alter_sparse_dense_layout(_attrs, inputs, _tinfos, _out_type):
     sparse_dense implementation for one that operates on a padded matrix. We
     also padd the matrix.
     """
-    # TODO(ANSHUMAN87): Handle for sparse_data case too
+    # TODO(ANSHUMAN87): Handle for sparse_lhs case too
     if (
         isinstance(inputs[1], relay.Constant)
         and isinstance(inputs[2], relay.Constant)

python/tvm/topi/nn/sparse.py

@@ -166,7 +166,7 @@ def f(i, row):
 
 
 def _sparse_dense_bsrmm_v1(data_data, data_indices, data_indptr, weight):
-    (k, _) = get_const_tuple(weight.shape)
+    (m, _) = get_const_tuple(weight.shape)
     (_, bs_r, bs_c) = get_const_tuple(data_data.shape)
     (num_blocks_plus_1,) = get_const_tuple(data_indptr.shape)
     num_blocks = num_blocks_plus_1 - 1
@@ -187,11 +187,11 @@ def _compute_block(nb_j, j, i):
     idxm = tvm.tir.indexmod
 
     bsrmm_block = te.compute(
-        (num_blocks, bs_r, k), _compute_block, tag="sparse_dense_bsrmm_block_v1"
+        (num_blocks, bs_r, m), _compute_block, tag="sparse_dense_bsrmm_block_v1"
     )
     return te.compute(
-        (num_blocks * bs_r, k),
-        lambda m, n: bsrmm_block[idxd(n, bs_r), idxm(n, bs_r), m],
+        (num_blocks * bs_r, m),
+        lambda m, n: bsrmm_block[idxd(m, bs_r), idxm(m, bs_r), n],
         tag="sparse_dense_bsrmm_v1",
     )
 

tests/python/topi/python/test_topi_sparse.py

@@ -272,6 +272,31 @@ def test_sparse_dense_csr():
     tvm.testing.assert_allclose(Y_tvm.asnumpy(), Y_np, atol=1e-4, rtol=1e-4)
 
 
+def test_sparse_dense_csr_reverse():
+    M, N, K, density = 1, 17, 47, 0.2
+    X_np = np.random.randn(M, K).astype("float32")
+    W_sp_np = sp.random(N, K, density=density, format="csr", dtype="float32")
+    W_np = W_sp_np.todense()
+    Y_np = W_np.dot(X_np.T)
+
+    W_data = te.placeholder(shape=W_sp_np.data.shape, dtype=str(W_sp_np.data.dtype))
+    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))
+    Y = topi.nn.sparse_dense(X, W_data, W_indices, W_indptr, sparse_lhs=True)
+    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=Y_np.dtype))
+    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, atol=1e-4, rtol=1e-4)
+
+
 def test_sparse_transpose_csr():
     N, density = 1023, 0.3
 
@@ -368,6 +393,31 @@ def test_sparse_dense_bsr_relu(ctx, target):
     verify_sparse_dense_bsr(M, N, K, BS_R, BS_C, density, False, ctx, target)
 
 
+def test_sparse_dense_bsr_reverse():
+    M, N, K, BS_R, BS_C, density = 1, 64, 128, 8, 16, 0.9
+    X_np = np.random.randn(M, K).astype("float32")
+    W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
+    W_np = W_sp_np.todense()
+    Y_np = W_np.dot(X_np.T)
+
+    W_data = te.placeholder(shape=W_sp_np.data.shape, dtype=str(W_sp_np.data.dtype))
+    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))
+    Y = topi.nn.sparse_dense(X, W_data, W_indices, W_indptr, sparse_lhs=True)
+    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=Y_np.dtype))
+    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, atol=1e-4, rtol=1e-4)
+
+
 @tvm.testing.uses_gpu
 def test_sparse_dense_bsr_randomized():
     for _ in range(20):
@@ -480,3 +530,5 @@ def test_sparse_dense_padded_alter_op():
     test_sparse_transpose_csr()
     test_sparse_dense_padded_cuda()
     test_sparse_dense_padded_alter_op()
+    test_sparse_dense_csr_reverse()
+    test_sparse_dense_bsr_reverse()