rusthrust

Rust 实现的 GPU 并行算法库，灵感来自 Nvidia Thrust。

基于 cust CUDA 驱动 API 绑定，提供 DeviceVec 容器和一系列光栅并行算法。

架构

rusthrust/
├── build.rs                      # 编译期内核工厂（PTX 预编译）
├── src/
│   ├── lib.rs                    # CUDA 上下文初始化 + vector_add
│   ├── device_vec.rs             # DeviceVec<T> GPU 容器
│   ├── reduce.rs                 # 并行规约（树形 reduce）
│   ├── scan.rs                   # 独占扫描（Blelloch prefix sum）
│   ├── sort.rs                   # 基数排序（8 轮 LSB）
│   ├── fusion/                   # 编译期内核融合
│   │   ├── mod.rs                # FusionKind, Fusable trait
│   │   ├── dispatch.rs           # 运行时分发 → Stream::launch
│   │   └── generated.rs          # include!(build.rs 生成)
│   └── iterator/                 # JIT 迭代器（运行时融合）
│       ├── mod.rs                # DeviceIterator trait + 组合子
│       ├── source.rs             # DeviceVecIter, ConstantIterator
│       ├── counting.rs           # CountingIterator
│       ├── transform.rs          # TransformIterator, BinaryTF + ops
│       ├── kernel_gen.rs         # CUDA 源码生成
│       └── launch.rs             # nvcc JIT + SHA256 缓存 + launch
├── examples/
│   ├── vector_add.rs             # 基本向量加法
│   ├── iterator_fusion.rs        # JIT 迭代器融合
│   ├── reduce.rs                 # 并行规约
│   ├── scan.rs                   # 独占扫描
│   ├── sort.rs                   # 基数排序
│   └── fusion.rs                 # 编译期融合（map + reduce）
└── kernels/
    ├── vector_add.cu             # 参考 CUDA 源码（手写）
    └── vector_add.ptx            # 参考 PTX（已不再使用）

核心概念

两种内核融合路径

特性	JIT 迭代器 (iterator/)	编译期融合 (fusion/)
融合时机	运行时	compile time（build.rs）
实现方式	walk chain → 生成 CUDA 字串 → nvcc 编译	build.rs 枚举模式 → nvcc -ptx
首次延迟	~0.2s（nvcc 编译）	零
灵活度	任意组合	仅预定义模式
缓存	SHA256 到 target/fused_kernels_cache/	嵌入二进制 include_str!
对标	通用回路	Thrust 的模板展开

标量作为内核参数

编译期融合把标量（乘数、加数）作为 CUDA 内核参数传入，而非嵌入源码——无需为不同标量值编译重复变体。

模块指南

1. DeviceVec — GPU 容器

对标 thrust::device_vector。实现 Deref / DerefMut 以透明访问底层 DeviceBuffer。

let v = DeviceVec::from_host(&[1.0, 2.0, 3.0])?;
let host = v.to_host()?;  // → vec![1.0, 2.0, 3.0]

2. 并行规约 (reduce.rs)

基于共享内存的树形规约：每线程加载 2 个元素，for (s=B/2; s>0; s>>=1) 折半相加。多级递归直至单个结果。

let sum = rusthrust::reduce::reduce(&data)?;

3. 独占扫描 (scan.rs)

Blelloch 算法：upsweep（树形求和）→ downsweep（前缀传播）。多块 pipeline：

tile_scan (块内 Blelloch) → scan_sums (递归) → add_prefix (块前缀加回)

let output = rusthrust::scan::scan(&input)?;  // exclusive scan

4. 基数排序 (sort.rs)

8 轮 LSB 基数排序，每轮 4 bits。float → uint 通过 sign-bit XOR 转换以保留负数排序。

• radix_histogram：每块统计 16 bins 的局部计数（共享内存 atomicAdd）
• CPU 端扫描 histogram 得到全局偏移
• radix_scatter：warp shuffle 前缀和计算线程在 digit 内的确定排名，无需原子操作即可稳定散列

rusthrust::sort::sort(&mut data)?;

5. JIT 迭代器 (iterator/)

运行时内核融合。链式 transform / add / mul_scalar 仅构建表达式树，.collect() 时生成单一 CUDA 内核。

let result = a.iter()
    .add(b.iter())
    .mul_scalar(3.0)
    .collect()?;

过程：

1. walk_chain() 遍历表达式树 → CUDA 源码字符串
2. SHA256 哈希 → 查缓存
3. 未命中 → nvcc -ptx 编译 → 存缓存
4. Module::from_ptx → Stream::launch → 返回 DeviceVec

6. 编译期融合 (fusion/)

build.rs 在编译时用 CUDA 模板生成 8 个融合内核的 PTX，嵌入到二进制。运行时通过 FusionKind 枚举选择预编译内核。

支持的 fusion 模式：

模式	表达式	输入	标量
MapAdd	a[i] + b[i]	2 DeviceVec	0
MapAddScalar	a[i] + s	1	1
MapMulScalar	a[i] * s	1	1
MapAddMulScalar	(a[i] + b[i]) * s	2	1
ReduceSum	sum a[i]	1	0
MapAddReduce	sum a[i] + b[i]	2	0
MapMulScalarReduce	sum a[i] * s	1	1

use rusthrust::fusion::{fused_collect, fused_reduce};

// 编译期 map（一次内核启动）
let c = fused_collect(a.iter().add(b.iter()).mul_scalar(3.0))?;

// 编译期 transform+reduce 融合（一次内核启动，省去中间缓冲区 + 多一次内核）
let sum = fused_reduce(a.iter().add(b.iter()))?;

示例运行

# 所有示例
cargo run --example vector_add
cargo run --example iterator_fusion
cargo run --example reduce
cargo run --example scan
cargo run --example sort
cargo run --example fusion

# 或一次性构建
cargo build --examples

依赖

crate	用途
cust 0.3.2	CUDA 驱动 API
anyhow 1.x	错误处理
sha2 0.10	JIT 缓存哈希

运行时依赖：CUDA 工具包 (nvcc 须在 PATH 中)。

编译期：build.rs 在编译时调用 nvcc -arch=sm_86 -ptx 生成融合内核。

GPU 要求

NVIDIA GPU，计算能力 ≥ 2.0（需要 atomicAdd、__shfl_up_sync、__syncthreads），推荐 ≥ 6.0。

开发环境：RTX 3050 (CC 8.6), CUDA 13.2, driver 595.

设计原则

1. 零外部 PTX 文件：所有内核源码以内联 Rust 字符串或 build.rs 生成的模板形式存在——无需手动管理 .ptx 文件
2. JIT 核融合：对于模式表中没有的任意迭代器链，运行时 JIT 编译仍然有效
3. 稳定的 sort：通过 warp 级前缀和代替 atomicAdd 确保了 LSD 基数排序的稳定性
4. 融合 reduce：单次内核中同时加载 + 变换 + 规约，消除了中间缓冲区