Rust FFI 调用 CUDA/OpenCL：GPU 高性能计算实践

你好！我是你们的“赛博朋克”老伙计，码农阿强。今天咱们来聊点硬核的，聊聊怎么用 Rust 这把“瑞士军刀”撬开 GPU 的大门，让你的程序像脱缰的野马一样在并行计算的世界里狂奔。

为什么选择 Rust + GPU？

你可能要问，GPU 编程不是 C/C++ 的地盘吗？Rust 来凑什么热闹？

别急，听我慢慢道来。Rust 的优势在于：

• 内存安全：Rust 的所有权和借用机制，让你在编译时就能消灭内存泄漏和数据竞争这些“老大难”问题。GPU 编程可是个“精细活”，一点点内存错误就可能导致程序崩溃，甚至硬件损坏。Rust 的安全性，让你能更专注于算法逻辑，而不是整天跟 bug 斗智斗勇。
• 零成本抽象：Rust 的很多高级特性，比如泛型、trait、迭代器等等，在编译后都能被优化成高效的机器码，不会带来额外的性能开销。这意味着你可以用更优雅、更安全的方式编写代码，同时还能保持与 C/C++ 媲美的性能。
• 强大的 FFI（外部函数接口）：Rust 可以轻松地调用 C/C++ 编写的库，反之亦然。这意味着你可以无缝地集成现有的 GPU 编程库，比如 CUDA 和 OpenCL，不需要“另起炉灶”。
• 活跃的社区和丰富的生态：Rust 社区非常活跃，有很多优秀的 GPU 编程库可供选择，比如 rust-cuda、ocl 等等。这些库大大降低了 GPU 编程的门槛，让你能更快地上手。

简单来说，Rust 就是想在保证安全和开发效率的同时，还能让你“榨干”GPU 的每一滴性能。

FFI：Rust 与 GPU 之间的桥梁

要让 Rust 调用 CUDA 或 OpenCL，我们需要用到 FFI（Foreign Function Interface，外部函数接口）。FFI 是一种机制，允许一种编程语言调用另一种编程语言编写的代码。Rust 提供了强大的 FFI 支持，可以让我们轻松地与 C/C++ 代码进行交互。

CUDA 示例

咱们先来看一个简单的 CUDA 示例，演示如何使用 Rust 调用 CUDA 核函数。

首先，我们需要一个 CUDA 核函数（kernel），这是一个在 GPU 上运行的函数。咱们写一个简单的向量加法核函数：

// add.cu
__global__ void add(int n, float *x, float *y, float *z) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index < n) {
    z[index] = x[index] + y[index];
  }
}

这个核函数接受三个浮点数数组 x、y 和 z，以及数组的长度 n。它将 x 和 y 对应位置的元素相加，并将结果存储在 z 中。

接下来，我们需要一个 C/C++ 的头文件，声明这个核函数：

// add.h
#ifndef ADD_H
#define ADD_H
 
extern "C" {
    void add(int n, float *x, float *y, float *z);
}
 
#endif

注意，我们需要使用 extern "C" 来告诉编译器，这个函数使用 C 语言的调用约定。这是因为 Rust 的 FFI 默认与 C 语言兼容。

现在，我们可以使用 nvcc（NVIDIA CUDA Compiler）将 CUDA 代码编译成一个动态链接库：

nvcc -shared -o libadd.so add.cu

接下来，是 Rust 部分的代码：

// main.rs
#[link(name = "add")]
extern "C" {
    fn add(n: i32, x: *const f32, y: *const f32, z: *mut f32);
}
 
fn main() {
    let n = 1024;
    let x: Vec<f32> = vec![1.0; n];
    let y: Vec<f32> = vec![2.0; n];
    let mut z: Vec<f32> = vec![0.0; n];
 
    unsafe {
        add(n as i32, x.as_ptr(), y.as_ptr(), z.as_mut_ptr());
    }
 
    println!("{:?}", &z[0..10]); // 打印 z 的前 10 个元素
}

在 Rust 代码中，我们使用 #[link(name = "add")] 告诉编译器，我们要链接名为 add 的库（也就是我们之前编译的 libadd.so）。然后，我们使用 extern "C" 声明了 add 函数，这与 C/C++ 头文件中的声明相对应。注意，我们需要使用指针类型（*const f32 和 *mut f32）来表示 C/C++ 中的数组。

由于 FFI 调用是不安全的（Rust 编译器无法保证外部代码的安全性），我们需要将 FFI 调用放在 unsafe 块中。

最后，我们创建了三个向量 x、y 和 z，并调用了 add 函数。然后，我们打印 z 的前 10 个元素，以验证结果是否正确。

要运行这段代码，我们需要先安装 CUDA Toolkit，然后使用 cargo 构建并运行 Rust 项目：

cargo run

OpenCL 示例

OpenCL 的示例与 CUDA 类似，只是我们需要使用不同的库和 API。这里我们使用 ocl crate 来简化 OpenCL 编程。

首先，我们需要在 Cargo.toml 中添加 ocl 依赖：

[dependencies]
ocl = "0.19"

然后，我们可以编写 Rust 代码：

// main.rs
use ocl::{ProQue, Buffer, Kernel};
 
const SRC: &'static str = r#"
    __kernel void add(__global float* x, __global float* y, __global float* z) {
        int gid = get_global_id(0);
        z[gid] = x[gid] + y[gid];
    }
"#;
 
fn main() -> ocl::Result<()> {
    let n = 1024;
    let x: Vec<f32> = vec![1.0; n];
    let y: Vec<f32> = vec![2.0; n];
    let mut z: Vec<f32> = vec![0.0; n];
 
    let pro_que = ProQue::builder()
        .src(SRC)
        .dims(n)
        .build()?;
 
    let x_buffer = Buffer::new(&pro_que, Some(&x), (n,), 0)?;
    let y_buffer = Buffer::new(&pro_que, Some(&y), (n,), 0)?;
    let mut z_buffer = Buffer::new(&pro_que, None::<&[f32]>, (n,), 0)?;
 
    let kernel = Kernel::new(&pro_que, "add", ())
        .arg_buf(&x_buffer)
        .arg_buf(&y_buffer)
        .arg_buf(&z_buffer);
 
    unsafe {
        kernel.enq()?;
    }
 
    z_buffer.read(&mut z).enq()?;
 
    println!("{:?}", &z[0..10]);
 
    Ok(())
}

这段代码首先定义了一个 OpenCL 核函数（SRC），然后创建了一个 ProQue 对象，这是 ocl crate 中用于管理 OpenCL 程序和队列的对象。接着，我们创建了三个缓冲区（Buffer），分别对应 x、y 和 z 数组。然后，我们创建了一个内核对象（Kernel），并设置了内核参数。最后，我们使用 enq() 方法将内核入队执行，并使用 read() 方法将结果从 z_buffer 读回到 z 向量中。

与 CUDA 示例类似，我们也需要将内核入队操作放在 unsafe 块中。

性能优化技巧

使用 FFI 调用 GPU 只是第一步，要充分发挥 GPU 的性能，我们还需要进行一些优化。

• 数据传输优化：GPU 和 CPU 之间的数据传输是性能瓶颈之一。尽量减少数据传输的次数和数据量。例如，可以将多个小数组合并成一个大数组，或者使用零拷贝（zero-copy）技术。
• 内核优化：优化 CUDA 或 OpenCL 核函数的代码。例如，使用共享内存（shared memory）来减少对全局内存（global memory）的访问，或者使用向量化指令来提高计算效率。
• 并行度优化：调整 CUDA 或 OpenCL 的线程块（block）和线程（thread）的数量，以充分利用 GPU 的并行计算能力。通常，线程块的数量应该与 GPU 的流多处理器（SM）数量相匹配，线程的数量应该与每个 SM 的最大线程数相匹配。
• 异步执行：使用 CUDA 或 OpenCL 的异步 API 来实现 CPU 和 GPU 的并行执行。例如，可以使用 CUDA streams 或 OpenCL command queues 来将多个内核和数据传输操作排队，让它们并行执行。
• 选择合适的GPU编程模型: CUDA适用于NVIDIA GPU, OpenCL则更具通用性, 支持多种GPU和CPU。根据你的硬件和需求, 选择最合适的模型。

总结

Rust 的 FFI 机制为我们提供了一种安全、高效的方式来调用 CUDA 和 OpenCL 等 GPU 编程库。通过结合 Rust 的内存安全、零成本抽象和强大的 FFI，我们可以编写出既安全又高效的 GPU 程序。当然，GPU 编程本身就是一个复杂的主题，需要我们不断学习和实践。希望这篇文章能帮助你入门 Rust GPU 编程，让你在高性能计算的世界里“如虎添翼”！

如果你对 Rust GPU 编程还有其他问题，或者想了解更多关于 Rust 的知识，欢迎在评论区留言，我会尽力解答。下次再见！