Repository: dmlc/dgl

Issue:
[Sparse] Support sparse-sparse and sparse-diag matmul
## 🔨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:
This issue has been automatically marked as stale due to lack of activity. It will be closed if no further activity occurs. Thank you

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
@@ -4,11 +4,11 @@
 
 import torch
 
-from .diag_matrix import DiagMatrix
+from .diag_matrix import diag, DiagMatrix
 
 from .sparse_matrix import SparseMatrix
 
-__all__ = ["spmm"]
+__all__ = ["spmm", "spspmm", "mm"]
 
 
 def spmm(A: Union[SparseMatrix, DiagMatrix], X: torch.Tensor) -> torch.Tensor:
@@ -53,51 +53,138 @@ def spmm(A: Union[SparseMatrix, DiagMatrix], X: torch.Tensor) -> torch.Tensor:
     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.
+def _diag_diag_mm(A1: DiagMatrix, A2: DiagMatrix) -> DiagMatrix:
+    """Internal function for multiplying a diagonal matrix by a diagonal matrix
 
     Parameters
     ----------
-    A1 : SparseMatrix
+    A1 : DiagMatrix
         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).
+    A2 : DiagMatrix
+        Matrix of shape (M, P), with values of shape (nnz2)
 
     Returns
     -------
-    torch.Tensor or SparseMatrix
+    DiagMatrix
         The result of multiplication.
+    """
+    M, N = A1.shape
+    N, P = A2.shape
+    common_diag_len = min(M, N, P)
+    new_diag_len = min(M, P)
+    diag_val = torch.zeros(new_diag_len)
+    diag_val[:common_diag_len] = (
+        A1.val[:common_diag_len] * A2.val[:common_diag_len]
+    )
+    return diag(diag_val.to(A1.device), (M, P))
+
+
+def spspmm(
+    A1: Union[SparseMatrix, DiagMatrix], A2: Union[SparseMatrix, DiagMatrix]
+) -> Union[SparseMatrix, DiagMatrix]:
+    """Multiply a sparse matrix by a sparse matrix. The non-zero values of the
+    two sparse matrices must be 1D.
+
+    Parameters
+    ----------
+    A1 : SparseMatrix or DiagMatrix
+        Sparse matrix of shape (N, M) with values of shape (nnz)
+    A2 : SparseMatrix or DiagMatrix
+        Sparse matrix of shape (M, P) with values of shape (nnz)
+
+    Returns
+    -------
+    SparseMatrix or DiagMatrix
+        The result of multiplication. It is a DiagMatrix object if both matrices
+        are DiagMatrix objects. It is a SparseMatrix object otherwise.
+
+    Examples
+    --------
+
+    >>> row1 = torch.tensor([0, 1, 1])
+    >>> col1 = torch.tensor([1, 0, 1])
+    >>> val1 = torch.ones(len(row1))
+    >>> A1 = create_from_coo(row1, col1, val1)
+
+    >>> row2 = torch.tensor([0, 1, 1])
+    >>> col2 = torch.tensor([0, 2, 1])
+    >>> val2 = torch.ones(len(row2))
+    >>> A2 = create_from_coo(row2, col2, val2)
+    >>> result = dgl.sparse.spspmm(A1, A2)
+    >>> print(result)
+    SparseMatrix(indices=tensor([[0, 0, 1, 1, 1],
+                                 [1, 2, 0, 1, 2]]),
+                 values=tensor([1., 1., 1., 1., 1.]),
+                 shape=(2, 3), nnz=5)
+    """
+    assert isinstance(
+        A1, (SparseMatrix, DiagMatrix)
+    ), f"Expect A1 to be a SparseMatrix or DiagMatrix object, got {type(A1)}"
+    assert isinstance(
+        A2, (SparseMatrix, DiagMatrix)
+    ), f"Expect A2 to be a SparseMatrix or DiagMatrix object, got {type(A2)}"
+
+    if isinstance(A1, DiagMatrix) and isinstance(A2, DiagMatrix):
+        return _diag_diag_mm(A1, A2)
+    if isinstance(A1, DiagMatrix):
+        A1 = A1.as_sparse()
+    if isinstance(A2, DiagMatrix):
+        A2 = A2.as_sparse()
+    return SparseMatrix(
+        torch.ops.dgl_sparse.spspmm(A1.c_sparse_matrix, A2.c_sparse_matrix)
+    )
+
+
+def mm(
+    A1: Union[SparseMatrix, DiagMatrix],
+    A2: Union[torch.Tensor, SparseMatrix, DiagMatrix],
+) -> Union[torch.Tensor, SparseMatrix, DiagMatrix]:
+    """Multiply a sparse/diagonal matrix by a dense/sparse/diagonal matrix.
+    If an input is a SparseMatrix or DiagMatrix, its non-zero values should
+    be 1-D.
+
+    Parameters
+    ----------
+    A1 : SparseMatrix or DiagMatrix
+        Matrix of shape (N, M), with values of shape (nnz1)
+    A2 : torch.Tensor, SparseMatrix, or DiagMatrix
+        Matrix of shape (M, P). If it is a SparseMatrix or DiagMatrix,
+        it should have values of shape (nnz2).
+
+    Returns
+    -------
+    torch.Tensor or DiagMatrix or SparseMatrix
+        The result of multiplication of shape (N, P)
 
         * It is a dense torch tensor if :attr:`A2` is so.
+        * It is a DiagMatrix object if both :attr:`A1` and :attr:`A2` are 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
+    >>> val = torch.randn(3)
+    >>> A1 = diag(val)
+    >>> A2 = torch.randn(3, 2)
+    >>> result = dgl.sparse.mm(A1, A2)
     >>> print(type(result))
     <class 'torch.Tensor'>
     >>> print(result.shape)
-    torch.Size([2, 3])
+    torch.Size([3, 2])
     """
+    assert isinstance(
+        A1, (SparseMatrix, DiagMatrix)
+    ), f"Expect arg1 to be a SparseMatrix, or DiagMatrix object, got {type(A1)}."
     assert isinstance(A2, (torch.Tensor, SparseMatrix, DiagMatrix)), (
-        f"Expect arg2 to be a torch Tensor, SparseMatrix, or DiagMatrix object,"
-        f"got {type(A2)}"
+        f"Expect arg2 to be a torch Tensor, SparseMatrix, or DiagMatrix"
+        f"object, got {type(A2)}."
     )
-
     if isinstance(A2, torch.Tensor):
         return spmm(A1, A2)
-    else:
-        raise NotImplementedError
+    if isinstance(A1, DiagMatrix) and isinstance(A2, DiagMatrix):
+        return _diag_diag_mm(A1, A2)
+    return spspmm(A1, A2)
 
 
-SparseMatrix.__matmul__ = mm_sp
+SparseMatrix.__matmul__ = mm
+DiagMatrix.__matmul__ = mm
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
@@ -11,8 +11,8 @@ def sddmm(
 ) -> 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``
+    ``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:
@@ -20,19 +20,23 @@ def sddmm(
     .. math::
         out = (mat1 @ mat2) * A
 
+    In particular, :attr:`mat1` and :attr:`mat2` can be 1-D, then ``mat1 @
+    mat2`` becomes the out-product of the two vector (which results in a
+    matrix).
+
     Parameters
     ----------
     A : SparseMatrix
-        Sparse matrix of shape `(M, N)`.
+        Sparse matrix of shape ``(M, N)``.
     mat1 : Tensor
-        Dense matrix of shape `(M, K)`
+        Dense matrix of shape ``(M, K)`` or ``(M,)``
     mat2 : Tensor
-        Dense matrix of shape `(K, N)`
+        Dense matrix of shape ``(K, N)`` or ``(N,)``
 
     Returns
     -------
     SparseMatrix
-        Sparse matrix of shape `(M, N)`.
+        Sparse matrix of shape ``(M, N)``.
 
     Examples
     --------


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
@@ -51,3 +51,41 @@ def test_spmm(create_func, shape, nnz, out_dim):
         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)])
+@pytest.mark.parametrize("shape_k", [3, 4])
+@pytest.mark.parametrize("nnz1", [1, 10])
+@pytest.mark.parametrize("nnz2", [1, 10])
+def test_sparse_sparse_mm(
+    create_func1, create_func2, shape_n_m, shape_k, nnz1, nnz2
+):
+    dev = F.ctx()
+    shape1 = shape_n_m
+    shape2 = (shape_n_m[1], shape_k)
+    A1 = create_func1(shape1, nnz1, dev)
+    A2 = create_func2(shape2, nnz2, dev)
+    A3 = A1 @ A2
+    grad = torch.randn_like(A3.val)
+    A3.val.backward(grad)
+
+    torch_A1 = sparse_matrix_to_torch_sparse(A1)
+    torch_A2 = sparse_matrix_to_torch_sparse(A2)
+    torch_A3 = torch.sparse.mm(torch_A1, torch_A2)
+    torch_A3_grad = sparse_matrix_to_torch_sparse(A3, grad)
+    torch_A3.backward(torch_A3_grad)
+
+    with torch.no_grad():
+        assert torch.allclose(A3.dense(), torch_A3.to_dense(), atol=1e-05)
+        assert torch.allclose(
+            val_like(A1, A1.val.grad).dense(),
+            torch_A1.grad.to_dense(),
+            atol=1e-05,
+        )
+        assert torch.allclose(
+            val_like(A2, A2.val.grad).dense(),
+            torch_A2.grad.to_dense(),
+            atol=1e-05,
+        )
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
@@ -13,7 +13,7 @@
 
 
 @pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
-@pytest.mark.parametrize("shape", [(2, 3), (5, 2)])
+@pytest.mark.parametrize("shape", [(5, 5), (5, 4)])
 @pytest.mark.parametrize("nnz", [2, 10])
 @pytest.mark.parametrize("hidden", [1, 5])
 def test_sddmm(create_func, shape, nnz, hidden):
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
@@ -1,4 +1,6 @@
+import numpy as np
 import torch
+
 from dgl.mock_sparse2 import (
     create_from_coo,
     create_from_csc,
@@ -6,6 +8,9 @@
     SparseMatrix,
 )
 
+np.random.seed(42)
+torch.random.manual_seed(42)
+
 
 def clone_detach_and_grad(t):
     t = t.clone().detach()
@@ -14,33 +19,80 @@ def clone_detach_and_grad(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)
+    # 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)
     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)
+    # 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)
     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
+    row_sorted, row_sorted_idx = torch.sort(row)
+    indices = col[row_sorted_idx]
     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)
+    # 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)
+    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)
+    col_sorted, col_sorted_idx = torch.sort(col)
+    indices = row[col_sorted_idx]
+    return create_from_csc(indptr, indices, val, shape=shape)
+
+
+def rand_coo_uncoalesced(shape, nnz, dev):
+    # Create a sparse matrix with possible duplicate entries.
+    row = torch.randint(shape[0], (nnz,), device=dev)
+    col = torch.randint(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_uncoalesced(shape, nnz, dev):
+    # Create a sparse matrix with possible duplicate entries.
+    row = torch.randint(shape[0], (nnz,), device=dev)
+    col = torch.randint(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)
+    row_sorted, row_sorted_idx = torch.sort(row)
+    indices = col[row_sorted_idx]
+    return create_from_csr(indptr, indices, val, shape=shape)
+
+
+def rand_csc_uncoalesced(shape, nnz, dev):
+    # Create a sparse matrix with possible duplicate entries.
+    row = torch.randint(shape[0], (nnz,), device=dev)
+    col = torch.randint(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
+    col_sorted, col_sorted_idx = torch.sort(col)
+    indices = row[col_sorted_idx]
     return create_from_csc(indptr, indices, val, shape=shape)
 
 
@@ -50,11 +102,13 @@ def sparse_matrix_to_dense(A: SparseMatrix):
     return dense
 
 
-def sparse_matrix_to_torch_sparse(A: SparseMatrix):
+def sparse_matrix_to_torch_sparse(A: SparseMatrix, val=None):
     row, col = A.coo()
     edge_index = torch.cat((row.unsqueeze(0), col.unsqueeze(0)), 0)
     shape = A.shape
-    val = A.val.clone().detach()
+    if val is None:
+        val = A.val
+    val = val.clone().detach()
     if len(A.val.shape) > 1:
         shape += (A.val.shape[-1],)
     ret = torch.sparse_coo_tensor(edge_index, val, shape).coalesce()