Rust宏库cubecl-macros的使用：高效代码生成与过程宏开发指南

概述

CubeCL是一个多平台高性能计算语言扩展，用于在Rust中编写GPU程序，利用零成本抽象开发可维护、灵活且高效的计算内核。CubeCL还提供了优化运行时，用于管理内存和执行。

支持平台

示例代码

以下是使用CubeCL宏库的完整示例：

use cubecl::prelude::*;

#[cube(launch_unchecked)]
/// 一个[Line]表示一系列连续元素，可能支持SIMD操作
/// 运行时将自动使用SIMD指令以提高性能
fn gelu_array<F: Float>(input: &Array<Line<F>>, output: &mut Array<Line<F>>) {
    if ABSOLUTE_POS < input.len() {
        output[ABSOLUTE_POS] = gelu_scalar(input[ABSOLUTE_POS]);
    }
}

#[cube]
fn gelu_scalar<F: Float>(x: Line<F>) -> Line<F> {
    // 在编译时执行sqrt函数
    let sqrt2 = F::new(comptime!(2.0f32.sqrt()));
    let tmp = x / Line::new(sqrt2);

    x * (Line::erf(tmp) + 1.0) / 2.0
}

然后可以使用自动生成的gelu_array::launch_unchecked函数启动内核：

pub fn launch<R: Runtime>(device: &R::Device) {
    let client = R::client(device);
    let input = &[-1., 0., 1., 5.];
    let vectorization = 4;
    let output_handle = client.empty(input.len() * core::mem::size_of::<f32>());
    let input_handle = client.create(f32::as_bytes(input));

    unsafe {
        gelu_array::launch_unchecked::<f极客时间, R>(
            &client,
            CubeCount::Static(1, 1, 1),
            CubeDim::new(input.len() as u32 / vectorization, 1, 1),
            ArrayArg::from_raw_parts::<f32>(&input_handle, input.len(), vectorization as u8),
            ArrayArg::from_raw_parts::<f32>(&output_handle, input.len(), vectorization as u8),
        )
    };

    let bytes = client.read_one(output_handle.binding());
    let output = f32::from_bytes(&bytes);

    // 结果应该是[-0.1587, 0.0000, 0.8413, 5.0000]
    println!("Executed gelu with runtime {:?} => {output:?}", R::name());
}

要运行这个GELU示例，可以使用以下命令：

cargo run --example gelu --features cuda  # CUDA运行时
cargo run --example gelu --features wgpu  # WGPU运行时

设计原理

CubeCL围绕"立方体"概念设计，其拓扑结构如下：

特殊功能

自动向量化

高性能内核应尽可能使用SIMD指令，CubeCL允许你在启动内核时指定每个输入变量的向量化因子。

编译时计算(Comptime)

CubeCL不仅是一种新的计算语言，它允许你编写编译器插件并完全自定义。Comptime是一种在首次编译内核时修改编译器IR的方法。

自动调优(Autotuning)

自动调优通过运行小型基准测试来确定在当前硬件上运行的最佳内核和配置，大大简化了内核选择。

安装

在项目目录中运行以下Cargo命令：

cargo add cubecl-macros

或者在Cargo.toml中添加：

cubecl-macros = "0.6.0"

完整示例

// 引入必要的模块
use cubecl::prelude::*;

// 定义向量化的GELU激活函数数组操作
#[cube(launch_unchecked)]
fn gelu_array<F: Float>(input: &Array<Line<F>>, output: &mut Array<Line<F>>) {
    // 检查当前绝对位置是否在输入范围内
    if ABSOLUTE_POS < input.len() {
        // 对每个元素应用GELU函数
        output[ABSOLUTE_POS] = gelu_scalar(input[ABSOLUTE_POS]);
    }
}

// 定义标量GELU激活函数
#[cube]
fn gelu_scalar<F: Float>(x: Line<F>) -> Line<F> {
    // 编译时计算2的平方根
    let sqrt2 = F::new(comptime!(2.0f32.sqrt()));
    let tmp = x / Line::new(sqrt2);

    // GELU公式: x * (1 + erf(x/√2)) / 2
    x * (Line::erf(tmp) + 1.0) / 2.0
}

// 主函数
fn main() {
    // 初始化CUDA运行时
    let device = cubecl::cuda::Device::new(0).unwrap();
    
    // 准备输入数据
    let input = [-1.0f32, 0.0, 1.0, 5.0];
    let vectorization = 4; // 向量化因子
    
    // 创建客户端
    let client = cubecl::cuda::Client::new(&device);
    
    // 分配内存
    let output_handle = client.empty(input.len() * std::mem::size_of::<f32>());
    let input_handle = client.create(f32::as_bytes(&input));
    
    // 启动内核
    unsafe {
        gelu_array::launch_unchecked::<f32, cubecl::cuda::Runtime>(
            &client,
            CubeCount::Static(1, 1, 1),
            CubeDim::new(input.len() as u32 / vectorization, 1, 1),
            ArrayArg::from_raw_parts::<f32>(&input_handle, input.len(), vectorization as u8),
            ArrayArg::from_raw_parts::<f32>(&output_handle, input.len(), vectorization as u8),
        );
    }
    
    // 读取结果
    let bytes = client.read_one(output_handle.binding());
    let output = f32::from_bytes(&bytes);
    
    println!("GELU计算结果: {:?}", output);
}

许可证

MIT OR Apache-2.0