burn_wgpu/
lib.rs

1#![cfg_attr(docsrs, feature(doc_cfg))]
2
3extern crate alloc;
4
5#[cfg(feature = "template")]
6pub use burn_cubecl::{
7    kernel::{KernelMetadata, into_contiguous},
8    kernel_source,
9    template::{KernelSource, SourceKernel, SourceTemplate, build_info},
10};
11
12pub use burn_cubecl::{BoolElement, FloatElement, IntElement};
13pub use burn_cubecl::{CubeBackend, tensor::CubeTensor};
14pub use cubecl::CubeDim;
15pub use cubecl::flex32;
16use cubecl::throughput::{ThroughputKey, ThroughputValue};
17
18#[cfg(feature = "metal")]
19use cubecl::wgpu::MslCompiler;
20#[cfg(feature = "vulkan")]
21use cubecl::wgpu::SpirvCompiler;
22#[cfg(feature = "webgpu")]
23use cubecl::wgpu::WgslCompiler;
24
25pub use cubecl::wgpu::{
26    AutoCompiler, MemoryConfiguration, RuntimeOptions, WgpuDevice, WgpuResource, WgpuRuntime,
27    WgpuSetup, WgpuStorage, init_device, init_setup, init_setup_async,
28};
29// Vulkan and WebGpu would have conflicting type names
30pub mod graphics {
31    pub use cubecl::wgpu::{AutoGraphicsApi, Dx12, GraphicsApi, Metal, OpenGl, Vulkan, WebGpu};
32}
33
34#[cfg(feature = "fusion")]
35type WgpuInner<C> = burn_fusion::Fusion<CubeBackend<cubecl::wgpu::WgpuRuntime<C>>>;
36
37#[cfg(not(feature = "fusion"))]
38type WgpuInner<C> = CubeBackend<cubecl::wgpu::WgpuRuntime<C>>;
39
40/// Tensor backend that uses the wgpu crate for executing GPU compute shaders.
41///
42/// This backend can target multiple graphics APIs, including:
43///   - [Vulkan][crate::graphics::Vulkan] on Linux, Windows, and Android.
44///   - [OpenGL](crate::graphics::OpenGl) on Linux, Windows, and Android.
45///   - [DirectX 12](crate::graphics::Dx12) on Windows.
46///   - [Metal][crate::graphics::Metal] on Apple hardware.
47///   - [WebGPU](crate::graphics::WebGpu) on supported browsers and `wasm` runtimes.
48///
49/// The selected graphics API is chosen automatically at runtime, and the appropriate shader
50/// compiler (WGSL, SPIR-V or MSL) is dispatched via [`AutoCompiler`]. When the target API is
51/// known ahead of time, prefer the dedicated `Vulkan`, `WebGpu` or `Metal` backend aliases
52/// (enabled by their respective Cargo features), which lock the compiler at compile time and
53/// avoid the runtime dispatch.
54///
55/// To configure the wgpu backend, eg. to select what graphics API to use or what memory strategy to use,
56/// you have to manually initialize the runtime. For example:
57///
58/// ```rust, ignore
59/// fn custom_init() {
60///     let device = Default::default();
61///     burn::backend::wgpu::init_setup::<burn::backend::wgpu::graphics::Vulkan>(
62///         &device,
63///         Default::default(),
64///     );
65/// }
66/// ```
67/// will mean the given device (in this case the default) will be initialized to use Vulkan as the graphics API.
68/// It's also possible to use an existing wgpu device, by using `init_device`.
69///
70/// # Notes
71///
72/// When the `fusion` feature flag is enabled (the default), this backend uses [burn_fusion] to
73/// compile and optimize streams of tensor operations for improved performance. You can disable
74/// the `fusion` feature flag to remove that functionality, which might be necessary on `wasm`
75/// for now.
76pub type Wgpu = WgpuInner<AutoCompiler>;
77
78/// Measure peak throughput on a wgpu `device` for each of the given `keys`.
79///
80/// Uses the auto-selected shader compiler, matching the default [`Wgpu`] backend.
81pub fn device_throughput(
82    device: &WgpuDevice,
83    keys: &[ThroughputKey],
84) -> alloc::vec::Vec<ThroughputValue> {
85    cubecl::std::throughput::device_throughput::<cubecl::wgpu::WgpuRuntime<AutoCompiler>>(
86        device, keys,
87    )
88}
89
90/// Tensor backend that leverages the Vulkan graphics API to execute GPU compute shaders compiled to SPIR-V.
91///
92/// This is a specialization of [`Wgpu`] that pins the shader compiler to SPIR-V at compile time,
93/// removing the runtime [`AutoCompiler`] dispatch. Enable the `vulkan` feature to use it.
94/// Multiple wgpu backend aliases (`Vulkan`, `WebGpu`, `Metal`) can be enabled simultaneously
95/// since each is a distinct type parameterized by its own compiler.
96#[cfg(feature = "vulkan")]
97pub type Vulkan = WgpuInner<SpirvCompiler>;
98
99/// Tensor backend that uses the wgpu crate to execute GPU compute shaders written in WGSL.
100///
101/// This is a specialization of [`Wgpu`] that pins the shader compiler to WGSL at compile time,
102/// removing the runtime [`AutoCompiler`] dispatch. Enable the `webgpu` feature to use it.
103#[cfg(feature = "webgpu")]
104pub type WebGpu = WgpuInner<WgslCompiler>;
105
106/// Tensor backend that leverages the Metal graphics API to execute GPU compute shaders compiled to MSL.
107///
108/// This is a specialization of [`Wgpu`] that pins the shader compiler to MSL at compile time,
109/// removing the runtime [`AutoCompiler`] dispatch. Enable the `metal` feature to use it.
110#[cfg(feature = "metal")]
111pub type Metal = WgpuInner<MslCompiler>;
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use burn_backend::{Backend, BoolStore, DType, DeviceOps};
117
118    #[test]
119    fn should_support_dtypes() {
120        type B = Wgpu;
121        let device = WgpuDevice::default();
122        let scheme = device.defaults().quantization.scheme;
123
124        assert!(B::supports_dtype(&device, DType::F32));
125        assert!(B::supports_dtype(&device, DType::I64));
126        assert!(B::supports_dtype(&device, DType::I32));
127        assert!(B::supports_dtype(&device, DType::U64));
128        assert!(B::supports_dtype(&device, DType::U32));
129        assert!(B::supports_dtype(&device, DType::QFloat(scheme)));
130        assert!(!B::supports_dtype(&device, DType::Bool(BoolStore::Native)));
131
132        #[cfg(feature = "vulkan")]
133        {
134            assert!(B::supports_dtype(&device, DType::F16));
135            assert!(B::supports_dtype(&device, DType::I16));
136            assert!(B::supports_dtype(&device, DType::I8));
137            assert!(B::supports_dtype(&device, DType::U16));
138            assert!(B::supports_dtype(&device, DType::U8));
139
140            // NOTE: F64 is not part of the default types, but is supported based on `shader_float64` feature
141            assert!(B::supports_dtype(&device, DType::F64));
142            assert!(!B::supports_dtype(&device, DType::Flex32));
143            // Not supported for any arithmetics, but buffer, conversion and possibly matmul (hw dependent)
144            assert!(!B::supports_dtype(&device, DType::BF16));
145        }
146
147        #[cfg(feature = "metal")]
148        {
149            assert!(B::supports_dtype(&device, DType::F16));
150            assert!(B::supports_dtype(&device, DType::I16));
151            assert!(B::supports_dtype(&device, DType::I8));
152            assert!(B::supports_dtype(&device, DType::U16));
153            assert!(B::supports_dtype(&device, DType::U8));
154
155            assert!(!B::supports_dtype(&device, DType::F64));
156            assert!(!B::supports_dtype(&device, DType::BF16));
157            assert!(!B::supports_dtype(&device, DType::Flex32));
158        }
159
160        // On macOS without the `metal` feature, wgpu still uses Metal at runtime,
161        // which doesn't support F64 or BF16.
162        #[cfg(all(not(any(feature = "vulkan", feature = "metal")), target_os = "macos"))]
163        {
164            assert!(B::supports_dtype(&device, DType::Flex32));
165            assert!(B::supports_dtype(&device, DType::F16));
166
167            assert!(!B::supports_dtype(&device, DType::F64));
168            assert!(!B::supports_dtype(&device, DType::BF16));
169            assert!(!B::supports_dtype(&device, DType::I16));
170            assert!(!B::supports_dtype(&device, DType::I8));
171            assert!(!B::supports_dtype(&device, DType::U16));
172            assert!(!B::supports_dtype(&device, DType::U8));
173        }
174
175        #[cfg(not(any(feature = "vulkan", feature = "metal", target_os = "macos")))]
176        {
177            assert!(B::supports_dtype(&device, DType::F64));
178            assert!(B::supports_dtype(&device, DType::Flex32));
179            assert!(B::supports_dtype(&device, DType::F16));
180
181            assert!(!B::supports_dtype(&device, DType::BF16));
182            assert!(!B::supports_dtype(&device, DType::I16));
183            assert!(!B::supports_dtype(&device, DType::I8));
184            assert!(!B::supports_dtype(&device, DType::U16));
185            assert!(!B::supports_dtype(&device, DType::U8));
186        }
187    }
188}
burn_wgpu/lib.rs

burn_wgpu/
lib.rs
