Repository: dmlc/dgl

Issue:
[Sparse] Support sparse-dense matmul and SDDMM
## 🔨Work Item


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:
This issue has been automatically marked as stale due to lack of activity. It will be closed if no further activity occurs. Thank you
#4999 is reverted due to design flaw. Will need two PRs to finish this task:
- [ ] Expose DGL's `SpMMCsr`, `SpMMCoo`, `SDDMMCsr` and `SDDMMCoo` to as include headers.
- [ ] Add a new `transpose_A` flag to the SpMM signature in the new sparse library to support implicit format materialization.
Sorry to merge it.

Resolving patch:
diff --git a/python/dgl/mock_sparse2/__init__.py b/python/dgl/mock_sparse2/__init__.py
--- a/python/dgl/mock_sparse2/__init__.py
+++ b/python/dgl/mock_sparse2/__init__.py
@@ -10,6 +10,8 @@
 from .sparse_matrix import *
 from .unary_op_diag import *
 from .unary_op_sp import *
+from .matmul import *
+from .sddmm import *
 
 
 def load_dgl_sparse():
diff --git a/python/dgl/mock_sparse2/matmul.py b/python/dgl/mock_sparse2/matmul.py
new file mode 100644
--- /dev/null
+++ b/python/dgl/mock_sparse2/matmul.py
@@ -0,0 +1,103 @@
+"""Matmul ops for SparseMatrix"""
+# pylint: disable=invalid-name
+from typing import Union
+
+import torch
+
+from .diag_matrix import DiagMatrix
+
+from .sparse_matrix import SparseMatrix
+
+__all__ = ["spmm"]
+
+
+def spmm(A: Union[SparseMatrix, DiagMatrix], X: torch.Tensor) -> torch.Tensor:
+    """Multiply a sparse matrix by a dense matrix.
+
+    Parameters
+    ----------
+    A : SparseMatrix or DiagMatrix
+        Sparse matrix of shape (N, M) with values of shape (nnz)
+    X : torch.Tensor
+        Dense tensor of shape (M, F) or (M)
+
+    Returns
+    -------
+    torch.Tensor
+        The multiplication result of shape (N, F) or (N)
+
+    Examples
+    --------
+
+    >>> row = torch.tensor([0, 1, 1])
+    >>> col = torch.tensor([1, 0, 1])
+    >>> val = torch.randn(len(row))
+    >>> A = create_from_coo(row, col, val)
+    >>> X = torch.randn(2, 3)
+    >>> result = dgl.sparse.spmm(A, X)
+    >>> print(type(result))
+    <class 'torch.Tensor'>
+    >>> print(result.shape)
+    torch.Size([2, 3])
+    """
+    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)}"
+
+    # The input is a DiagMatrix. Cast it to SparseMatrix
+    if not isinstance(A, SparseMatrix):
+        A = A.as_sparse()
+    return torch.ops.dgl_sparse.spmm(A.c_sparse_matrix, X)
+
+
+def mm_sp(
+    A1: SparseMatrix, A2: Union[torch.Tensor, SparseMatrix, DiagMatrix]
+) -> Union[torch.Tensor, SparseMatrix]:
+    """Internal function for multiplying a sparse matrix by
+    a dense/sparse/diagonal matrix.
+
+    Parameters
+    ----------
+    A1 : SparseMatrix
+        Matrix of shape (N, M), with values of shape (nnz1)
+    A2 : torch.Tensor, SparseMatrix, or DiagMatrix
+        If A2 is a dense tensor, it can have shapes of (M, P) or (M, ).
+        Otherwise it must have a shape of (M, P).
+
+    Returns
+    -------
+    torch.Tensor or SparseMatrix
+        The result of multiplication.
+
+        * It is a dense torch tensor if :attr:`A2` is so.
+        * It is a SparseMatrix object otherwise.
+
+    Examples
+    --------
+
+    >>> row = torch.tensor([0, 1, 1])
+    >>> col = torch.tensor([1, 0, 1])
+    >>> val = torch.randn(len(row))
+    >>> A1 = create_from_coo(row, col, val)
+    >>> A2 = torch.randn(2, 3)
+    >>> result = A1 @ A2
+    >>> print(type(result))
+    <class 'torch.Tensor'>
+    >>> print(result.shape)
+    torch.Size([2, 3])
+    """
+    assert isinstance(A2, (torch.Tensor, SparseMatrix, DiagMatrix)), (
+        f"Expect arg2 to be a torch Tensor, SparseMatrix, or DiagMatrix object,"
+        f"got {type(A2)}"
+    )
+
+    if isinstance(A2, torch.Tensor):
+        return spmm(A1, A2)
+    else:
+        raise NotImplementedError
+
+
+SparseMatrix.__matmul__ = mm_sp
diff --git a/python/dgl/mock_sparse2/sddmm.py b/python/dgl/mock_sparse2/sddmm.py
new file mode 100644
--- /dev/null
+++ b/python/dgl/mock_sparse2/sddmm.py
@@ -0,0 +1,54 @@
+"""Sampled Dense-Dense Matrix Multiplication (SDDMM) operator module."""
+import torch
+
+from .sparse_matrix import SparseMatrix
+
+__all__ = ["sddmm"]
+
+
+def sddmm(
+    A: SparseMatrix, mat1: torch.Tensor, mat2: torch.Tensor
+) -> SparseMatrix:
+    r"""Sampled-Dense-Dense Matrix Multiplication (SDDMM).
+
+    ``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
+
+    Parameters
+    ----------
+    A : SparseMatrix
+        Sparse matrix of shape `(M, N)`.
+    mat1 : Tensor
+        Dense matrix of shape `(M, K)`
+    mat2 : Tensor
+        Dense matrix of shape `(K, N)`
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix of shape `(M, N)`.
+
+    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.randn(3, 5)
+    >>> mat2 = torch.randn(5, 4)
+    >>> dgl.mock_sparse.sddmm(A, mat1, mat2)
+    SparseMatrix(indices=tensor([[1, 1, 2],
+            [2, 3, 3]]),
+    values=tensor([ 1.3097, -1.0977,  1.6953]),
+    shape=(3, 4), nnz=3)
+    """
+    return SparseMatrix(
+        torch.ops.dgl_sparse.sddmm(A.c_sparse_matrix, mat1, mat2)
+    )


Test patch:
diff --git a/tests/pytorch/mock_sparse2/test_matmul.py b/tests/pytorch/mock_sparse2/test_matmul.py
new file mode 100644
--- /dev/null
+++ b/tests/pytorch/mock_sparse2/test_matmul.py
@@ -0,0 +1,52 @@
+import sys
+
+import backend as F
+import pytest
+import torch
+
+from dgl.mock_sparse2 import val_like
+
+from .utils import (
+    clone_detach_and_grad,
+    rand_coo,
+    rand_csc,
+    rand_csr,
+    sparse_matrix_to_dense,
+    sparse_matrix_to_torch_sparse,
+)
+
+# TODO(#4818): Skipping tests on win.
+if not sys.platform.startswith("linux"):
+    pytest.skip("skipping tests on win", allow_module_level=True)
+
+
+@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])
+@pytest.mark.parametrize("out_dim", [None, 10])
+def test_spmm(create_func, shape, nnz, out_dim):
+    dev = F.ctx()
+    A = create_func(shape, nnz, dev)
+    if out_dim is not None:
+        X = torch.randn(shape[1], out_dim, requires_grad=True, device=dev)
+    else:
+        X = torch.randn(shape[1], requires_grad=True, device=dev)
+
+    sparse_result = A @ X
+    grad = torch.randn_like(sparse_result)
+    sparse_result.backward(grad)
+
+    adj = sparse_matrix_to_torch_sparse(A)
+    XX = clone_detach_and_grad(X)
+    torch_sparse_result = torch.sparse.mm(
+        adj, XX.view(-1, 1) if out_dim is None else XX
+    )
+    if out_dim is None:
+        torch_sparse_result = torch_sparse_result.view(-1)
+    torch_sparse_result.backward(grad)
+    assert torch.allclose(sparse_result, torch_sparse_result)
+    assert torch.allclose(X.grad, XX.grad)
+    assert torch.allclose(
+        adj.grad.coalesce().to_dense(),
+        sparse_matrix_to_dense(val_like(A, A.val.grad)),
+    )
diff --git a/tests/pytorch/mock_sparse2/test_sddmm.py b/tests/pytorch/mock_sparse2/test_sddmm.py
new file mode 100644
--- /dev/null
+++ b/tests/pytorch/mock_sparse2/test_sddmm.py
@@ -0,0 +1,50 @@
+import sys
+
+import backend as F
+import pytest
+import torch
+from dgl.mock_sparse2 import sddmm
+
+from .utils import clone_detach_and_grad, rand_coo, rand_csc, rand_csr
+
+# TODO(#4818): Skipping tests on win.
+if not sys.platform.startswith("linux"):
+    pytest.skip("skipping tests on win", allow_module_level=True)
+
+
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("shape", [(2, 3), (5, 2)])
+@pytest.mark.parametrize("nnz", [2, 10])
+@pytest.mark.parametrize("hidden", [1, 5])
+def test_sddmm(create_func, shape, nnz, hidden):
+    dev = F.ctx()
+    A = create_func(shape, nnz, dev)
+    if hidden > 1:
+        B = torch.rand(shape[0], hidden, requires_grad=True, device=dev)
+        C = torch.rand(hidden, shape[1], requires_grad=True, device=dev)
+    else:
+        B = torch.rand(shape[0], requires_grad=True, device=dev)
+        C = torch.rand(shape[1], 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 = sddmm(A, B, C)
+
+    grad = torch.rand_like(sparse_result.val)
+    sparse_result.val.backward(grad)
+
+    if hidden == 1:
+        dense_result = dense_B.view(-1, 1) @ dense_C.view(1, -1)
+    else:
+        dense_result = dense_B @ dense_C
+
+    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)
+    assert torch.allclose(dense_C.grad, C.grad)
+    assert torch.allclose(dense_B.grad, B.grad)
+    assert torch.allclose(A_val_clone.grad, A.val.grad)
diff --git a/tests/pytorch/mock_sparse2/utils.py b/tests/pytorch/mock_sparse2/utils.py
new file mode 100644
--- /dev/null
+++ b/tests/pytorch/mock_sparse2/utils.py
@@ -0,0 +1,62 @@
+import torch
+from dgl.mock_sparse2 import (
+    create_from_coo,
+    create_from_csc,
+    create_from_csr,
+    SparseMatrix,
+)
+
+
+def clone_detach_and_grad(t):
+    t = t.clone().detach()
+    t.requires_grad_()
+    return t
+
+
+def rand_coo(shape, nnz, dev):
+    row = torch.randint(0, shape[0], (nnz,), device=dev)
+    col = torch.randint(0, shape[1], (nnz,), device=dev)
+    val = torch.randn(nnz, device=dev, requires_grad=True)
+    return create_from_coo(row, col, val, shape)
+
+
+def rand_csr(shape, nnz, dev):
+    row = torch.randint(0, shape[0], (nnz,), device=dev)
+    col = torch.randint(0, shape[1], (nnz,), device=dev)
+    val = torch.randn(nnz, 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
+    indptr = torch.cumsum(indptr, 0)
+    indices = col
+    return create_from_csr(indptr, indices, val, shape=shape)
+
+
+def rand_csc(shape, nnz, dev):
+    row = torch.randint(0, shape[0], (nnz,), device=dev)
+    col = torch.randint(0, shape[1], (nnz,), device=dev)
+    val = torch.randn(nnz, 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
+    indptr = torch.cumsum(indptr, 0)
+    indices = row
+    return create_from_csc(indptr, indices, val, shape=shape)
+
+
+def sparse_matrix_to_dense(A: SparseMatrix):
+    dense = A.dense()
+    dense.requires_grad_()
+    return dense
+
+
+def sparse_matrix_to_torch_sparse(A: SparseMatrix):
+    row, col = A.coo()
+    edge_index = torch.cat((row.unsqueeze(0), col.unsqueeze(0)), 0)
+    shape = A.shape
+    val = A.val.clone().detach()
+    if len(A.val.shape) > 1:
+        shape += (A.val.shape[-1],)
+    ret = torch.sparse_coo_tensor(edge_index, val, shape).coalesce()
+    ret.requires_grad_()
+    return ret