# 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

这是兼容原生 CUDA 与 Torch-FL / 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 实现。
原生 CUDA 环境继续使用 `cuda` 设备名；FlagOS 环境使用 `flagos` 设备名。
Hygon 后端默认启用可选的 C++ 主机扩展，负责缓冲区分配、执行计划缓存、kernel 启动和部分
Autograd 调度封装；GPU 数值计算仍由 Triton kernel 完成。

文档与项目链接：

- [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`，按设备和数据类型选择 Triton 或 PyTorch sparse 路径
- 图操作：随机游走、邻居采样、SAINT、重标号等
- 可选 METIS 图划分：`SparseTensor.partition`

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

## 环境要求

- Python >= 3.9
- 已安装与目标硬件软件栈匹配的 PyTorch 和 Triton
- 原生 CUDA：使用匹配的 PyTorch CUDA 环境和驱动，设备名使用 `cuda`
- Torch-FL / FlagOS：预先配置 Torch-FL、FlagGems 和 FlagTree；程序中先 `import torch_fl`，再导入 `torch`，设备名使用 `flagos`（例如 `flagos:0`）
- `scipy`
- 运行 benchmark 还需要 `wget`
- Hygon C++ 主机加速构建：需要 C++17 编译器、Python 开发头文件和当前 PyTorch 的头文件/库
- 可选 METIS 构建需要 C/C++ 编译器、PyTorch C++ headers 和 64 位索引的 METIS

## 安装

进入本目录：

```bash
cd /path_to_dir/torch_sparse
```

### 按硬件后端安装

源码安装或构建 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（默认编译 C++ 主机扩展）
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
```

Hygon 源码安装默认编译并安装 C++ 主机扩展，负责主机端分配、调度和启动优化，GPU 数值计算仍由
FlagTree Triton 完成。先激活目标 Torch-FL 环境，准备 C++17 编译器（可用 `CXX=g++` 指定）、
Python 开发头文件和当前 PyTorch 的头文件/库。上述安装命令会完成 C++ 编译；编译失败则安装报错。
运行时直接加载已编译的扩展，不再在首次使用时编译 C++ 或写入主机扩展缓存。

编译该主机扩展不需要 `nvcc`/`hipcc`，GPU 运行仍需要匹配的 HIP/DCU 运行时和驱动；
`hipModuleLaunchKernel` 是用于启动 Triton 编译产物的运行时 API。
Triton kernel 本身仍可能在首次调用新配置时进行 GPU JIT 编译，因此性能测试需要预热。
主机扩展未安装或无法加载时会发出 warning，并回退到 Python/Triton 路径。

两个开关默认启用；设为 `0` 关闭，取消变量或设为 `1` 恢复。运行时开关须在启动 Python 前设置：

| 环境变量 | 设为 `0` 的作用 |
| --- | --- |
| `TORCH_SPARSE_HYGON_HOST` | 安装/构建时跳过 C++ 主机扩展编译；运行时禁用扩展加载，使用 Python/Triton 路径。不影响可选 METIS 构建；未编译扩展时，需重新安装才能启用。 |
| `TORCH_SPARSE_HYGON_NATIVE_LAUNCH` | 运行时仅关闭直接 HIP 启动优化，保留 C++ 主机加速，改用 Triton 原有 launcher。 |

例如，临时关闭 C++ 主机加速运行程序：

```bash
TORCH_SPARSE_HYGON_HOST=0 python your_program.py
```

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

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

构建 wheel 时使用相同的变量（Hygon 默认将已编译的扩展打入 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
```

Hygon 二进制 wheel 需匹配目标平台、Python 和 PyTorch ABI；安装匹配的 wheel 无需本机编译器。
切换 PyTorch 环境后应重新构建扩展。开发源码推荐使用 editable 安装；普通安装的验证应在仓库目录外进行，
以免本地源码遮蔽已安装的包。METIS 是独立的可选 CPU 扩展，见下一节。

检查已安装的后端：

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

Hygon 安装后，在已激活的 Torch-FL 环境中检查 C++ 主机扩展是否成功加载：

```python
import torch_fl  # 必须先于 torch 导入
import torch_sparse
from torch_sparse._hygon.host import extension

host = extension()
print(torch_sparse.__version__, torch_sparse.backend_name())
print("C++ host enabled:", host is not None)
if host is not None:
    print(host.__file__)
```

## 可选 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()
```

### Torch-FL / FlagOS

在已配置好的 Torch-FL / FlagOS 环境中，先导入 `torch_fl`，并将稀疏索引、稀疏值和稠密矩阵放到同一个 `flagos` 设备上：

```python
import torch_fl  # noqa: F401 — 必须先于 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` 沿用输入张量的设备执行计算。FlagOS 下请使用 `flagos` 设备名，不要将 FlagOS 张量与 `cuda` 设备名混用于同一次计算；原生 CUDA 环境仍使用 `cuda`。

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

### 稀疏-稀疏矩阵乘

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

后端按设备和数据类型选择 Triton 或 PyTorch sparse 实现。half/bfloat16 等 dtype 的
支持范围取决于安装的 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/
```

测试覆盖转换、稀疏结构操作、前向和反向传播等公共接口。部分 SpSpMM dtype、加速设备不可用
或未启用 METIS 时，相应测试会按条件跳过；具体原因可通过 `-rs` 查看，数量随环境而变化。

在 Torch-FL / FlagOS 环境中，先注册设备，再使用同一批通用测试：

```python
import torch_fl  # 必须先于 torch 导入
import torch
import pytest
import torch_sparse.testing as testing

testing.devices[:] = [torch.device("cpu"), torch.device("flagos:0")]
raise SystemExit(pytest.main([
    "-q", "-rs", "test/", "--ignore=test/test_triton.py",
]))
```

上述命令应在仓库根目录运行；`test_triton.py` 使用原生 CUDA 接口，在 FlagOS 通用测试中单独排除。

只测试 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。

## 全公开接口性能测试

```bash
cd benchmark
python public_api_benchmark.py
```

默认四种浮点精度、七类输入，直接逐行打印延迟和校验结果。筛选、计时口径与覆盖范围见 [benchmark/public_api_benchmark.md](benchmark/public_api_benchmark.md)。