Repository: pyg-team/pytorch_geometric

Issue:
Batching sparse tensors is very slow
### 🐛 Describe the bug

Hi,

I've been storing some sparse matrices(different dtypes - float32, int64, complex64) in my Data object which I batch block-diagonally. I've noticed that my training times have increased ~3x since then. I've narrowed it down to the batching of these sparse matrices. 

I've made a simplified example that reprodouces the behaviour (I've left out the different dtypes and the block diagonal batching). 

```python
from typing import Tuple

import numpy as np
import scipy
import torch
from scipy.sparse import csr_matrix
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
from tqdm import tqdm


def csr_to_sparse_tensor(csr: csr_matrix, size: Tuple[int, ...]):
    """
    Convert a scipy csr matrix to a torch csr tensor.

    :param csr: csr matrix
    :param size: Size.
    :return: csr tensor.
    """
    return torch.sparse_csr_tensor(crow_indices=torch.LongTensor(csr.indptr),
                                   col_indices=torch.LongTensor(csr.indices),
                                   values=torch.tensor(csr.data),
                                   size=size)


def main():
    num_rows, num_columns = 42, 123
    sparse = True

    data_list = []
    for i in range(1000):
        random_csr_tensor = csr_to_sparse_tensor(csr=scipy.sparse.random(num_rows, num_columns, format='csr', dtype=np.float32),
                                                 size=(num_rows, num_columns))

        data_list.append(
            Data(
                sparse_or_dense_matrix=random_csr_tensor if sparse else random_csr_tensor.to_dense()
            )
        )

    loader = DataLoader(data_list, batch_size=128, shuffle=True)

    for epoch in tqdm(range(500), ncols=75, ascii=True, desc="Execution time"):
        for batch in loader:
            pass


if __name__ == "__main__":
    main()

```

The output of the script with `sparse = False` is: `Execution time: 100%|###################| 500/500 [00:01<00:00, 254.06it/s]`,

with `sparse = True`: `Execution time: 100%|####################| 500/500 [00:51<00:00,  9.74it/s]`

### Versions

I'm on Windows so I couldn't run the above command, but my pip freeze is:
```
pip freeze            
antlr4-python3-runtime==4.9.3
appdirs==1.4.4
certifi==2023.7.22
charset-normalizer==3.3.0
click==8.1.7
colorama==0.4.6
contourpy==1.1.1
cycler==0.12.1
deepdiff==6.6.0
docker-pycreds==0.4.0
filelock==3.12.4
fonttools==4.43.1
fsspec==2023.9.2
gitdb==4.0.10
GitPython==3.1.37
hydra-core==1.3.2
idna==3.4
importlib-resources==6.1.0
Jinja2==3.1.2
joblib==1.3.2
kiwisolver==1.4.5
lightning-utilities==0.10.0
MarkupSafe==2.1.3
matplotlib==3.7.1
mpmath==1.3.0
networkx==3.1
numpy==1.24.3
omegaconf==2.3.0
ordered-set==4.1.0
packaging==23.2
pandapower==2.12.1
pandas==1.4.4
pathtools==0.1.2
Pillow==10.0.1
protobuf==4.24.4
psutil==5.9.5
pyg-nightly==2.4.0.dev20231028
pyparsing==3.1.1
PYPOWER==5.1.16
python-dateutil==2.8.2
pytz==2023.3.post1
PyYAML==6.0.1
requests==2.31.0
scikit-learn==1.2.2
scipy==1.10.1
sentry-sdk==1.32.0
setproctitle==1.3.3
six==1.16.0
smmap==5.0.1
sympy==1.12
threadpoolctl==3.2.0
torch==2.1.0+cu121
torch-cluster==1.6.3+pt21cu121
torch-scatter==2.1.2+pt21cu121
torch-sparse==0.6.18+pt21cu121
torch-spline-conv==1.2.2+pt21cu121
torch_geometric==2.4.0
torchaudio==2.1.0+cu121
torchmetrics==1.2.1
torchvision==0.16.0+cu121
tqdm==4.65.0
typing_extensions==4.8.0
urllib3==2.0.6
wandb==0.15.12
zipp==3.17.0

```

Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index a180ccbd0..f2065ca7b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -43,6 +43,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Accelerated mini-batching of `torch.sparse` tensors ([#8670](https://github.com/pyg-team/pytorch_geometric/pull/8670))
 - Fixed RPC timeout due to worker closing in `DistLoader` with `atexit` not executed correctly in `worker_init_fn` ([#8605](https://github.com/pyg-team/pytorch_geometric/pull/8605))
 - `ExplainerDataset` will now contain node labels for any motif generator ([#8519](https://github.com/pyg-team/pytorch_geometric/pull/8519))
 - Made `utils.softmax` faster via `softmax_csr` ([#8399](https://github.com/pyg-team/pytorch_geometric/pull/8399))
diff --git a/torch_geometric/utils/sparse.py b/torch_geometric/utils/sparse.py
index d29f75029..ad68ce8c3 100644
--- a/torch_geometric/utils/sparse.py
+++ b/torch_geometric/utils/sparse.py
@@ -493,37 +493,195 @@ def index2ptr(index: Tensor, size: int) -> Tensor:
         index, size, out_int32=index.dtype == torch.int32)
 
 
-def cat(tensors: List[Tensor], dim: Union[int, Tuple[int, int]]) -> Tensor:
-    # TODO (matthias) We can make this more efficient by directly operating on
-    # the individual sparse tensor layouts.
+def cat_coo(tensors: List[Tensor], dim: Union[int, Tuple[int, int]]) -> Tensor:
     assert dim in {0, 1, (0, 1)}
+    assert tensors[0].layout == torch.sparse_coo
+
+    indices, values = [], []
+    num_rows = num_cols = 0
+
+    if dim == 0:
+        for i, tensor in enumerate(tensors):
+            if i == 0:
+                indices.append(tensor.indices())
+            else:
+                offset = torch.tensor([[num_rows], [0]], device=tensor.device)
+                indices.append(tensor.indices() + offset)
+            values.append(tensor.values())
+            num_rows += tensor.size(0)
+            num_cols = max(num_cols, tensor.size(1))
+
+    elif dim == 1:
+        for i, tensor in enumerate(tensors):
+            if i == 0:
+                indices.append(tensor.indices())
+            else:
+                offset = torch.tensor([[0], [num_cols]], device=tensor.device)
+                indices.append(tensor.indices() + offset)
+            values.append(tensor.values())
+            num_rows = max(num_rows, tensor.size(0))
+            num_cols += tensor.size(1)
 
-    size = (0, 0)
-    edge_indices = []
-    edge_attrs = []
-    for tensor in tensors:
-        assert is_torch_sparse_tensor(tensor)
-        edge_index, edge_attr = to_edge_index(tensor)
-        edge_index = edge_index.clone()
-
-        if dim == 0:
-            edge_index[0] += size[0]
-            size = (size[0] + tensor.size(0), max(size[1], tensor.size(1)))
-        elif dim == 1:
-            edge_index[1] += size[1]
-            size = (max(size[0], tensor.size(0)), size[1] + tensor.size(1))
-        else:
-            edge_index[0] += size[0]
-            edge_index[1] += size[1]
-            size = (size[0] + tensor.size(0), size[1] + tensor.size(1))
-
-        edge_indices.append(edge_index)
-        edge_attrs.append(edge_attr)
-
-    return to_torch_sparse_tensor(
-        edge_index=torch.cat(edge_indices, dim=1),
-        edge_attr=torch.cat(edge_attrs, dim=0),
-        size=size,
-        is_coalesced=dim == (0, 1),
-        layout=tensors[0].layout,
+    else:
+        for i, tensor in enumerate(tensors):
+            if i == 0:
+                indices.append(tensor.indices())
+            else:
+                offset = torch.tensor([[num_rows], [num_cols]],
+                                      device=tensor.device)
+                indices.append(tensor.indices() + offset)
+            values.append(tensor.values())
+            num_rows += tensor.size(0)
+            num_cols += tensor.size(1)
+
+    return torch.sparse_coo_tensor(
+        indices=torch.cat(indices, dim=-1),
+        values=torch.cat(values),
+        size=(num_rows, num_cols) + values[-1].size()[1:],
+        device=tensor.device,
     )
+
+
+def cat_csr(tensors: List[Tensor], dim: Union[int, Tuple[int, int]]) -> Tensor:
+    assert dim in {0, 1, (0, 1)}
+    assert tensors[0].layout == torch.sparse_csr
+
+    rows, cols, values = [], [], []
+    num_rows = num_cols = nnz = 0
+
+    if dim == 0:
+        for i, tensor in enumerate(tensors):
+            if i == 0:
+                rows.append(tensor.crow_indices())
+            else:
+                rows.append(tensor.crow_indices()[1:] + nnz)
+            cols.append(tensor.col_indices())
+            values.append(tensor.values())
+            num_rows += tensor.size(0)
+            num_cols = max(num_cols, tensor.size(1))
+            nnz += cols[-1].numel()
+
+        return torch.sparse_csr_tensor(
+            crow_indices=torch.cat(rows),
+            col_indices=torch.cat(cols),
+            values=torch.cat(values),
+            size=(num_rows, num_cols) + values[-1].size()[1:],
+            device=tensor.device,
+        )
+
+    elif dim == 1:
+        for i, tensor in enumerate(tensors):
+            rows.append(ptr2index(tensor.crow_indices()))
+            if i == 0:
+                cols.append(tensor.col_indices())
+            else:
+                cols.append(tensor.col_indices() + num_cols)
+            values.append(tensor.values())
+            num_rows = max(num_rows, tensor.size(0))
+            num_cols += tensor.size(1)
+
+        return torch.sparse_coo_tensor(
+            indices=torch.stack((torch.cat(rows), torch.cat(cols)), 0),
+            values=torch.cat(values),
+            size=(num_rows, num_cols) + values[-1].size()[1:],
+            device=tensor.device,
+        )
+
+    else:
+        for i, tensor in enumerate(tensors):
+            if i == 0:
+                rows.append(tensor.crow_indices())
+                cols.append(tensor.col_indices())
+            else:
+                rows.append(tensor.crow_indices()[1:] + nnz)
+                cols.append(tensor.col_indices() + num_cols)
+            values.append(tensor.values())
+            num_rows += tensor.size(0)
+            num_cols += tensor.size(1)
+            nnz += cols[-1].numel()
+
+        return torch.sparse_csr_tensor(
+            crow_indices=torch.cat(rows),
+            col_indices=torch.cat(cols),
+            values=torch.cat(values),
+            size=(num_rows, num_cols) + values[-1].size()[1:],
+            device=tensor.device,
+        )
+
+
+def cat_csc(tensors: List[Tensor], dim: Union[int, Tuple[int, int]]) -> Tensor:
+    assert dim in {0, 1, (0, 1)}
+    assert tensors[0].layout == torch.sparse_csc
+
+    rows, cols, values = [], [], []
+    num_rows = num_cols = nnz = 0
+
+    if dim == 0:
+        for i, tensor in enumerate(tensors):
+            cols.append(ptr2index(tensor.ccol_indices()))
+            if i == 0:
+                rows.append(tensor.row_indices())
+            else:
+                rows.append(tensor.row_indices() + num_rows)
+            values.append(tensor.values())
+            num_rows += tensor.size(0)
+            num_cols = max(num_cols, tensor.size(1))
+
+        return torch.sparse_coo_tensor(
+            indices=torch.stack((torch.cat(rows), torch.cat(cols)), 0),
+            values=torch.cat(values),
+            size=(num_rows, num_cols) + values[-1].size()[1:],
+            device=tensor.device,
+        )
+
+    elif dim == 1:
+        for i, tensor in enumerate(tensors):
+            if i == 0:
+                cols.append(tensor.ccol_indices())
+            else:
+                cols.append(tensor.ccol_indices()[1:] + nnz)
+            rows.append(tensor.row_indices())
+            values.append(tensor.values())
+            num_rows = max(num_rows, tensor.size(0))
+            num_cols += tensor.size(1)
+            nnz += rows[-1].numel()
+
+        return torch.sparse_csc_tensor(
+            row_indices=torch.cat(rows),
+            ccol_indices=torch.cat(cols),
+            values=torch.cat(values),
+            size=(num_rows, num_cols) + values[-1].size()[1:],
+            device=tensor.device,
+        )
+
+    else:
+        for i, tensor in enumerate(tensors):
+            if i == 0:
+                rows.append(tensor.row_indices())
+                cols.append(tensor.ccol_indices())
+            else:
+                rows.append(tensor.row_indices() + num_rows)
+                cols.append(tensor.ccol_indices()[1:] + nnz)
+            values.append(tensor.values())
+            num_rows += tensor.size(0)
+            num_cols += tensor.size(1)
+            nnz += rows[-1].numel()
+
+        return torch.sparse_csc_tensor(
+            row_indices=torch.cat(rows),
+            ccol_indices=torch.cat(cols),
+            values=torch.cat(values),
+            size=(num_rows, num_cols) + values[-1].size()[1:],
+            device=tensor.device,
+        )
+
+
+def cat(tensors: List[Tensor], dim: Union[int, Tuple[int, int]]) -> Tensor:
+    assert is_torch_sparse_tensor(tensors[0])
+
+    if tensors[0].layout == torch.sparse_coo:
+        return cat_coo(tensors, dim)
+    elif tensors[0].layout == torch.sparse_csr:
+        return cat_csr(tensors, dim)
+    else:
+        return cat_csc(tensors, dim)


Test patch:
diff --git a/test/data/test_batch.py b/test/data/test_batch.py
index e382ca950..5ab69df63 100644
--- a/test/data/test_batch.py
+++ b/test/data/test_batch.py
@@ -509,7 +509,7 @@ def test_torch_sparse_batch(layout):
     batch = Batch.from_data_list([data, data])
 
     assert batch.x.size() == (6, 4)
-    assert batch.x.layout == layout
+    assert batch.x.layout in {torch.sparse_coo, torch.sparse_csr}
     assert torch.equal(batch.x.to_dense(), torch.cat([x_dense, x_dense], 0))
 
     assert batch.adj.size() == (6, 6)
diff --git a/test/utils/test_sparse.py b/test/utils/test_sparse.py
index b8e61316b..5bf9c9c6f 100644
--- a/test/utils/test_sparse.py
+++ b/test/utils/test_sparse.py
@@ -1,10 +1,11 @@
 import os.path as osp
 
+import pytest
 import torch
 
 import torch_geometric.typing
 from torch_geometric.profile import benchmark
-from torch_geometric.testing import is_full_test, withPackage
+from torch_geometric.testing import is_full_test, withCUDA, withPackage
 from torch_geometric.typing import SparseTensor
 from torch_geometric.utils import (
     dense_to_sparse,
@@ -14,7 +15,9 @@ from torch_geometric.utils import (
     to_torch_coo_tensor,
     to_torch_csc_tensor,
     to_torch_csr_tensor,
+    to_torch_sparse_tensor,
 )
+from torch_geometric.utils.sparse import cat
 
 
 def test_dense_to_sparse():
@@ -246,6 +249,47 @@ def test_to_edge_index():
         assert edge_attr.tolist() == [1., 1., 1., 1., 1., 1.]
 
 
+@withCUDA
+@pytest.mark.parametrize(
+    'layout',
+    [torch.sparse_coo, torch.sparse_csr, torch.sparse_csc],
+)
+@pytest.mark.parametrize('dim', [0, 1, (0, 1)])
+def test_cat(layout, dim, device):
+    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], device=device)
+    if torch_geometric.typing.WITH_PT20:
+        edge_weight = torch.rand(4, 2, device=device)
+    else:
+        edge_weight = torch.rand(4, device=device)
+
+    adj = to_torch_sparse_tensor(edge_index, edge_weight, layout=layout)
+
+    out = cat([adj, adj], dim=dim)
+    edge_index, edge_weight = to_edge_index(out.to_sparse_csr())
+
+    if dim == 0:
+        if torch_geometric.typing.WITH_PT20:
+            assert out.size() == (6, 3, 2)
+        else:
+            assert out.size() == (6, 3)
+        assert edge_index[0].tolist() == [0, 1, 1, 2, 3, 4, 4, 5]
+        assert edge_index[1].tolist() == [1, 0, 2, 1, 1, 0, 2, 1]
+    elif dim == 1:
+        if torch_geometric.typing.WITH_PT20:
+            assert out.size() == (3, 6, 2)
+        else:
+            assert out.size() == (3, 6)
+        assert edge_index[0].tolist() == [0, 0, 1, 1, 1, 1, 2, 2]
+        assert edge_index[1].tolist() == [1, 4, 0, 2, 3, 5, 1, 4]
+    else:
+        if torch_geometric.typing.WITH_PT20:
+            assert out.size() == (6, 6, 2)
+        else:
+            assert out.size() == (6, 6)
+        assert edge_index[0].tolist() == [0, 1, 1, 2, 3, 4, 4, 5]
+        assert edge_index[1].tolist() == [1, 0, 2, 1, 4, 3, 5, 4]
+
+
 if __name__ == '__main__':
     import argparse