You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2020/11/16 18:00:42 UTC
[GitHub] [incubator-tvm] tkonolige commented on a change in pull request #6889: [TOPI] sparse_dense Op sparse_data input added
tkonolige commented on a change in pull request #6889:
URL: https://github.com/apache/incubator-tvm/pull/6889#discussion_r524456070
##########
File path: src/relay/op/nn/sparse.cc
##########
@@ -85,10 +117,11 @@ RELAY_REGISTER_OP("nn.sparse_dense")
)code" TVM_ADD_FILELINE)
.set_attrs_type<SparseDenseAttrs>()
.set_num_inputs(4)
- .add_argument("data", "nD Tensor", "Input data.")
- .add_argument("weight_data", "1D Tensor", "Weight data matrix.")
- .add_argument("weight_indices", "1D Tensor", "Weight indices matrix.")
- .add_argument("weight_indptr", "1D Tensor", "Weight indptr matrix.")
+ .add_argument("input_tensor1", "nD Tensor",
Review comment:
Please switch these back to their original names (data, weight_data, etc) or something descriptive.
##########
File path: python/tvm/topi/nn/sparse.py
##########
@@ -52,13 +52,120 @@ def sparse_dense(data, weight_data, weight_indices, weight_indptr):
"""
assert len(weight_data.shape) in (1, 3)
if len(weight_data.shape) == 1:
- func = _sparse_dense_csrmm
+ func = _sparse_dense_csrmm_v2
if len(weight_data.shape) == 3:
- func = _sparse_dense_bsrmm
+ func = _sparse_dense_bsrmm_v2
return func(data, weight_data, weight_indices, weight_indptr)
-def _sparse_dense_csrmm(data, weight_data, weight_indices, weight_indptr):
+def sparse_dense_v1(data_data, data_indices, data_indptr, weight):
+ """
+ Computes sparse-dense matrix multiplication of
+ `(data_data, data_indices, data_indptr)` and `weight.T`
+
+ Parameters
+ ----------
+ data_data:
+ 1-D with shape [nnz] (CSR) or
+ 3-D with shape [num_blocks, bs_r, bs_c] (BSR)
+
+ data_indices:
+ 1-D with shape [nnz] (CSR) or
+ 1-D with shape [num_blocks] (BSR)
+
+ data_indptr:
+ 1-D with shape [M + 1] (CSR) or
+ 1-D with shape [(M + 1) // bs_r] (BSR)
+
+ weight:
+ 2-D with shape [N, K], float32 when weight is dense
+
+ Returns
+ -------
+ output : tvm.te.Tensor
+ 2-D with shape [M, N]
+ """
+ assert len(data_data.shape) in (1, 3)
+ if len(data_data.shape) == 1:
+ func = _sparse_dense_csrmm_v1
+ if len(data_data.shape) == 3:
+ func = _sparse_dense_bsrmm_v1
+ return func(data_data, data_indices, data_indptr, weight)
+
+
+# pylint: disable=no-else-return,inconsistent-return-statements
+def sparse_dense(input_1, input_2, input_3, input_4, sparse_data=False):
+ """
+ Computes sparse-dense matrix multiplication of `data` and
+ `(weight_data, weight_indices, weight_indptr).T`
+ or
+ Computes sparse-dense matrix multiplication of
+ `(data_data, data_indices, data_indptr)` and `weight.T`
+
+ Parameters
+ ----------
+ input_1 : tvm.te.Tensor
Review comment:
Please name these something more descriptive. Maybe sparse_data, sparse_indices, sparse_indptr?
##########
File path: python/tvm/relay/op/nn/nn.py
##########
@@ -1993,7 +1993,8 @@ def batch_matmul(x, y):
return _make.batch_matmul(x, y)
-def sparse_dense(data, weight):
+# pylint: disable=no-else-return,inconsistent-return-statements
+def sparse_dense(data, weight, sparse_data=False):
Review comment:
Please document `sparse_data` here. Also update the doc comment with the new input types.
##########
File path: python/tvm/topi/nn/sparse.py
##########
@@ -52,13 +52,120 @@ def sparse_dense(data, weight_data, weight_indices, weight_indptr):
"""
assert len(weight_data.shape) in (1, 3)
if len(weight_data.shape) == 1:
- func = _sparse_dense_csrmm
+ func = _sparse_dense_csrmm_v2
if len(weight_data.shape) == 3:
- func = _sparse_dense_bsrmm
+ func = _sparse_dense_bsrmm_v2
return func(data, weight_data, weight_indices, weight_indptr)
-def _sparse_dense_csrmm(data, weight_data, weight_indices, weight_indptr):
+def sparse_dense_v1(data_data, data_indices, data_indptr, weight):
+ """
+ Computes sparse-dense matrix multiplication of
+ `(data_data, data_indices, data_indptr)` and `weight.T`
+
+ Parameters
+ ----------
+ data_data:
+ 1-D with shape [nnz] (CSR) or
+ 3-D with shape [num_blocks, bs_r, bs_c] (BSR)
+
+ data_indices:
+ 1-D with shape [nnz] (CSR) or
+ 1-D with shape [num_blocks] (BSR)
+
+ data_indptr:
+ 1-D with shape [M + 1] (CSR) or
+ 1-D with shape [(M + 1) // bs_r] (BSR)
+
+ weight:
+ 2-D with shape [N, K], float32 when weight is dense
+
+ Returns
+ -------
+ output : tvm.te.Tensor
+ 2-D with shape [M, N]
+ """
+ assert len(data_data.shape) in (1, 3)
+ if len(data_data.shape) == 1:
+ func = _sparse_dense_csrmm_v1
+ if len(data_data.shape) == 3:
+ func = _sparse_dense_bsrmm_v1
+ return func(data_data, data_indices, data_indptr, weight)
+
+
+# pylint: disable=no-else-return,inconsistent-return-statements
+def sparse_dense(input_1, input_2, input_3, input_4, sparse_data=False):
Review comment:
Please document `sparse_data`.
----------------------------------------------------------------
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:
users@infra.apache.org