From c11d23de43c969642c20be5eb9be06b7342e9f85 Mon Sep 17 00:00:00 2001 From: ANSHUMAN TRIPATHY Date: Mon, 9 Nov 2020 21:56:28 +0530 Subject: [PATCH 1/9] [TOPI] sparse_dense op sparse_data input added --- include/tvm/relay/attrs/nn.h | 6 +- python/tvm/relay/frontend/tensorflow.py | 27 +++- python/tvm/relay/op/nn/_nn.py | 2 +- python/tvm/relay/op/nn/nn.py | 11 +- python/tvm/relay/op/strategy/generic.py | 2 +- python/tvm/topi/cuda/sparse.py | 6 +- python/tvm/topi/nn/sparse.py | 146 ++++++++++++++++-- src/relay/op/nn/sparse.cc | 45 +++++- .../frontend/tensorflow/test_forward.py | 12 +- 9 files changed, 224 insertions(+), 33 deletions(-) diff --git a/include/tvm/relay/attrs/nn.h b/include/tvm/relay/attrs/nn.h index b2555de6d35e..08dfc9368a68 100644 --- a/include/tvm/relay/attrs/nn.h +++ b/include/tvm/relay/attrs/nn.h @@ -937,7 +937,11 @@ struct DenseAttrs : public tvm::AttrsNode { /*! \brief Attributes for sparse_dense operator */ struct SparseDenseAttrs : public tvm::AttrsNode { - TVM_DECLARE_ATTRS(SparseDenseAttrs, "relay.attrs.SparseDenseAttrs") {} + bool sparse_data; + + TVM_DECLARE_ATTRS(SparseDenseAttrs, "relay.attrs.SparseDenseAttrs") { + TVM_ATTR_FIELD(sparse_data).set_default(false).describe("Indicate whether data or weight is sparse. True if data is sparse"); + } }; /*! \brief Attributes for sparse_transpose operator */ diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index c6079b4535c4..4656a3f8a470 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -916,6 +916,13 @@ def _impl(inputs, attr, params, mod): data = inputs[3] + # By default, in tensorflow the first input ,i.e., data is sparse + sparse_data = True + + # If both are true means First input was dense and second was sparse + if attr.get("adjoint_a") and attr.get("adjoint_b"): + sparse_data = False + rows = [x[0] for x in indices_tensor] cols = [x[1] for x in indices_tensor] @@ -923,21 +930,29 @@ def _impl(inputs, attr, params, mod): weight_sp = csr_matrix( (values_tensor, (rows, cols)), shape=tuple(dense_shape_tensor.tolist()) ) - weight_sp = csr_matrix(weight_sp.transpose()) + + if sparse_data: + data = _op.transpose(data) + else: + weight_sp = csr_matrix(weight_sp.transpose()) weight_data = _expr.const(weight_sp.data, weight_sp.data.dtype) weight_indptrs = _expr.const(weight_sp.indptr, weight_sp.indptr.dtype) weight_indices = _expr.const(weight_sp.indices, weight_sp.indices.dtype) - ret = _op.nn.sparse_dense(data, [weight_data, weight_indices, weight_indptrs]) + if sparse_data: + ret = _op.nn.sparse_dense([weight_data, weight_indices, weight_indptrs], data, sparse_data=True) + else: + ret = _op.nn.sparse_dense(data, [weight_data, weight_indices, weight_indptrs]) + ret = _op.transpose(ret) - # If both are true means First input was dense and second was sparse + # Case 1. If both are true means first input was dense and second was sparse + # Case 2. If both are false means first input was sparse and second was dense # TODO(ANSHUMAN87): Support other adjoint option too - if attr.get("adjoint_a") and attr.get("adjoint_b"): - ret = _op.transpose(ret) - else: + if not ((attr.get("adjoint_a") and attr.get("adjoint_b")) or ((not attr.get("adjoint_a")) and (not attr.get("adjoint_b")))): raise tvm.error.OpAttributeUnImplemented( "Only tf.sparse.sparse_dense_matmul() with adjoint_a=True and adjoint_b=True" + "or with adjoint_a=False and adjoint_b=False" " is supported, but adjoint_a={} and adjoint_b={} was supplied.".format( attr.get("adjoint_a"), attr.get("adjoint_b") ) diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py index c235f87d1e99..f92808f4cf0c 100644 --- a/python/tvm/relay/op/nn/_nn.py +++ b/python/tvm/relay/op/nn/_nn.py @@ -69,7 +69,7 @@ def compute_fifo_buffer(attrs, inputs, out_type): @reg.register_compute("nn.sparse_dense") def compute_sparse_dense(attrs, inputs, out_type): """Compute definition of sparse_dense""" - return [topi.nn.sparse_dense(inputs[0], inputs[1], inputs[2], inputs[3])] + return [topi.nn.sparse_dense(inputs[0], inputs[1], inputs[2], inputs[3], attrs["sparse_data"])] reg.register_strategy("nn.sparse_dense", strategy.sparse_dense_strategy) diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py index 4810bdc35bbd..6d2de1e47be2 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/python/tvm/relay/op/nn/nn.py @@ -1993,7 +1993,7 @@ def batch_matmul(x, y): return _make.batch_matmul(x, y) -def sparse_dense(data, weight): +def sparse_dense(data, weight, sparse_data=False): r""" Computes the matrix multiplication of `data` and `weight`, where `data` is a dense matrix and `weight` is a sparse (either BSR or CSR) namedtuple with @@ -2025,8 +2025,13 @@ def sparse_dense(data, weight): The computed result. """ if hasattr(weight, "indices"): - return _make.sparse_dense(data, weight.data, weight.indices, weight.indptr) - return _make.sparse_dense(data, weight[0], weight[1], weight[2]) + return _make.sparse_dense(data, weight.data, weight.indices, weight.indptr, sparse_data) + elif isinstance(weight, (tuple, list)): + return _make.sparse_dense(data, weight[0], weight[1], weight[2], sparse_data) + elif hasattr(data, "indices"): + return _make.sparse_dense(data.data, data.indices, data.indptr, weight, sparse_data) + else: + return _make.sparse_dense(data[0], data[1], data[2], weight, sparse_data) def sparse_transpose(x): diff --git a/python/tvm/relay/op/strategy/generic.py b/python/tvm/relay/op/strategy/generic.py index bdefbcb79009..ce747f3b674a 100644 --- a/python/tvm/relay/op/strategy/generic.py +++ b/python/tvm/relay/op/strategy/generic.py @@ -706,7 +706,7 @@ def wrap_compute_sparse_dense(topi_compute): """wrap sparse dense topi compute""" def _compute_sparse_dense(attrs, inputs, out_type): - return [topi_compute(inputs[0], inputs[1], inputs[2], inputs[3])] + return [topi_compute(inputs[0], inputs[1], inputs[2], inputs[3], attrs["sparse_data"])] return _compute_sparse_dense diff --git a/python/tvm/topi/cuda/sparse.py b/python/tvm/topi/cuda/sparse.py index ebac5517d46c..8f1a26869c07 100644 --- a/python/tvm/topi/cuda/sparse.py +++ b/python/tvm/topi/cuda/sparse.py @@ -65,10 +65,11 @@ def schedule_sparse_dense(outs): # pylint:disable=invalid-name s = te.create_schedule([x.op for x in outs]) + # TODO(ANSHUMAN87): Add for sparse_dense_bsrmm_v1 also def _callback(op): - if op.tag == "sparse_dense_bsrmm": + if op.tag == "sparse_dense_bsrmm_v2": y_bsrmm = op.input_tensors[0] - assert y_bsrmm.op.tag == "sparse_dense_bsrmm_block" + assert y_bsrmm.op.tag == "sparse_dense_bsrmm_block_v2" out = s.outputs[0].output(0) if op not in s.outputs: @@ -362,6 +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 if ( isinstance(inputs[1], relay.Constant) and isinstance(inputs[2], relay.Constant) diff --git a/python/tvm/topi/nn/sparse.py b/python/tvm/topi/nn/sparse.py index 55b3e6a7d1e5..b57a347ba684 100644 --- a/python/tvm/topi/nn/sparse.py +++ b/python/tvm/topi/nn/sparse.py @@ -23,7 +23,7 @@ from ..utils import get_const_tuple -def sparse_dense(data, weight_data, weight_indices, weight_indptr): +def sparse_dense_v2(data, weight_data, weight_indices, weight_indptr): """ Computes sparse-dense matrix multiplication of `data` and `(weight_data, weight_indices, weight_indptr).T` @@ -52,13 +52,115 @@ 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_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) + +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 + data: + 2-D with shape [M, K], float32 when input data is dense or + + data_data: + 1-D with shape [nnz] (CSR) or + 3-D with shape [num_blocks, bs_r, bs_c] (BSR) + + input_2 : tvm.te.Tensor + weight_data: + 1-D with shape [nnz] (CSR) or + 3-D with shape [num_blocks, bs_r, bs_c] (BSR) or + + data_indices: + 1-D with shape [nnz] (CSR) or + 1-D with shape [num_blocks] (BSR) + + input_3 : tvm.te.Tensor + weight_indices: + 1-D with shape [nnz] (CSR) or + 1-D with shape [num_blocks] (BSR) or -def _sparse_dense_csrmm(data, weight_data, weight_indices, weight_indptr): + data_indptr: + 1-D with shape [M + 1] (CSR) or + 1-D with shape [(M + 1) // bs_r] (BSR) + + input_4 : tvm.te.Tensor + weight_indptr: + 1-D with shape [N + 1] (CSR) or + 1-D with shape [(N + 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] + """ + if sparse_data: + return sparse_dense_v1(input_1, input_2, input_3, input_4) + else: + return sparse_dense_v2(input_1, input_2, input_3, input_4) + +def _sparse_dense_csrmm_v1(data_data, data_indices, data_indptr, weight): + oshape = (get_const_tuple(data_indptr.shape)[0] - 1, get_const_tuple(weight.shape)[0]) + + def f(row, i): + row_start = data_indptr[row] + row_end = data_indptr[row + 1] + row_elems = row_end - row_start + elem_idx = te.reduce_axis((0, row_elems), name="elem_idx") + elem = row_start + elem_idx + a_val = data_data[elem] + weight_val = weight[i, data_indices[elem]] + return te.sum(a_val * weight_val, axis=elem_idx) + + return te.compute(oshape, f, tag="sparse_dense_csrmm_v1") + +def _sparse_dense_csrmm_v2(data, weight_data, weight_indices, weight_indptr): oshape = (get_const_tuple(data.shape)[0], get_const_tuple(weight_indptr.shape)[0] - 1) def f(i, row): @@ -71,10 +173,37 @@ def f(i, row): weight_val = data[i, weight_indices[elem]] return te.sum(a_val * weight_val, axis=elem_idx) - return te.compute(oshape, f, tag="sparse_dense_csrmm") + return te.compute(oshape, f, tag="sparse_dense_csrmm_v2") + +def _sparse_dense_bsrmm_v1(data_data, data_indices, data_indptr, weight): + (k, m) = get_const_tuple(weight.shape) + (l, 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 + + def _compute_block(nb_j, j, i): + row_start = data_indptr[nb_j] + row_end = data_indptr[nb_j + 1] + row_elems = row_end - row_start + elem_idx = te.reduce_axis((0, row_elems), name="elem_idx") + block_offset = row_start + elem_idx + c = te.reduce_axis((0, bs_c), name="c") + block_j = data_indices[block_offset] + block_ij_val = data_data[block_offset][j][c] + x_val = weight[i, bs_c * block_j + c] + return te.sum(block_ij_val * x_val, axis=[elem_idx, c]) + + idxd = tvm.tir.indexdiv + idxm = tvm.tir.indexmod + bsrmm_block = te.compute((num_blocks, bs_r, k), _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], + tag="sparse_dense_bsrmm_v1", + ) -def _sparse_dense_bsrmm(data, weight_data, weight_indices, weight_indptr): +def _sparse_dense_bsrmm_v2(data, weight_data, weight_indices, weight_indptr): (m, _) = get_const_tuple(data.shape) (_, bs_r, bs_c) = get_const_tuple(weight_data.shape) (num_blocks_plus_1,) = get_const_tuple(weight_indptr.shape) @@ -95,14 +224,13 @@ def _compute_block(i, nb_j, j): idxd = tvm.tir.indexdiv idxm = tvm.tir.indexmod - bsrmm_block = te.compute((m, num_blocks, bs_r), _compute_block, tag="sparse_dense_bsrmm_block") + bsrmm_block = te.compute((m, num_blocks, bs_r), _compute_block, tag="sparse_dense_bsrmm_block_v2") return te.compute( (m, num_blocks * bs_r), lambda m, n: bsrmm_block[m, idxd(n, bs_r), idxm(n, bs_r)], - tag="sparse_dense_bsrmm", + tag="sparse_dense_bsrmm_v2", ) - def sparse_transpose(sparse_data, sparse_indices, sparse_indptr): """ Transpose a square sparse matrix, diff --git a/src/relay/op/nn/sparse.cc b/src/relay/op/nn/sparse.cc index 09dca09a82de..7456668a3e09 100644 --- a/src/relay/op/nn/sparse.cc +++ b/src/relay/op/nn/sparse.cc @@ -39,6 +39,35 @@ TVM_REGISTER_NODE_TYPE(SparseDenseAttrs); bool SparseDenseRel(const Array& types, int num_inputs, const Attrs& attrs, const TypeReporter& reporter) { ICHECK_EQ(types.size(), 5); + const auto* param = attrs.as(); + ICHECK(param != nullptr); + + if (param->sparse_data) { + const auto* weight = types[3].as(); + const auto* data_data = types[0].as(); + ICHECK(data_data->shape.size() == 1 || data_data->shape.size() == 3); + const auto* data_indptr = types[2].as(); + if (weight == nullptr) return false; + + if (data_data->shape.size() == 1) { + // CSR case. + Array oshape({data_indptr->shape[0] - 1, weight->shape[0]}); + reporter->Assign(types[4], TensorType(oshape, weight->dtype)); + return true; + } + + if (data_data->shape.size() == 3) { + // BSR case. + Array oshape( + {(data_indptr->shape[0] - 1) * data_data->shape[1], weight->shape[0]}); + reporter->Assign(types[4], TensorType(oshape, weight->dtype)); + return true; + } + LOG(FATAL) << "Unknown data ndim for nn.sparse_dense, should be 1 (CSR) or 3 (BSR)"; + return false; + + }else { + const auto* data = types[0].as(); const auto* weight_data = types[1].as(); ICHECK(weight_data->shape.size() == 1 || weight_data->shape.size() == 3); @@ -61,22 +90,24 @@ bool SparseDenseRel(const Array& types, int num_inputs, const Attrs& attrs } LOG(FATAL) << "Unknown weight ndim for nn.sparse_dense, should be 1 (CSR) or 3 (BSR)"; return false; + } } // Positional relay function to create dense operator used by frontend FFI. -Expr MakeSparseDense(Expr data, Expr weight_data, Expr weight_indices, Expr weight_indptr) { +Expr MakeSparseDense(Expr data, Expr weight_data, Expr weight_indices, Expr weight_indptr, bool sparse_data) { auto attrs = make_object(); + attrs->sparse_data = std::move(sparse_data); static const Op& op = Op::Get("nn.sparse_dense"); return Call(op, {data, weight_data, weight_indices, weight_indptr}, Attrs(attrs), {}); } TVM_REGISTER_GLOBAL("relay.op.nn._make.sparse_dense") .set_body([](const TVMArgs& args, TVMRetValue* rv) { - runtime::detail::unpack_call(MakeSparseDense, args, rv); + runtime::detail::unpack_call(MakeSparseDense, args, rv); }); RELAY_REGISTER_OP("nn.sparse_dense") - .describe(R"code(Applies a sparse linear transformation: :math:`Y = XW^T` with W sparse. + .describe(R"code(Applies a sparse linear transformation: :math:`Y = XW^T` with either X or W sparse. - **data**: `(x1, x2, ..., xn, input_dim)` - **weight**: `(units, input_dim)` @@ -85,10 +116,10 @@ RELAY_REGISTER_OP("nn.sparse_dense") )code" TVM_ADD_FILELINE) .set_attrs_type() .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", "Input data if dense, otherwise data_data matrix if sparse.") + .add_argument("input_tensor2", "nD Tensor", "Weight_data matrix or data_indices matrix.") + .add_argument("input_tensor3", "nD Tensor", "Weight_indices matrix or data_indptr matrix.") + .add_argument("input_tensor4", "nD Tensor", "Weight_indptr matrix or weight matrix if dense.") .set_support_level(1) .add_type_rel("SparseDense", SparseDenseRel); diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index b79bd8bbba52..d1ad56aed20c 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -1794,9 +1794,15 @@ def test_forward_sparse_dense_matmul(): # # ------------------------------------------------------------------ - # TODO(ANSHUMAN87): False case for flip need to be supported - # _test_sparse_dense_matmul([[0, 0], [1, 2]], [4.0, 8.0], [3, 4], [4, 3], "float32") - _test_sparse_dense_matmul([[0, 0], [1, 2]], [4.0, 8.0], [3, 5], [4, 3], "float32", True) + _test_sparse_dense_matmul([[0, 0], [1, 2]], [4.0, 8.0], [3, 4], [4, 3], "float32") + _test_sparse_dense_matmul([[0, 0], [1, 2]], [4.0, 8.0], [3, 3], [3, 3], "float32") + _test_sparse_dense_matmul( + [[0, 0], [1, 3], [4, 3]], [3.0, 6.0, 9.0], [5, 5], [5, 5], "float32" + ) + _test_sparse_dense_matmul( + [[0, 0], [1, 3], [4, 3]], [3.0, 6.0, 9.0], [7, 9], [9, 5], "float32" + ) + _test_sparse_dense_matmul([[0, 0], [1, 2]], [4.0, 8.0], [4, 3], [3, 4], "float32", True) _test_sparse_dense_matmul([[0, 0], [1, 2]], [4.0, 8.0], [3, 3], [3, 3], "float32", True) _test_sparse_dense_matmul( [[0, 0], [1, 3], [4, 3]], [3.0, 6.0, 9.0], [5, 5], [5, 5], "float32", True From 0b866349012fdd711967eebfe6badf2d4fc9dde0 Mon Sep 17 00:00:00 2001 From: ANSHUMAN TRIPATHY Date: Mon, 9 Nov 2020 21:58:21 +0530 Subject: [PATCH 2/9] [1] clang issue resolved --- include/tvm/relay/attrs/nn.h | 6 ++- src/relay/op/nn/sparse.cc | 102 ++++++++++++++++++----------------- 2 files changed, 56 insertions(+), 52 deletions(-) diff --git a/include/tvm/relay/attrs/nn.h b/include/tvm/relay/attrs/nn.h index 08dfc9368a68..a064facee9b4 100644 --- a/include/tvm/relay/attrs/nn.h +++ b/include/tvm/relay/attrs/nn.h @@ -940,8 +940,10 @@ struct SparseDenseAttrs : public tvm::AttrsNode { bool sparse_data; TVM_DECLARE_ATTRS(SparseDenseAttrs, "relay.attrs.SparseDenseAttrs") { - TVM_ATTR_FIELD(sparse_data).set_default(false).describe("Indicate whether data or weight is sparse. True if data is sparse"); - } + TVM_ATTR_FIELD(sparse_data) + .set_default(false) + .describe("Indicate whether data or weight is sparse. True if data is sparse"); + } }; /*! \brief Attributes for sparse_transpose operator */ diff --git a/src/relay/op/nn/sparse.cc b/src/relay/op/nn/sparse.cc index 7456668a3e09..bbb72b9e5e9e 100644 --- a/src/relay/op/nn/sparse.cc +++ b/src/relay/op/nn/sparse.cc @@ -43,58 +43,58 @@ bool SparseDenseRel(const Array& types, int num_inputs, const Attrs& attrs ICHECK(param != nullptr); if (param->sparse_data) { - const auto* weight = types[3].as(); - const auto* data_data = types[0].as(); - ICHECK(data_data->shape.size() == 1 || data_data->shape.size() == 3); - const auto* data_indptr = types[2].as(); - if (weight == nullptr) return false; - - if (data_data->shape.size() == 1) { - // CSR case. - Array oshape({data_indptr->shape[0] - 1, weight->shape[0]}); - reporter->Assign(types[4], TensorType(oshape, weight->dtype)); - return true; - } - - if (data_data->shape.size() == 3) { - // BSR case. - Array oshape( - {(data_indptr->shape[0] - 1) * data_data->shape[1], weight->shape[0]}); - reporter->Assign(types[4], TensorType(oshape, weight->dtype)); - return true; - } - LOG(FATAL) << "Unknown data ndim for nn.sparse_dense, should be 1 (CSR) or 3 (BSR)"; - return false; - - }else { - - const auto* data = types[0].as(); - const auto* weight_data = types[1].as(); - ICHECK(weight_data->shape.size() == 1 || weight_data->shape.size() == 3); - const auto* weight_indptr = types[3].as(); - if (data == nullptr) return false; - - if (weight_data->shape.size() == 1) { - // CSR case. - Array oshape({data->shape[0], weight_indptr->shape[0] - 1}); - reporter->Assign(types[4], TensorType(oshape, data->dtype)); - return true; - } - - if (weight_data->shape.size() == 3) { - // BSR case. - Array oshape( - {data->shape[0], (weight_indptr->shape[0] - 1) * weight_data->shape[1]}); - reporter->Assign(types[4], TensorType(oshape, data->dtype)); - return true; - } - LOG(FATAL) << "Unknown weight ndim for nn.sparse_dense, should be 1 (CSR) or 3 (BSR)"; - return false; + const auto* weight = types[3].as(); + const auto* data_data = types[0].as(); + ICHECK(data_data->shape.size() == 1 || data_data->shape.size() == 3); + const auto* data_indptr = types[2].as(); + if (weight == nullptr) return false; + + if (data_data->shape.size() == 1) { + // CSR case. + Array oshape({data_indptr->shape[0] - 1, weight->shape[0]}); + reporter->Assign(types[4], TensorType(oshape, weight->dtype)); + return true; + } + + if (data_data->shape.size() == 3) { + // BSR case. + Array oshape( + {(data_indptr->shape[0] - 1) * data_data->shape[1], weight->shape[0]}); + reporter->Assign(types[4], TensorType(oshape, weight->dtype)); + return true; + } + LOG(FATAL) << "Unknown data ndim for nn.sparse_dense, should be 1 (CSR) or 3 (BSR)"; + return false; + + } else { + const auto* data = types[0].as(); + const auto* weight_data = types[1].as(); + ICHECK(weight_data->shape.size() == 1 || weight_data->shape.size() == 3); + const auto* weight_indptr = types[3].as(); + if (data == nullptr) return false; + + if (weight_data->shape.size() == 1) { + // CSR case. + Array oshape({data->shape[0], weight_indptr->shape[0] - 1}); + reporter->Assign(types[4], TensorType(oshape, data->dtype)); + return true; + } + + if (weight_data->shape.size() == 3) { + // BSR case. + Array oshape( + {data->shape[0], (weight_indptr->shape[0] - 1) * weight_data->shape[1]}); + reporter->Assign(types[4], TensorType(oshape, data->dtype)); + return true; + } + LOG(FATAL) << "Unknown weight ndim for nn.sparse_dense, should be 1 (CSR) or 3 (BSR)"; + return false; } } // Positional relay function to create dense operator used by frontend FFI. -Expr MakeSparseDense(Expr data, Expr weight_data, Expr weight_indices, Expr weight_indptr, bool sparse_data) { +Expr MakeSparseDense(Expr data, Expr weight_data, Expr weight_indices, Expr weight_indptr, + bool sparse_data) { auto attrs = make_object(); attrs->sparse_data = std::move(sparse_data); static const Op& op = Op::Get("nn.sparse_dense"); @@ -107,7 +107,8 @@ TVM_REGISTER_GLOBAL("relay.op.nn._make.sparse_dense") }); RELAY_REGISTER_OP("nn.sparse_dense") - .describe(R"code(Applies a sparse linear transformation: :math:`Y = XW^T` with either X or W sparse. + .describe( + R"code(Applies a sparse linear transformation: :math:`Y = XW^T` with either X or W sparse. - **data**: `(x1, x2, ..., xn, input_dim)` - **weight**: `(units, input_dim)` @@ -116,7 +117,8 @@ RELAY_REGISTER_OP("nn.sparse_dense") )code" TVM_ADD_FILELINE) .set_attrs_type() .set_num_inputs(4) - .add_argument("input_tensor1", "nD Tensor", "Input data if dense, otherwise data_data matrix if sparse.") + .add_argument("input_tensor1", "nD Tensor", + "Input data if dense, otherwise data_data matrix if sparse.") .add_argument("input_tensor2", "nD Tensor", "Weight_data matrix or data_indices matrix.") .add_argument("input_tensor3", "nD Tensor", "Weight_indices matrix or data_indptr matrix.") .add_argument("input_tensor4", "nD Tensor", "Weight_indptr matrix or weight matrix if dense.") From bb733b4b3b2e48ff210d78c39adefcace418fed0 Mon Sep 17 00:00:00 2001 From: ANSHUMAN TRIPATHY Date: Mon, 9 Nov 2020 22:10:06 +0530 Subject: [PATCH 3/9] [2] python format resolved --- python/tvm/relay/frontend/tensorflow.py | 9 +++++++-- python/tvm/topi/nn/sparse.py | 15 +++++++++++++-- tests/python/frontend/tensorflow/test_forward.py | 8 ++------ 3 files changed, 22 insertions(+), 10 deletions(-) diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index 4656a3f8a470..f5b12d7b3e34 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -941,7 +941,9 @@ def _impl(inputs, attr, params, mod): weight_indices = _expr.const(weight_sp.indices, weight_sp.indices.dtype) if sparse_data: - ret = _op.nn.sparse_dense([weight_data, weight_indices, weight_indptrs], data, sparse_data=True) + ret = _op.nn.sparse_dense( + [weight_data, weight_indices, weight_indptrs], data, sparse_data=True + ) else: ret = _op.nn.sparse_dense(data, [weight_data, weight_indices, weight_indptrs]) ret = _op.transpose(ret) @@ -949,7 +951,10 @@ def _impl(inputs, attr, params, mod): # Case 1. If both are true means first input was dense and second was sparse # Case 2. If both are false means first input was sparse and second was dense # TODO(ANSHUMAN87): Support other adjoint option too - if not ((attr.get("adjoint_a") and attr.get("adjoint_b")) or ((not attr.get("adjoint_a")) and (not attr.get("adjoint_b")))): + if not ( + (attr.get("adjoint_a") and attr.get("adjoint_b")) + or ((not attr.get("adjoint_a")) and (not attr.get("adjoint_b"))) + ): raise tvm.error.OpAttributeUnImplemented( "Only tf.sparse.sparse_dense_matmul() with adjoint_a=True and adjoint_b=True" "or with adjoint_a=False and adjoint_b=False" diff --git a/python/tvm/topi/nn/sparse.py b/python/tvm/topi/nn/sparse.py index b57a347ba684..7ef804250fd2 100644 --- a/python/tvm/topi/nn/sparse.py +++ b/python/tvm/topi/nn/sparse.py @@ -57,6 +57,7 @@ def sparse_dense_v2(data, weight_data, weight_indices, weight_indptr): func = _sparse_dense_bsrmm_v2 return func(data, weight_data, weight_indices, weight_indptr) + def sparse_dense_v1(data_data, data_indices, data_indptr, weight): """ Computes sparse-dense matrix multiplication of @@ -91,6 +92,7 @@ def sparse_dense_v1(data_data, data_indices, data_indptr, weight): func = _sparse_dense_bsrmm_v1 return func(data_data, data_indices, data_indptr, weight) + def sparse_dense(input_1, input_2, input_3, input_4, sparse_data=False): """ Computes sparse-dense matrix multiplication of `data` and @@ -145,6 +147,7 @@ def sparse_dense(input_1, input_2, input_3, input_4, sparse_data=False): else: return sparse_dense_v2(input_1, input_2, input_3, input_4) + def _sparse_dense_csrmm_v1(data_data, data_indices, data_indptr, weight): oshape = (get_const_tuple(data_indptr.shape)[0] - 1, get_const_tuple(weight.shape)[0]) @@ -160,6 +163,7 @@ def f(row, i): return te.compute(oshape, f, tag="sparse_dense_csrmm_v1") + def _sparse_dense_csrmm_v2(data, weight_data, weight_indices, weight_indptr): oshape = (get_const_tuple(data.shape)[0], get_const_tuple(weight_indptr.shape)[0] - 1) @@ -175,6 +179,7 @@ def f(i, row): return te.compute(oshape, f, tag="sparse_dense_csrmm_v2") + def _sparse_dense_bsrmm_v1(data_data, data_indices, data_indptr, weight): (k, m) = get_const_tuple(weight.shape) (l, bs_r, bs_c) = get_const_tuple(data_data.shape) @@ -196,13 +201,16 @@ def _compute_block(nb_j, j, i): idxd = tvm.tir.indexdiv idxm = tvm.tir.indexmod - bsrmm_block = te.compute((num_blocks, bs_r, k), _compute_block, tag="sparse_dense_bsrmm_block_v1") + bsrmm_block = te.compute( + (num_blocks, bs_r, k), _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], tag="sparse_dense_bsrmm_v1", ) + def _sparse_dense_bsrmm_v2(data, weight_data, weight_indices, weight_indptr): (m, _) = get_const_tuple(data.shape) (_, bs_r, bs_c) = get_const_tuple(weight_data.shape) @@ -224,13 +232,16 @@ def _compute_block(i, nb_j, j): idxd = tvm.tir.indexdiv idxm = tvm.tir.indexmod - bsrmm_block = te.compute((m, num_blocks, bs_r), _compute_block, tag="sparse_dense_bsrmm_block_v2") + bsrmm_block = te.compute( + (m, num_blocks, bs_r), _compute_block, tag="sparse_dense_bsrmm_block_v2" + ) return te.compute( (m, num_blocks * bs_r), lambda m, n: bsrmm_block[m, idxd(n, bs_r), idxm(n, bs_r)], tag="sparse_dense_bsrmm_v2", ) + def sparse_transpose(sparse_data, sparse_indices, sparse_indptr): """ Transpose a square sparse matrix, diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index d1ad56aed20c..6720c2e13bfe 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -1796,12 +1796,8 @@ def test_forward_sparse_dense_matmul(): _test_sparse_dense_matmul([[0, 0], [1, 2]], [4.0, 8.0], [3, 4], [4, 3], "float32") _test_sparse_dense_matmul([[0, 0], [1, 2]], [4.0, 8.0], [3, 3], [3, 3], "float32") - _test_sparse_dense_matmul( - [[0, 0], [1, 3], [4, 3]], [3.0, 6.0, 9.0], [5, 5], [5, 5], "float32" - ) - _test_sparse_dense_matmul( - [[0, 0], [1, 3], [4, 3]], [3.0, 6.0, 9.0], [7, 9], [9, 5], "float32" - ) + _test_sparse_dense_matmul([[0, 0], [1, 3], [4, 3]], [3.0, 6.0, 9.0], [5, 5], [5, 5], "float32") + _test_sparse_dense_matmul([[0, 0], [1, 3], [4, 3]], [3.0, 6.0, 9.0], [7, 9], [9, 5], "float32") _test_sparse_dense_matmul([[0, 0], [1, 2]], [4.0, 8.0], [4, 3], [3, 4], "float32", True) _test_sparse_dense_matmul([[0, 0], [1, 2]], [4.0, 8.0], [3, 3], [3, 3], "float32", True) _test_sparse_dense_matmul( From f7b8d98b0604dbf289d079f822b69eadda2a72d1 Mon Sep 17 00:00:00 2001 From: ANSHUMAN TRIPATHY Date: Mon, 9 Nov 2020 22:21:15 +0530 Subject: [PATCH 4/9] [3] lint error resolved --- python/tvm/relay/op/nn/nn.py | 1 + python/tvm/topi/nn/sparse.py | 5 +++-- src/relay/transforms/convert_sparse_dense.cc | 4 +++- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py index 6d2de1e47be2..eac972771c12 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/python/tvm/relay/op/nn/nn.py @@ -1993,6 +1993,7 @@ def batch_matmul(x, y): return _make.batch_matmul(x, y) +# pylint: disable=no-else-return,inconsistent-return-statements def sparse_dense(data, weight, sparse_data=False): r""" Computes the matrix multiplication of `data` and `weight`, where `data` is diff --git a/python/tvm/topi/nn/sparse.py b/python/tvm/topi/nn/sparse.py index 7ef804250fd2..cb664179680f 100644 --- a/python/tvm/topi/nn/sparse.py +++ b/python/tvm/topi/nn/sparse.py @@ -93,6 +93,7 @@ def sparse_dense_v1(data_data, data_indices, data_indptr, weight): 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 @@ -181,8 +182,8 @@ def f(i, row): def _sparse_dense_bsrmm_v1(data_data, data_indices, data_indptr, weight): - (k, m) = get_const_tuple(weight.shape) - (l, bs_r, bs_c) = get_const_tuple(data_data.shape) + (k, _) = 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 diff --git a/src/relay/transforms/convert_sparse_dense.cc b/src/relay/transforms/convert_sparse_dense.cc index 5f4dbe642c3d..26a4d487196d 100644 --- a/src/relay/transforms/convert_sparse_dense.cc +++ b/src/relay/transforms/convert_sparse_dense.cc @@ -103,8 +103,10 @@ class DenseToSparseDenseMutator : public ExprRewriter { Var weight_data(prefix + ".data", ws_data_type); Var weight_indices(prefix + ".indices", ws_indices_type); Var weight_indptr(prefix + ".indptr", ws_indptr_type); + auto attrs = make_object(); - return Call(sparse_dense_op_, {data, weight_data, weight_indices, weight_indptr}); + return Call(sparse_dense_op_, {data, weight_data, weight_indices, weight_indptr}, + Attrs(attrs)); } } } From 98444f2082a32bb3870ab5e33830511d78568b40 Mon Sep 17 00:00:00 2001 From: ANSHUMAN TRIPATHY Date: Fri, 20 Nov 2020 02:11:36 +0530 Subject: [PATCH 5/9] [4] Review comments handled --- include/tvm/relay/attrs/nn.h | 6 ++--- python/tvm/relay/frontend/tensorflow.py | 17 ++++++------ python/tvm/relay/op/nn/_nn.py | 2 +- python/tvm/relay/op/nn/nn.py | 13 +++++---- python/tvm/relay/op/strategy/generic.py | 2 +- python/tvm/topi/nn/sparse.py | 35 +++++++------------------ src/relay/op/nn/sparse.cc | 21 +++++++-------- 7 files changed, 39 insertions(+), 57 deletions(-) diff --git a/include/tvm/relay/attrs/nn.h b/include/tvm/relay/attrs/nn.h index a064facee9b4..e6b07078b40b 100644 --- a/include/tvm/relay/attrs/nn.h +++ b/include/tvm/relay/attrs/nn.h @@ -937,12 +937,12 @@ struct DenseAttrs : public tvm::AttrsNode { /*! \brief Attributes for sparse_dense operator */ struct SparseDenseAttrs : public tvm::AttrsNode { - bool sparse_data; + bool sparse_lhs; TVM_DECLARE_ATTRS(SparseDenseAttrs, "relay.attrs.SparseDenseAttrs") { - TVM_ATTR_FIELD(sparse_data) + TVM_ATTR_FIELD(sparse_lhs) .set_default(false) - .describe("Indicate whether data or weight is sparse. True if data is sparse"); + .describe("Indicate whether lhs or rhs matrix is sparse. True if lhs matrix is sparse"); } }; diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index f5b12d7b3e34..3c8fb8ecbeb6 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -917,11 +917,11 @@ def _impl(inputs, attr, params, mod): data = inputs[3] # By default, in tensorflow the first input ,i.e., data is sparse - sparse_data = True + sparse_lhs = True # If both are true means First input was dense and second was sparse if attr.get("adjoint_a") and attr.get("adjoint_b"): - sparse_data = False + sparse_lhs = False rows = [x[0] for x in indices_tensor] cols = [x[1] for x in indices_tensor] @@ -931,7 +931,7 @@ def _impl(inputs, attr, params, mod): (values_tensor, (rows, cols)), shape=tuple(dense_shape_tensor.tolist()) ) - if sparse_data: + if sparse_lhs: data = _op.transpose(data) else: weight_sp = csr_matrix(weight_sp.transpose()) @@ -940,12 +940,11 @@ def _impl(inputs, attr, params, mod): weight_indptrs = _expr.const(weight_sp.indptr, weight_sp.indptr.dtype) weight_indices = _expr.const(weight_sp.indices, weight_sp.indices.dtype) - if sparse_data: - ret = _op.nn.sparse_dense( - [weight_data, weight_indices, weight_indptrs], data, sparse_data=True - ) - else: - ret = _op.nn.sparse_dense(data, [weight_data, weight_indices, weight_indptrs]) + ret = _op.nn.sparse_dense( + data, [weight_data, weight_indices, weight_indptrs], sparse_lhs + ) + + if not sparse_lhs: ret = _op.transpose(ret) # Case 1. If both are true means first input was dense and second was sparse diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py index f92808f4cf0c..93149b5fa1f4 100644 --- a/python/tvm/relay/op/nn/_nn.py +++ b/python/tvm/relay/op/nn/_nn.py @@ -69,7 +69,7 @@ def compute_fifo_buffer(attrs, inputs, out_type): @reg.register_compute("nn.sparse_dense") def compute_sparse_dense(attrs, inputs, out_type): """Compute definition of sparse_dense""" - return [topi.nn.sparse_dense(inputs[0], inputs[1], inputs[2], inputs[3], attrs["sparse_data"])] + return [topi.nn.sparse_dense(inputs[0], inputs[1], inputs[2], inputs[3], attrs["sparse_lhs"])] reg.register_strategy("nn.sparse_dense", strategy.sparse_dense_strategy) diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py index eac972771c12..100a3691dd84 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/python/tvm/relay/op/nn/nn.py @@ -1994,7 +1994,7 @@ def batch_matmul(x, y): # pylint: disable=no-else-return,inconsistent-return-statements -def sparse_dense(data, weight, sparse_data=False): +def sparse_dense(data, weight, sparse_lhs=False): r""" Computes the matrix multiplication of `data` and `weight`, where `data` is a dense matrix and `weight` is a sparse (either BSR or CSR) namedtuple with @@ -2020,19 +2020,18 @@ def sparse_dense(data, weight, sparse_data=False): weight : Union[namedtuple, Tuple[ndarray, ndarray, ndarray]]. The sparse weight matrix for the matrix multiplication. + sparse_lhs : bool, optional + Indicates whether lhs or rhs matrix is sparse. + Returns ------- result: tvm.relay.Expr The computed result. """ if hasattr(weight, "indices"): - return _make.sparse_dense(data, weight.data, weight.indices, weight.indptr, sparse_data) - elif isinstance(weight, (tuple, list)): - return _make.sparse_dense(data, weight[0], weight[1], weight[2], sparse_data) - elif hasattr(data, "indices"): - return _make.sparse_dense(data.data, data.indices, data.indptr, weight, sparse_data) + return _make.sparse_dense(data, weight.data, weight.indices, weight.indptr, sparse_lhs) else: - return _make.sparse_dense(data[0], data[1], data[2], weight, sparse_data) + return _make.sparse_dense(data, weight[0], weight[1], weight[2], sparse_lhs) def sparse_transpose(x): diff --git a/python/tvm/relay/op/strategy/generic.py b/python/tvm/relay/op/strategy/generic.py index ce747f3b674a..7b2b3e35e077 100644 --- a/python/tvm/relay/op/strategy/generic.py +++ b/python/tvm/relay/op/strategy/generic.py @@ -706,7 +706,7 @@ def wrap_compute_sparse_dense(topi_compute): """wrap sparse dense topi compute""" def _compute_sparse_dense(attrs, inputs, out_type): - return [topi_compute(inputs[0], inputs[1], inputs[2], inputs[3], attrs["sparse_data"])] + return [topi_compute(inputs[0], inputs[1], inputs[2], inputs[3], attrs["sparse_lhs"])] return _compute_sparse_dense diff --git a/python/tvm/topi/nn/sparse.py b/python/tvm/topi/nn/sparse.py index cb664179680f..efbca7fd48d7 100644 --- a/python/tvm/topi/nn/sparse.py +++ b/python/tvm/topi/nn/sparse.py @@ -94,7 +94,7 @@ def sparse_dense_v1(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): +def sparse_dense(dense_data, sparse_data, sparse_indices, sparse_indptr, sparse_lhs=False): """ Computes sparse-dense matrix multiplication of `data` and `(weight_data, weight_indices, weight_indptr).T` @@ -104,49 +104,34 @@ def sparse_dense(input_1, input_2, input_3, input_4, sparse_data=False): Parameters ---------- - input_1 : tvm.te.Tensor + dense_data : tvm.te.Tensor data: 2-D with shape [M, K], float32 when input data is dense or - data_data: - 1-D with shape [nnz] (CSR) or - 3-D with shape [num_blocks, bs_r, bs_c] (BSR) - - input_2 : tvm.te.Tensor - weight_data: + sparse_data : tvm.te.Tensor 1-D with shape [nnz] (CSR) or 3-D with shape [num_blocks, bs_r, bs_c] (BSR) or - data_indices: - 1-D with shape [nnz] (CSR) or - 1-D with shape [num_blocks] (BSR) - - input_3 : tvm.te.Tensor - weight_indices: + sparse_indices : tvm.te.Tensor 1-D with shape [nnz] (CSR) or 1-D with shape [num_blocks] (BSR) or - data_indptr: - 1-D with shape [M + 1] (CSR) or - 1-D with shape [(M + 1) // bs_r] (BSR) - - input_4 : tvm.te.Tensor - weight_indptr: + sparse_indptr : tvm.te.Tensor 1-D with shape [N + 1] (CSR) or 1-D with shape [(N + 1) // bs_r] (BSR) - weight: - 2-D with shape [N, K], float32 when weight is dense + sparse_lhs : bool, optional + Indicates whether lhs or rhs matrix is sparse. Returns ------- output : tvm.te.Tensor 2-D with shape [M, N] """ - if sparse_data: - return sparse_dense_v1(input_1, input_2, input_3, input_4) + if sparse_lhs: + return sparse_dense_v1(sparse_data, sparse_indices, sparse_indptr, dense_data) else: - return sparse_dense_v2(input_1, input_2, input_3, input_4) + return sparse_dense_v2(dense_data, sparse_data, sparse_indices, sparse_indptr) def _sparse_dense_csrmm_v1(data_data, data_indices, data_indptr, weight): diff --git a/src/relay/op/nn/sparse.cc b/src/relay/op/nn/sparse.cc index bbb72b9e5e9e..e9073730641d 100644 --- a/src/relay/op/nn/sparse.cc +++ b/src/relay/op/nn/sparse.cc @@ -42,11 +42,11 @@ bool SparseDenseRel(const Array& types, int num_inputs, const Attrs& attrs const auto* param = attrs.as(); ICHECK(param != nullptr); - if (param->sparse_data) { - const auto* weight = types[3].as(); - const auto* data_data = types[0].as(); + if (param->sparse_lhs) { + const auto* weight = types[0].as(); + const auto* data_data = types[1].as(); ICHECK(data_data->shape.size() == 1 || data_data->shape.size() == 3); - const auto* data_indptr = types[2].as(); + const auto* data_indptr = types[3].as(); if (weight == nullptr) return false; if (data_data->shape.size() == 1) { @@ -94,9 +94,9 @@ bool SparseDenseRel(const Array& types, int num_inputs, const Attrs& attrs // Positional relay function to create dense operator used by frontend FFI. Expr MakeSparseDense(Expr data, Expr weight_data, Expr weight_indices, Expr weight_indptr, - bool sparse_data) { + bool sparse_lhs) { auto attrs = make_object(); - attrs->sparse_data = std::move(sparse_data); + attrs->sparse_lhs = std::move(sparse_lhs); static const Op& op = Op::Get("nn.sparse_dense"); return Call(op, {data, weight_data, weight_indices, weight_indptr}, Attrs(attrs), {}); } @@ -117,11 +117,10 @@ RELAY_REGISTER_OP("nn.sparse_dense") )code" TVM_ADD_FILELINE) .set_attrs_type() .set_num_inputs(4) - .add_argument("input_tensor1", "nD Tensor", - "Input data if dense, otherwise data_data matrix if sparse.") - .add_argument("input_tensor2", "nD Tensor", "Weight_data matrix or data_indices matrix.") - .add_argument("input_tensor3", "nD Tensor", "Weight_indices matrix or data_indptr matrix.") - .add_argument("input_tensor4", "nD Tensor", "Weight_indptr matrix or weight matrix if dense.") + .add_argument("dense_data", "nD Tensor", "Input dense data.") + .add_argument("sparse_data", "1D or 3D Tensor", "Sparse data matrix.") + .add_argument("sparse_indices", "1D Tensor", "Sparse indices matrix.") + .add_argument("sparse_indptr", "1D Tensor", "Sparse indptr matrix.") .set_support_level(1) .add_type_rel("SparseDense", SparseDenseRel); From 64fa8ae18deaa5a1e5d76dc172e6eb8394e5b4a1 Mon Sep 17 00:00:00 2001 From: ANSHUMAN TRIPATHY Date: Fri, 20 Nov 2020 02:20:13 +0530 Subject: [PATCH 6/9] [5] Lint error resolved --- python/tvm/relay/frontend/tensorflow.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index 3c8fb8ecbeb6..aa87c2284697 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -940,9 +940,7 @@ def _impl(inputs, attr, params, mod): weight_indptrs = _expr.const(weight_sp.indptr, weight_sp.indptr.dtype) weight_indices = _expr.const(weight_sp.indices, weight_sp.indices.dtype) - ret = _op.nn.sparse_dense( - data, [weight_data, weight_indices, weight_indptrs], sparse_lhs - ) + ret = _op.nn.sparse_dense(data, [weight_data, weight_indices, weight_indptrs], sparse_lhs) if not sparse_lhs: ret = _op.transpose(ret) From 68252223bf30558dbaec9eddcdb4eca037fd073d Mon Sep 17 00:00:00 2001 From: ANSHUMAN TRIPATHY Date: Fri, 20 Nov 2020 03:04:19 +0530 Subject: [PATCH 7/9] [6] Review comments handled --- include/tvm/relay/attrs/nn.h | 4 +++- python/tvm/relay/op/nn/nn.py | 2 +- python/tvm/topi/nn/sparse.py | 13 ++++++------- 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/include/tvm/relay/attrs/nn.h b/include/tvm/relay/attrs/nn.h index e6b07078b40b..3fdbd7205a9d 100644 --- a/include/tvm/relay/attrs/nn.h +++ b/include/tvm/relay/attrs/nn.h @@ -942,7 +942,9 @@ struct SparseDenseAttrs : public tvm::AttrsNode { TVM_DECLARE_ATTRS(SparseDenseAttrs, "relay.attrs.SparseDenseAttrs") { TVM_ATTR_FIELD(sparse_lhs) .set_default(false) - .describe("Indicate whether lhs or rhs matrix is sparse. True if lhs matrix is sparse"); + .describe( + "Indicate whether sparse matrix is multiplied on the right or the left. If true, then " + "the operation is S * D^T (D dense, S sparse). If false, the operation is D * S^T"); } }; diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py index 100a3691dd84..25b4cc19b1fa 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/python/tvm/relay/op/nn/nn.py @@ -2021,7 +2021,7 @@ def sparse_dense(data, weight, sparse_lhs=False): The sparse weight matrix for the matrix multiplication. sparse_lhs : bool, optional - Indicates whether lhs or rhs matrix is sparse. + Indicates whether lhs or rhs matrix is sparse. Default value is False. Returns ------- diff --git a/python/tvm/topi/nn/sparse.py b/python/tvm/topi/nn/sparse.py index efbca7fd48d7..2b21320917a5 100644 --- a/python/tvm/topi/nn/sparse.py +++ b/python/tvm/topi/nn/sparse.py @@ -97,31 +97,30 @@ def sparse_dense_v1(data_data, data_indices, data_indptr, weight): def sparse_dense(dense_data, sparse_data, sparse_indices, sparse_indptr, sparse_lhs=False): """ Computes sparse-dense matrix multiplication of `data` and - `(weight_data, weight_indices, weight_indptr).T` + `(weight_data, weight_indices, weight_indptr).T`, if sparse_lhs=False or Computes sparse-dense matrix multiplication of - `(data_data, data_indices, data_indptr)` and `weight.T` + `(data_data, data_indices, data_indptr)` and `weight.T`, if sparse_lhs=True Parameters ---------- dense_data : tvm.te.Tensor - data: - 2-D with shape [M, K], float32 when input data is dense or + 2-D with shape [M, K], float32 sparse_data : tvm.te.Tensor 1-D with shape [nnz] (CSR) or - 3-D with shape [num_blocks, bs_r, bs_c] (BSR) or + 3-D with shape [num_blocks, bs_r, bs_c] (BSR) sparse_indices : tvm.te.Tensor 1-D with shape [nnz] (CSR) or - 1-D with shape [num_blocks] (BSR) or + 1-D with shape [num_blocks] (BSR) sparse_indptr : tvm.te.Tensor 1-D with shape [N + 1] (CSR) or 1-D with shape [(N + 1) // bs_r] (BSR) sparse_lhs : bool, optional - Indicates whether lhs or rhs matrix is sparse. + Indicates whether lhs or rhs matrix is sparse. Default value is False. Returns ------- From 4623979e51c9ad37f6b56adfb2722f8a83c04edd Mon Sep 17 00:00:00 2001 From: ANSHUMAN TRIPATHY Date: Fri, 20 Nov 2020 10:11:17 +0530 Subject: [PATCH 8/9] [7] Review comments handled --- python/tvm/relay/op/nn/nn.py | 37 ++++++++++++++++++++++++------------ python/tvm/topi/nn/sparse.py | 2 +- 2 files changed, 26 insertions(+), 13 deletions(-) diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py index 25b4cc19b1fa..920aaa7da1c8 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/python/tvm/relay/op/nn/nn.py @@ -1994,17 +1994,26 @@ def batch_matmul(x, y): # pylint: disable=no-else-return,inconsistent-return-statements -def sparse_dense(data, weight, sparse_lhs=False): +def sparse_dense(dense_mat, sparse_mat, sparse_lhs=False): r""" - Computes the matrix multiplication of `data` and `weight`, where `data` is - a dense matrix and `weight` is a sparse (either BSR or CSR) namedtuple with + Computes the matrix multiplication of `dense_mat` and `sparse_mat`, where `dense_mat` is + a dense matrix and `sparse_mat` is a sparse (either BSR or CSR) namedtuple with fields `data`, `indices`, and `indptr`. .. math:: - \mbox{sparse_dense}(data, weight)[m, n] = \mbox{matmul}(x, \mbox{as_dense}(weight)^T)[m, n] + if sparse_lhs=True - where `as_dense` returns dense equivalent of the given sparse matrix. + \mbox{sparse_dense}(dense_mat, sparse_mat)[m, n] + = \mbox{matmul}(D, \mbox{as_dense}(S)^T)[m, n] + + if sparse_lhs=False + + \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). See https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_matrix.html @@ -2014,11 +2023,11 @@ def sparse_dense(data, weight, sparse_lhs=False): Parameters ---------- - data : tvm.relay.Expr - The input data for the matrix multiplication + dense_mat : tvm.relay.Expr + The input dense matrix for the matrix multiplication - weight : Union[namedtuple, Tuple[ndarray, ndarray, ndarray]]. - The sparse weight matrix for the matrix multiplication. + sparse_mat : Union[namedtuple, Tuple[ndarray, ndarray, ndarray]]. + The input sparse matrix for the matrix multiplication. sparse_lhs : bool, optional Indicates whether lhs or rhs matrix is sparse. Default value is False. @@ -2028,10 +2037,14 @@ def sparse_dense(data, weight, sparse_lhs=False): result: tvm.relay.Expr The computed result. """ - if hasattr(weight, "indices"): - return _make.sparse_dense(data, weight.data, weight.indices, weight.indptr, sparse_lhs) + if hasattr(sparse_mat, "indices"): + return _make.sparse_dense( + dense_mat, sparse_mat.data, sparse_mat.indices, sparse_mat.indptr, sparse_lhs + ) else: - return _make.sparse_dense(data, weight[0], weight[1], weight[2], sparse_lhs) + return _make.sparse_dense( + dense_mat, sparse_mat[0], sparse_mat[1], sparse_mat[2], sparse_lhs + ) def sparse_transpose(x): diff --git a/python/tvm/topi/nn/sparse.py b/python/tvm/topi/nn/sparse.py index 2b21320917a5..706971e887ec 100644 --- a/python/tvm/topi/nn/sparse.py +++ b/python/tvm/topi/nn/sparse.py @@ -78,7 +78,7 @@ def sparse_dense_v1(data_data, data_indices, data_indptr, weight): 1-D with shape [(M + 1) // bs_r] (BSR) weight: - 2-D with shape [N, K], float32 when weight is dense + 2-D with shape [N, K], float32 Returns ------- From a4ed84fd2d496fab03db3c6731bf213bd812b363 Mon Sep 17 00:00:00 2001 From: ANSHUMAN TRIPATHY Date: Tue, 8 Dec 2020 21:42:29 +0530 Subject: [PATCH 9/9] [8] Review comments handled --- python/tvm/relay/op/nn/nn.py | 16 +++--- python/tvm/topi/cuda/sparse.py | 2 +- python/tvm/topi/nn/sparse.py | 8 +-- tests/python/topi/python/test_topi_sparse.py | 52 ++++++++++++++++++++ 4 files changed, 65 insertions(+), 13 deletions(-) diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py index 920aaa7da1c8..eb6ff45d2f12 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/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). diff --git a/python/tvm/topi/cuda/sparse.py b/python/tvm/topi/cuda/sparse.py index 8f1a26869c07..c59e6887d47e 100644 --- a/python/tvm/topi/cuda/sparse.py +++ b/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) diff --git a/python/tvm/topi/nn/sparse.py b/python/tvm/topi/nn/sparse.py index 706971e887ec..94d6d9a16330 100644 --- a/python/tvm/topi/nn/sparse.py +++ b/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", ) diff --git a/tests/python/topi/python/test_topi_sparse.py b/tests/python/topi/python/test_topi_sparse.py index 62f49e21418f..e47bfddbf7fc 100644 --- a/tests/python/topi/python/test_topi_sparse.py +++ b/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()