15 releases (9 breaking)
Uses new Rust 2024
0.17.0 | Apr 23, 2025 |
---|---|
0.16.0 | Jan 14, 2025 |
0.15.0 | Oct 28, 2024 |
0.13.2 | May 3, 2024 |
0.8.0 | Jul 25, 2023 |
#1220 in Machine learning
17,312 downloads per month
Used in 41 crates
(6 directly)
2MB
41K
SLoC
Burn WGPU Backend
Burn WGPU backend
This crate provides a WGPU backend for Burn using the wgpu.
The backend supports Vulkan, Metal, DirectX11/12, OpenGL, WebGPU.
Usage Example
#[cfg(feature = "wgpu")]
mod wgpu {
use burn_autodiff::Autodiff;
use burn_wgpu::{Wgpu, WgpuDevice};
use mnist::training;
pub fn run() {
let device = WgpuDevice::default();
training::run::<Autodiff<Wgpu<f32, i32>>>(device);
}
}
⚠️ Warning
When using one of thewgpu
backends, you may encounter compilation errors related to recursive type evaluation. This is due to complex type nesting within thewgpu
dependency chain.
To resolve this issue, add the following line at the top of yourmain.rs
orlib.rs
file:#![recursion_limit = "256"]
The default recursion limit (128) is often just below the required depth (typically 130-150) due to deeply nested associated types and trait bounds.
Configuration
You can set BURN_WGPU_MAX_TASKS
to a positive integer that determines how many computing tasks are
submitted in batches to the graphics API.
Alternative SPIR-V backend
When targeting Vulkan, the spirv
feature flag can be enabled to enable the SPIR-V compiler
backend, which performs significantly better than WGSL. This is especially true for matrix
multiplication, where SPIR-V can make use of TensorCores and run at f16
precision. This isn't
currently supported by WGSL. The compiler can also be selected at runtime by setting the
corresponding generic parameter to either SpirV
or Wgsl
.
Platform Support
Option | CPU | GPU | Linux | MacOS | Windows | Android | iOS | WASM |
---|---|---|---|---|---|---|---|---|
Metal | No | Yes | No | Yes | No | No | Yes | No |
Vulkan | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No |
OpenGL | No | Yes | Yes | Yes | Yes | Yes | Yes | No |
WebGpu | No | Yes | No | No | No | No | No | Yes |
Dx11/Dx12 | No | Yes | No | No | Yes | No | No | No |
Dependencies
~30–63MB
~1M SLoC