# torch_sparse

**Repository Path**: onescience-ai/torch_sparse

## Basic Information

- **Project Name**: torch_sparse
- **Description**: 本项目为基于flagos统一中间层实现的torch_sparse库
- **Primary Language**: Python
- **License**: MIT
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-08-25
- **Last Updated**: 2026-11-12

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# torch_sparse

This Triton-based refactoring of `torch_sparse 0.6.18` supports both native
CUDA and Torch-FL / FlagOS devices. It
preserves the public import name, Python API, and `torch.ops.torch_sparse`
schemas of the upstream `torch_sparse` package. COO/CSR conversion, diagonal
operations, CSR SpMM, and random walks on accelerator devices use FlagTree
Triton kernels, while CPU paths and graph sampling are implemented with
PyTorch/Python. Native CUDA environments continue to use the `cuda` device
name; FlagOS environments use `flagos`.

Documentation and project links:

- [Upstream torch-sparse project](https://github.com/rusty1s/pytorch_sparse)
- [torch-sparse API documentation](https://github.com/rusty1s/pytorch_sparse#functions)
- [PyTorch Sparse documentation](https://pytorch.org/docs/stable/sparse.html)

## Feature Overview

This library provides sparse matrix operations with Autograd support. Sparse
structures can be represented using COO indices or `SparseTensor`. Major
features include:

- COO/CSR conversion: `ind2ptr`, `ptr2ind`
- Sparse matrix manipulation: `coalesce`, `transpose`, and diagonal insertion
  and removal
- Sparse-dense matrix multiplication: `spmm` and `SparseTensor.matmul`, with
  support for `sum`, `add`, `mean`, `min`, and `max`
- Sparse-sparse matrix multiplication: `spspmm`, currently implemented using
  the PyTorch sparse path
- Graph operations: random walks, neighbor sampling, SAINT, relabeling, and more
- Optional METIS graph partitioning: `SparseTensor.partition`

Index tensors contain discrete data and do not participate in gradient
computation. Sparse values and dense matrices support Autograd.

## Requirements

- Python >= 3.9
- PyTorch and Triton compatible with the target hardware software stack
- Native CUDA: a matching PyTorch CUDA environment and driver; use the `cuda` device name
- Torch-FL / FlagOS: configure Torch-FL, FlagGems, and FlagTree, import `torch_fl` before `torch`, and use the `flagos` device name (for example, `flagos:0`)
- `scipy`
- `wget` for running benchmarks
- Optional METIS builds require a C/C++ compiler, PyTorch C++ headers, and
  METIS built with 64-bit indices

## Default Installation

When METIS is not required, install directly from this directory. This path
does not compile any C++ extensions:

```bash
cd /path_to_dir/torch_sparse
python -m pip install -e . --no-deps --no-build-isolation
```

Build and install a pure Python/Triton wheel:

```bash
python -m pip wheel . --no-deps --no-build-isolation --wheel-dir dist
python -m pip install --force-reinstall --no-deps \
  dist/torch_sparse-*-py3-none-any.whl
```

Verify the installation:

```bash
python -c "import torch_sparse; print(torch_sparse.__version__); print(torch_sparse.__file__)"
```

New Triton configurations are JIT-compiled on first use, so the first call may
be slower.

### Installing for a Hardware Backend

Use the generic `FLAGOS_BACKEND` environment variable when installing from source or building a
wheel:

```bash
# NVIDIA (default; used when the variable is not set)
python -m pip install . --no-deps --no-build-isolation

# NVIDIA/GPU
FLAGOS_BACKEND=nvidia python -m pip install . --no-deps --no-build-isolation

# Hygon/DCU
FLAGOS_BACKEND=hygon python -m pip install . --no-deps --no-build-isolation

# MThreads/MUSA
FLAGOS_BACKEND=mthreads python -m pip install . --no-deps --no-build-isolation
```

Use the same variable for editable development installs and wheel builds:

```bash
# Hygon is shown as an example; replace hygon with the target hardware backend
FLAGOS_BACKEND=hygon python -m pip install -e . --no-deps --no-build-isolation
FLAGOS_BACKEND=hygon python -m pip wheel . --no-deps --no-build-isolation --wheel-dir dist
```

The only accepted values are `nvidia`, `hygon`, and `mthreads`; the default is `nvidia`. The
variable is read only during installation or wheel building, and changing it afterwards does not
dynamically switch the backend. Wheels carry a backend-specific version marker such as
`0.6.18+triton.hygon` and can be distributed directly:

```bash
python -m pip install --force-reinstall --no-deps dist/torch_sparse-0.6.18+triton.hygon-*.whl
python -c "import torch_sparse; print(torch_sparse.__version__); print(torch_sparse.backend_name())"
```

### Optimizing Individual Backends

Backend entry points live in `csrc/_nvidia`, `csrc/_hygon`, and `csrc/_mthreads`; all currently
reuse the implementation in `csrc/_triton`. Add backend-specific kernels, launch configurations,
or device handling in the corresponding directory while preserving the `spmm_forward` interface.
The dispatch layer loads the entry point selected by the install-time marker, so kernels should not
branch at runtime on `FLAGOS_BACKEND`. Run the full tests, Triton-specific tests, and benchmarks on
each target device, and compare numerics with the shared Triton implementation. Before using the
MThreads entry on real MUSA hardware, verify that the selected PyTorch/Triton stack supports its
device context and compilation target.

## Optional METIS Support

METIS is an optional CPU graph-partitioning extension that implements
`partition` and `partition2`. It does not affect other Python/Triton operators.
METIS must be built with a 64-bit `idx_t` (`IDXTYPEWIDTH=64`). To install with
METIS support:

```bash
cd /path_to_dir/torch_sparse
export METIS_ROOT=/opt/metis-5.1.0-idx64
export LD_LIBRARY_PATH="$METIS_ROOT/lib:$METIS_ROOT/lib64:$LD_LIBRARY_PATH"
TORCH_SPARSE_WITH_METIS=1 \
python -m pip install --force-reinstall . --no-deps --no-build-isolation
```

`WITH_METIS=1` is also supported as a compatibility option. Verify the build:

```bash
python -c "import torch_sparse; print(torch_sparse.has_metis)"
pytest -q -rs /path_to_dir/torch_sparse/test/test_metis.py
```

Graph partitioning always runs on the CPU. For a `SparseTensor` on an
accelerator, the graph structure is temporarily moved to the CPU and the
results are moved back to the original device afterward.

MT-METIS is a separate multithreaded CPU implementation that depends on
`libmtmetis` and `libwildriver`. This version can be built with
`TORCH_SPARSE_WITH_MTMETIS=1` (or the compatible `WITH_MTMETIS=1` option). Use
`MTMETIS_ROOT` to specify the installation prefix for both libraries, or set
`WILDRIVER_ROOT` separately for WildRiver when needed. When it is not enabled,
`mt_partition` reports that support was not compiled. This feature is not part
of the default installation.

```bash
export MTMETIS_ROOT=/opt/mt-metis
export LD_LIBRARY_PATH="$MTMETIS_ROOT/lib:$LD_LIBRARY_PATH"
TORCH_SPARSE_WITH_MTMETIS=1 \
python -m pip install --force-reinstall . --no-deps --no-build-isolation
```

## Basic Usage

### Coalesce

```text
torch_sparse.coalesce(index, value, m, n, op="add")
    -> (torch.LongTensor, torch.Tensor)
```

`coalesce` sorts indices by row and merges duplicate entries. Duplicate values
are combined using the reduction specified by `op`.

```python
import torch
from torch_sparse import coalesce

index = torch.tensor([[1, 0, 1, 0, 2, 1],
                      [0, 1, 1, 1, 0, 0]])
value = torch.tensor([[1., 2.], [2., 3.], [3., 4.],
                      [4., 5.], [5., 6.], [6., 7.]])

index, value = coalesce(index, value, m=3, n=2)
```

Result:

```text
index = tensor([[0, 1, 1, 2],
                [1, 0, 1, 0]])
value = tensor([[6., 8.],
                [7., 9.],
                [3., 4.],
                [5., 6.]])
```

### Transpose

```text
torch_sparse.transpose(index, value, m, n, coalesced=True)
    -> (torch.LongTensor, torch.Tensor)
```

This operator swaps the row and column dimensions of a sparse matrix:

```python
from torch_sparse import transpose

out_index, out_value = transpose(index, value, 3, 2)
```

### Sparse-Dense Matrix Multiplication

```text
torch_sparse.spmm(index, value, m, n, matrix) -> torch.Tensor
```

```python
import torch
from torch_sparse import spmm

index = torch.tensor([[0, 0, 1, 2, 2],
                      [0, 2, 1, 0, 1]], device="cuda")
value = torch.tensor([1., 2., 4., 1., 3.], device="cuda",
                     requires_grad=True)
matrix = torch.tensor([[1., 4.], [2., 5.], [3., 6.]],
                      device="cuda", requires_grad=True)

out = spmm(index, value, 3, 3, matrix)
# tensor([[ 7., 16.],
#         [ 8., 20.],
#         [ 7., 19.]], device="cuda:0")
out.sum().backward()
```

### Torch-FL / FlagOS

In a configured Torch-FL / FlagOS environment, import `torch_fl` first and place the sparse indices, sparse values, and dense matrix on the same `flagos` device:

```python
import torch_fl  # noqa: F401 — must be imported before torch
import torch
from torch_sparse import spmm

device = torch.device("flagos:0")
index = torch.tensor([[0, 0, 1, 2, 2],
                      [0, 2, 1, 0, 1]], device=device)
value = torch.tensor([1., 2., 4., 1., 3.], device=device,
                     requires_grad=True)
matrix = torch.tensor([[1., 4.], [2., 5.], [3., 6.]],
                      device=device, requires_grad=True)

out = spmm(index, value, 3, 3, matrix)
# tensor([[ 7., 16.],
#         [ 8., 20.],
#         [ 7., 19.]], device='flagos:0')
out.sum().backward()
```

`torch_sparse` follows the device of its input tensors. In FlagOS programs, use `flagos` rather than `cuda`; native CUDA programs continue to use `cuda`.

`SparseTensor.matmul` also supports the `sum`, `add`, `mean`, `min`, and `max`
reductions.

### Sparse-Sparse Matrix Multiplication

```text
torch_sparse.spspmm(indexA, valueA, indexB, valueB, m, k, n,
                    coalesced=False)
    -> (torch.LongTensor, torch.Tensor)
```

Both input matrices must be coalesced. The current implementation uses
`torch.sparse.mm`; support for half/bfloat16 Sparse SpSpMM depends on the
installed PyTorch version and device backend.

```python
import torch
from torch_sparse import spspmm

indexA = torch.tensor([[0, 0, 1, 2, 2], [1, 2, 0, 0, 1]])
valueA = torch.tensor([1., 2., 3., 4., 5.])
indexB = torch.tensor([[0, 2], [1, 0]])
valueB = torch.tensor([2., 4.])

indexC, valueC = spspmm(indexA, valueA, indexB, valueB, 3, 3, 2)
# indexC = tensor([[0, 1, 2], [0, 1, 1]])
# valueC = tensor([8., 6., 8.])
```

### METIS Graph Partitioning

```python
import torch
from torch_sparse import SparseTensor

row = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3])
col = torch.tensor([1, 3, 0, 2, 1, 3, 0, 2])
graph = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))

partitioned, partptr, perm = graph.partition(
    num_parts=2,
    recursive=False,
    weighted=False,
)
```

`partptr` describes the range of each partition, and `perm` is the node
permutation sorted by partition ID. Different METIS versions or default random
states may produce different but equally valid partition labels. Validation
should therefore compare partition validity, permutation, and edge-cut quality
rather than fixed label IDs.

## Correctness Tests

Run the complete test suite:

```bash
cd /path_to_dir/torch_sparse
pytest -q -rs test/
```

The 12 skipped `half/bfloat16 SpSpMM` cases in the test results are expected
(the total is normally 12 when METIS is enabled; additional METIS cases are
skipped when it is disabled). The current `spspmm` implementation uses
`torch.sparse.mm`, and PyTorch does not implement these two dtypes for this
sparse-sparse path on the current backend, so the tests explicitly skip them.
This does not indicate a correctness failure in Triton SpMM and does not affect
the test results for supported dtypes.

Run only the Triton-accelerated tests:

```bash
pytest -q test/test_triton.py
```

Run only the optional METIS tests:

```bash
pytest -q -rs test/test_metis.py
```

## Benchmarks

By default, the benchmark reads the `StocF-1465`, `ldoor`,
`citationCiteseer`, and `web-Stanford` datasets. Missing dataset files are
downloaded automatically.

```bash
cd /path_to_dir/torch_sparse/benchmark
python main.py
```

CPU mode:

```bash
python main.py --device cpu
```

Benchmark results are affected by the device model, graph size, sparsity
distribution, cache state, and Triton JIT compilation on first use. Compare
results using the same environment and parameters.

Benchmark COO row-index to CSR row-pointer conversion:

```bash
cd /path_to_dir/torch_sparse/benchmark

# Test both the CPU and available CUDA/DCU devices
python ptr2ind.py --device both

# Test only the CPU with a shorter trial duration
python ptr2ind.py --device cpu --duration 0.05 --warmup 0.02

# Test only the accelerator device
python ptr2ind.py --device cuda
```

### Real-Model Input Benchmark

`mattergen_mp20_sparse_batch256.npz` comes from the OneScience MatterGen MP-20
cache. It contains GemNet adjacency inputs for 256 structures, including the
`row`, `col`, `edge_batch`, and `num_atoms` fields, rather than synthetically
generated shapes.

`model_benchmark.py` measures the `SparseTensor` row queries and triplet
filtering required by GemNet `get_triplets`, as well as SpMM forward and
backward passes. By default, it uses 32 structures and 512 channels, with 10
warm-up iterations and 50 timed iterations:

```bash
cd /path_to_dir/torch_sparse/benchmark
python model_benchmark.py --device cuda --warmup 10 --iters 50
```

Use `--input` to specify another `.npz` snapshot with the same fields,
`--num-structures 0` to use all structures, and `--compare-torch` to
additionally benchmark PyTorch COO SpMM on the same graph.