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:

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,7 @@
 from .sparse_matrix import *
 from .unary_op_diag import *
 from .unary_op_sp import *
+from .matmul 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,107 @@
+"""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 result of multiplication
+
+    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 = 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)}"
+    assert (
+        A.shape[1] == X.shape[0]
+    ), f"Expect arg1.shape[1] == arg2.shape[0], got {A.shape[1]} and {X.shape[0]}"
+    val_dim = len(A.val.shape)
+    assert val_dim == 1, f"Expect arg1.val to be a 1D tensor, got {val_dim}D"
+    val_dim = len(X.shape)
+    assert val_dim <= 2, f"Expect arg2 to be a 1D/2D tensor, got {val_dim}D"
+
+    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, got {type(A2)}"
+
+    if isinstance(A2, torch.Tensor):
+        return spmm(A1, A2)
+    else:
+        raise NotImplementedError
+
+
+SparseMatrix.__matmul__ = mm_sp


Test patch:
diff --git a/tests/pytorch/mock_sparse2/test_mm.py b/tests/pytorch/mock_sparse2/test_mm.py
new file mode 100644
--- /dev/null
+++ b/tests/pytorch/mock_sparse2/test_mm.py
@@ -0,0 +1,114 @@
+import sys
+
+import backend as F
+import pytest
+import torch
+
+from dgl.mock_sparse2 import create_from_coo, create_from_csc, create_from_csr
+
+# TODO(#4818): Skipping tests on win.
+if not sys.platform.startswith("linux"):
+    pytest.skip("skipping tests on win", allow_module_level=True)
+
+
+def get_adj(A):
+    row, col = A.coo()
+    edge_index = torch.cat((row.unsqueeze(0), col.unsqueeze(0)), 0)
+    shape = A.shape
+    val = A.val.detach()
+    if len(A.val.shape) > 1:
+        shape += (A.val.shape[-1],)
+    return torch.sparse_coo_tensor(edge_index, val, shape).coalesce()
+
+
+def test_spmm_coo():
+    dev = F.ctx()
+    # A: shape (N, M), X: shape (M, F)
+    row = torch.tensor([0, 1, 1, 1]).to(dev)
+    col = torch.tensor([1, 0, 1, 2]).to(dev)
+    val = torch.randn(len(row), requires_grad=True, device=dev)
+    A = create_from_coo(row, col, val)
+    X = torch.randn(3, 4, requires_grad=True, device=dev)
+    sparse_result = A @ X
+    grad = torch.randn_like(sparse_result)
+    sparse_result.backward(grad)
+
+    adj = get_adj(A)
+    adj.requires_grad_()
+    XX = X.clone().detach()
+    XX.requires_grad_()
+    dense_result = torch.sparse.mm(adj, XX)
+    dense_result.backward(grad)
+    assert torch.allclose(sparse_result, dense_result)
+    assert torch.allclose(X.grad, XX.grad)
+    assert torch.allclose(adj.grad.coalesce().values(), val.grad)
+
+
+def test_spmm_coo_one_dim_rhs():
+    dev = F.ctx()
+    # A: shape (N, M), X: shape (M,)
+    row = torch.tensor([0, 1, 1, 1]).to(dev)
+    col = torch.tensor([1, 0, 1, 2]).to(dev)
+    val = torch.randn(len(row), requires_grad=True, device=dev)
+    A = create_from_coo(row, col, val)
+    X = torch.randn(3, requires_grad=True, device=dev)
+    sparse_result = A @ X
+    grad = torch.randn_like(sparse_result)
+    sparse_result.backward(grad)
+
+    adj = get_adj(A)
+    adj.requires_grad_()
+    XX = X.clone().detach()
+    XX.requires_grad_()
+    dense_result = torch.sparse.mm(adj, XX.view(-1, 1))
+    dense_result = dense_result.view(-1)
+    dense_result.backward(grad)
+    assert torch.allclose(sparse_result, dense_result)
+    assert torch.allclose(X.grad, XX.grad)
+    assert torch.allclose(adj.grad.coalesce().values(), val.grad)
+
+
+def test_spmm_csr():
+    dev = F.ctx()
+    # A: shape (N, M), X: shape (M, F)
+    indptr = torch.tensor([0, 1, 4]).to(dev)
+    indices = torch.tensor([1, 0, 1, 2]).to(dev)
+    val = torch.randn(len(indices), requires_grad=True, device=dev)
+    A = create_from_csr(indptr, indices, val, shape=(2, 3))
+    X = torch.randn(3, 4, requires_grad=True, device=dev)
+    sparse_result = A @ X
+    grad = torch.randn_like(sparse_result)
+    sparse_result.backward(grad)
+
+    adj = get_adj(A)
+    adj.requires_grad_()
+    XX = X.clone().detach()
+    XX.requires_grad_()
+    dense_result = torch.sparse.mm(adj, XX)
+    dense_result.backward(grad)
+    assert torch.allclose(sparse_result, dense_result)
+    assert torch.allclose(X.grad, XX.grad)
+    assert torch.allclose(adj.grad.coalesce().values(), val.grad)
+
+
+def test_spmm_csc():
+    dev = F.ctx()
+    # A: shape (N, M), X: shape (M, F)
+    indptr = torch.tensor([0, 1, 3, 4]).to(dev)
+    indices = torch.tensor([0, 0, 1, 1]).to(dev)
+    val = torch.randn(len(indices), requires_grad=True, device=dev)
+    A = create_from_csc(indptr, indices, val, shape=(2, 3))
+    X = torch.randn(3, 4, requires_grad=True, device=dev)
+    sparse_result = A @ X
+    grad = torch.randn_like(sparse_result)
+    sparse_result.backward(grad)
+
+    adj = get_adj(A)
+    adj.requires_grad_()
+    XX = X.clone().detach()
+    XX.requires_grad_()
+    dense_result = torch.sparse.mm(adj, XX)
+    dense_result.backward(grad)
+    assert torch.allclose(sparse_result, dense_result)
+    assert torch.allclose(X.grad, XX.grad)
+    assert torch.allclose(adj.grad.coalesce().values(), val.grad)