Rust线性代数库cubecl-linalg的使用：高性能矩阵运算与向量计算库

CubeCL是一个多平台高性能计算语言扩展，用于在Rust中编写GPU计算内核。它支持函数、泛型和结构体，并部分支持特征、方法和类型推断。

示例代码

use cubecl::prelude::*;

#[cube(launch_unchecked)]
/// 代表一个连续的元素序列，其中可能可用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
}

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::<f32, 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());
}

完整示例

以下是一个完整的Rust线性代数库cubecl-linalg使用示例：

use cubecl::prelude::*;
use cubecl_linalg::Matrix;

// 定义一个矩阵乘法内核
#[cube(launch_unchecked)]
fn matrix_multiply<F: Float>(
    a: &Matrix<F>,
    b: &Matrix<F>,
    output: &mut Matrix<F>,
) {
    let row = ABSOLUTE_POS_X;
    let col = ABSOLUTE_POS_Y;
    
    if row < a.rows() && col < b.cols() {
        let mut sum = F::zero();
        for k in 0..a.cols() {
            sum += a[[row, k]] * b[[k, col]];
        }
        output[[row, col]] = sum;
    }
}

fn main() {
    // 初始化CUDA运行时
    let device = cubecl::runtime::Cuda::default_device();
    
    // 创建两个随机矩阵
    let a = Matrix::random(device, 1024, 1024);
    let b = Matrix::random(device, 1024, 1024);
    
    // 创建输出矩阵
    let mut output = Matrix::zeros(device, 1024, 1024);
    
    // 启动矩阵乘法内核
    unsafe {
        matrix_multiply::launch_unchecked::<f32, _>(
            &a, &b, &mut output,
            CubeCount::Static(1, 1, 1),
            CubeDim::new(32, 32, 1),
        );
    }
    
    // 验证结果
    let result = output.to_host();
    println!("Matrix multiplication result: {:?}", result);
}

运行环境

CubeCL支持以下GPU运行时：

• WGPU 用于跨平台GPU支持(Vulkan, Metal, DirectX, WebGPU)
• CUDA 用于NVIDIA GPU支持
• ROCm/HIP 用于AMD GPU支持(开发中)

特殊功能

1. 自动向量化：指定输入变量的向量化因子，运行时能够使用最佳指令
2. 编译时计算(Comptime)：在运行时修改编译器IR，实现多种优化
3. 自动调优(Autotune)：运行时运行小基准测试来确定最佳内核配置

安装

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

cargo add cubecl-linalg

或者在Cargo.toml中添加：

cubecl-linalg = "0.5.0"

拓扑设计

CubeCL基于立方体(cuboids)设计，所有计算API都需要映射到硬件，这些硬件可以通过3D表示访问。