Metadata-Version: 2.2
Name: torch_cluster
Version: 1.6.3+triton.hygon
Summary: FlagTree Triton implementation of torch-cluster
Home-page: https://github.com/rusty1s/pytorch_cluster
Author: FlagOS operator team; based on work by Matthias Fey
Keywords: pytorch geometric-deep-learning graph-neural-networks cluster-algorithms triton flagtree
Classifier: Development Status :: 4 - Beta
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3 :: Only
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch
Requires-Dist: triton
Requires-Dist: scipy
Provides-Extra: test
Requires-Dist: pytest; extra == "test"
Requires-Dist: pytest-cov; extra == "test"
Dynamic: author
Dynamic: home-page
Dynamic: keywords
Dynamic: provides-extra
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# PyTorch Cluster

--------------------------------------------------------------------------------

这是面向 FlagOS 的 `torch_cluster` 实现，保留上游
[`pytorch_cluster`](https://github.com/rusty1s/pytorch_cluster) 的公开导入名和核心
API。加速设备上的受支持输入默认使用 FlagTree Triton 内核，CPU 输入或当前 Triton
内核尚不支持的输入使用纯 PyTorch 参考实现。本版本不包含 C++、CUDA/HIP 扩展源码或
预编译共享库。

## 算子库介绍

本库提供图构建、点云采样和图聚类常用算子：

- **邻域查询**：`radius`、`radius_graph`、`knn`、`knn_graph` 和 `nearest`；
- **点云与网格**：`fps` 和 `grid_cluster`；
- **图采样与聚类**：`random_walk`、`neighbor_sampler` 和 `graclus_cluster`。

`radius_graph` 和 `knn_graph` 直接生成 PyG 常用的 `[2, num_edges]`
`edge_index`；批输入通过有序的 `batch` 向量隔离不同样本。`fps` 用于最远点采样，
`random_walk` 和 `neighbor_sampler` 接收 CSR 图结构，`graclus_cluster` 用于图的粗化
与聚类。

## 环境要求

- Python >= 3.9
- 与目标加速设备及其软件栈匹配的 PyTorch 和 FlagTree Triton
- `scipy`
- 运行测试需要 `pytest`，运行 Triton 测试和性能测试需要可用的加速设备

不同加速设备可能通过 PyTorch 的兼容设备接口对外提供计算能力。以下示例中的
`device="cuda"` 表示 PyTorch 设备接口，不限定底层物理加速卡的厂商或型号。

## 安装

进入本目录：

```bash
cd /path_to_dir/torch_cluster
```

### 按硬件后端安装

源码安装或构建 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 会带有对应
的版本标记（例如 `1.6.3+triton.hygon`），可以直接安装，安装时无需再次设置变量：

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

检查已安装的硬件后端：

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

首次使用新的 Triton 配置时会进行即时编译，第一次调用可能较慢。

## 基本用法

下面示例为两组点分别构造半径图。相同 `batch` 编号的点之间才会建立边：

```python
import torch
from torch_cluster import radius_graph

device = torch.device("cuda")
x = torch.tensor(
    [[0.0, 0.0], [0.5, 0.0], [2.0, 0.0],
     [0.0, 0.0], [0.0, 0.5], [0.0, 2.0]],
    device=device,
)
batch = torch.tensor([0, 0, 0, 1, 1, 1], device=device)

edge_index = radius_graph(
    x,
    r=0.8,
    batch=batch,
    loop=False,
    max_num_neighbors=32,
)
print(edge_index.shape)  # [2, num_edges]
```

## 算子示例

下面示例展示 k 近邻图和显式 Triton 后端。`triton` 模式遇到 CPU 输入或当前内核
不支持的输入会直接报错，便于发现模型中意外回退到参考实现：

```python
import torch
from torch_cluster import knn_graph, use_backend

device = torch.device("cuda")
x = torch.randn(1024, 3, device=device)

with use_backend("triton"):
    edge_index = knn_graph(x, k=16, loop=False)

assert edge_index.dtype == torch.long
assert edge_index.shape[0] == 2
```

## 运行时实现选择

`FLAGOS_BACKEND` 决定安装包面向的硬件；`TORCH_CLUSTER_BACKEND` 则决定算子运行时使用
Triton 内核还是 PyTorch 参考实现。可通过以下环境变量选择运行时实现：

```bash
export TORCH_CLUSTER_BACKEND=auto     # 默认：受支持的加速输入走 Triton，其余走参考实现
export TORCH_CLUSTER_BACKEND=triton   # 强制 Triton；不支持时立即报错
export TORCH_CLUSTER_BACKEND=native   # 强制使用纯 PyTorch 参考实现
```

也可以使用 `torch_cluster.set_backend(...)` 或上下文管理器
`torch_cluster.use_backend(...)`。这里的 `native` 为兼容已有测试和性能测试参数而
保留，不表示平台专用编译扩展。

## 正确性测试

在本目录执行完整测试：

```bash
cd /path_to_dir/torch_cluster
pytest -q
```

只运行 Triton 专项和真实模型输入测试：

```bash
TORCH_CLUSTER_BACKEND=triton pytest -q \
  test/test_flagtree_triton.py \
  test/real_models
```

测试覆盖前向接口、空输入、不同数据类型、批次隔离、邻居上限、随机算子约束、
PyTorch 参考实现/Triton 差分，以及 OneScience DiffDock 的真实调用快照。对于输出顺序不属于
API 契约的图算子，测试会规范化边集合后比较；随机算子使用合法性、唯一性和统计约束。

## 性能测试

在仓库根目录、同一设备和同一软件栈下运行 PyTorch 参考实现/Triton 对比：

```bash
python benchmark/all_ops.py --warmup 10 --iterations 100
```

该入口覆盖全部底层算子。结果写入：

- `benchmark/results/benchmark_all_ops.csv`；
- `benchmark/results/benchmark_all_ops.json`。

脚本对每个样本计时前后执行设备同步，并输出首次调用和稳态延迟。性能结果受设备型号、
输入分布、缓存和首次即时编译影响，应执行预热，并只比较同一环境下生成的结果。这里的
参考路径是纯 PyTorch 实现，不能作为历史平台专用编译扩展的性能基线。

### 模型实际输入测例

`data/onescience_diffdock/` 保存了从 OneScience DiffDock 采集的 `6o5u` 和
`1a46` 输入快照，包含 batch 1 和 batch 10 的 `radius`、`radius_graph`、`knn_graph`
调用，并附带 SHA256 校验文件。运行真实调用 benchmark：

```bash
python benchmark/diffdock.py --warmup 10 --iterations 100
```

结果写入 `benchmark/results/benchmark_diffdock.csv` 和 `.json`。以下入口复用真实
DiffDock 坐标、batch 和图边，为其余算子构造代表性工作负载：

```bash
python benchmark/diffdock_derived.py --warmup 10 --iterations 100
```

`diffdock_derived.py` 的输入形状来自真实科学模型，但不是 DiffDock 中直接观测到的算子
调用；报告中会保留来源和新增参数，不能将其表述为 DiffDock 原生覆盖。
