burn_wgpu/
lib.rs

1#![cfg_attr(docsrs, feature(doc_cfg))]
2
3//! The [wgpu](https://github.com/gfx-rs/wgpu) runtime for [Burn](https://github.com/tracel-ai/burn)'s
4//! CubeCL backend: Vulkan, Metal, DirectX 12, OpenGL and WebGPU through one API.
5//!
6//! Applications enable Burn's `wgpu` feature and create a device with `Device::wgpu(..)`, which
7//! picks a graphics API for the platform, or configure it first with `Device::wgpu_options()`.
8//! The `vulkan`, `metal` and `webgpu` features add devices pinned to that API. Applications do not
9//! need this crate directly.
10//!
11//! This crate also re-exports the wgpu runtime types ([`WgpuDevice`], [`WgpuSetup`],
12//! [`MemoryConfiguration`], ...) for sharing an existing wgpu device with Burn, and, with the
13//! `template` feature, the API for launching hand-written WGSL kernels.
14//!
15//! # Feature flags
16//!
17//! - `std` (default): standard library support.
18//! - `fusion` (default): kernel fusion.
19//! - `autotune` (default): benchmark kernel variants at runtime and keep the fastest.
20//! - `vulkan`, `metal`, `webgpu`: devices pinned to that graphics API, using its native shader
21//!   compiler where available.
22//! - `template`: launch hand-written kernels.
23//! - `exclusive-memory-only`: never share a memory page between allocations (always the case on
24//!   wasm).
25//! - `tracing`: instrument operations with the `tracing` crate.
26
27extern crate alloc;
28
29#[cfg(feature = "template")]
30pub use burn_cubecl::{
31    kernel::{KernelMetadata, into_contiguous},
32    kernel_source,
33    template::{KernelSource, SourceKernel, SourceTemplate, build_info},
34};
35
36pub use burn_cubecl::{BoolElement, FloatElement, IntElement};
37pub use burn_cubecl::{CubeBackend, tensor::CubeTensor};
38pub use cubecl::CubeDim;
39pub use cubecl::flex32;
40
41#[cfg(feature = "metal")]
42pub use cubecl::wgpu::MslCompiler;
43#[cfg(not(target_family = "wasm"))]
44pub use cubecl::wgpu::try_init_setup;
45pub use cubecl::wgpu::{
46    AutoCompiler, MemoryConfiguration, RuntimeOptions, WgpuBackend, WgpuDevice, WgpuInitError,
47    WgpuResource, WgpuRuntime, WgpuSetup, WgpuStorage, init_device, init_device_with_api,
48    init_setup, init_setup_async, try_init_device, try_init_device_with_api, try_init_setup_async,
49    wgpu,
50};
51// Vulkan and WebGpu would have conflicting type names
52pub mod graphics {
53    pub use cubecl::wgpu::{AutoGraphicsApi, Dx12, GraphicsApi, Metal, OpenGl, Vulkan, WebGpu};
54}
55
56#[cfg(feature = "fusion")]
57type WgpuInner = burn_fusion::Fusion<CubeBackend>;
58
59#[cfg(not(feature = "fusion"))]
60type WgpuInner = CubeBackend;
61
62/// Tensor backend that uses the wgpu crate for executing GPU compute shaders.
63///
64/// This backend can target multiple graphics APIs, including:
65///   - [Vulkan][crate::graphics::Vulkan] on Linux, Windows, and Android.
66///   - [OpenGL](crate::graphics::OpenGl) on Linux, Windows, and Android.
67///   - [DirectX 12](crate::graphics::Dx12) on Windows.
68///   - [Metal][crate::graphics::Metal] on Apple hardware.
69///   - [WebGPU](crate::graphics::WebGpu) on supported browsers and `wasm` runtimes.
70///
71/// Automatic devices use [`AutoCompiler`] to select WGSL, SPIR-V or MSL, with WGSL fallback
72/// when a native compiler is unavailable. Enabling `metal` makes MSL available without changing
73/// the compiler used by unrelated devices.
74///
75/// With `metal` enabled, an explicitly selected Metal device, created with Burn's
76/// `burn::tensor::Device::metal` or [`cubecl::Device::metal_msl`], requires native MSL and panics
77/// during initialization if it is unavailable. The same applies to [`init_setup`] or
78/// [`init_setup_async`] with [`graphics::Metal`], and to [`init_device_with_api`] with
79/// [`graphics::Metal`]. Importing through [`init_device`] retains automatic compiler selection
80/// and fallback. The corresponding `try_*` functions return initialization errors instead
81/// of panicking.
82///
83/// Multiple backend features can be enabled together. The deprecated `Vulkan`, `WebGpu` and
84/// `Metal` aliases name the same backend; use [`Wgpu`] with an explicit device constructor to
85/// select the graphics API. Compiler selection follows the device and enabled features.
86///
87/// Application code configures and initializes the runtime with `Device::wgpu_options`:
88///
89/// ```rust, ignore
90/// use burn::tensor::{Device, wgpu::WgpuBackend};
91/// let device = Device::wgpu_options()
92///     .graphics_api(WgpuBackend::Vulkan)
93///     .tasks_max(32)
94///     .init()?;
95/// ```
96/// Use `init_async().await` for native async applications and browsers, or
97/// `.setup(existing_setup).init()` to share an application's existing device and queue.
98/// Initialize once, then clone the Burn device to share its runtime. The low-level
99/// [`try_init_setup_async`] and [`try_init_device`] APIs remain available for backend authors.
100///
101/// # Notes
102///
103/// When the `fusion` feature flag is enabled (the default), this backend uses [burn_fusion] to
104/// compile and optimize streams of tensor operations for improved performance. You can disable
105/// the `fusion` feature flag to remove that functionality, which might be necessary on `wasm`
106/// for now.
107pub type Wgpu = WgpuInner;
108
109/// Deprecated alias of [`Wgpu`].
110///
111/// Select Vulkan explicitly with Burn's `burn::tensor::Device::vulkan` or
112/// [`cubecl::Device::vulkan`]. Compiler selection follows the device; see [`Wgpu`].
113#[cfg(feature = "vulkan")]
114#[deprecated(
115    since = "0.22.0",
116    note = "Use `Wgpu` with `burn::tensor::Device::vulkan` to select Vulkan explicitly. This alias is identical to `Wgpu` and does not select a graphics API."
117)]
118pub type Vulkan = WgpuInner;
119
120/// Deprecated alias of [`Wgpu`].
121///
122/// Select WebGPU explicitly with Burn's `burn::tensor::Device::webgpu` or
123/// [`cubecl::Device::webgpu`]. Compiler selection follows the device; see [`Wgpu`].
124#[cfg(feature = "webgpu")]
125#[deprecated(
126    since = "0.22.0",
127    note = "Use `Wgpu` with `burn::tensor::Device::webgpu` to select WebGPU explicitly. This alias is identical to `Wgpu` and does not select a graphics API."
128)]
129pub type WebGpu = WgpuInner;
130
131/// Deprecated alias of [`Wgpu`].
132///
133/// Use Burn's `burn::tensor::Device::metal` or
134/// [`cubecl::Device::metal_msl`] to select Metal explicitly. Both select
135/// [`WgpuBackend::Metal`], which requires native MSL support with `metal` enabled and panics
136/// during initialization if it is unavailable. Automatic devices retain WGSL fallback.
137///
138/// To import an existing Metal setup with the same requirement, use
139/// `init_device_with_api::<graphics::Metal>(setup, RuntimeOptions::default())`.
140#[cfg(feature = "metal")]
141#[deprecated(
142    since = "0.22.0",
143    note = "Use `Wgpu` with `burn::tensor::Device::metal` to select Metal explicitly. This alias is identical to `Wgpu` and does not select a graphics API."
144)]
145pub type Metal = WgpuInner;
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150    use burn_backend::{Backend, BoolStore, DType, DeviceOps};
151
152    fn assert_common_dtypes(device: &cubecl::Device) {
153        // Metal and Vulkan alias Wgpu; the device selects the compiler.
154        type B = Wgpu;
155        let defaults = device.defaults();
156        let scheme = defaults.quantization.scheme;
157
158        assert!(B::supports_dtype(device, DType::F32));
159        assert!(B::supports_dtype(device, DType::F16));
160        assert!(B::supports_dtype(device, DType::I64));
161        assert!(B::supports_dtype(device, DType::I32));
162        assert!(B::supports_dtype(device, DType::U64));
163        assert!(B::supports_dtype(device, DType::U32));
164        assert!(B::supports_dtype(device, DType::QFloat(scheme)));
165        assert!(!B::supports_dtype(device, DType::Bool(BoolStore::Native)));
166        assert!(B::supports_dtype(device, defaults.bool_dtype.into()));
167    }
168
169    #[cfg(any(
170        all(feature = "vulkan", not(target_family = "wasm")),
171        all(feature = "metal", target_vendor = "apple")
172    ))]
173    fn assert_fp4_dtypes(device: &cubecl::Device) {
174        use burn_backend::quantization::{QuantScheme, QuantStore, QuantValue, ScaleDtype};
175
176        // FP8 block scale storage and software conversion support NVFP4 and MXFP4.
177        let fp4 = QuantScheme::default()
178            .with_value(QuantValue::E2M1)
179            .with_store(QuantStore::PackedU32(0));
180        let nvfp4 = fp4
181            .per_block([16], ScaleDtype::UE4M3)
182            .per_tensor(ScaleDtype::F32);
183        let mxfp4 = fp4.per_block([32], ScaleDtype::UE8M0);
184        assert!(Wgpu::supports_dtype(device, DType::QFloat(nvfp4)));
185        assert!(Wgpu::supports_dtype(device, DType::QFloat(mxfp4)));
186    }
187
188    #[test]
189    fn should_support_dtypes() {
190        let device = cubecl::Device::Wgpu(WgpuDevice::default());
191        assert_common_dtypes(&device);
192    }
193
194    #[cfg(all(feature = "vulkan", not(target_family = "wasm")))]
195    #[test]
196    fn should_support_vulkan_dtypes() {
197        type B = Wgpu;
198        let device = cubecl::Device::Wgpu(WgpuDevice::default().on(WgpuBackend::Vulkan));
199        assert_common_dtypes(&device);
200
201        assert!(B::supports_dtype(&device, DType::I16));
202        assert!(B::supports_dtype(&device, DType::I8));
203        assert!(B::supports_dtype(&device, DType::U16));
204        assert!(B::supports_dtype(&device, DType::U8));
205
206        // F64 is supported through the shader_float64 feature.
207        assert!(B::supports_dtype(&device, DType::F64));
208        assert!(!B::supports_dtype(&device, DType::Flex32));
209        // BF16 supports storage and conversion, but not all scalar arithmetic operations.
210        assert!(!B::supports_dtype(&device, DType::BF16));
211
212        assert_fp4_dtypes(&device);
213    }
214
215    #[cfg(all(feature = "metal", target_vendor = "apple"))]
216    #[test]
217    fn should_support_metal_dtypes() {
218        type B = Wgpu;
219        let device = cubecl::Device::Wgpu(WgpuDevice::default().on(WgpuBackend::Metal));
220        assert_common_dtypes(&device);
221
222        assert!(B::supports_dtype(&device, DType::I16));
223        assert!(B::supports_dtype(&device, DType::I8));
224        assert!(B::supports_dtype(&device, DType::U16));
225        assert!(B::supports_dtype(&device, DType::U8));
226
227        assert!(!B::supports_dtype(&device, DType::F64));
228        // MSL 3.2 carries bfloat natively.
229        assert!(B::supports_dtype(&device, DType::BF16));
230        assert!(!B::supports_dtype(&device, DType::Flex32));
231
232        assert_fp4_dtypes(&device);
233    }
234
235    #[cfg(not(any(feature = "vulkan", feature = "metal")))]
236    #[test]
237    fn should_support_wgsl_dtypes() {
238        type B = Wgpu;
239        let device = cubecl::Device::Wgpu(WgpuDevice::default());
240        assert_common_dtypes(&device);
241
242        assert!(B::supports_dtype(&device, DType::Flex32));
243        #[cfg(target_os = "macos")]
244        assert!(!B::supports_dtype(&device, DType::F64));
245        #[cfg(not(target_os = "macos"))]
246        assert!(B::supports_dtype(&device, DType::F64));
247        assert!(!B::supports_dtype(&device, DType::BF16));
248        assert!(!B::supports_dtype(&device, DType::I16));
249        assert!(!B::supports_dtype(&device, DType::I8));
250        assert!(!B::supports_dtype(&device, DType::U16));
251        assert!(!B::supports_dtype(&device, DType::U8));
252    }
253}
burn_wgpu/lib.rs

burn_wgpu/
lib.rs
