Rust threads on the GPU

发表于 更新于 温馨的阅读时光

原文网址

Rust threads on the GPU

作者完成了一项重大的成就：把为CPU设计的线程映射到了GPU上

执行模型

CPU 和 GPU 执行程序时从根本上不同，CPU程序从一个单独的线程出发，然后按需要派生出若干额外线程，每个线程彼此独立地运行，由程序员管理并发

GPU 程序由一个或多个kernel组成，每个kernel启动时会并行地伴随着若干instances，并发不是需要程序员显式声明的，而是GPU程序的固有属性

graph TD
    %% 定义顶部入口
    Start("launch_by_cpu()")
    
    %% 定义中间的 kernel 节点
    K0["kernel()"]
    K1["kernel()"]
    K2["kernel()"]
    K3["kernel()"]
    KN["kernel()"]
    
    %% 定义底部的 Warp 节点
    W0["Warp 0"]
    W1["Warp 1"]
    W2["Warp 2"]
    W3["Warp 3"]
    WN["Warp N"]
    
    %% 定义结束节点
    End("exit()")
    
    %% 建立连接关系 (从上到下)
    Start --> K0
    Start --> K1
    Start --> K2
    Start --> K3
    Start --> KN
    
    K0 --> W0
    K1 --> W1
    K2 --> W2
    K3 --> W3
    KN --> WN
    
    W0 --> End
    W1 --> End
    W2 --> End
    W3 --> End
    WN --> End

对于矩阵乘法、图像处理和图形渲染这样，每个warp(线程束)对不同的数据做相同的事，的任务，刚才的模型可以工作的很好

随着GPU 程序变得越来越复杂，开发者使用warp specialization(线程束特化)来并发地激活一个程序中的不同部分

graph TD
    %% 定义顶部入口
    Start("launch_by_cpu()")
    
    %% 定义中间的 kernel 节点
    K0["kernel()"]
    K1["kernel()"]
    K2["kernel()"]
    K3["kernel()"]
    KN["kernel()"]
    
    %% 中间层
    N0["a0"]
    N1["b0"]
    N2["c0"]
    N3["d0"]
    N4["n0"]

    %% 定义底部的 Warp 节点
    W0["Warp 0"]
    W1["Warp 1"]
    W2["Warp 2"]
    W3["Warp 3"]
    WN["Warp N"]
    
    %% 定义结束节点
    End("exit()")
    
    %% 建立连接关系 (从上到下)
    Start --> K0
    Start --> K1
    Start --> K2
    Start --> K3
    Start --> KN
    
    K0 --> N0 --> W0
    K1 --> N1 --> W1
    K2 --> N2 --> W2
    K3 --> N3 --> W3
    KN --> N4 --> WN
    
    W0 --> End
    W1 --> End
    W2 --> End
    W3 --> End
    WN --> End

函数即程序

大部分CPU程序以一个确定的main函数作为程序的入口，显然可以称为“函数即程序”

而大部分GPU程序模型也用一个函数作为入口点，程序员写下这个函数就好像它只会执行一次，但是实际上硬件会把他并发启动成千上万次。GPU的kernels看起来和正常的CPU函数相似，但表现得很不同：

// CUDA C kernel
__global__ void scale(float* data) {
    // This code is executed thousands of times in parallel,
    // each with the same function parameters, but different values for the global
    // indexes such as `blockIdx` and `threadIdx`.
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    data[i] *= 2.0f;
}

这种看起来和实际运行的不匹配就是造成GPU编程困难的主要原因，无论是程序员还是编译器都无法只通过代码轻易推断出程序的行为

在实践中，这就需要程序员为其负责，手动保证索引正确、避免竞争

GPU程序在rust里也是同样的模式：

use core::arch::nvptx::*;
 
pub unsafe extern "ptx-kernel" fn scale(data: *mut f32) {
    // This code is executed thousands of times in parallel,
    // each with the same function parameters, but different values
    // for the global indexes such as `_block_idx_x` and `_thread_idx_x`.
    let i = (_block_idx_x() * _block_dim_x() + _thread_idx_x()) as usize;
    *data.add(i) *= 2.0;
}

因为这个函数会被执行无数次，每次只有几个全局变量不同，所以rust的所有权系统没法工作了，需要打上unsafe的标记，同时无可避免地使用裸指针

然而理想的情况是rust的安全保障也能拓展适配到GPU上

第一个方案是引入新的类型和用法去匹配GPU具体的语义，但是这会创造一个新的和原有rust分离的编程模型，导致原先的rust生态不适配，不是作者想要的

第二个方案是教会rust编译器分析GPU行为，这是一个很长很重的工程，由rust编译器团队逐步推进，但是作者现在就想写safe的GPU代码

第三个方案是用CPU模拟运行，用miri一类的工具去考虑可能的UB，但是总有些情况是模拟不到的，专业的说叫leaky abstraction。理想的最终状态是GPU程序表现地就像CPU程序一样，这样rust编译器可以如常地分析其行为

单函数kernel 入口点由于并发是隐式的使得实现很难，而显式并发模型对于程序员和编译器都更容易分析，threads 线程就是这样的模型

为什么要支持std::thread于GPU

rust主要用两种并发模型，异步future和threads

作者的这一篇博客演示了如何实现async/await在GPU上

然而很多rust生态是围绕着thread建立的，比如常用的有线程池rayon，异步运行时tokio和一些并行库等等，支持threads可以释放当前生态的很大潜力

为什么不只解决把std::thread映射到 GPU threads

在warp(线程束)里，GPU有很多个threads，又称为lanes(通道)，一个很明显的方法是吗std::thread映射到这些lanes之一

然而这两者不是一回事，一个GPU thread更像是一个CPU的SIMD 通道而不是一个可以独立执行的上下文

一个CPU 线程有自己的栈、PC并且可以被独立调度，但是GPU的lanes在一个warp里同步运行，直接映射有违语义

同时也会变慢，当一个warp里的lanes选择了不同的分支时，GPU硬件会掩盖不活跃的lanes并且把它们顺序执行，称为divergence，于是并发地优势荡然无存

世界第一：GPU上的std::thread

终于来到了本文的重点，作者完成的工作，将std::thread搬到了GPU上，其用法和正常的CPU程序完全一样：

use std::thread;
 
fn main() {
    let a = thread::spawn(|| {
        let mut sum = 0u64;
        for i in 0..1000 {
            sum += i;
        }
        sum
    });
 
    let b = thread::spawn(|| {
        let mut product = 1u64;
        for i in 1..20 {
            product *= i;
        }
        product
    });
 
    let sum = a.join().unwrap();
    let factorial = b.join().unwrap();
 
    println!("sum: {sum}, factorial: {factorial}");
}

一个线程算累加，一个算阶乘，主线程等算完输出，看起来平平无奇，作者接下来放出来一个运行演示，依旧看起来和CPU程序没什么两样

实现

之所以std::thread能在GPU上被支持主要源于三个关键的观察

  1. Warps 可以表现得和CPU线程相似：每一个warp都有自己的PC、自己的寄存器文件，同时可以和其它warp独立执行。GPU的warp调度器在warp之间的转换和操作系统调度器在线程之间的切换很像，在CPU上thread能做什么，warp也能做到

  2. 一个GPU的kernel不必在启动时就是并发的。默认的GPU执行模型同时启动所有warp，但是其实并不需要他们一开始就做有用的工作。通过从一个warp开始，让其它warp按序启用，就可以恢复到CPU的执行模型

  3. 线程束特化(warp specialization)是不需要语言支持的手动线程分区。GPU开发者已经把不同的任务分配给不同的warp了，std::thread是同样的概念，额外内置了语言提供的所有权、类型检查、生命周期保证的接口

于是作者把每个std::thread映射到一个GPUwarp，当kernel启动时，只有Warp 0运行main，其余休眠。通过调用thread::spawn()唤醒一个休眠的warp来运行一个派生的闭包，调用thread::join()阻塞一个父warp直到其子warp结束

随着warp-as-thread模型就位，剩余部分的std::thread也就自然而然了

thread::current()返回一个当前warp的标识符，thread::sleep()使得当前warp休眠，thread::yield_now()把这个warp返回到调度器，线程名、id和构建器模式都按照预期地构建

作者的实现针对于英伟达GPU，但是这套方法并非是对CUDA特有的.Vulkan 有subgroups的概念，HIP/ROCm 也有wavefronts，这两者都可以被用作实现同样的warp-as-thread映射

优点

最大的优点就是程序员不需要去思考warps blocks 或是grids了，可以就像在CPU上用使用线程一样使用GPU，把GPU的工作原理抽象出去了

另一个优点则是架构上避免了warp内串行导致的divergence。thread::spawn()把一个闭包映射到一个warp。每个lane在这个warp里都运行相同的代码，所以绝不会产生分支。最坏的情况下一个warp里只有一个lane在工作，但即使这样也比divergent好，因为闲着只会浪费空间而分支会导致串行。需要跨lane执行的代码仍然可以使用warp_shuffle_idx一类的函数显式地进行

更重要的是，这样一来rust的借用检查和生命周期就能工作了！

缺点

warps是有限的资源，分出太多线程会用尽可用的warp，接着要么排队要么崩溃。然而实践中这并不是太大的问题，因为大部分rust代码已经对硬件做了适应，可以查询可用的warp/thread数量

线程同步的代价变得更加高昂。锁、条件变量和其它阻塞原语需要warp水平的调度决策，这会降低吞吐量

每个thread会消耗一整个warp，如果生成的代码没有用完lanes，GPU就是没有充分利用的，这对于复杂的任务是可以接受的，但是对于简单的并行任务是一种浪费，数据并行可能会更合适

空闲的warp没有任务，也是一种浪费

最后，每个线程的栈内存需要从GPU的内存中分配，这比CPU内存更受限，递归栈变深会耗尽可用的内存，而默认的CUDA栈对于复杂应用来说又太小必须调整

其它

这篇博客发表在VectorWare公司的官网上，作者是他的创始人LegNeato，rust生态里的又一位神仙人物

VectorWare这家公司主打GPU原生，进入首页上来就是We are at the start of a new software industry.

作者并没有提一个warp里，如何让thread利用上若干lanes，大概就是默认只用一个lane，作者说对于复杂任务是可以接受的，因为对于一个GPU有上万个lanes，512个warp，而CPU如何能达到512 core，所以依然是可以接受的。

同时关于串行，因为所有lanes是完全同步的，所以分支就是完全的浪费，分支越多越浪费，完全可以丢给另外的warp，而warp是相当充裕的，同时还有cache问题

嗯，还有很多东西并不能完全理解。看来GPU领域值我得学习的内容还有很多

# rust# GPU