跳到主要内容

2 篇博文 含有标签「ROCm」

查看所有标签

概述​

Burn 0.15.0 带来了显著的性能改进,特别是在矩阵乘法和卷积操作方面。

此外,此版本还引入了以下重要更新:

  • 实验性支持:新增 ROCm/HIP 和 SPIR-V 支持,通过 CubeCL 运行时实现。
  • 多后端兼容性:奠定多后端支持的基础。
  • 新特性:增加了量化操作支持。
  • ONNX 支持扩展:包括更多的算子支持和错误修复,以提升覆盖率。

除此之外,Burn 0.15.0 还包含多项错误修复、性能优化、新的张量操作,以及改进的文档支持。


模块与张量相关更新​

  • 移除:对常量泛型模块的拷贝限制。
  • 新增:deform_conv2d(实现于 torchvision)、Softmin、round、floor、ceil 等浮点操作。
  • 增强:为张量同步增加支持,添加 tensor.one_hot 整数操作。
  • 更改:LR 调度器调整为首次调用 .step() 时返回初始学习率。

ONNX 支持扩展​

  • 支持多维索引的 gather 操作。
  • 增强张量形状跟踪能力。
  • 新增 ConvTranspose1d 和 trilu 操作支持。
  • 修复 where 操作在标量输入下的行为。

后端改进​

  • 支持 CudaDevice 和 MetalDevice,避免重复创建设备。
  • 新增 SPIR-V 编译器支持 (burn-wgpu) 和 HIP 支持 (burn-hip)。
  • 引入 BackendRouter,为分布式后端处理铺路。
  • 修复自动微分相关的内存泄漏和 NaN 问题。

文档与示例​

  • 新增自定义 cubecl 内核的文档。
  • 改进了回归任务的示例和 burn-tch 文档。
  • 修复了多个 Burn Book 的链接及 Raspberry Pi 示例的编译问题。

性能与优化​

  • 性能提升:增强了切片内核的性能,改进了 conv2d 和 conv_transpose2d 的自动调优。
  • 数据局部性优化:为隐式 GEMM 提供更好的性能支持,并新增边界检查以支持任意输入形状。

Miscellaneous 更新​

  • 工具链:更新了 CI 工作流及工具,修复编译器设置的多处问题。
  • 兼容性:确保最小支持 Rust 版本为 1.81。

参考​


通过 Burn 0.15.0,深度学习开发者可以更高效地利用 GPU 加速和量化技术,同时享受多后端支持带来的灵活性。欢迎尝试新版本并加入我们的社区,共同推动 Rust 生态的技术进步!

鱼雪

使用 CubeCL,您可以通过 Rust 编写 GPU 程序,借助零成本抽象开发出可维护、灵活且高效的计算内核。CubeCL 目前完全支持函数、泛型和结构体,并部分支持 trait、方法和类型推导。随着项目的演进,我们预计将支持更多 Rust 语言特性,同时保持最佳性能。

示例​

只需用 cube 属性标注函数,即可指定其运行在 GPU 上:

use cubecl::prelude::*;

#[cube(launch_unchecked)]
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> {
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::<f32, R>(
&client,
CubeCount::Static(1, 1, 1),
CubeDim::new(input.len() as u32 / vectorization, 1, 1),
ArrayArg::from_raw_parts(&input_handle, input.len(), vectorization as u8),
ArrayArg::from_raw_parts(&output_handle, input.len(), vectorization as u8),
)
};

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

println!("Executed gelu with runtime {:?} => {output:?}", R::name());
}

运行以下命令即可体验 GELU 示例:

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

运行时支持​

支持以下 GPU 运行时:

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

未来还计划开发一个使用 SIMD 指令的优化 JIT CPU 运行时, 基于 Cranelift。

项目动机​

CubeCL 的目标是简化高性能、跨硬件可移植计算内核的开发。 目前,要在不同硬件上实现最佳性能,通常需要使用不同语言(如 CUDA、Metal 或 ROCm)编写定制内核。 这种繁琐流程激发了我们开发 CubeCL 的初衷。

CubeCL 采用了以下核心特性:

  1. 自动向量化:在编译时自动使用最优 SIMD 指令。
  2. Comptime:在编译 GPU 内核时动态修改 IR。
  3. 自动调优:在运行时选择最佳内核配置。

这些特性不仅提升了性能,还增强了代码的可组合性、可重用性和可维护性。

我们的愿景不仅是提供优化的计算语言,还包括构建一个 Rust 高性能计算生态系统。

CubeCL 已经提供了线性代数组件,并计划支持卷积、随机数生成、快速傅里叶变换等算法。

特性概览​

自动向量化​

通过 CubeCL,可以为输入变量指定向量化因子,运行时将自动使用最佳指令集。 内核代码始终保持简单,向量化处理由 CubeCL 自动完成。

Comptime 优化​

CubeCL 允许在编译时动态修改 IR,实现指令特化、循环展开、形状特化等优化。 这样可以避免为不同硬件手写多个内核变种。

自动调优​

自动调优通过运行小型基准测试,自动选择性能最佳的内核配置,并缓存结果以优化后续运行时间。

学习资源​

目前的学习资源较少,但可以参考 线性代数库 了解 CubeCL 的实际使用。

声明与历史​

CubeCL 当前处于 alpha 阶段。

最初 CubeCL 仅作为 Burn 的 WebGPU 后端。 随着优化需求的增加,开发了中间表示(IR),并支持 CUDA 编译目标。

通过 Rust 的 proc 宏,创建了一个易用的前端,最终形成了 CubeCL。

通过 CubeCL,释放 Rust 在高性能计算中的潜力,欢迎您的参与和贡献!

链接​

鱼雪