Repository: dmlc/dgl

Issue:
[Sparse] Batched SpMM and SDDMM
## 🔨Work Item

**IMPORTANT:**
* This template is only for dev team to track project progress. For feature request or bug report, please use the corresponding issue templates.
* DO NOT create a new work item if the purpose is to fix an existing issue or feature request. We will directly use the issue in the project tracker.

Project tracker: https://github.com/orgs/dmlc/projects/2

## Description

<!-- short description of the work item -->

## Depending work items or issues

<!-- what must be done before this -->



Discussion:
Implement the mock first https://github.com/dmlc/dgl/issues/5055 to unblock #5056 
Let's wait for the example get drafted and see how well the API is, before we implement the release version of BSDDMM.
Extend this task to include bspmm which is needed by the GAT example.

Resolving patch:
diff --git a/python/dgl/mock_sparse2/matmul.py b/python/dgl/mock_sparse2/matmul.py
--- a/python/dgl/mock_sparse2/matmul.py
+++ b/python/dgl/mock_sparse2/matmul.py
@@ -8,7 +8,7 @@
 
 from .sparse_matrix import SparseMatrix
 
-__all__ = ["spmm", "spspmm", "mm"]
+__all__ = ["spmm", "bspmm", "spspmm", "mm"]
 
 
 def spmm(A: Union[SparseMatrix, DiagMatrix], X: torch.Tensor) -> torch.Tensor:
@@ -53,6 +53,44 @@ def spmm(A: Union[SparseMatrix, DiagMatrix], X: torch.Tensor) -> torch.Tensor:
     return torch.ops.dgl_sparse.spmm(A.c_sparse_matrix, X)
 
 
+def bspmm(A: Union[SparseMatrix, DiagMatrix], X: torch.Tensor) -> torch.Tensor:
+    """Multiply a sparse matrix by a dense matrix by batches.
+
+    Parameters
+    ----------
+    A : SparseMatrix or DiagMatrix
+        Sparse matrix of shape (N, M, B) with values of shape (nnz)
+    X : torch.Tensor
+        Dense tensor of shape (M, F, B)
+
+    Returns
+    -------
+    torch.Tensor
+        The multiplication result of shape (N, F, B)
+
+    Examples
+    --------
+
+    >>> row = torch.tensor([0, 1, 1])
+    >>> col = torch.tensor([1, 0, 2])
+    >>> val = torch.randn(len(row), 2)
+    >>> A = create_from_coo(row, col, val, shape=(3, 3))
+    >>> X = torch.randn(3, 3, 2)
+    >>> result = dgl.sparse.bspmm(A, X)
+    >>> print(type(result))
+    <class 'torch.Tensor'>
+    >>> print(result.shape)
+    torch.Size([3, 3, 2])
+    """
+    assert isinstance(
+        A, (SparseMatrix, DiagMatrix)
+    ), f"Expect arg1 to be a SparseMatrix or DiagMatrix object, got {type(A)}"
+    assert isinstance(
+        X, torch.Tensor
+    ), f"Expect arg2 to be a torch.Tensor, got {type(X)}"
+    return spmm(A, X)
+
+
 def _diag_diag_mm(A1: DiagMatrix, A2: DiagMatrix) -> DiagMatrix:
     """Internal function for multiplying a diagonal matrix by a diagonal matrix
 
diff --git a/python/dgl/mock_sparse2/sddmm.py b/python/dgl/mock_sparse2/sddmm.py
--- a/python/dgl/mock_sparse2/sddmm.py
+++ b/python/dgl/mock_sparse2/sddmm.py
@@ -3,7 +3,7 @@
 
 from .sparse_matrix import SparseMatrix
 
-__all__ = ["sddmm"]
+__all__ = ["sddmm", "bsddmm"]
 
 
 def sddmm(
@@ -47,7 +47,7 @@ def sddmm(
     >>> A = create_from_coo(row, col, val, (3, 4))
     >>> mat1 = torch.randn(3, 5)
     >>> mat2 = torch.randn(5, 4)
-    >>> dgl.mock_sparse.sddmm(A, mat1, mat2)
+    >>> dgl.sparse.sddmm(A, mat1, mat2)
     SparseMatrix(indices=tensor([[1, 1, 2],
             [2, 3, 3]]),
     values=tensor([ 1.3097, -1.0977,  1.6953]),
@@ -56,3 +56,55 @@ def sddmm(
     return SparseMatrix(
         torch.ops.dgl_sparse.sddmm(A.c_sparse_matrix, mat1, mat2)
     )
+
+
+def bsddmm(
+    A: SparseMatrix, mat1: torch.Tensor, mat2: torch.Tensor
+) -> SparseMatrix:
+    r"""Sampled-Dense-Dense Matrix Multiplication (SDDMM) by batches.
+
+    ``sddmm`` multiplies two dense matrices :attr:`mat1` and :attr:`mat2`
+    at the nonzero locations of sparse matrix :attr:`A`. Values of :attr:`A`
+    is not considered during the computation.
+
+    Mathematically ``sddmm`` is formulated as:
+
+    .. math::
+        out = (mat1 @ mat2) * A
+
+    The batch dimension is the last dimension for input matrices. In particular,
+    if the sparse matrix has scalar non-zero values, it will be broadcasted
+    for bsddmm.
+
+    Parameters
+    ----------
+    A : SparseMatrix
+        Sparse matrix of shape ``(M, N)`` or ``(M, N, B)``.
+    mat1 : Tensor
+        Dense matrix of shape ``(M, K, B)``
+    mat2 : Tensor
+        Dense matrix of shape ``(K, N, B)``
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix of shape ``(M, N, B)``.
+
+    Examples
+    --------
+
+    >>> row = torch.tensor([1, 1, 2])
+    >>> col = torch.tensor([2, 3, 3])
+    >>> val = torch.arange(1, 4).float()
+    >>> A = create_from_coo(row, col, val, (3, 4))
+    >>> mat1 = torch.arange(0, 3 * 5 * 2).view(3, 5, 2).float()
+    >>> mat2 = torch.arange(0, 5 * 4 * 2).view(5, 4, 2).float()
+    >>> dgl.sparse.bsddmm(A, mat1, mat2)
+    SparseMatrix(indices=tensor([[1, 1, 2],
+            [2, 3, 3]]),
+    values=tensor([[1560., 1735.],
+            [3400., 3770.],
+            [8400., 9105.]]),
+    shape=(3, 4), nnz=3)
+    """
+    return sddmm(A, mat1, mat2)


Test patch:
diff --git a/tests/pytorch/mock_sparse2/test_matmul.py b/tests/pytorch/mock_sparse2/test_matmul.py
--- a/tests/pytorch/mock_sparse2/test_matmul.py
+++ b/tests/pytorch/mock_sparse2/test_matmul.py
@@ -4,10 +4,11 @@
 import pytest
 import torch
 
-from dgl.mock_sparse2 import create_from_coo, val_like
+from dgl.mock_sparse2 import bspmm, create_from_coo, val_like
 
 from .utils import (
     clone_detach_and_grad,
+    dense_mask,
     rand_coo,
     rand_csc,
     rand_csr,
@@ -53,6 +54,34 @@ def test_spmm(create_func, shape, nnz, out_dim):
     )
 
 
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("shape", [(2, 7), (5, 2)])
+@pytest.mark.parametrize("nnz", [1, 10])
+def test_bspmm(create_func, shape, nnz):
+    dev = F.ctx()
+    A = create_func(shape, nnz, dev, 2)
+    X = torch.randn(shape[1], 10, 2, requires_grad=True, device=dev)
+
+    sparse_result = bspmm(A, X)
+    grad = torch.randn_like(sparse_result)
+    sparse_result.backward(grad)
+
+    XX = clone_detach_and_grad(X)
+    torch_A = A.dense().clone().detach().requires_grad_()
+    torch_result = torch_A.permute(2, 0, 1) @ XX.permute(2, 0, 1)
+
+    torch_result.backward(grad.permute(2, 0, 1))
+    assert torch.allclose(
+        sparse_result.permute(2, 0, 1), torch_result, atol=1e-05
+    )
+    assert torch.allclose(X.grad, XX.grad, atol=1e-05)
+    assert torch.allclose(
+        dense_mask(torch_A.grad, A),
+        sparse_matrix_to_dense(val_like(A, A.val.grad)),
+        atol=1e-05,
+    )
+
+
 @pytest.mark.parametrize("create_func1", [rand_coo, rand_csr, rand_csc])
 @pytest.mark.parametrize("create_func2", [rand_coo, rand_csr, rand_csc])
 @pytest.mark.parametrize("shape_n_m", [(5, 5), (5, 6)])
diff --git a/tests/pytorch/mock_sparse2/test_sddmm.py b/tests/pytorch/mock_sparse2/test_sddmm.py
--- a/tests/pytorch/mock_sparse2/test_sddmm.py
+++ b/tests/pytorch/mock_sparse2/test_sddmm.py
@@ -3,7 +3,7 @@
 import backend as F
 import pytest
 import torch
-from dgl.mock_sparse2 import sddmm
+from dgl.mock_sparse2 import bsddmm, sddmm
 
 from .utils import clone_detach_and_grad, rand_coo, rand_csc, rand_csr
 
@@ -48,3 +48,36 @@ def test_sddmm(create_func, shape, nnz, hidden):
     assert torch.allclose(dense_C.grad, C.grad, atol=1e-05)
     assert torch.allclose(dense_B.grad, B.grad, atol=1e-05)
     assert torch.allclose(A_val_clone.grad, A.val.grad, atol=1e-05)
+
+
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("shape", [(5, 5), (5, 4)])
+@pytest.mark.parametrize("nnz", [2, 10])
+@pytest.mark.parametrize("nz_dim", [2, 10])
+def test_bsddmm(create_func, shape, nnz, nz_dim):
+    dev = F.ctx()
+    hidden = 2
+    A = create_func(shape, nnz, dev, nz_dim)
+    B = torch.rand(shape[0], hidden, nz_dim, requires_grad=True, device=dev)
+    C = torch.rand(hidden, shape[1], nz_dim, requires_grad=True, device=dev)
+
+    A_val_clone = clone_detach_and_grad(A.val)
+    dense_B = clone_detach_and_grad(B)
+    dense_C = clone_detach_and_grad(C)
+
+    sparse_result = bsddmm(A, B, C)
+
+    grad = torch.rand_like(sparse_result.val)
+    sparse_result.val.backward(grad)
+
+    dense_result = dense_B.permute(2, 0, 1) @ dense_C.permute(2, 0, 1)
+    dense_result = dense_result.permute(1, 2, 0)
+
+    row, col = A.coo()
+    dense_val = dense_result[row, col] * A_val_clone
+    dense_val.backward(grad)
+
+    assert torch.allclose(dense_val, sparse_result.val, atol=1e-05)
+    assert torch.allclose(dense_C.grad, C.grad, atol=1e-05)
+    assert torch.allclose(dense_B.grad, B.grad, atol=1e-05)
+    assert torch.allclose(A_val_clone.grad, A.val.grad, atol=1e-05)
diff --git a/tests/pytorch/mock_sparse2/utils.py b/tests/pytorch/mock_sparse2/utils.py
--- a/tests/pytorch/mock_sparse2/utils.py
+++ b/tests/pytorch/mock_sparse2/utils.py
@@ -18,23 +18,29 @@ def clone_detach_and_grad(t):
     return t
 
 
-def rand_coo(shape, nnz, dev):
+def rand_coo(shape, nnz, dev, nz_dim=None):
     # Create a sparse matrix without duplicate entries.
     nnzid = np.random.choice(shape[0] * shape[1], nnz, replace=False)
     nnzid = torch.tensor(nnzid, device=dev).long()
     row = torch.div(nnzid, shape[1], rounding_mode="floor")
     col = nnzid % shape[1]
-    val = torch.randn(nnz, device=dev, requires_grad=True)
+    if nz_dim is None:
+        val = torch.randn(nnz, device=dev, requires_grad=True)
+    else:
+        val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True)
     return create_from_coo(row, col, val, shape)
 
 
-def rand_csr(shape, nnz, dev):
+def rand_csr(shape, nnz, dev, nz_dim=None):
     # Create a sparse matrix without duplicate entries.
     nnzid = np.random.choice(shape[0] * shape[1], nnz, replace=False)
     nnzid = torch.tensor(nnzid, device=dev).long()
     row = torch.div(nnzid, shape[1], rounding_mode="floor")
     col = nnzid % shape[1]
-    val = torch.randn(nnz, device=dev, requires_grad=True)
+    if nz_dim is None:
+        val = torch.randn(nnz, device=dev, requires_grad=True)
+    else:
+        val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True)
     indptr = torch.zeros(shape[0] + 1, device=dev, dtype=torch.int64)
     for r in row.tolist():
         indptr[r + 1] += 1
@@ -44,13 +50,16 @@ def rand_csr(shape, nnz, dev):
     return create_from_csr(indptr, indices, val, shape=shape)
 
 
-def rand_csc(shape, nnz, dev):
+def rand_csc(shape, nnz, dev, nz_dim=None):
     # Create a sparse matrix without duplicate entries.
     nnzid = np.random.choice(shape[0] * shape[1], nnz, replace=False)
     nnzid = torch.tensor(nnzid, device=dev).long()
     row = torch.div(nnzid, shape[1], rounding_mode="floor")
     col = nnzid % shape[1]
-    val = torch.randn(nnz, device=dev, requires_grad=True)
+    if nz_dim is None:
+        val = torch.randn(nnz, device=dev, requires_grad=True)
+    else:
+        val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True)
     indptr = torch.zeros(shape[1] + 1, device=dev, dtype=torch.int64)
     for c in col.tolist():
         indptr[c + 1] += 1
@@ -114,3 +123,11 @@ def sparse_matrix_to_torch_sparse(A: SparseMatrix, val=None):
     ret = torch.sparse_coo_tensor(edge_index, val, shape).coalesce()
     ret.requires_grad_()
     return ret
+
+
+def dense_mask(dense, sparse):
+    ret = torch.zeros_like(dense)
+    row, col = sparse.coo()
+    for r, c in zip(row, col):
+        ret[r, c] = dense[r, c]
+    return ret