CubeCL基于Rust的跨平台现代GPU科学计算的生态系统CubeCL基于Rust的跨平台现代GPU科学计算的生态系统

CubeCL基于Rust的跨平台现代GPU科学计算的生态系统

CubeCL 是一个新推出的 Rust 库，专注于提供简洁而强大的接口，用于与 OpenCL（Open Computing Language）进行交互。OpenCL 是一个框架，用于编写在异构平台上运行的程序，包括中央处理器（CPU）、图形处理器（GPU）和其他处理单元。

与目前的 PyCUDA 相比，CubeCL 专注于利用 Rust 的安全性和性能优势来提供高效的 GPU 计算接口。虽然 PyCUDA 是一个非常成熟且广泛使用的 Python 库，它主要用于 CUDA 平台，而 CubeCL 则针对 OpenCL 平台，为开发者提供了更广泛的硬件支持选择。

CubeCL 库旨在简化 OpenCL 的使用，使开发者能够更轻松地编写高性能计算代码，而不必深入了解 OpenCL 的复杂细节。通过使用 Rust 的安全性和并发性优势，CubeCL 提供了一种更安全、更高效的解决方案。

CubeCL 库的关键特性：

简化的 API：CubeCL 提供了一个简化的 API，使得与 OpenCL 进行交互变得更加直观和容易。
内存安全：利用 Rust 的所有权系统，CubeCL 确保了内存访问的安全性，减少了内存泄漏和数据竞争的风险。
高性能：通过高效的内存管理和并行计算能力，CubeCL 能够充分发挥硬件的性能优势。
跨平台支持：CubeCL 支持在多种硬件和操作系统平台上运行，包括 Windows、Linux 和 macOS。

快速入门

use cubecl::prelude::*;
use cubecl::buffer::Array;
use cubecl::types::cl_float;

#[cube(launch)]
fn gelu_array<F: Float>(input: &Array<F>, output: &mut Array<F>) {
    if ABSOLUTE_POS < input.len() {
        output[ABSOLUTE_POS] = gelu_scalar::<F>(input[ABSOLUTE_POS]);
    }
}

#[cube]
fn gelu_scalar<F: Float>(x: F) -> F {
    x * (F::erf(x / F::sqrt(2.0.into())) + 1.0) / 2.0
}

fn main() {
    // 初始化 OpenCL 上下文
    let context = Context::default().unwrap();

    // 创建命令队列
    let queue = Queue::new(&context).unwrap();

    // 输入数据
    let input: Vec<cl_float> = vec![1.0, 2.0, 3.0, 4.0];
    let mut output: Vec<cl_float> = vec![0.0; input.len()];

    // 创建输入和输出数组
    let input_array = Array::new(&context, &input, cubecl::flags::MEM_READ_ONLY).unwrap();
    let mut output_array = Array::new(&context, &output, cubecl::flags::MEM_WRITE_ONLY).unwrap();

    // 启动内核
    queue.enqueue_kernel_with_args(gelu_array, (&input_array, &mut output_array), input.len()).unwrap();

    // 从设备读取结果
    queue.read_buffer(&output_array, &mut output).unwrap();

    // 打印结果
    println!("Output: {:?}", output);
}

#[cube(launch)] 是一个属性宏，用于标记这个函数将作为 OpenCL 内核来执行。

gelu_array 函数接受一个输入数组 input 和一个输出数组 output，将 input 中的每个元素经过 GELU 函数处理后存储到 output 中。

ABSOLUTE_POS 是 CubeCL 提供的一个特殊变量，表示当前 OpenCL 计算单元的全局索引。

看到它在行动中，运行工作GELU例如与以下命令：

cargo run --example gelu --features cuda # cuda runtime
cargo run --example gelu --features wgpu # wgpu runtime

如何工作的

CubeCL 利用 Rust 的过程宏系统，采用独特的两步过程：

解析：过程宏使用 syn crate 解析 GPU 内核代码。扩展：宏并不会立即生成中间表示（IR），而是生成一个新的 Rust 函数。生成的函数在语义上与原函数相似，并且在调用时负责创建 IR。这种方法不同于传统编译器，它们通常在解析后直接生成 IR。我们的方法实现了几个关键功能：

编译时计算（Comptime）：由于不对原代码进行转换，因此集成编译时优化变得非常容易。自动向量化：通过简单地向量化 CubeCL 函数的输入，我们可以在扩展期间确定每个中间变量的向量化因子。 Rust 集成：生成的代码仍然是有效的 Rust 代码，可以不依赖特定运行时进行打包。

CubeCL是怎么兼容多个平台的

CubeCL 设计围绕立方体！更具体地说，它基于长方体，因为并非所有的轴都是相同的大小。由于所有计算 API 都需要映射到硬件，而硬件是可以使用三维表示访问的网格，我们的拓扑结构可以轻松映射到其他 API 的概念。

CubeCL - 拓扑结构

一个立方体由单元组成，因此一个 3x3x3 的立方体有 27 个单元，可以通过它们在 x、y 和 z 轴上的位置进行访问。类似地，超立方体由立方体组成，就像立方体由单元组成一样。每个超立方体中的立方体可以通过其相对于超立方体在 x、y 和 z 轴上的位置进行访问。因此，一个 3x3x3 的超立方体将有 27 个立方体。在这个例子中，总的工作单元数将是 27 x 27 = 729。

更多内容可以查看 burn.dev/docs/burn/