burn_tensor/wgpu.rs
1//! Initialize wgpu through [`Device::wgpu_options`].
2//!
3//! In browsers, call `Device::wgpu_options().init_async().await?` before creating
4//! models or tensors. Use the returned [`Device`] for subsequent operations.
5//! Native applications can use `.init()` instead.
6//!
7//! ```no_run
8//! # async fn example() -> Result<(), burn_tensor::wgpu::WgpuInitError> {
9//! use burn_tensor::{Device, Tensor};
10//!
11//! let device = Device::wgpu_options().tasks_max(32).init_async().await?;
12//! let input = Tensor::<2>::ones([2, 3], &device);
13//! # Ok(())
14//! # }
15//! ```
16//!
17//! Supply a [`WgpuSetup`](crate::wgpu::WgpuSetup) with
18//! [`.setup(...)`](crate::wgpu::WgpuOptions::setup) to use
19//! an existing application's device and queue. Register it once, then pass or
20//! clone the returned Burn device into models. Repeated registration of the same
21//! handles is not detected. Registering a setup does not
22//! change which runtime [`Device::default`] selects.
23//!
24//! [`wgpu`](crate::wgpu::wgpu) exposes the exact wgpu version whose handles this runtime accepts.
25//! Sharing handles does not convert tensors into buffers or wait for GPU work:
26//! flush pending Burn work before submitting external consumers on the shared
27//! queue, and keep shared allocations alive until those consumers finish.
28
29use burn_dispatch::devices::wgpu as runtime;
30pub use runtime::{MemoryConfiguration, WgpuBackend, WgpuInitError, WgpuSetup, wgpu};
31
32use crate::{Device, DeviceKind, device::wgpu_device};
33
34impl Device {
35 /// Options for initializing Burn's wgpu runtime or registering an existing setup.
36 ///
37 /// Starts with the default device selector, automatic graphics API, and runtime
38 /// defaults. Call `init` or `init_async` before creating tensors on the same selector.
39 /// An initialized runtime cannot be reconfigured. Clone its device to reuse it.
40 pub fn wgpu_options() -> WgpuOptions {
41 WgpuOptions::default()
42 }
43}
44
45#[derive(Default)]
46struct Options {
47 memory: Option<MemoryConfiguration>,
48 tasks_max: Option<usize>,
49}
50
51impl Options {
52 fn resolve(self) -> Result<runtime::RuntimeOptions, WgpuInitError> {
53 let mut options = runtime::RuntimeOptions::with_tasks_max(self.tasks_max)?;
54 if let Some(memory) = self.memory {
55 options.memory_config = memory;
56 }
57 Ok(options)
58 }
59}
60
61/// Options for initializing Burn's wgpu runtime.
62///
63/// Construct with [`Device::wgpu_options`] or [`Default::default`]. Options cover
64/// device acquisition and runtime settings such as memory management and task batching.
65/// Setters do not acquire a GPU. Initialization registers the selected runtime for reuse.
66/// Dropping the returned device does not deregister that runtime.
67#[must_use = "wgpu options do nothing until initialized"]
68pub struct WgpuOptions {
69 kind: DeviceKind,
70 api: WgpuBackend,
71 selection_explicit: bool,
72 options: Options,
73}
74
75impl Default for WgpuOptions {
76 fn default() -> Self {
77 Self {
78 kind: DeviceKind::DefaultDevice,
79 api: WgpuBackend::Auto,
80 selection_explicit: false,
81 options: Options::default(),
82 }
83 }
84}
85
86impl WgpuOptions {
87 /// Select an adapter kind. `Existing` identifiers cannot be initialized here;
88 /// register an existing setup with [`Self::setup`] instead.
89 ///
90 /// Browser selection follows the runtime's power preferences; browsers do not
91 /// expose native adapter indices. The default prefers a high-performance GPU.
92 pub fn device_kind(mut self, kind: DeviceKind) -> Self {
93 self.kind = kind;
94 self.selection_explicit = true;
95 self
96 }
97
98 /// Pin the graphics API, or let [`WgpuBackend::Auto`] choose it.
99 ///
100 /// An explicit API never falls back to another API. This selects the graphics
101 /// API, not a shader compiler; the runtime chooses a supported compiler.
102 /// With the `metal` feature enabled, an explicit Metal selection requires native
103 /// MSL support. Initialization returns an error if it is unavailable; automatic
104 /// selection permits WGSL fallback.
105 pub fn graphics_api(mut self, api: WgpuBackend) -> Self {
106 self.api = api;
107 self.selection_explicit = true;
108 self
109 }
110
111 /// Set the runtime's memory allocation policy.
112 pub fn memory_config(mut self, config: MemoryConfiguration) -> Self {
113 self.options.memory = Some(config);
114 self
115 }
116
117 /// Set the maximum tasks aggregated into a GPU command. Must be nonzero.
118 /// Overrides `CUBECL_WGPU_MAX_TASKS`, including a malformed environment value.
119 pub fn tasks_max(mut self, tasks: usize) -> Self {
120 self.options.tasks_max = Some(tasks);
121 self
122 }
123
124 /// Register existing wgpu handles instead of acquiring another device.
125 ///
126 /// Memory and task options carry over. Explicit calls to [`Self::device_kind`]
127 /// or [`Self::graphics_api`] conflict with a supplied setup, even when those
128 /// options name defaults; [`WgpuSetupOptions::init`] returns an error.
129 ///
130 /// Handles must belong to one coherent setup. Compiler selection follows the
131 /// same backend and adapter checks as acquired setups. The supplied device
132 /// must enable the selected compiler's required features, including native
133 /// extensions for SPIR-V or MSL; registration does not validate all of them.
134 pub fn setup(self, setup: WgpuSetup) -> WgpuSetupOptions {
135 WgpuSetupOptions {
136 setup,
137 options: self.options,
138 selection_explicit: self.selection_explicit,
139 }
140 }
141
142 fn into_request(self) -> Result<(runtime::WgpuDevice, runtime::RuntimeOptions), WgpuInitError> {
143 if matches!(self.kind, DeviceKind::Existing(_)) {
144 return Err(WgpuInitError::InvalidConfiguration {
145 message: "Existing is a runtime identity, not an adapter selector; supply a setup or reuse its device".into(),
146 });
147 }
148 Ok((wgpu_device(self.kind, self.api), self.options.resolve()?))
149 }
150
151 /// Initialize the device synchronously on native targets.
152 ///
153 /// Use [`Self::init_async`] when running in a browser or an async application.
154 ///
155 /// # Errors
156 ///
157 /// Returns acquisition, configuration, or registration errors.
158 ///
159 /// Returns an error if the runtime for the selected device kind and graphics
160 /// API has already been initialized, either explicitly or by a tensor operation.
161 /// Runtime options must be set before first use.
162 #[cfg(not(target_family = "wasm"))]
163 pub fn init(self) -> Result<Device, WgpuInitError> {
164 self.init_with_setup().map(|(device, _)| device)
165 }
166
167 /// Initialize the device asynchronously on native or browser targets.
168 ///
169 /// # Errors
170 ///
171 /// Returns acquisition, configuration, or registration errors.
172 ///
173 /// Returns an error if the runtime for the selected device kind and graphics
174 /// API has already been initialized, either explicitly or by a tensor operation.
175 /// Runtime options must be set before first use.
176 pub async fn init_async(self) -> Result<Device, WgpuInitError> {
177 self.init_with_setup_async().await.map(|(device, _)| device)
178 }
179
180 /// Initialize synchronously and return the exact handles used by this runtime.
181 ///
182 /// Share these handles with another wgpu component. Do not register the
183 /// returned setup again; clone the accompanying Burn device to reuse it.
184 ///
185 /// # Errors
186 ///
187 /// Returns acquisition, configuration, or registration errors.
188 ///
189 /// Returns an error if the runtime for the selected device kind and graphics
190 /// API has already been initialized, either explicitly or by a tensor operation.
191 /// Runtime options must be set before first use.
192 #[cfg(not(target_family = "wasm"))]
193 pub fn init_with_setup(self) -> Result<(Device, WgpuSetup), WgpuInitError> {
194 let (device, options) = self.into_request()?;
195 let setup = runtime::try_init_setup::<runtime::AutoGraphicsApi>(&device, options)?;
196 Ok((Device::from(device), setup))
197 }
198
199 /// Initialize asynchronously and return the exact handles used by this runtime.
200 ///
201 /// Available on native and browser targets. The returned device is ready for
202 /// tensor operations; the setup is ready to share with another wgpu component.
203 ///
204 /// # Errors
205 ///
206 /// Returns acquisition, configuration, or registration errors.
207 ///
208 /// Returns an error if the runtime for the selected device kind and graphics
209 /// API has already been initialized, either explicitly or by a tensor operation.
210 /// Runtime options must be set before first use.
211 pub async fn init_with_setup_async(self) -> Result<(Device, WgpuSetup), WgpuInitError> {
212 let (device, options) = self.into_request()?;
213 let setup =
214 runtime::try_init_setup_async::<runtime::AutoGraphicsApi>(&device, options).await?;
215 Ok((Device::from(device), setup))
216 }
217}
218
219/// Runtime options for registering an existing wgpu setup.
220///
221/// Created by [`WgpuOptions::setup`]. Registration is synchronous even in
222/// a browser, since the application has already acquired the device and queue.
223#[must_use = "wgpu setup options do nothing until initialized"]
224pub struct WgpuSetupOptions {
225 setup: WgpuSetup,
226 options: Options,
227 selection_explicit: bool,
228}
229
230impl WgpuSetupOptions {
231 /// Set the runtime's memory allocation policy.
232 pub fn memory_config(mut self, config: MemoryConfiguration) -> Self {
233 self.options.memory = Some(config);
234 self
235 }
236
237 /// Set a nonzero task count, overriding `CUBECL_WGPU_MAX_TASKS`.
238 pub fn tasks_max(mut self, tasks: usize) -> Self {
239 self.options.tasks_max = Some(tasks);
240 self
241 }
242
243 /// Register the setup and return its initialized Burn device.
244 ///
245 /// Each call creates a new runtime identity; repeated handles are not detected.
246 /// Failed initialization does not change an existing registration.
247 ///
248 /// # Errors
249 ///
250 /// Returns an error for conflicting selection options, invalid runtime options,
251 /// incompatible capabilities, or a failure to register the runtime.
252 pub fn init(self) -> Result<Device, WgpuInitError> {
253 if self.selection_explicit {
254 return Err(WgpuInitError::InvalidConfiguration {
255 message: "device_kind and graphics_api cannot be combined with an existing setup"
256 .into(),
257 });
258 }
259 runtime::try_init_device(self.setup, self.options.resolve()?).map(Device::from)
260 }
261}
262
263#[cfg(test)]
264mod tests {
265 use super::*;
266
267 #[test]
268 fn existing_identity_is_not_an_adapter_to_create() {
269 assert!(matches!(
270 Device::wgpu_options()
271 .device_kind(DeviceKind::Existing(42))
272 .into_request(),
273 Err(WgpuInitError::InvalidConfiguration { .. })
274 ));
275 }
276
277 #[test]
278 fn zero_tasks_are_rejected_before_acquiring_a_gpu() {
279 assert!(matches!(
280 WgpuOptions::default().tasks_max(0).into_request(),
281 Err(WgpuInitError::InvalidConfiguration { .. })
282 ));
283 }
284
285 #[cfg(all(feature = "std", not(target_family = "wasm")))]
286 #[test]
287 fn explicit_tasks_override_a_malformed_environment_default() {
288 // Isolate environment changes from other tests and runtime initialization.
289 const CHILD: &str = "BURN_TEST_WGPU_TASKS_CHILD";
290 if std::env::var_os(CHILD).is_some() {
291 assert!(matches!(
292 Device::wgpu_options().into_request(),
293 Err(WgpuInitError::InvalidConfiguration { .. })
294 ));
295 assert_eq!(
296 Device::wgpu_options()
297 .tasks_max(7)
298 .into_request()
299 .unwrap()
300 .1
301 .tasks_max,
302 7
303 );
304 return;
305 }
306 let status = std::process::Command::new(std::env::current_exe().unwrap())
307 .args([
308 "--exact",
309 "wgpu::tests::explicit_tasks_override_a_malformed_environment_default",
310 ])
311 .env(CHILD, "1")
312 .env("CUBECL_WGPU_MAX_TASKS", "invalid")
313 .status()
314 .unwrap();
315 assert!(status.success());
316 }
317}