# 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-09-30

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# torch_sparse

这是面向 FlagOS 的 `torch_sparse 0.6.18` Triton 重构版本，保留上游
`torch_sparse` 的公开导入名、Python API 和 `torch.ops.torch_sparse` schema。
加速设备上的 COO/CSR 转换、对角线处理、CSR SpMM 和随机游走使用 FlagTree
Triton kernel；CPU 路径和图采样功能使用 PyTorch/Python 实现。

文档与项目链接：

- [torch-sparse 上游项目](https://github.com/rusty1s/pytorch_sparse)
- [torch-sparse API 文档](https://github.com/rusty1s/pytorch_sparse#functions)
- [PyTorch Sparse 文档](https://pytorch.org/docs/stable/sparse.html)

## 功能概览

本库提供带 Autograd 支持的稀疏矩阵操作，稀疏结构可以使用 COO 索引或
`SparseTensor` 表示。主要功能包括：

- COO/CSR 转换：`ind2ptr`、`ptr2ind`
- 稀疏矩阵整理：`coalesce`、`transpose`、对角线插入与删除
- 稀疏-稠密矩阵乘：`spmm` 和 `SparseTensor.matmul`，支持 `sum`、`add`、
  `mean`、`min`、`max`
- 稀疏-稀疏矩阵乘：`spspmm`，当前复用 PyTorch sparse 路径
- 图操作：随机游走、邻居采样、SAINT、重标号等
- 可选 METIS 图划分：`SparseTensor.partition`

索引张量是离散数据，不参与求导；稀疏值和稠密矩阵支持 Autograd。

## 环境要求

- Python >= 3.9
- 已安装与当前 FlagOS/DCU 软件栈匹配的 PyTorch 和 Triton
- `scipy`
- 运行 benchmark 还需要 `wget`
- 可选 METIS 构建需要 C/C++ 编译器、PyTorch C++ headers 和 64 位索引的 METIS

## 默认安装

不需要 METIS 时，直接进入本目录安装。该路径不编译 C++ 扩展：

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

构建并安装纯 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
```

验证安装：

```bash
python -c "import torch_sparse; print(torch_sparse.__version__); print(torch_sparse.__file__)"
```
首次使用新的 Triton 配置时会进行 JIT编译，因此首次调用可能较慢。

### 按硬件后端安装

源码安装或构建 wheel 时，可以通过通用环境变量 `FLAGOS_BACKEND` 选择目标硬件：

```bash
# NVIDIA（默认值；不设置变量时使用）
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
```

开发和调试推荐使用 editable 安装：

```bash
# 以下以 Hygon 为例；请根据目标硬件替换 hygon
FLAGOS_BACKEND=hygon python -m pip install -e . --no-deps --no-build-isolation
```

构建 wheel 时使用相同的变量:

```bash
FLAGOS_BACKEND=hygon python -m pip wheel . --no-deps --no-build-isolation --wheel-dir dist
```

可选值只有 `nvidia`、`hygon` 和 `mthreads`，默认是 `nvidia`。变量只在安装/构建阶段
读取，安装完成后修改它不会动态切换后端。生成的 wheel 带有后端版本标记，例如
`0.6.18+triton.hygon`，可以直接分发和安装：

```bash
python -m pip install --force-reinstall --no-deps dist/torch_sparse-0.6.18+triton.hygon-*.whl
```

检查已安装的后端：
```bash
python -c "import torch_sparse; print(torch_sparse.__version__); print(torch_sparse.backend_name())"
```


## 可选 METIS 支持

METIS 是可选的 CPU 图划分扩展，实现 `partition` 和 `partition2`，不影响其他
Python/Triton 算子。METIS 必须使用 64 位 `idx_t`（`IDXTYPEWIDTH=64`）。安装后：

```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` 也可作为兼容开关。验证：

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

图划分始终在 CPU 执行；加速器上的 `SparseTensor` 会暂时把图结构移到 CPU，完成
后再把结果移回原设备。

MT-METIS 是独立的多线程 CPU 实现，依赖 `libmtmetis` 和 `libwildriver`。当前版本
支持通过 `TORCH_SPARSE_WITH_MTMETIS=1`（兼容 `WITH_MTMETIS=1`）构建；使用
`MTMETIS_ROOT` 指定两个库的安装前缀，必要时用 `WILDRIVER_ROOT` 单独指定
WildRiver。未启用时，`mt_partition` 会报告未编译支持；该功能不在默认安装中。

```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
```

## 基本用法

### Coalesce

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

`coalesce` 按行排序索引并合并重复位置。重复项使用 `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)
```

结果：

```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)
```

该算子交换稀疏矩阵的行列维度：

```python
from torch_sparse import transpose

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

### 稀疏-稠密矩阵乘

```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()
```

`SparseTensor.matmul` 还支持 `sum`、`add`、`mean`、`min` 和 `max` 归约。

### 稀疏-稀疏矩阵乘

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

两个输入矩阵需要是 coalesced 状态。当前实现复用 `torch.sparse.mm`，其
half/bfloat16 Sparse SpSpMM 支持取决于安装的 PyTorch 和设备后端。

```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 图划分

```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` 描述各分区范围，`perm` 是按 partition id 排序后的节点排列。METIS
不同版本或默认随机状态可能给出不同但同样合法的 partition 标签，因此验证时应比较
分区合法性、排列和切边质量，而不是固定标签编号。

## 正确性测试

完整测试：

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

测试结果中的 12 个 `half/bfloat16 SpSpMM` skipped 是预期的（启用 METIS 时通常总数为
12；未启用 METIS 时还会增加 METIS 用例的 skipped）：当前
`spspmm` 复用 `torch.sparse.mm`，而 PyTorch 对该稀疏-稀疏路径的这两种 dtype
在当前后端未实现；测试因此显式跳过。这不表示 Triton 的 SpMM 正确性失败，
也不影响已覆盖 dtype 的测试结果。

只测试 Triton 加速路径：

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

只测试可选 METIS：

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

## 性能测试

benchmark 默认读取 `StocF-1465`、`ldoor`、`citationCiteseer` 和
`web-Stanford` 数据集；文件不存在时会自动下载。

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

CPU 模式：

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

benchmark 结果受设备型号、图规模、稀疏分布、缓存和 Triton 首次 JIT 编译影响，
应在同一环境与参数下比较。

COO 行索引转 CSR 行指针的 benchmark：

```bash
cd /path_to_dir/torch_sparse/benchmark

# CPU 和可用的 CUDA/DCU 设备都测试
python ptr2ind.py --device both

# 只测试 CPU；缩短试跑时间
python ptr2ind.py --device cpu --duration 0.05 --warmup 0.02

# 只测试加速设备
python ptr2ind.py --device cuda
```

### 模型实际输入测例

`mattergen_mp20_sparse_batch256.npz` 来自 OneScience MatterGen 的 MP-20 cache，包含
256 个结构的 GemNet 邻接输入（`row`、`col`、`edge_batch`、`num_atoms` 等字段），
不是随机构造的形状。

`model_benchmark.py` 测量 GemNet `get_triplets` 所需的 `SparseTensor` 行查询、
triplet 过滤及 SpMM 前向/反向；默认使用 32 个结构、512 个通道、预热 10 次并迭代
50 次：

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

可通过 `--input` 指定其他同字段的 `.npz` 快照，`--num-structures 0` 使用全部结构，
`--compare-torch` 额外测试相同图上的 PyTorch COO SpMM。