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;
16
17pub use cubecl::wgpu::{
18    AutoCompiler, MemoryConfiguration, RuntimeOptions, WgpuDevice, WgpuResource, WgpuRuntime,
19    WgpuSetup, WgpuStorage, init_device, init_setup, init_setup_async,
20};
21// Vulkan and WebGpu would have conflicting type names
22pub mod graphics {
23    pub use cubecl::wgpu::{AutoGraphicsApi, Dx12, GraphicsApi, Metal, OpenGl, Vulkan, WebGpu};
24}
25
26#[cfg(feature = "cubecl-wgsl")]
27pub use cubecl::wgpu::WgslCompiler;
28#[cfg(feature = "cubecl-spirv")]
29pub use cubecl::wgpu::vulkan::VkSpirvCompiler;
30
31#[cfg(feature = "fusion")]
32/// Tensor backend that uses the wgpu crate for executing GPU compute shaders.
33///
34/// This backend can target multiple graphics APIs, including:
35///   - [Vulkan][crate::graphics::Vulkan] on Linux, Windows, and Android.
36///   - [OpenGL](crate::graphics::OpenGl) on Linux, Windows, and Android.
37///   - [DirectX 12](crate::graphics::Dx12) on Windows.
38///   - [Metal][crate::graphics::Metal] on Apple hardware.
39///   - [WebGPU](crate::graphics::WebGpu) on supported browsers and `wasm` runtimes.
40///
41/// To configure the wgpu backend, eg. to select what graphics API to use or what memory strategy to use,
42/// you have to manually initialize the runtime. For example:
43///
44/// ```rust, ignore
45/// fn custom_init() {
46///     let device = Default::default();
47///     burn::backend::wgpu::init_setup::<burn::backend::wgpu::graphics::Vulkan>(
48///         &device,
49///         Default::default(),
50///     );
51/// }
52/// ```
53/// will mean the given device (in this case the default) will be initialized to use Vulkan as the graphics API.
54/// It's also possible to use an existing wgpu device, by using `init_device`.
55///
56/// # Notes
57///
58/// This version of the wgpu backend uses [burn_fusion] to compile and optimize streams of tensor
59/// operations for improved performance.
60///
61/// You can disable the `fusion` feature flag to remove that functionality, which might be
62/// necessary on `wasm` for now.
63pub type Wgpu<F = f32, I = i32, B = u32> =
64    burn_fusion::Fusion<CubeBackend<cubecl::wgpu::WgpuRuntime, F, I, B>>;
65
66#[cfg(not(feature = "fusion"))]
67/// Tensor backend that uses the wgpu crate for executing GPU compute shaders.
68///
69/// This backend can target multiple graphics APIs, including:
70///   - [Vulkan] on Linux, Windows, and Android.
71///   - [OpenGL](crate::OpenGl) on Linux, Windows, and Android.
72///   - [DirectX 12](crate::Dx12) on Windows.
73///   - [Metal] on Apple hardware.
74///   - [WebGPU](crate::WebGpu) on supported browsers and `wasm` runtimes.
75///
76/// To configure the wgpu backend, eg. to select what graphics API to use or what memory strategy to use,
77/// you have to manually initialize the runtime. For example:
78///
79/// ```rust, ignore
80/// fn custom_init() {
81///     let device = Default::default();
82///     burn::backend::wgpu::init_setup::<burn::backend::wgpu::graphics::Vulkan>(
83///         &device,
84///         Default::default(),
85///     );
86/// }
87/// ```
88/// will mean the given device (in this case the default) will be initialized to use Vulkan as the graphics API.
89/// It's also possible to use an existing wgpu device, by using `init_device`.
90///
91/// # Notes
92///
93/// This version of the wgpu backend doesn't use [burn_fusion] to compile and optimize streams of tensor
94/// operations.
95///
96/// You can enable the `fusion` feature flag to add that functionality, which might improve
97/// performance.
98pub type Wgpu<F = f32, I = i32, B = u32> = CubeBackend<cubecl::wgpu::WgpuRuntime, F, I, B>;
99
100#[cfg(feature = "vulkan")]
101/// Tensor backend that leverages the Vulkan graphics API to execute GPU compute shaders compiled to SPIR-V.
102pub type Vulkan<F = f32, I = i32, B = u8> = Wgpu<F, I, B>;
103
104#[cfg(feature = "webgpu")]
105/// Tensor backend that uses the wgpu crate to execute GPU compute shaders written in WGSL.
106pub type WebGpu<F = f32, I = i32, B = u32> = Wgpu<F, I, B>;
107
108#[cfg(feature = "metal")]
109/// Tensor backend that leverages the Metal graphics API to execute GPU compute shaders compiled to MSL.
110pub type Metal<F = f32, I = i32, B = u8> = Wgpu<F, I, B>;
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use burn_backend::{Backend, BoolStore, DType, QTensorPrimitive};
116
117    #[test]
118    fn should_support_dtypes() {
119        type B = Wgpu;
120        let device = Default::default();
121
122        assert!(B::supports_dtype(&device, DType::F32));
123        assert!(B::supports_dtype(&device, DType::I64));
124        assert!(B::supports_dtype(&device, DType::I32));
125        assert!(B::supports_dtype(&device, DType::U64));
126        assert!(B::supports_dtype(&device, DType::U32));
127        assert!(B::supports_dtype(
128            &device,
129            DType::QFloat(CubeTensor::<WgpuRuntime>::default_scheme())
130        ));
131        // Registered as supported type but we don't actually use it?
132        assert!(B::supports_dtype(&device, DType::Bool(BoolStore::Native)));
133
134        #[cfg(feature = "vulkan")]
135        {
136            assert!(B::supports_dtype(&device, DType::F16));
137            assert!(B::supports_dtype(&device, DType::I16));
138            assert!(B::supports_dtype(&device, DType::I8));
139            assert!(B::supports_dtype(&device, DType::U16));
140            assert!(B::supports_dtype(&device, DType::U8));
141
142            assert!(!B::supports_dtype(&device, DType::F64));
143            assert!(!B::supports_dtype(&device, DType::Flex32));
144            // Not supported for any arithmetics, but buffer, conversion and possibly matmul (hw dependent)
145            assert!(!B::supports_dtype(&device, DType::BF16));
146        }
147
148        #[cfg(feature = "metal")]
149        {
150            assert!(B::supports_dtype(&device, DType::F16));
151            assert!(B::supports_dtype(&device, DType::I16));
152            assert!(B::supports_dtype(&device, DType::I8));
153            assert!(B::supports_dtype(&device, DType::U16));
154            assert!(B::supports_dtype(&device, DType::U8));
155
156            assert!(!B::supports_dtype(&device, DType::F64));
157            assert!(!B::supports_dtype(&device, DType::BF16));
158            assert!(!B::supports_dtype(&device, DType::Flex32));
159        }
160
161        // On macOS without the `metal` feature, wgpu still uses Metal at runtime,
162        // which doesn't support F64 or BF16.
163        #[cfg(all(not(any(feature = "vulkan", feature = "metal")), target_os = "macos"))]
164        {
165            assert!(B::supports_dtype(&device, DType::Flex32));
166            assert!(B::supports_dtype(&device, DType::F16));
167
168            assert!(!B::supports_dtype(&device, DType::F64));
169            assert!(!B::supports_dtype(&device, DType::BF16));
170            assert!(!B::supports_dtype(&device, DType::I16));
171            assert!(!B::supports_dtype(&device, DType::I8));
172            assert!(!B::supports_dtype(&device, DType::U16));
173            assert!(!B::supports_dtype(&device, DType::U8));
174        }
175
176        #[cfg(not(any(feature = "vulkan", feature = "metal", target_os = "macos")))]
177        {
178            assert!(B::supports_dtype(&device, DType::F64));
179            assert!(B::supports_dtype(&device, DType::Flex32));
180            assert!(B::supports_dtype(&device, DType::F16));
181
182            assert!(!B::supports_dtype(&device, DType::BF16));
183            assert!(!B::supports_dtype(&device, DType::I16));
184            assert!(!B::supports_dtype(&device, DType::I8));
185            assert!(!B::supports_dtype(&device, DType::U16));
186            assert!(!B::supports_dtype(&device, DType::U8));
187        }
188    }
189}
burn_wgpu/lib.rs

burn_wgpu/
lib.rs
