burn_tensor/device.rs
1pub use burn_std::{
2 DeviceError, DeviceSettings, ExecutionError, backtrace::BackTrace, device::DeviceId,
3};
4
5#[cfg(feature = "cubecl")]
6pub use burn_backend::cubecl::{
7 AdapterLuid, DeviceIdentity, MemoryAccess, PciAddress, PciVendor, PhysicalDevice,
8 ThroughputError, ThroughputKey, ThroughputMode, ThroughputValue,
9};
10use burn_backend::{Backend, DeviceOps};
11pub use burn_backend::{
12 MemoryPoolUsage, ProfileDuration, ProfileOptions, ProfileTicks, SlicedPoolReport, TimingMethod,
13};
14#[allow(unused)]
15use burn_dispatch::DispatchDeviceId;
16#[cfg(feature = "autodiff")]
17use burn_dispatch::GradientCheckpointingStrategy;
18#[cfg(any(feature = "cpu", feature = "cuda", feature = "rocm", feature = "wgpu"))]
19use burn_dispatch::devices::RuntimeId;
20#[cfg(any(feature = "metal", feature = "vulkan", feature = "webgpu"))]
21use burn_dispatch::devices::WgpuBackend;
22use burn_dispatch::{Dispatch, DispatchDevice};
23use burn_std::{BoolDType, FloatDType, IntDType, TensorData};
24
25#[cfg(feature = "capture")]
26pub use burn_dispatch::backends::capture::{
27 CaptureError, CaptureScope, CapturedGraph, CompletedCaptureScope, TensorId,
28};
29
30#[cfg(any(feature = "cpu", feature = "cuda", feature = "rocm", feature = "wgpu"))]
31use alloc::string::String;
32use alloc::vec;
33use alloc::vec::Vec;
34
35/// A high-level device handle for tensor operations.
36///
37/// [`Device`] provides a unified interface to interact with the underlying compute backend.
38///
39/// Autodiff support is configured on the device rather than through a separate type parameter.
40/// Tensors inherit that context when created and can later change it independently.
41#[cfg_attr(
42 feature = "autodiff",
43 doc = "Wrap a device with [`.autodiff()`](Device::autodiff) to enable automatic differentiation with the device."
44)]
45#[cfg_attr(
46 not(feature = "autodiff"),
47 doc = "Enable the `autodiff` feature to add automatic differentiation support to devices."
48)]
49///
50/// # Backend selection
51///
52/// Default Cargo features do not enable an execution backend. Select one explicitly,
53/// for example with the `wgpu` or `flex` feature; there is no implicit CPU fallback.
54/// Backend-free builds can expose tensor/model APIs, but cannot create an execution device.
55///
56/// [`Device::default()`] selects the first enabled backend in this order:
57/// CUDA, Metal, ROCm, Vulkan, WebGPU, wgpu, CPU, LibTorch, Flex, Remote, NdArray.
58/// Use an explicit factory method when the choice must be independent of Cargo feature unification.
59/// Without an execution backend, `Device::default()` panics with configuration guidance.
60///
61/// Enable the desired backend via Cargo feature flags, then call the
62/// corresponding factory method:
63///
64/// ```rust,ignore
65/// // Default CUDA device (requires the `cuda` feature).
66/// let device = Device::cuda(DeviceIndex::Default);
67///
68/// // CUDA device at hardware index 1.
69/// let device = Device::cuda(1);
70///
71/// // WGPU with explicit selector (requires `wgpu`/`vulkan`/`metal`/`webgpu`).
72/// let device = Device::wgpu(DeviceKind::DiscreteGpu(0));
73///
74/// // Default device for whichever backend is enabled.
75/// let device = Default::default();
76/// ```
77///
78/// Available factory methods (each gated by its matching Cargo feature):
79/// `Device::cpu`, `Device::cuda` / `Device::rocm` / `Device::libtorch_cuda`
80/// (take an integer index or a [`DeviceIndex`]), `Device::wgpu` /
81/// `Device::vulkan` / `Device::metal` / `Device::webgpu` (take a
82/// [`DeviceKind`]), `Device::flex`, `Device::ndarray`, `Device::libtorch`,
83/// `Device::libtorch_mps`, `Device::libtorch_vulkan`, `Device::capture`.
84///
85/// # Autodiff
86///
87/// Requires `autodiff` feature.
88///
89/// Gradient computation is opt-in for a device:
90///
91/// ```rust,ignore
92/// let device = Device::default().autodiff();
93///
94/// // Tensors created on this device can participate in an autodiff graph.
95/// let x = Tensor::<1>::from_floats([1.0, 2.0, 3.0], &device).require_grad();
96/// ```
97pub struct Device {
98 blob: device_opaque::Opaque,
99}
100
101// Aligned, type-erased storage for `DispatchDevice`. See `crate::macros` for
102// why this indirection exists (it keeps the dispatch type tree out of
103// downstream MIR).
104burn_std::obfuscate!(
105 type: DispatchDevice,
106 module: device_opaque,
107 derives: [Send, Sync]
108);
109
110impl Clone for Device {
111 fn clone(&self) -> Self {
112 Self::new(self.as_dispatch().clone())
113 }
114}
115
116impl Default for Device {
117 fn default() -> Self {
118 Self::new(DispatchDevice::default())
119 }
120}
121
122impl core::fmt::Debug for Device {
123 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
124 write!(f, "Device<{:?}>", self.as_dispatch())
125 }
126}
127
128// Manually implement both `eq` and `ne` to add documentation on equality.
129#[allow(clippy::partialeq_ne_impl)]
130impl PartialEq for Device {
131 /// Compares devices based on hardware identity.
132 ///
133 /// Returns `true` if both devices represent the same compute resource.
134 /// Note that this comparison ignores autodiff and checkpointing settings. Inspect
135 /// [`Device::is_autodiff`] and `Device::gradient_checkpointing_strategy()` when execution
136 /// context also matters.
137 fn eq(&self, other: &Self) -> bool {
138 self.as_dispatch() == other.as_dispatch()
139 }
140
141 /// Compares devices based on hardware identity.
142 ///
143 /// Returns `false` if both devices represent the same compute resource,
144 /// even if one has autodiff enabled and the other does not.
145 fn ne(&self, other: &Self) -> bool {
146 !self.eq(other)
147 }
148}
149
150impl Eq for Device {}
151
152impl Device {
153 /// Wrap a backend-specific device in a unified [`Device`].
154 ///
155 /// Used by:
156 /// - the backend-specific factory methods below (`Device::cuda`, etc.)
157 /// — these are the recommended entry points for downstream code;
158 /// - burn-tensor's bridge ops, which already hold a [`DispatchDevice`]
159 /// and just need to wrap it;
160 /// - direct callers (tests, type-erased helpers) that have a concrete
161 /// backend device type at hand.
162 ///
163 /// Anything convertible into [`DispatchDevice`] is accepted, including
164 /// `DispatchDevice` itself.
165 pub fn new(device: impl Into<DispatchDevice>) -> Self {
166 Self {
167 blob: device_opaque::Opaque::new(device.into()),
168 }
169 }
170
171 /// Borrow the underlying [`DispatchDevice`].
172 ///
173 /// The inverse of [`Device::new`]. Useful to backend-extension authors who need to dispatch on
174 /// the concrete backend variant (e.g. matching `DispatchDevice::Remote(_)`).
175 pub fn as_dispatch(&self) -> &DispatchDevice {
176 self.blob.as_ref()
177 }
178
179 /// Crate-internal owning extraction of the underlying dispatch device.
180 pub(crate) fn into_dispatch(self) -> DispatchDevice {
181 self.blob.into_inner()
182 }
183}
184
185impl<D: Into<DispatchDevice>> From<D> for Device {
186 fn from(device: D) -> Self {
187 Self::new(device)
188 }
189}
190
191/// Selector for the hardware index of a backend whose devices are simply
192/// indexed (e.g. CUDA, ROCm).
193///
194/// Backend factory methods that take an index (`Device::cuda`, `Device::rocm`,
195/// `Device::libtorch_cuda`) accept `impl Into<DeviceIndex>`, so the common
196/// shorthand is to pass a plain integer literal:
197///
198/// ```rust,ignore
199/// Device::cuda(0); // hardware index 0
200/// Device::cuda(DeviceIndex::Default); // backend-chosen default
201/// ```
202#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Default)]
203pub enum DeviceIndex {
204 /// Target a specific hardware device by its index.
205 Specified(usize),
206 /// Let the backend pick its default device (typically index `0`).
207 #[default]
208 Default,
209}
210
211impl DeviceIndex {
212 /// Construct a [`DeviceIndex::Specified`] from anything convertible into
213 /// a `usize`.
214 pub fn new(index: impl Into<usize>) -> Self {
215 Self::Specified(index.into())
216 }
217
218 /// Resolve to a concrete hardware index, defaulting to `0` for
219 /// [`DeviceIndex::Default`]. Backend factory methods are each gated by a
220 /// Cargo feature, so this looks dead when none of them are enabled.
221 #[allow(dead_code)]
222 pub(crate) fn resolve(self) -> usize {
223 match self {
224 DeviceIndex::Specified(i) => i,
225 DeviceIndex::Default => 0,
226 }
227 }
228}
229
230impl From<usize> for DeviceIndex {
231 fn from(i: usize) -> Self {
232 Self::Specified(i)
233 }
234}
235
236impl From<u32> for DeviceIndex {
237 fn from(i: u32) -> Self {
238 Self::Specified(i as usize)
239 }
240}
241
242impl From<u64> for DeviceIndex {
243 fn from(i: u64) -> Self {
244 Self::Specified(i as usize)
245 }
246}
247
248impl From<i32> for DeviceIndex {
249 fn from(i: i32) -> Self {
250 Self::Specified(usize::try_from(i).expect("device index must be non-negative"))
251 }
252}
253
254impl From<i64> for DeviceIndex {
255 fn from(i: i64) -> Self {
256 Self::Specified(usize::try_from(i).expect("device index must be non-negative"))
257 }
258}
259
260/// Selector for the more flexible backends whose device handle is a tagged
261/// enum (e.g. WGPU, which can target a discrete/integrated/virtual GPU, a CPU
262/// adapter, an externally-created wgpu setup, or just "best available").
263///
264/// The variants mirror `WgpuDeviceKind` from cubecl so the mapping is direct, but
265/// it is kept as a burn-owned enum so callers don't have to depend on cubecl.
266#[derive(Clone, Debug, Hash, PartialEq, Eq, Default)]
267pub enum DeviceKind {
268 /// Discrete GPU with the given index. The index is the index of the discrete GPU in the list
269 /// of all discrete GPUs found on the system.
270 DiscreteGpu(usize),
271
272 /// Integrated GPU with the given index. The index is the index of the integrated GPU in the
273 /// list of all integrated GPUs found on the system.
274 IntegratedGpu(usize),
275
276 /// Virtual GPU with the given index. The index is the index of the virtual GPU in the list of
277 /// all virtual GPUs found on the system.
278 VirtualGpu(usize),
279
280 /// CPU.
281 Cpu,
282
283 /// The best available device found with the current graphics API.
284 ///
285 /// This will prioritize GPUs wgpu recognizes as "high power". Additionally, you can override this using
286 /// the `CUBECL_WGPU_DEFAULT_DEVICE` environment variable. This variable is spelled as if i was a `WgpuDevice`,
287 /// so for example `CUBECL_WGPU_DEFAULT_DEVICE=IntegratedGpu(1)` or `CUBECL_WGPU_DEFAULT_DEVICE=Cpu`
288 #[default]
289 DefaultDevice,
290
291 /// Use an externally created, existing, wgpu setup. This is helpful when using `CubeCL` in conjunction
292 /// with some existing wgpu setup (eg. egui or bevy), as resources can be transferred in & out of `CubeCL`.
293 ///
294 /// # Notes
295 ///
296 /// This identifies an already registered runtime. To register an application's
297 /// setup, use `Device::wgpu_options().setup(setup).init()` and retain the
298 /// returned device. Clone that device to share it; do not invent an existing ID.
299 Existing(u32),
300}
301
302impl Device {
303 /// Create a reusable graph-capture device.
304 ///
305 /// Operations on tensors moved to this device are recorded rather than executed. Use
306 /// [`Device::capture_scope`] to delimit each capture and declare its graph boundaries.
307 #[cfg(feature = "capture")]
308 pub fn capture() -> Self {
309 Self::new(DispatchDevice::capture())
310 }
311
312 /// Capture the operations performed by `capture` on this device.
313 ///
314 /// The closure receives a [`CaptureScope`] and must return the token produced by
315 /// [`CaptureScope::complete`], containing the ordered runtime input and output tensor IDs.
316 /// Requiring this return value prevents a capture from being finalized without an explicit
317 /// boundary declaration. Completing the scope immediately rejects further tensor operations;
318 /// the device can then be reused for later, independent scopes after the closure returns.
319 ///
320 /// Returns [`CaptureError::InvalidDevice`] if this is not a capture device, and
321 /// [`CaptureError::AlreadyActive`] if another scope is active on the same device.
322 #[cfg(feature = "capture")]
323 pub fn capture_scope(
324 &self,
325 capture: impl FnOnce(CaptureScope) -> CompletedCaptureScope,
326 ) -> Result<CapturedGraph, CaptureError> {
327 match self.as_dispatch() {
328 DispatchDevice::Capture(device) => device.capture_scope(capture),
329 #[allow(unreachable_patterns)] // Capture can be the only enabled backend.
330 _ => Err(CaptureError::InvalidDevice),
331 }
332 }
333
334 /// Default CPU device backed by CubeCL's CPU backend.
335 #[cfg(feature = "cpu")]
336 pub fn cpu() -> Self {
337 Self::new(burn_dispatch::devices::CpuDevice::default())
338 }
339
340 /// CUDA device at the given hardware index.
341 ///
342 /// Accepts a plain integer (e.g. `Device::cuda(0)`) or a
343 /// [`DeviceIndex`] — use [`DeviceIndex::Default`] to let the backend
344 /// pick.
345 #[cfg(feature = "cuda")]
346 pub fn cuda(index: impl Into<DeviceIndex>) -> Self {
347 Self::new(burn_dispatch::devices::CudaDevice::new(
348 index.into().resolve(),
349 ))
350 }
351
352 /// ROCm/HIP device at the given hardware index.
353 ///
354 /// Same selector semantics as [`Device::cuda`].
355 #[cfg(feature = "rocm")]
356 pub fn rocm(index: impl Into<DeviceIndex>) -> Self {
357 Self::new(burn_dispatch::devices::RocmDevice::new(
358 index.into().resolve(),
359 ))
360 }
361
362 /// Flex backend device.
363 #[cfg(feature = "flex")]
364 pub fn flex() -> Self {
365 Self::new(burn_dispatch::devices::FlexDevice)
366 }
367
368 /// Default NdArray (CPU) device.
369 #[cfg(feature = "ndarray")]
370 #[deprecated(
371 since = "0.22.0",
372 note = "burn-ndarray is deprecated and will be removed in a future release. Use `Device::flex()` for pure-Rust CPU execution instead."
373 )]
374 #[allow(deprecated)] // constructing the deprecated device is this constructor's job
375 pub fn ndarray() -> Self {
376 Self::new(burn_dispatch::devices::NdArrayDevice::default())
377 }
378
379 /// LibTorch CPU device.
380 #[cfg(feature = "tch")]
381 #[deprecated(
382 since = "0.22.0",
383 note = "burn-tch is deprecated and will be removed in a future release. Use a CubeCL backend (`Device::cuda`, `Device::rocm`, `Device::metal`, `Device::vulkan`, `Device::cpu`) or `Device::flex()` instead."
384 )]
385 #[allow(deprecated)] // constructing the deprecated device is this constructor's job
386 pub fn libtorch() -> Self {
387 Self::new(burn_dispatch::devices::LibTorchDevice::Cpu)
388 }
389
390 /// LibTorch CUDA device at the given hardware index.
391 #[cfg(feature = "tch")]
392 #[deprecated(
393 since = "0.22.0",
394 note = "burn-tch is deprecated and will be removed in a future release. Use a CubeCL backend (`Device::cuda`, `Device::rocm`, `Device::metal`, `Device::vulkan`, `Device::cpu`) or `Device::flex()` instead."
395 )]
396 #[allow(deprecated)] // constructing the deprecated device is this constructor's job
397 pub fn libtorch_cuda(index: impl Into<DeviceIndex>) -> Self {
398 Self::new(burn_dispatch::devices::LibTorchDevice::Cuda(
399 index.into().resolve(),
400 ))
401 }
402
403 /// LibTorch Metal Performance Shaders (MPS) device.
404 #[cfg(feature = "tch")]
405 #[deprecated(
406 since = "0.22.0",
407 note = "burn-tch is deprecated and will be removed in a future release. Use a CubeCL backend (`Device::cuda`, `Device::rocm`, `Device::metal`, `Device::vulkan`, `Device::cpu`) or `Device::flex()` instead."
408 )]
409 #[allow(deprecated)] // constructing the deprecated device is this constructor's job
410 pub fn libtorch_mps() -> Self {
411 Self::new(burn_dispatch::devices::LibTorchDevice::Mps)
412 }
413
414 /// LibTorch Vulkan device.
415 #[cfg(feature = "tch")]
416 #[deprecated(
417 since = "0.22.0",
418 note = "burn-tch is deprecated and will be removed in a future release. Use a CubeCL backend (`Device::cuda`, `Device::rocm`, `Device::metal`, `Device::vulkan`, `Device::cpu`) or `Device::flex()` instead."
419 )]
420 #[allow(deprecated)] // constructing the deprecated device is this constructor's job
421 pub fn libtorch_vulkan() -> Self {
422 Self::new(burn_dispatch::devices::LibTorchDevice::Vulkan)
423 }
424
425 /// WGPU device, selected via [`DeviceKind`].
426 ///
427 /// [`AutoCompiler`](burn_dispatch::devices::AutoCompiler) dispatches to the most
428 /// appropriate shader language (WGSL, SPIR-V, or MSL) based on the enabled features.
429 ///
430 /// For [`DeviceKind::DefaultDevice`], the adapter is picked by `wgpu`'s
431 /// selection heuristics (high-power GPU preferred, or overridden by
432 /// `CUBECL_WGPU_DEFAULT_DEVICE`).
433 ///
434 /// The graphics API — and so the shader compiler — is the runtime's to settle, from the
435 /// enabled features and what the machine offers. `Device::vulkan`, `Device::metal` and
436 /// `Device::webgpu` pin one instead: the same adapter on two APIs is two devices.
437 ///
438 /// This constructor is lazy. Use [`Device::wgpu_options`] to initialize with runtime
439 /// options or share an existing wgpu setup. In a browser, initialize asynchronously
440 /// with [`init_async`](crate::wgpu::WgpuOptions::init_async) before creating tensors.
441 #[cfg(feature = "wgpu")]
442 pub fn wgpu(device_kind: DeviceKind) -> Self {
443 Self::new(wgpu_device(
444 device_kind,
445 burn_dispatch::devices::WgpuBackend::Auto,
446 ))
447 }
448
449 /// Vulkan-backed WGPU device, selected via [`DeviceKind`] — and so `SPIR-V`, where the
450 /// build and the adapter allow it.
451 ///
452 /// Pinned to Vulkan: the device comes up there or not at all, rather than quietly on
453 /// whichever API [`Device::wgpu`] would settle on.
454 #[cfg(feature = "vulkan")]
455 pub fn vulkan(device_kind: DeviceKind) -> Self {
456 Self::new(wgpu_device(
457 device_kind,
458 burn_dispatch::devices::WgpuBackend::Vulkan,
459 ))
460 }
461
462 /// Metal-backed WGPU device, selected via [`DeviceKind`], requiring native MSL support.
463 ///
464 /// Pinned to Metal, as `Device::vulkan` is to Vulkan. Runtime initialization panics if
465 /// native MSL is unavailable on the selected device. Use [`Device::wgpu`] for automatic
466 /// compiler selection with WGSL fallback.
467 #[cfg(feature = "metal")]
468 pub fn metal(device_kind: DeviceKind) -> Self {
469 Self::new(wgpu_device(
470 device_kind,
471 burn_dispatch::devices::WgpuBackend::Metal,
472 ))
473 }
474
475 /// WebGPU-backed device, selected via [`DeviceKind`] — the browser's own. Pinned to
476 /// WebGPU, as `Device::vulkan` is to Vulkan.
477 #[cfg(feature = "webgpu")]
478 pub fn webgpu(device_kind: DeviceKind) -> Self {
479 Self::new(wgpu_device(
480 device_kind,
481 burn_dispatch::devices::WgpuBackend::WebGpu,
482 ))
483 }
484
485 /// Enables autodiff on this device.
486 ///
487 /// Tensors created on the returned device inherit its autodiff context. A tensor can later
488 /// change that context independently with [`Tensor::autodiff`](crate::Tensor::autodiff) or
489 /// [`Tensor::without_autodiff`](crate::Tensor::without_autodiff).
490 ///
491 /// Calling this method on a device that already has autodiff enabled returns it unchanged,
492 /// preserving its gradient-checkpointing strategy. This operation is idempotent.
493 /// Calling it repeatedly doesn't enable higher-order differentiation; only first-order
494 /// autodiff is supported.
495 ///
496 /// # Example
497 ///
498 /// ```rust,ignore
499 /// let device = Device::default().autodiff();
500 /// let x = Tensor::<1>::from_floats([1.0, 2.0, 3.0], &device).require_grad();
501 /// let gradients = x.backward();
502 /// ```
503 ///
504 #[cfg(feature = "autodiff")]
505 #[must_use]
506 pub fn autodiff(self) -> Self {
507 match self.into_dispatch() {
508 device @ DispatchDevice::Autodiff(_) => Self::new(device),
509 other => Self::new(DispatchDevice::autodiff(other)),
510 }
511 }
512
513 /// Returns this device's gradient-checkpointing strategy when autodiff is enabled.
514 #[cfg(feature = "autodiff")]
515 pub fn gradient_checkpointing_strategy(&self) -> Option<GradientCheckpointingStrategy> {
516 match self.as_dispatch() {
517 DispatchDevice::Autodiff(device) => Some(device.gradient_checkpointing_strategy()),
518 _ => None,
519 }
520 }
521
522 /// Enables gradient checkpointing on the autodiff device.
523 ///
524 /// Gradient checkpointing recomputes activations during backpropagation for operations
525 /// marked as memory-bound, while compute-bound operations still cache their
526 /// output. This reduces peak memory usage at the cost of additional computation
527 /// for memory-bound ops.
528 ///
529 /// # Example
530 ///
531 /// ```rust,ignore
532 /// let device = Device::default().autodiff().gradient_checkpointing();
533 /// ```
534 ///
535 /// # Panics
536 ///
537 /// Panics if autodiff is not enabled on this device.
538 #[cfg(feature = "autodiff")]
539 #[must_use]
540 pub fn gradient_checkpointing(self) -> Self {
541 self.with_gradient_checkpointing_strategy(GradientCheckpointingStrategy::Balanced)
542 }
543
544 /// Sets the gradient-checkpointing strategy on this autodiff device.
545 #[cfg(feature = "autodiff")]
546 #[must_use]
547 pub(crate) fn with_gradient_checkpointing_strategy(
548 self,
549 strategy: GradientCheckpointingStrategy,
550 ) -> Self {
551 match self.into_dispatch() {
552 DispatchDevice::Autodiff(device) => Self::new(
553 DispatchDevice::autodiff_with_gradient_checkpointing(device.inner(), strategy),
554 ),
555 _ => panic!(
556 "Gradient checkpointing requires autodiff; use Device::autodiff().gradient_checkpointing()"
557 ),
558 }
559 }
560
561 /// Returns this device without its autodiff association.
562 ///
563 /// If autodiff is not enabled, the device is returned unchanged. This operation is idempotent.
564 ///
565 /// # Example
566 ///
567 /// ```rust,ignore
568 /// let device = Device::default().autodiff();
569 /// let inference_device = device.without_autodiff();
570 ///
571 /// assert!(!inference_device.is_autodiff());
572 /// ```
573 #[must_use]
574 pub fn without_autodiff(self) -> Self {
575 if self.is_autodiff() {
576 Self::new(self.into_dispatch().inner())
577 } else {
578 self
579 }
580 }
581
582 /// Applies `source`'s autodiff association and gradient-checkpointing strategy to this device,
583 /// discarding this device's own.
584 #[must_use]
585 #[doc(hidden)]
586 pub fn with_autodiff_context_from(self, source: &Self) -> Self {
587 #[cfg(feature = "autodiff")]
588 {
589 match source.gradient_checkpointing_strategy() {
590 Some(strategy) => self
591 .autodiff()
592 .with_gradient_checkpointing_strategy(strategy),
593 None => self.without_autodiff(),
594 }
595 }
596
597 #[cfg(not(feature = "autodiff"))]
598 {
599 let _ = source;
600 self
601 }
602 }
603
604 /// Returns this device without its autodiff association.
605 ///
606 /// This is equivalent to [`without_autodiff`](Self::without_autodiff). `inner` reflects the
607 /// historical backend-decorator model, while `without_autodiff` describes the device's runtime
608 /// property directly.
609 #[must_use]
610 pub fn inner(self) -> Self {
611 self.without_autodiff()
612 }
613
614 /// Synchronize the device, waiting for all pending operations to complete.
615 ///
616 /// # Errors
617 ///
618 /// Returns an [`ExecutionError`] if an operation failed to execute.
619 pub fn sync(&self) -> Result<(), ExecutionError> {
620 Dispatch::sync(self.as_dispatch())
621 }
622
623 /// Flush the device's pending operations, handing them off for execution without waiting
624 /// for them to complete.
625 ///
626 /// Backends that buffer work hold registered operations in a local queue until enough
627 /// accumulate: the fusion backend batches ops to build optimizations, and the remote backend
628 /// batches them before sending them over the network. `flush` forces that queue out now — the
629 /// fusion backend processes its pending optimizations and the remote backend sends its batch to
630 /// the server.
631 ///
632 /// Unlike [`sync`](Self::sync), this does not block on results — it only ensures buffered
633 /// operations are dispatched instead of sitting idle. Eager backends, which execute each
634 /// operation as it is registered, have nothing buffered and treat this as a no-op.
635 ///
636 /// # Errors
637 ///
638 /// Returns an [`ExecutionError`] when the buffered operations cannot be dispatched, e.g. on a
639 /// device that is poisoned.
640 pub fn flush(&self) -> Result<(), ExecutionError> {
641 Dispatch::flush(self.as_dispatch())
642 }
643
644 /// Measure how long this device spends on the work `func` puts on it, in
645 /// device time.
646 ///
647 /// The measurement is a [`ProfileDuration`]: a future the device answers
648 /// once it has run both ends of the window, so nothing here waits on the
649 /// device, and windows nest — an inner `profile` costs the outer one
650 /// nothing. Collect them and [`resolve`](ProfileDuration::resolve) once
651 /// the run is over.
652 ///
653 /// The window spans the stream from the call to `func`'s return. Work the
654 /// stream still owed from before falls in; work a backend queues past the
655 /// end falls out — the fusion backend holds a closure's last operations
656 /// back to batch them, so a window over lazy work alone can read as
657 /// empty. Ending the closure with a read, or
658 /// [`profile_with`](Self::profile_with) and
659 /// [`ProfileOptions::flush`], closes the window over all of it. A window
660 /// that nothing ran in reads as no time.
661 ///
662 /// A backend with no device clock (ndarray, LibTorch, a remote device
663 /// whose server has none) measures wall-clock time between two syncs
664 /// instead: that one waits, and an inner window's syncs are charged to
665 /// the outer.
666 ///
667 /// ```rust,ignore
668 /// let (output, duration) = device.profile(|| model.forward(input))?;
669 /// // Later, once the run is over:
670 /// let ticks = duration.resolve().await.expect("the window carried work");
671 /// println!("forward: {:?}", ticks.duration());
672 /// ```
673 ///
674 /// # Errors
675 ///
676 /// Returns an [`ExecutionError`] when the device refuses to open or close
677 /// the window — a remote device does from a browser thread, which cannot
678 /// wait on the server.
679 pub fn profile<O: Send + 'static>(
680 &self,
681 func: impl FnOnce() -> O + Send,
682 ) -> Result<(O, ProfileDuration), ExecutionError> {
683 self.profile_with(ProfileOptions::default(), func)
684 }
685
686 /// [`profile`](Self::profile) with [`ProfileOptions`].
687 pub fn profile_with<O: Send + 'static>(
688 &self,
689 options: ProfileOptions,
690 func: impl FnOnce() -> O + Send,
691 ) -> Result<(O, ProfileDuration), ExecutionError> {
692 Dispatch::profile(self.as_dispatch(), options, func)
693 }
694
695 /// Seeds the random number generator for this device.
696 ///
697 /// Seeding before tensor operations that involve randomness (e.g. [`Tensor::random`](crate::Tensor::random))
698 /// makes those operations reproducible in a single-threaded program.
699 ///
700 /// # Note
701 ///
702 /// Depending on the backend, the seed may be applied globally rather than scoped
703 /// to this specific device. It is guaranteed that at least this device will be seeded.
704 ///
705 /// # Example
706 ///
707 /// ```rust,ignore
708 /// let device = Default::default();
709 /// device.seed(42);
710 /// let t = Tensor::<1>::random([8], Distribution::Default, &device);
711 /// ```
712 pub fn seed(&self, seed: u64) {
713 Dispatch::seed(self.as_dispatch(), seed)
714 }
715
716 /// Returns whether this device is associated with autodiff.
717 ///
718 /// This is device context inherited by newly created tensors, not a statement that any tensor
719 /// participates in a graph or retains gradients.
720 ///
721 /// # Example
722 ///
723 /// ```rust,ignore
724 /// let device = Default::default();
725 /// assert!(!device.is_autodiff());
726 ///
727 /// let ad_device = device.autodiff();
728 /// assert!(ad_device.is_autodiff());
729 /// ```
730 pub fn is_autodiff(&self) -> bool {
731 Dispatch::ad_enabled(self.as_dispatch())
732 }
733
734 /// Returns `true` if this device supports `dtype` for general computation:
735 /// storage, conversion, *and* arithmetic.
736 ///
737 /// A type can be less than generally supported — bf16 on a Vulkan device,
738 /// for example, is often storable and convertible but has no arithmetic
739 /// (SPIR-V's `SPV_KHR_bfloat16` permits only conversions, dot products,
740 /// and cooperative-matrix use). Computing in such a type produces
741 /// backend-dependent garbage, so check before selecting a reduced
742 /// precision:
743 ///
744 /// ```rust,ignore
745 /// let dtype = if device.supports_dtype(FloatDType::BF16) {
746 /// FloatDType::BF16
747 /// } else {
748 /// FloatDType::F32
749 /// };
750 /// ```
751 pub fn supports_dtype(&self, dtype: impl Into<burn_std::DType>) -> bool {
752 Dispatch::supports_dtype(self.as_dispatch(), dtype.into())
753 }
754
755 /// Sets the current allocation mode to persistent.
756 pub fn memory_persistent_allocations<
757 Output: Send,
758 Input: Send,
759 Func: Fn(Input) -> Output + Send,
760 >(
761 &self,
762 input: Input,
763 func: Func,
764 ) -> Output {
765 Dispatch::memory_persistent_allocations(self.as_dispatch(), input, func)
766 }
767
768 /// Triggers a memory cleanup on this device.
769 ///
770 /// The amount of memory reclaimed depends on the allocator implementation.
771 /// Calling this method does not guarantee that any memory will be freed.
772 pub fn memory_cleanup(&self) {
773 Dispatch::memory_cleanup(self.as_dispatch());
774 }
775
776 /// This device's dynamic pools, in the order they were installed. `None` on
777 /// a backend that does not report them.
778 pub fn memory_pool_report(&self) -> Option<Vec<SlicedPoolReport>> {
779 Dispatch::memory_pool_report(self.as_dispatch())
780 }
781
782 /// What this device's allocator currently holds. `None` on a backend that
783 /// does not report it.
784 pub fn memory_pool_usage(&self) -> Option<MemoryPoolUsage> {
785 Dispatch::memory_pool_usage(self.as_dispatch())
786 }
787
788 /// Prepares the given data for transfer between the CPU and accelerator devices such as GPUs.
789 ///
790 /// Depending on the backend, the data may be transferred to pinned memory
791 /// or another transfer-optimized format to improve transfer performance.
792 pub fn staging<'a, Iter>(&self, data: Iter)
793 where
794 Iter: Iterator<Item = &'a mut TensorData>,
795 {
796 Dispatch::staging(data, self.as_dispatch());
797 }
798
799 /// Returns the [`DeviceSettings`] for this device.
800 ///
801 /// Settings include the default float, integer, and boolean data types used when creating
802 /// tensors on this device.
803 ///
804 /// Before initialization, returns a snapshot of the backend defaults without locking them.
805 /// Another thread may configure the device afterward, so subsequent tensor operations may use
806 /// different settings. [Configure the device](Device::configure) before relying on its defaults.
807 pub fn settings(&self) -> DeviceSettings {
808 burn_backend::get_device_settings::<Dispatch>(self.as_dispatch())
809 }
810
811 /// Returns settings for tensor operations, initializing them to defaults if unset.
812 pub(crate) fn get_or_init_settings(&self) -> DeviceSettings {
813 burn_backend::get_or_init_device_settings::<Dispatch>(self.as_dispatch())
814 }
815
816 /// Configures the [settings](DeviceSettings) for this device.
817 ///
818 /// This configures the dtype used when no explicit type is specified at tensor
819 /// creation time.
820 ///
821 /// Settings can only be initialized once per device. Configure defaults before creating
822 /// tensors or initializing model parameters: tensor creation locks them to the backend's
823 /// defaults, even with an explicit dtype, and any later call returns
824 /// [`DeviceError::AlreadyInitialized`]. Querying [`settings`](Device::settings) does not lock them.
825 ///
826 /// Individual tensors can still use an explicit supported dtype at creation or be converted
827 /// with [`Tensor::cast`](crate::Tensor::cast); neither changes the defaults.
828 ///
829 /// # Errors
830 ///
831 /// Returns [`DeviceError::UnsupportedDType`] if a requested dtype is unsupported.
832 /// Returns [`DeviceError::AlreadyInitialized`] if settings have already been initialized
833 /// for this device, either by a prior call or by a tensor operation.
834 ///
835 /// # Example
836 ///
837 /// ```rust,ignore
838 /// use burn_tensor::{Device, FloatDType, Int, IntDType, Tensor};
839 ///
840 /// let mut device = Device::cuda(0);
841 /// device.configure((FloatDType::F16, IntDType::I32))?;
842 ///
843 /// // Float tensors will now use F16
844 /// let floats = Tensor::<2>::zeros([2, 3], &device);
845 /// // Int tensors will now use I32
846 /// let ints = Tensor::<2, Int>::zeros([2, 3], &device);
847 /// ```
848 pub fn configure(&mut self, config: impl Into<DeviceConfig>) -> Result<(), DeviceError> {
849 let mut config = config.into();
850
851 let defaults = self.as_dispatch().defaults();
852
853 let float_dtype = config.float_dtype.take().unwrap_or(defaults.float_dtype);
854 let int_dtype = config.int_dtype.take().unwrap_or(defaults.int_dtype);
855 let bool_dtype = config.bool_dtype.take().unwrap_or(defaults.bool_dtype);
856
857 burn_backend::set_default_dtypes::<Dispatch>(
858 self.as_dispatch(),
859 float_dtype,
860 int_dtype,
861 bool_dtype,
862 )
863 }
864
865 /// Retrieves all available [`Device`]s that match the given [`DeviceType`] filter.
866 ///
867 /// Backends enumerate the hardware found on this machine, and `DeviceType::Remote` every
868 /// device a remote server hosts:
869 ///
870 /// ```rust,ignore
871 /// // Every CUDA device on this machine.
872 /// let local = Device::enumerate(DeviceType::Cuda);
873 ///
874 /// // Filters combine with `|`.
875 /// let both = Device::enumerate(DeviceType::Cuda | DeviceType::Remote(host));
876 /// ```
877 ///
878 /// # Panics
879 ///
880 /// Where `RemoteHost::devices` returns an error for a `DeviceType::Remote`, and on wasm for any
881 /// `DeviceType::Remote`, which a browser can only list with `RemoteHost::devices_async`.
882 pub fn enumerate(filter: impl Into<DeviceFilter>) -> Devices {
883 #[allow(unused)]
884 let mut devices = Vec::new();
885
886 #[allow(clippy::never_loop)] // at least one backend is expected to be enabled.
887 for device_type in filter.into() {
888 #[allow(unused)]
889 let type_id = match device_type {
890 // The cubecl runtimes are one dispatch backend, so `DispatchDeviceId::Cube`
891 // would list every runtime's devices at once. These name the runtime the
892 // caller asked for instead. `Metal`, `Vulkan` and `WebGpu` pin wgpu to the
893 // requested graphics API; `Wgpu` leaves the API to auto-selection.
894 #[cfg(feature = "cpu")]
895 DeviceType::Cpu => {
896 push_cube(&mut devices, RuntimeId::Cpu);
897 continue;
898 }
899 #[cfg(feature = "cuda")]
900 DeviceType::Cuda => {
901 push_cube(&mut devices, RuntimeId::Cuda);
902 continue;
903 }
904 #[cfg(feature = "rocm")]
905 DeviceType::Rocm => {
906 push_cube(&mut devices, RuntimeId::Hip);
907 continue;
908 }
909 #[cfg(feature = "wgpu")]
910 DeviceType::Wgpu => {
911 push_cube(&mut devices, RuntimeId::Wgpu);
912 continue;
913 }
914 #[cfg(feature = "metal")]
915 DeviceType::Metal => {
916 push_wgpu(&mut devices, WgpuBackend::Metal);
917 continue;
918 }
919 #[cfg(feature = "vulkan")]
920 DeviceType::Vulkan => {
921 push_wgpu(&mut devices, WgpuBackend::Vulkan);
922 continue;
923 }
924 #[cfg(feature = "webgpu")]
925 DeviceType::WebGpu => {
926 push_wgpu(&mut devices, WgpuBackend::WebGpu);
927 continue;
928 }
929 #[cfg(feature = "flex")]
930 DeviceType::Flex => DispatchDeviceId::Flex,
931 #[cfg(feature = "ndarray")]
932 DeviceType::NdArray => DispatchDeviceId::NdArray,
933 #[cfg(feature = "tch")]
934 DeviceType::LibTorch => DispatchDeviceId::LibTorch,
935 #[cfg(feature = "remote")]
936 DeviceType::Remote(host) => {
937 devices.extend(host.enumerate());
938 continue;
939 }
940 };
941
942 #[allow(unreachable_code)] // need to have one backend enabled, so it is reachable
943 for device in Dispatch::enumerate(type_id) {
944 devices.push(Device::new(device))
945 }
946 }
947
948 Devices(devices)
949 }
950
951 /// Who this device is: the part's name, what its kernels are compiled for and, for a card
952 /// on a bus, the card itself. `None` for a backend that does not report one. Opens the
953 /// device.
954 #[cfg(feature = "cubecl")]
955 pub fn identity(&self) -> Option<DeviceIdentity> {
956 self.as_dispatch().identity()
957 }
958
959 /// Every card this build can reach, once each, with the devices that reach it.
960 ///
961 /// A card visible to two runtimes (an NVIDIA GPU under CUDA and under Vulkan) is one entry
962 /// holding both devices, so a caller placing work on cards never puts two stages on one
963 /// card by another name. Opens every device of every runtime, so call it once and keep
964 /// the answer.
965 #[cfg(any(feature = "cpu", feature = "cuda", feature = "rocm", feature = "wgpu"))]
966 pub fn enumerate_physical() -> Vec<PhysicalGpu> {
967 let mut gpus: Vec<PhysicalGpu> = Vec::new();
968 for device in Dispatch::enumerate(DispatchDeviceId::Cube) {
969 let device = Device::new(device);
970 let Some(DeviceIdentity {
971 name,
972 physical: Some(physical),
973 ..
974 }) = device.identity()
975 else {
976 continue;
977 };
978 match gpus
979 .iter_mut()
980 .find(|gpu| gpu.physical.is_same_card(&physical))
981 {
982 Some(gpu) => {
983 gpu.physical.fill_from(&physical);
984 gpu.devices.0.push(device);
985 }
986 None => gpus.push(PhysicalGpu {
987 physical,
988 name,
989 devices: Devices(vec![device]),
990 }),
991 }
992 }
993 gpus
994 }
995
996 /// Measure peak compute and memory throughput for this device.
997 ///
998 /// Runs cubecl-std's throughput benchmarks for each [`ThroughputKey`],
999 /// returning one [`ThroughputStat`] per key (in the same order). A key the
1000 /// device has no peak for keeps its place, carrying the [`ThroughputError`]
1001 /// saying why. Only cubecl-backed devices (cuda, wgpu, ...) report
1002 /// measurements; other backends return an empty vector.
1003 #[cfg(feature = "cubecl")]
1004 pub fn performance_stats(&self, keys: &[ThroughputKey]) -> Vec<ThroughputStat> {
1005 self.as_dispatch()
1006 .performance_stats(keys)
1007 .into_iter()
1008 .zip(keys.iter().copied())
1009 .map(|(value, key)| ThroughputStat { key, value })
1010 .collect()
1011 }
1012}
1013
1014/// One card and every device that reaches it, from [`Device::enumerate_physical`].
1015#[cfg(any(feature = "cpu", feature = "cuda", feature = "rocm", feature = "wgpu"))]
1016#[derive(Debug, Clone)]
1017pub struct PhysicalGpu {
1018 /// The card's PCI address, Windows LUID and vendor, from whichever of its runtimes reported
1019 /// each.
1020 pub physical: PhysicalDevice,
1021 /// The part as the first runtime that reached it names it.
1022 pub name: String,
1023 /// Every device that runs on this card, one per runtime that reaches it.
1024 pub devices: Devices,
1025}
1026
1027/// A single peak-throughput measurement produced by [`Device::performance_stats`].
1028#[cfg(feature = "cubecl")]
1029#[derive(Debug, Clone, Copy, PartialEq)]
1030pub struct ThroughputStat {
1031 /// The measurement key (mode + dtype) that was benchmarked.
1032 pub key: ThroughputKey,
1033 /// The measured throughput for that key, or why the device has none.
1034 pub value: Result<ThroughputValue, ThroughputError>,
1035}
1036
1037/// Short, column-friendly name for a throughput mode.
1038#[cfg(feature = "cubecl")]
1039fn mode_label(mode: &ThroughputMode) -> &'static str {
1040 match mode {
1041 ThroughputMode::ComputeDirect { .. } => "compute-direct",
1042 ThroughputMode::ComputeCmma { .. } => "compute-cmma",
1043 ThroughputMode::Memory(spec) => match spec.access {
1044 MemoryAccess::Copy => "memory",
1045 MemoryAccess::Read => "memory-read",
1046 MemoryAccess::Write => "memory-write",
1047 },
1048 ThroughputMode::Launch => "launch",
1049 }
1050}
1051
1052#[cfg(feature = "cubecl")]
1053impl core::fmt::Display for ThroughputStat {
1054 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1055 // Width/alignment flags are ignored on `ThroughputMode`/`ElemType` directly
1056 // (their fmt impls don't call `f.pad`), so render them to `String`s first —
1057 // `str`'s `Display` honors padding. The mode is labelled by hand rather than
1058 // derived through `Debug`: its variants carry payloads that would blow out the column.
1059 let mode = mode_label(&self.key.mode);
1060
1061 // `ThroughputKey::dtype()` reports f32 for the modes that don't compute with an
1062 // element type, so blank the column there rather than print a misleading type.
1063 let dtype = match self.key.mode {
1064 ThroughputMode::ComputeDirect { dtype } | ThroughputMode::ComputeCmma { dtype, .. } => {
1065 alloc::format!("{dtype}")
1066 }
1067 ThroughputMode::Memory(_) | ThroughputMode::Launch => alloc::string::String::new(),
1068 };
1069
1070 let value = match &self.value {
1071 Ok(value) => value.format(&self.key),
1072 Err(error) => alloc::format!("{error}"),
1073 };
1074
1075 write!(f, "{mode:<14} {dtype:<5} {value}")
1076 }
1077}
1078
1079/// Append every device of the cubecl `runtime` to `devices`, skipping any already listed.
1080///
1081/// A filter can name the same runtime twice (e.g. two `DeviceType::Cuda` entries),
1082/// and a caller enumerating hardware should see each device once.
1083#[cfg(any(feature = "cpu", feature = "cuda", feature = "rocm", feature = "wgpu"))]
1084fn push_cube(devices: &mut Vec<Device>, runtime: RuntimeId) {
1085 for device in Dispatch::enumerate_cube(runtime) {
1086 let device = Device::new(device);
1087 if !devices.contains(&device) {
1088 devices.push(device);
1089 }
1090 }
1091}
1092
1093/// Append every wgpu device on `backend`, skipping any already listed.
1094#[cfg(any(feature = "metal", feature = "vulkan", feature = "webgpu"))]
1095fn push_wgpu(devices: &mut Vec<Device>, backend: WgpuBackend) {
1096 push_wgpu_with(
1097 devices,
1098 backend,
1099 burn_dispatch::devices::CubeDevice::enumerate,
1100 );
1101}
1102
1103#[cfg(any(feature = "metal", feature = "vulkan", feature = "webgpu"))]
1104fn push_wgpu_with(
1105 devices: &mut Vec<Device>,
1106 backend: WgpuBackend,
1107 mut enumerate: impl FnMut(DeviceId) -> Vec<DeviceId>,
1108) {
1109 use burn_dispatch::devices::{CubeDevice, WgpuDevice, WgpuDeviceKind};
1110
1111 // Enumeration preserves the seed's kind and graphics API. Ask for each concrete
1112 // kind on the requested API: Auto may expose different kinds or fewer adapters.
1113 // DefaultDevice is an alias, and Existing devices cannot be discovered.
1114 for kind in [
1115 WgpuDeviceKind::DiscreteGpu(0),
1116 WgpuDeviceKind::IntegratedGpu(0),
1117 WgpuDeviceKind::VirtualGpu(0),
1118 WgpuDeviceKind::Cpu,
1119 WgpuDeviceKind::Other(0),
1120 ] {
1121 let seed = CubeDevice::Wgpu(WgpuDevice::new(kind).on(backend));
1122 for id in enumerate(seed.to_id()) {
1123 let device = Device::new(CubeDevice::from_id(id));
1124 if !devices.contains(&device) {
1125 devices.push(device);
1126 }
1127 }
1128 }
1129}
1130
1131/// Map our backend-agnostic [`DeviceKind`] onto cubecl's `WgpuDevice`, on `backend`.
1132///
1133/// Shared by [`Device::wgpu`], which leaves the graphics API to the runtime, and
1134/// [`Device::vulkan`], [`Device::metal`] and [`Device::webgpu`], which each pin theirs.
1135#[cfg(feature = "wgpu")]
1136pub(crate) fn wgpu_device(
1137 device_kind: DeviceKind,
1138 backend: burn_dispatch::devices::WgpuBackend,
1139) -> burn_dispatch::devices::WgpuDevice {
1140 use burn_dispatch::devices::{WgpuDevice, WgpuDeviceKind};
1141
1142 let kind = match device_kind {
1143 DeviceKind::DiscreteGpu(i) => WgpuDeviceKind::DiscreteGpu(i),
1144 DeviceKind::IntegratedGpu(i) => WgpuDeviceKind::IntegratedGpu(i),
1145 DeviceKind::VirtualGpu(i) => WgpuDeviceKind::VirtualGpu(i),
1146 DeviceKind::Cpu => WgpuDeviceKind::Cpu,
1147 DeviceKind::DefaultDevice => WgpuDeviceKind::DefaultDevice,
1148 DeviceKind::Existing(id) => WgpuDeviceKind::Existing(id),
1149 };
1150
1151 WgpuDevice::new(kind).on(backend)
1152}
1153
1154/// Represents the devices that can be used.
1155///
1156/// `DeviceType` is used to filter the available device types for [`Device::enumerate`]. Each
1157/// variant selects a backend's hardware on this machine, except `Remote`, which selects a remote
1158/// server's devices.
1159///
1160/// Variants combine into a [`DeviceFilter`] with the `|` operator, so a single
1161/// [`Device::enumerate`] call can span several backends and remote servers.
1162#[allow(missing_docs)]
1163#[derive(Debug, Clone, PartialEq, Eq)]
1164pub enum DeviceType {
1165 #[cfg(feature = "cpu")]
1166 Cpu,
1167 #[cfg(feature = "cuda")]
1168 Cuda,
1169 #[cfg(feature = "rocm")]
1170 Rocm,
1171 #[cfg(feature = "wgpu")]
1172 Wgpu,
1173 #[cfg(feature = "metal")]
1174 Metal,
1175 #[cfg(feature = "vulkan")]
1176 Vulkan,
1177 #[cfg(feature = "webgpu")]
1178 WebGpu,
1179 #[cfg(feature = "flex")]
1180 Flex,
1181 #[cfg(feature = "ndarray")]
1182 NdArray,
1183 #[cfg(feature = "tch")]
1184 LibTorch,
1185 /// Every device the remote server `host` hosts, connected on first use like any listed device.
1186 /// [`RemoteHost::devices`](crate::remote::RemoteHost::devices) does the same and returns an
1187 /// error where `enumerate` panics.
1188 #[cfg(feature = "remote")]
1189 Remote(crate::remote::RemoteHost),
1190}
1191
1192/// A set of [`DeviceType`]s passed to [`Device::enumerate`].
1193///
1194/// Built from a single [`DeviceType`], a `Vec<DeviceType>`, or by combining variants with the
1195/// `|` operator (`DeviceType::Cuda | DeviceType::Cpu`).
1196#[derive(Debug, Clone, Default)]
1197pub struct DeviceFilter(Vec<DeviceType>);
1198
1199impl DeviceFilter {
1200 /// Create an empty filter.
1201 pub fn new() -> Self {
1202 Self::default()
1203 }
1204
1205 /// Add a [`DeviceType`] to the filter.
1206 pub fn with(mut self, device_type: DeviceType) -> Self {
1207 self.0.push(device_type);
1208 self
1209 }
1210}
1211
1212impl From<DeviceType> for DeviceFilter {
1213 fn from(value: DeviceType) -> Self {
1214 DeviceFilter(vec![value])
1215 }
1216}
1217
1218impl From<Vec<DeviceType>> for DeviceFilter {
1219 fn from(value: Vec<DeviceType>) -> Self {
1220 DeviceFilter(value)
1221 }
1222}
1223
1224impl IntoIterator for DeviceFilter {
1225 type Item = DeviceType;
1226 type IntoIter = alloc::vec::IntoIter<DeviceType>;
1227 fn into_iter(self) -> Self::IntoIter {
1228 self.0.into_iter()
1229 }
1230}
1231
1232impl core::ops::BitOr for DeviceType {
1233 type Output = DeviceFilter;
1234 fn bitor(self, rhs: Self) -> DeviceFilter {
1235 DeviceFilter(vec![self, rhs])
1236 }
1237}
1238
1239impl core::ops::BitOr<DeviceType> for DeviceFilter {
1240 type Output = DeviceFilter;
1241 fn bitor(mut self, rhs: DeviceType) -> DeviceFilter {
1242 self.0.push(rhs);
1243 self
1244 }
1245}
1246
1247/// Configuration options used to initialize a device.
1248///
1249/// Unlike [`DeviceSettings`], this type represents partial user-provided
1250/// configuration and does not require all settings to be specified.
1251///
1252/// Any unspecified options will be resolved to device-specific defaults
1253/// when the device is initialized.
1254///
1255/// Use [`Device::configure`] to apply this configuration to a device.
1256#[derive(new, Debug, Clone, Default)]
1257pub struct DeviceConfig {
1258 /// Default floating-point data type.
1259 pub float_dtype: Option<FloatDType>,
1260
1261 /// Default integer data type.
1262 pub int_dtype: Option<IntDType>,
1263
1264 /// Default boolean data type.
1265 pub bool_dtype: Option<BoolDType>,
1266 // TODO: maybe quantization, but for now we keep this as device defaults
1267}
1268
1269impl DeviceConfig {
1270 /// Sets the default floating-point data type for tensors created on the device.
1271 pub fn float_dtype(mut self, dtype: impl Into<FloatDType>) -> Self {
1272 self.float_dtype = Some(dtype.into());
1273 self
1274 }
1275
1276 /// Sets the default integer data type for tensors created on the device.
1277 pub fn int_dtype(mut self, dtype: impl Into<IntDType>) -> Self {
1278 self.int_dtype = Some(dtype.into());
1279 self
1280 }
1281
1282 /// Sets the default boolean data type storage precision for tensors created on the device.
1283 pub fn bool_dtype(mut self, dtype: impl Into<BoolDType>) -> Self {
1284 self.bool_dtype = Some(dtype.into());
1285 self
1286 }
1287}
1288
1289impl From<FloatDType> for DeviceConfig {
1290 fn from(value: FloatDType) -> Self {
1291 DeviceConfig::new(Some(value), None, None)
1292 }
1293}
1294
1295impl From<IntDType> for DeviceConfig {
1296 fn from(value: IntDType) -> Self {
1297 DeviceConfig::new(None, Some(value), None)
1298 }
1299}
1300
1301impl From<BoolDType> for DeviceConfig {
1302 fn from(value: BoolDType) -> Self {
1303 DeviceConfig::new(None, None, Some(value))
1304 }
1305}
1306
1307impl From<(FloatDType, IntDType)> for DeviceConfig {
1308 fn from(value: (FloatDType, IntDType)) -> Self {
1309 DeviceConfig::new(Some(value.0), Some(value.1), None)
1310 }
1311}
1312
1313/// A collection of [`Device`]s returned by [`Device::enumerate`].
1314///
1315/// This type provides bulk operations and transformations over multiple
1316/// devices, such as enabling autodiff or configuring the device settings.
1317///
1318/// # Example
1319///
1320/// ```rust,ignore
1321/// let mut devices = Device::enumerate(DeviceType::Cuda)
1322/// .autodiff();
1323///
1324/// devices.configure(
1325/// DeviceConfig::default().float_dtype(FloatDType::F16),
1326/// )?;
1327/// ```
1328///
1329/// `Devices` dereferences to a slice of [`Device`], so it can be iterated,
1330/// indexed, and passed anywhere a `&[Device]` is expected.
1331#[derive(Debug, Clone)]
1332pub struct Devices(Vec<Device>);
1333
1334impl Devices {
1335 /// Enables autodiff across all contained devices.
1336 ///
1337 /// Devices that already have autodiff enabled are left unchanged, preserving their
1338 /// gradient-checkpointing strategy.
1339 ///
1340 /// See [`Device::autodiff`].
1341 #[cfg(feature = "autodiff")]
1342 #[must_use]
1343 pub fn autodiff(mut self) -> Self {
1344 for device in &mut self.0 {
1345 *device = core::mem::take(device).autodiff();
1346 }
1347
1348 self
1349 }
1350
1351 /// Removes the autodiff association from every contained device.
1352 ///
1353 /// See [`Device::without_autodiff`].
1354 #[must_use]
1355 pub fn without_autodiff(mut self) -> Self {
1356 for device in &mut self.0 {
1357 *device = core::mem::take(device).without_autodiff();
1358 }
1359
1360 self
1361 }
1362
1363 /// Configures the [settings](DeviceSettings) for all devices.
1364 ///
1365 /// This configures the dtype used when no explicit type is specified at tensor
1366 /// creation time.
1367 ///
1368 /// Settings can only be initialized once per device. Configure defaults before creating
1369 /// tensors or initializing model parameters; tensor creation locks them, even with an explicit
1370 /// dtype. Querying [`Device::settings`] does not lock them.
1371 ///
1372 /// Stops at the first error; devices configured before it keep their settings.
1373 ///
1374 /// See [`Device::configure`].
1375 pub fn configure(&mut self, config: impl Into<DeviceConfig>) -> Result<(), DeviceError> {
1376 let config = config.into();
1377 for device in &mut self.0 {
1378 device.configure(config.clone())?;
1379 }
1380 Ok(())
1381 }
1382
1383 /// Returns the `Vec` of [`Device`]s.
1384 pub fn into_vec(self) -> Vec<Device> {
1385 self.0
1386 }
1387}
1388
1389// Loop over `&Devices` or `Devices` seamlessly
1390impl FromIterator<Device> for Devices {
1391 fn from_iter<I: IntoIterator<Item = Device>>(devices: I) -> Self {
1392 Self(devices.into_iter().collect())
1393 }
1394}
1395
1396impl IntoIterator for Devices {
1397 type Item = Device;
1398 type IntoIter = alloc::vec::IntoIter<Device>;
1399 fn into_iter(self) -> Self::IntoIter {
1400 self.0.into_iter()
1401 }
1402}
1403
1404impl core::ops::Deref for Devices {
1405 type Target = [Device];
1406 fn deref(&self) -> &Self::Target {
1407 &self.0
1408 }
1409}
1410
1411#[cfg(all(test, feature = "capture"))]
1412mod capture_tests {
1413 use super::*;
1414
1415 #[cfg(all(feature = "flex", feature = "autodiff"))]
1416 #[test]
1417 #[should_panic(expected = "Cannot move a tensor to a capture device with autodiff enabled")]
1418 fn capture_transfer_rejects_autodiff_context() {
1419 let tensor = crate::Tensor::<1, crate::Int>::from_ints([1, 2], &Device::flex().autodiff());
1420 let capture = Device::capture();
1421 let _ = capture.capture_scope(|scope| {
1422 let _ = tensor.to_device(&capture);
1423 scope.complete([], [])
1424 });
1425 }
1426
1427 #[cfg(all(feature = "flex", feature = "autodiff"))]
1428 #[test]
1429 #[should_panic(expected = "Cannot move a tensor to a capture device with autodiff enabled")]
1430 fn capture_float_transfer_rejects_autodiff_context() {
1431 let tensor =
1432 crate::Tensor::<1>::from_floats([1.0, 2.0], &Device::flex().autodiff()).require_grad();
1433 let capture = Device::capture();
1434 let _ = capture.capture_scope(|scope| {
1435 let _ = tensor.to_device(&capture);
1436 scope.complete([], [])
1437 });
1438 }
1439
1440 #[test]
1441 fn user_facing_capture_device_supports_repeated_scopes() {
1442 let device = Device::capture();
1443
1444 let first = device
1445 .capture_scope(|scope| scope.complete([], []))
1446 .unwrap();
1447 let second = device
1448 .capture_scope(|scope| scope.complete([], []))
1449 .unwrap();
1450
1451 assert!(first.graph.operations.is_empty());
1452 assert!(second.graph.operations.is_empty());
1453 }
1454
1455 #[test]
1456 fn capture_scope_rejects_a_non_capture_device() {
1457 let device = Device::default();
1458
1459 let result = device.capture_scope(|scope| scope.complete([], []));
1460
1461 assert!(matches!(result, Err(CaptureError::InvalidDevice)));
1462 }
1463
1464 #[test]
1465 fn capture_device_reports_recordable_dtype_support() {
1466 let device = Device::capture();
1467
1468 let captured = device.capture_scope(|scope| {
1469 assert!(device.supports_dtype(FloatDType::F32));
1470 assert!(device.supports_dtype(FloatDType::F64));
1471 assert!(device.supports_dtype(FloatDType::BF16));
1472 assert!(device.supports_dtype(IntDType::I32));
1473 assert!(device.supports_dtype(BoolDType::Native));
1474 scope.complete([], [])
1475 });
1476
1477 assert!(captured.is_ok());
1478 }
1479}
1480
1481#[cfg(all(test, feature = "flex", feature = "autodiff"))]
1482mod autodiff_move_tests {
1483 use crate::{Bool, Device, GradientCheckpointingStrategy, Int, Tensor, TensorData};
1484
1485 // A tensor without autodiff (e.g. a gradient) can be moved onto an autodiff device; it lands
1486 // on the underlying hardware and stays outside autodiff.
1487 #[test]
1488 fn move_non_autodiff_float_tensor_to_autodiff_device() {
1489 let device = Device::default();
1490 let ad_device = device.clone().autodiff();
1491
1492 let t = Tensor::<2>::from_floats([[1.0, 2.0], [3.0, 4.0]], &device);
1493 let moved = t.to_device(&ad_device);
1494
1495 assert!(!moved.is_autodiff());
1496 assert!(!moved.is_tracked());
1497 assert_eq!(
1498 moved.try_into_vec_as::<f32>().unwrap(),
1499 vec![1.0, 2.0, 3.0, 4.0]
1500 );
1501 }
1502
1503 #[test]
1504 fn move_autodiff_tensor_preserves_association_strategy_and_graph() {
1505 let device = Device::default();
1506 let ad_device = device.clone().autodiff().gradient_checkpointing();
1507
1508 let tensor = Tensor::<1>::from_floats([1.0, 2.0], &ad_device).require_grad();
1509 let moved = tensor.clone().to_device(&device);
1510
1511 assert!(moved.is_autodiff());
1512 assert!(moved.is_tracked());
1513 // The transfer is a differentiable operation, so its result is a derived tensor whose own
1514 // gradient isn't retained by default.
1515 assert!(!moved.is_require_grad());
1516 assert_eq!(
1517 moved.gradient_checkpointing_strategy(),
1518 Some(GradientCheckpointingStrategy::Balanced)
1519 );
1520
1521 let grads = moved.sum().backward();
1522 tensor
1523 .grad(&grads)
1524 .expect("the transfer should preserve the graph")
1525 .into_data()
1526 .assert_eq(&TensorData::from([1.0f32, 1.0]), true);
1527 }
1528
1529 #[test]
1530 fn move_non_float_tensors_preserves_context() {
1531 let device = Device::default();
1532 let ad_device = device.clone().autodiff().gradient_checkpointing();
1533
1534 let plain_int = Tensor::<1, Int>::from_ints([1, 2], &device).to_device(&ad_device);
1535 assert!(!plain_int.is_autodiff());
1536
1537 let autodiff_bool =
1538 Tensor::<1, Bool>::from_bool([true, false], &ad_device).to_device(&device);
1539 assert!(autodiff_bool.is_autodiff());
1540 assert_eq!(
1541 autodiff_bool.gradient_checkpointing_strategy(),
1542 Some(GradientCheckpointingStrategy::Balanced)
1543 );
1544 }
1545
1546 #[test]
1547 #[should_panic(
1548 expected = "Gradient checkpointing requires autodiff; use Device::autodiff().gradient_checkpointing()"
1549 )]
1550 fn gradient_checkpointing_requires_autodiff() {
1551 let _ = Device::default().gradient_checkpointing();
1552 }
1553}
1554
1555#[cfg(all(test, feature = "cuda", feature = "wgpu"))]
1556mod tests {
1557 use super::*;
1558
1559 /// `cargo test -p burn-tensor --features cuda,wgpu a_card_reached_by_two_runtimes -- --ignored`
1560 #[test]
1561 #[ignore = "needs an NVIDIA card reachable through CUDA and Vulkan"]
1562 fn a_card_reached_by_two_runtimes_is_listed_once() {
1563 let gpus = Device::enumerate_physical();
1564 let nvidia: Vec<_> = gpus
1565 .iter()
1566 .filter(|gpu| gpu.physical.vendor == Some(PciVendor::Nvidia))
1567 .collect();
1568 assert!(!nvidia.is_empty(), "no NVIDIA card in {gpus:?}");
1569 for gpu in &nvidia {
1570 assert!(
1571 gpu.physical.pci_address.is_some(),
1572 "{} has no PCI address",
1573 gpu.name
1574 );
1575 assert!(
1576 gpu.devices.len() >= 2,
1577 "{} is reached by {:?} only",
1578 gpu.name,
1579 gpu.devices
1580 );
1581 }
1582 for (i, gpu) in gpus.iter().enumerate() {
1583 for other in &gpus[i + 1..] {
1584 assert!(
1585 !gpu.physical.is_same_card(&other.physical),
1586 "one card listed twice: {gpu:?} and {other:?}"
1587 );
1588 }
1589 }
1590 }
1591}
1592
1593#[cfg(all(test, any(feature = "metal", feature = "vulkan", feature = "webgpu")))]
1594mod enumerate_wgpu_tests {
1595 use super::*;
1596 use burn_dispatch::devices::{CubeDevice, WgpuDevice, WgpuDeviceKind};
1597
1598 fn cube_device(kind: WgpuDeviceKind, backend: WgpuBackend) -> CubeDevice {
1599 CubeDevice::Wgpu(WgpuDevice::new(kind).on(backend))
1600 }
1601
1602 #[test]
1603 fn enumerate_requested_api_independently_of_auto() {
1604 for backend in [WgpuBackend::Metal, WgpuBackend::Vulkan, WgpuBackend::WebGpu] {
1605 // An existing Auto device must not limit discovery on the requested API
1606 // or be deduplicated with a device pinned to that API.
1607 let auto = cube_device(WgpuDeviceKind::DiscreteGpu(0), WgpuBackend::Auto);
1608 let expected = [
1609 WgpuDeviceKind::DiscreteGpu(0),
1610 WgpuDeviceKind::DiscreteGpu(1),
1611 WgpuDeviceKind::IntegratedGpu(0),
1612 WgpuDeviceKind::VirtualGpu(0),
1613 WgpuDeviceKind::Cpu,
1614 WgpuDeviceKind::Other(0),
1615 ]
1616 .map(|kind| cube_device(kind, backend));
1617 let enumerate = |seed: DeviceId| {
1618 let CubeDevice::Wgpu(seed_device) = CubeDevice::from_id(seed) else {
1619 panic!("expected a wgpu seed");
1620 };
1621 assert_eq!(seed_device.backend, backend);
1622 expected
1623 .iter()
1624 .filter(|device| device.to_id().type_id == seed.type_id)
1625 .map(CubeDevice::to_id)
1626 .collect()
1627 };
1628
1629 let mut devices = vec![Device::new(auto.clone())];
1630 push_wgpu_with(&mut devices, backend, enumerate);
1631 let expected_devices: Vec<_> = core::iter::once(&auto)
1632 .chain(expected.iter())
1633 .cloned()
1634 .map(Device::new)
1635 .collect();
1636 assert_eq!(devices, expected_devices);
1637
1638 push_wgpu_with(&mut devices, backend, enumerate);
1639 assert_eq!(
1640 devices, expected_devices,
1641 "repeated filters must not duplicate devices"
1642 );
1643 }
1644 }
1645
1646 #[test]
1647 fn enumerate_unavailable_api_is_empty() {
1648 let mut devices = Vec::new();
1649 push_wgpu_with(&mut devices, WgpuBackend::Vulkan, |_| Vec::new());
1650 assert!(devices.is_empty());
1651 }
1652
1653 #[cfg(not(target_family = "wasm"))]
1654 #[test]
1655 fn enumerate_pins_the_graphics_api() {
1656 let cases = [
1657 #[cfg(feature = "metal")]
1658 (DeviceType::Metal, WgpuBackend::Metal),
1659 #[cfg(feature = "vulkan")]
1660 (DeviceType::Vulkan, WgpuBackend::Vulkan),
1661 #[cfg(feature = "webgpu")]
1662 (DeviceType::WebGpu, WgpuBackend::WebGpu),
1663 ];
1664 for (device_type, backend) in cases {
1665 for device in Device::enumerate(device_type).iter() {
1666 match device.as_dispatch() {
1667 DispatchDevice::Cube(CubeDevice::Wgpu(wgpu)) => {
1668 assert_eq!(wgpu.backend, backend);
1669 }
1670 other => panic!("expected a wgpu device, got {other:?}"),
1671 }
1672 }
1673 }
1674 }
1675}