< 返回版块

长琴 发表于 2024-07-21 09:14

Tags:rust,日报,wgpu,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(),如果有特定的编译需求,可以使用这个字段。
  • RenderPipelineDescriptorComputePipelineDescriptor 现在有一个 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计算中的三个主要挑战:

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

下面是GELU函数的例子:

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 关键字会自动生成一个函数来运行生成的内核:

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 日报小组 长琴

社区学习交流平台订阅:

评论区

写评论

还没有评论

1 共 0 条评论, 1 页