[TOPI] Sparse Add Op added

python/tvm/relay/op/strategy/generic.py

@@ -799,6 +799,34 @@ def sparse_dense_padded_strategy(attrs, inputs, out_type, target):
     raise NotImplementedError("sparse_dense_padded is only implemented for cuda")
 
 
+# sparse_add
+def wrap_compute_sparse_add(topi_compute):
+    """wrap sparse add topi compute"""
+
+    def _compute_sparse_add(attrs, inputs, out_type):
+        return [topi_compute(inputs[0], inputs[1], inputs[2], inputs[3])]
+
+    return _compute_sparse_add
+
+
+@override_native_generic_func("sparse_add_strategy")
+def sparse_add_strategy(attrs, inputs, out_type, target):
+    """sparse add generic strategy"""
+    logger.warning("sparse add is not optimized for this platform.")
+    strategy = _op.OpStrategy()
+    strategy.add_implementation(
+        wrap_compute_sparse_add(topi.nn.sparse_add),
+        wrap_topi_schedule(topi.generic.schedule_sparse_add),
+        name="sparse_add.generic",
+    )
+    return strategy
+
+@generic_func
+def schedule_sparse_add(attrs, outs, target):
+    """schedule sparse_add"""
+    with target:
+        return topi.generic.schedule_sparse_add(outs)
+
 # sparse_transpose
 @generic_func
 def schedule_sparse_transpose(attrs, outs, target):

python/tvm/topi/generic/nn.py

@@ -729,6 +729,21 @@ def schedule_sparse_transpose(outs):
     """
     return _default_schedule(outs, False)
 
+def schedule_sparse_add(outs):
+    """Schedule for sparse_add
+
+    Parameters
+    ----------
+    outs: Array of Tensor
+          The computation graph description of sparse_add
+          in the format of an array of tensors.
+
+    Returns
+    -------
+    sch: Schedule
+        The computation schedule for the op.
+    """
+    return _default_schedule(outs, False)
 
 def schedule_batch_matmul(outs):
     """Schedule for batch_matmul

python/tvm/topi/nn/sparse.py

@@ -356,3 +356,67 @@ def sparse_dense_alter_layout(_attrs, _inputs, _tinfos, _out_type):
     Unlike other TOPI functions, this function operates on both graph level and operator level.
     """
     return None
+
+def sparse_add(dense_data, sparse_data, sparse_indices, sparse_indptr):
+    """
+    Computes sparse-dense addition
+
+    Parameters
+    ----------
+    dense_data : tvm.te.Tensor
+        2-D with shape [M, K], float32
+
+    sparse_data : tvm.te.Tensor
+        1-D with shape [nnz] (CSR) or
+
+    sparse_indices : tvm.te.Tensor
+        1-D with shape [nnz] (CSR) or
+
+    sparse_indptr : tvm.te.Tensor
+        1-D with shape [N + 1] (CSR) or
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        2-D with shape [M, N]
+    """
+    #TODO(ANSHUMAN87): support BSR format too
+    assert len(sparse_data.shape) == 1
+    return _sparse_add_csr(dense_data, sparse_data, sparse_indices, sparse_indptr)
+
+def _sparse_add_csr(dense_data_inp, sparse_data_inp, sparse_indices_inp, sparse_indptr_inp):
+    oshape = get_const_tuple(dense_data_inp.shape)
+
+    def _csr_add_ir(dense_data, sparse_data, sparse_indices, sparse_indptr, out_data):
+        irb = tvm.tir.ir_builder.create()
+        dense_data_ptr = irb.buffer_ptr(dense_data)
+        sparse_data_ptr = irb.buffer_ptr(sparse_data)
+        sparse_indices_ptr = irb.buffer_ptr(sparse_indices)
+        sparse_indptr_ptr = irb.buffer_ptr(sparse_indptr)
+
+        out_data_ptr = irb.buffer_ptr(out_data)
+
+        with irb.for_range(0, oshape[0], kind="serial", name="row") as row:
+            with irb.for_range(0, oshape[1], kind="serial", name="col") as col:
+                out_data_ptr[row, col] = dense_data_ptr[row, col]
+
+        with irb.for_range(0, oshape[0], kind="serial", name="row") as row:
+            offset = sparse_indptr_ptr[row]
+            diff = sparse_indptr_ptr[row + 1] - sparse_indptr_ptr[row]
+            with irb.for_range(0, diff, kind="serial", name="idx") as idx:
+                real_idx = offset + idx
+                col = sparse_indices_ptr[real_idx]
+                out_data_ptr[row, col] = sparse_data_ptr[real_idx] + out_data_ptr[row, col]
+
+        return irb.get()
+
+    return te.extern(
+        shape=oshape,
+        inputs=[dense_data_inp, sparse_data_inp, sparse_indices_inp, sparse_indptr_inp],
+        fcompute=lambda ins, outs: _csr_add_ir(
+            ins[0], ins[1], ins[2], ins[3], outs[0]
+        ),
+        tag="sparse_add_csr",
+        dtype=["float32", "float32", "int32", "int32"],
+        name="out",
+    )

src/relay/op/nn/sparse.cc

@@ -196,5 +196,46 @@ RELAY_REGISTER_OP("nn.sparse_transpose")
     .set_support_level(1)
     .add_type_rel("SparseTranspose", SparseTransposeRel);
 
+// relay.nn.sparse_add
+bool SparseAddRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                        const TypeReporter& reporter) {
+  ICHECK_EQ(types.size(), 5);
+  const auto* dense_data = types[0].as<TensorTypeNode>();
+  const auto* sparse_data = types[1].as<TensorTypeNode>();
+  ICHECK_EQ(sparse_data->shape.size(), 1);
+  const auto* sparse_indices = types[2].as<TensorTypeNode>();
+  ICHECK_EQ(sparse_indices->shape.size(), 1);
+  const auto* sparse_indptr = types[3].as<TensorTypeNode>();
+
+  reporter->Assign(types[4], TensorType(dense_data->shape, dense_data->dtype));
+  return true;
+}
+
+Expr MakeSparseAdd(Expr dense_data, Expr sparse_data, Expr sparse_indices, Expr sparse_indptr) {
+  static const Op& op = Op::Get("nn.sparse_add");
+  return Call(op, {dense_data, sparse_data, sparse_indices, sparse_indptr}, Attrs(), {});
+}
+
+TVM_REGISTER_GLOBAL("relay.op.nn._make.sparse_add").set_body_typed(MakeSparseAdd);
+
+RELAY_REGISTER_OP("nn.sparse_add")
+    .describe(R"code(Add a dense matrix X with sparse matrix Y. Only support square sparse matrix
+
+- **dense**: `(N, N)`
+- **sparse**: `(N, N)`
+
+- **out**: `(N, N)`.
+
+)code" TVM_ADD_FILELINE)
+    .set_attrs_type<SparseTransposeAttrs>()
+    .set_num_inputs(4)
+    .add_argument("dense_data", "2D Tensor", "Dense data matrix.")
+    .add_argument("sparse_data", "1D Tensor", "Sparse data matrix.")
+    .add_argument("sparse_indices", "1D Tensor", "Sparse indices matrix.")
+    .add_argument("sparse_indptr", "1D Tensor", "Sparse index pointer matrix.")
+    .set_support_level(1)
+    .add_type_rel("SparseAdd", SparseAddRel);
+
+
 }  // namespace relay
 }  // namespace tvm

tests/python/topi/python/test_topi_sparse.py

@@ -525,6 +525,29 @@ def test_sparse_dense_padded_alter_op():
         with tvm.transform.PassContext(opt_level=3, required_pass="AlterOpLayout"):
             x = relay.build(tvm.IRModule.from_expr(f), target=tvm.target.Target("cuda"))
 
+def test_sparse_add_csr():
+    M, K, density = 3, 49, 0.2
+    X_np = np.random.randn(M, K).astype("float32")
+    Y_sp_np = sp.random(M, K, density=density, format="csr", dtype="float32")
+    Y_np = Y_sp_np.todense()
+    Z_np = X_np + Y_np
+
+    Y_data = te.placeholder(shape=Y_sp_np.data.shape, dtype=str(Y_sp_np.data.dtype))
+    Y_indices = te.placeholder(shape=Y_sp_np.indices.shape, dtype=str(Y_sp_np.indices.dtype))
+    Y_indptr = te.placeholder(shape=Y_sp_np.indptr.shape, dtype=str(Y_sp_np.indptr.dtype))
+    X = te.placeholder(shape=X_np.shape, dtype=str(X_np.dtype))
+    Z = topi.nn.sparse_add(X, Y_data, Y_indices, Y_indptr)
+    s = te.create_schedule(Z.op)
+    func = tvm.build(s, [X, Y_data, Y_indices, Y_indptr, Z])
+    Z_tvm = tvm.nd.array(np.zeros(Z_np.shape, dtype=Z_np.dtype))
+    func(
+        tvm.nd.array(X_np),
+        tvm.nd.array(Y_sp_np.data),
+        tvm.nd.array(Y_sp_np.indices),
+        tvm.nd.array(Y_sp_np.indptr),
+        Z_tvm,
+    )
+    tvm.testing.assert_allclose(Z_tvm.asnumpy(), Z_np, atol=1e-4, rtol=1e-4)
 
 if __name__ == "__main__":
     test_csrmv()
@@ -537,3 +560,4 @@ def test_sparse_dense_padded_alter_op():
     test_sparse_dense_padded_alter_op()
     test_sparse_dense_csr_reverse()
     test_sparse_dense_bsr_reverse()
+    test_sparse_add_csr()