【Rust日报】跨平台高性能计算语言扩展CubeCL

[rust is best]10 亿行挑战方案

作者尝试优化 Rust 中「10 亿行挑战」的解决方案。从原始的 5 分钟优化到了 9 秒。

作者的主要目标是创建一段简单、可维护且生产就绪的代码，而且没有使用不安全的操作。以下是一些关键经验：

使用 --release 优化构建
避免在关键路径中使用 println!；使用日志库进行调试
谨慎使用 FromIterator::collect()；它会触发新的分配
最小化不必要的分配，特别是避免使用 to_owned() 和 clone()
更换哈希函数，FxHashMap 比标准的 HashMap 稍微更快
对于大文件，优先使用缓冲读取而不是加载整个文件
当不需要 UTF-8 验证时，使用字节片段（[u8]）而不是字符串
只有在优化单线程性能后才进行并行化

作者采取了迭代式的方法，每个解决方案都作为一个单独的提交。

Blog: https://naveenaidu.dev/tackling-the-1-billion-row-challenge-in-rust-a-journey-from-5-minutes-to-9-seconds

GitHub: https://github.com/Naveenaidu/rust-1brc

[new ver] wgpu v22.0

主要更新内容如下：

所有与着色器相关的配置结构体现在都有一个 compilation_options 字段。目前只是将其设置为 Default::default()，如果有特定的编译需求，可以使用这个字段。
RenderPipelineDescriptor 和 ComputePipelineDescriptor 现在有一个 cache 字段。这允许在着色器编译过程中提供一个缓存来使用。这主要对 Android 设备有用，因为大多数桌面硬件/驱动程序提供了缓存。目前设置为 None。
DeviceDescriptor 现在有一个 memory_hint 字段。可以使用这个字段请求 GPU 优先考虑性能、内存使用情况，或允许请求自定义的内存块大小。不过，这些只是提示，硬件决定最终如何执行。目前设置为 Default::default()。

Wgpu是WebGPU API规范的Rust实现。WebGPU是由GPU for the Web社区组发布的规范，旨在以安全可靠的方式让Web代码访问GPU功能。它通过模仿Vulkan API，并将其转换为主机硬件正在使用的API（如DirectX、Metal、Vulkan）来实现这一目的。

Document: https://sotrh.github.io/learn-wgpu/news/22.0/

GitHub: https://github.com/sotrh/learn-wgpu

[new lib] CubeCL

CubeCL旨在现代化GPU计算，使编写最佳和可移植的内核更加容易。CubeCL允许使用Rust语法的子集编写GPU内核，并正在进行工作以支持更多语言特性。

CubeCL解决了GPU计算中的三个主要挑战：

可移植性：相同的代码库可以用来在任何GPU上进行编程，而不会降低性能。
可用性：无需使用新的着色器语言，只需在Rust代码顶部添加一个属性就可以在任何GPU上运行。
性能：通过创新的编译时系统生成细粒度的内核专用化，以利用最有效的指令。

下面是GELU函数的例子：

go 复制代码

use cubecl::prelude::*;

#[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
}

cube 属性中的 launch 关键字会自动生成一个函数来运行生成的内核：

go 复制代码

fn main() {
    type Runtime = cubecl::cuda::CudaRuntime;
    let device = Default::default();
    let client = Runtime::client(&device);
    let input = &[-1., 0., 1., 5.];
    let output_handle = client.empty(input.len() * core::mem::size_of::<f32>());
    let input_handle = client.create(f32::as_bytes(input));

    gelu_array::launch::<F32, Runtime>(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new(input.len() as u32, 1, 1),
        ArrayArg::new(&input_handle, input.len()),
        ArrayArg::new(&output_handle, input.len()),
    );

    let bytes = client.read(output_handle.binding());
    let output = f32::from_bytes(&bytes);
    // Should be [-0.1587,  0.0000,  0.8413,  5.0000]
    println!("Executed gelu with runtime {:?} => {output:?}", Runtime::name());
}

值得一提的是，这也是burn团队的一个项目。

GitHub: https://github.com/tracel-ai/cubecl

From 日报小组长琴

社区学习交流平台订阅：

Rustcc 论坛：支持 rss
微信公众号：Rust 语言中文社区