[Sparse] Add sparse matrix C++ implementation (#4773)

* [Sparse] Add sparse matrix C++ implementation

* Add documentation

* Update

* Minor fix

* Move Python code to dgl/mock_sparse2

* Move headers to include

* lint

* Update

* Add dgl_sparse directory

* Move src code to dgl_sparse

* Add __init__.py in tests to avoid naming conflict

* Add dgl sparse so in Jenkinsfile

* Complete docstring & SparseMatrix basic op

* lint

* Disable win tests

tests/pytorch/mock_sparse2/test_diag.py

+import pytest
+import torch
+import sys
+
+from dgl.mock_sparse2 import diag, identity, DiagMatrix
+
+# FIXME: Skipping tests on win.
+if not sys.platform.startswith("linux"):
+    pytest.skip("skipping tests on win", allow_module_level=True)
+
+
+@pytest.mark.parametrize("val_shape", [(3,), (3, 2)])
+@pytest.mark.parametrize("mat_shape", [None, (3, 5), (5, 3)])
+def test_diag(val_shape, mat_shape):
+    # creation
+    val = torch.randn(val_shape)
+    mat = diag(val, mat_shape)
+
+    # val, shape attributes
+    assert torch.allclose(mat.val, val)
+    if mat_shape is None:
+        mat_shape = (val_shape[0], val_shape[0])
+    assert mat.shape == mat_shape
+
+    # __call__
+    val = torch.randn(val_shape)
+    mat = mat(val)
+    assert torch.allclose(mat.val, val)
+
+    # nnz
+    assert mat.nnz == val.shape[0]
+    # dtype
+    assert mat.dtype == val.dtype
+    # device
+    assert mat.device == val.device
+
+    # as_sparse
+    sp_mat = mat.as_sparse()
+    # shape
+    assert tuple(sp_mat.shape) == mat_shape
+    # nnz
+    assert sp_mat.nnz == mat.nnz
+    # dtype
+    assert sp_mat.dtype == mat.dtype
+    # device
+    assert sp_mat.device == mat.device
+    # row, col, val
+    edge_index = torch.arange(len(val)).to(mat.device)
+    row, col, val = sp_mat.coo()
+    assert torch.allclose(row, edge_index)
+    assert torch.allclose(col, edge_index)
+    assert torch.allclose(val, val)
+
+
+@pytest.mark.parametrize("shape", [(3, 3), (3, 5), (5, 3)])
+@pytest.mark.parametrize("d", [None, 2])
+def test_identity(shape, d):
+    # creation
+    mat = identity(shape, d)
+    # type
+    assert isinstance(mat, DiagMatrix)
+    # shape
+    assert mat.shape == shape
+    # val
+    len_val = min(shape)
+    if d is None:
+        val_shape = len_val
+    else:
+        val_shape = (len_val, d)
+    val = torch.ones(val_shape)
+    assert torch.allclose(val, mat.val)

tests/pytorch/mock_sparse2/test_elementwise_op_diag.py

+import operator
+
+import numpy as np
+import pytest
+import torch
+import sys
+from dgl.mock_sparse import diag
+
+# FIXME: Skipping tests on win.
+if not sys.platform.startswith("linux"):
+    pytest.skip("skipping tests on win", allow_module_level=True)
+
+
+def all_close_sparse(A, B):
+    assert torch.allclose(A.indices(), B.indices())
+    assert torch.allclose(A.values(), B.values())
+    assert A.shape == B.shape
+
+
+@pytest.mark.parametrize(
+    "op", [operator.add, operator.sub, operator.mul, operator.truediv]
+)
+def test_diag_op_diag(op):
+    D1 = diag(torch.arange(1, 4))
+    D2 = diag(torch.arange(10, 13))
+    assert np.allclose(op(D1, D2).val, op(D1.val, D2.val), rtol=1e-4, atol=1e-4)
+
+
+@pytest.mark.parametrize("v_scalar", [2, 2.5])
+def test_diag_op_scalar(v_scalar):
+    D1 = diag(torch.arange(1, 50))
+    assert np.allclose(
+        D1.val * v_scalar, (D1 * v_scalar).val, rtol=1e-4, atol=1e-4
+    )
+    assert np.allclose(
+        v_scalar * D1.val, (D1 * v_scalar).val, rtol=1e-4, atol=1e-4
+    )
+    assert np.allclose(
+        D1.val / v_scalar, (D1 / v_scalar).val, rtol=1e-4, atol=1e-4
+    )
+    assert np.allclose(
+        pow(D1.val, v_scalar), pow(D1, v_scalar).val, rtol=1e-4, atol=1e-4
+    )

tests/pytorch/mock_sparse2/test_elementwise_op_sp.py

+import operator
+
+import numpy as np
+import pytest
+import torch
+import sys
+import dgl
+from dgl.mock_sparse2 import create_from_coo, diag
+
+# FIXME: Skipping tests on win.
+if not sys.platform.startswith("linux"):
+    pytest.skip("skipping tests on win", allow_module_level=True)
+
+def all_close_sparse(A, row, col, val, shape):
+    rowA, colA, valA = A.coo()
+    assert torch.allclose(rowA, row)
+    assert torch.allclose(colA, col)
+    assert torch.allclose(valA, val)
+    assert A.shape == shape
+
+
+@pytest.mark.parametrize("op", [operator.add])
+def test_sparse_op_sparse(op):
+    rowA = torch.tensor([1, 0, 2, 7, 1])
+    colA = torch.tensor([0, 49, 2, 1, 7])
+    valA = torch.rand(len(rowA))
+    A = create_from_coo(rowA, colA, valA, shape=(10, 50))
+    w = torch.rand(len(rowA))
+    A1 = create_from_coo(rowA, colA, w, shape=(10, 50))
+
+    def _test():
+        all_close_sparse(op(A, A1), rowA, colA, valA + w, (10, 50))
+
+    _test()
+
+
+@pytest.mark.skip(
+    reason="No way to test it because we does not element-wise op \
+    between matrices with different sparsity"
+)
+@pytest.mark.parametrize("op", [operator.add])
+def test_sparse_op_diag(op):
+    rowA = torch.tensor([1, 0, 2, 7, 1])
+    colA = torch.tensor([0, 49, 2, 1, 7])
+    valA = torch.rand(len(rowA))
+    A = create_from_coo(rowA, colA, valA, shape=(10, 50))
+    D = diag(torch.arange(2, 12), shape=A.shape)
+    D_sp = D.as_sparse()
+
+    def _test():
+        all_close_sparse(op(A, D), *D_sp.coo(), [10, 50])
+
+    _test()

tests/pytorch/mock_sparse2/test_sparse_matrix.py

+import pytest
+import torch
+import sys
+
+from dgl.mock_sparse2 import create_from_coo, create_from_csr, create_from_csc
+
+# FIXME: Skipping tests on win.
+if not sys.platform.startswith("linux"):
+    pytest.skip("skipping tests on win", allow_module_level=True)
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("row", [[0, 0, 1, 2], (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (3, 5), (5, 3)])
+def test_create_from_coo(dense_dim, row, col, shape):
+    # Skip invalid matrices
+    if shape is not None and (
+        max(row) >= shape[0] or max(col) >= shape[1]
+    ):
+        return
+
+    val_shape = (len(row),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    val = torch.randn(val_shape)
+    row = torch.tensor(row)
+    col = torch.tensor(col)
+    mat = create_from_coo(row, col, val, shape)
+
+    if shape is None:
+        shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
+    
+    mat_row, mat_col, mat_val = mat.coo()
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_row, row)
+    assert torch.allclose(mat_col, col)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [[0, 0, 1, 4], (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (3, 5)])
+def test_create_from_csr(dense_dim, indptr, indices, shape):
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    val = torch.randn(val_shape)
+    indptr = torch.tensor(indptr)
+    indices = torch.tensor(indices)
+    mat = create_from_csr(indptr, indices, val, shape)
+
+    if shape is None:
+        shape = (indptr.numel() - 1, torch.max(indices).item() + 1)
+
+    assert mat.device == val.device
+    assert mat.shape == shape
+    assert mat.nnz == indices.numel()
+    assert mat.dtype == val.dtype
+    mat_indptr, mat_indices, mat_val = mat.csr()
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+    assert torch.allclose(mat_val, val)
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [[0, 0, 1, 4], (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (5, 3)])
+def test_create_from_csc(dense_dim, indptr, indices, shape):
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    val = torch.randn(val_shape)
+    indptr = torch.tensor(indptr)
+    indices = torch.tensor(indices)
+    mat = create_from_csc(indptr, indices, val, shape)
+
+    if shape is None:
+        shape = (torch.max(indices).item() + 1, indptr.numel() - 1)
+
+    assert mat.device == val.device
+    assert mat.shape == shape
+    assert mat.nnz == indices.numel()
+    assert mat.dtype == val.dtype
+    mat_indptr, mat_indices, mat_val = mat.csc()
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+    assert torch.allclose(mat_val, val)

tests/scripts/build_dgl.sh

@@ -15,7 +15,7 @@ if [ "$1" != "cugraph" ]; then
     # We do not build pytorch for cugraph because currently building
     # pytorch against all the supported cugraph versions is not supported
     # See issue: https://github.com/rapidsai/cudf/issues/8510
-    CMAKE_VARS="$CMAKE_VARS -DBUILD_TORCH=ON -DTORCH_PYTHON_INTERPS=/opt/conda/envs/pytorch-ci/bin/python"
+    CMAKE_VARS="$CMAKE_VARS -DBUILD_TORCH=ON -DBUILD_SPARSE=ON -DTORCH_PYTHON_INTERPS=/opt/conda/envs/pytorch-ci/bin/python"
 fi
 
 #This is implemented to detect underlying architecture and enable arch specific optimization.

tests/scripts/task_lint.sh

@@ -3,6 +3,7 @@
 # cpplint
 echo 'Checking code style of C++ codes...'
 python3 tests/lint/lint.py dgl cpp include src || exit 1
+python3 tests/lint/lint.py dgl_sparse cpp dgl_sparse/include dgl_sparse/src || exit 1
 
 # pylint
 echo 'Checking code style of python codes...'