burn_dispatch/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2#![warn(missing_docs)]
3#![cfg_attr(docsrs, feature(doc_cfg))]
4// Without a backend, generated dispatch bodies diverge and their arguments/imports
5// are unused. Keep the API available for libraries that let consumers select a backend.
6#![cfg_attr(
7    not(backend_enabled),
8    allow(
9        unused_imports,
10        unused_variables,
11        unused_mut,
12        unused_macros,
13        unused_assignments,
14        dead_code,
15        irrefutable_let_patterns,
16        unreachable_code
17    )
18)]
19
20//! Burn multi-backend dispatch.
21//!
22//! # Available Backends
23//!
24//! The dispatch backend supports the following variants, each enabled via cargo features:
25//!
26//! | Backend    | Feature    | Description |
27//! |------------|------------|-------------|
28//! | `Cube`     | `cpu`, `cuda`, `metal`, `rocm`, `vulkan`, `webgpu`, `wgpu` | Every cubecl runtime. One backend: the features decide which runtimes are compiled in, and a tensor's device says which one it runs on |
29//! | `Flex`     | `flex`     | Pure Rust CPU backend using `burn-flex` |
30//! | `Remote`   | `remote`   | Devices hosted by another process or machine, through `burn-remote` |
31//! | `Capture`  | `capture`  | Records operation graphs instead of executing them, through `burn-capture` |
32//! | `Group`    | `group`    | Experimental: tensors split over a group of Cube or Flex devices, through `burn-group` |
33//! | `Autodiff` | `autodiff` | Autodiff-enabled backend (used in combination with any of the backends above) |
34//!
35//! **Note:** The features can be combined freely. The cubecl-backed ones all
36//! select the same backend, so they share the one `DispatchDevice::Cube`
37//! variant — enabling several compiles several runtimes in, and the device a
38//! tensor carries is what picks between them.
39
40#[macro_use]
41mod macros;
42
43/// Dispatch backend module.
44pub mod backend;
45/// Dispatch device module.
46pub mod device;
47#[cfg(feature = "group")]
48mod group;
49mod ops;
50/// Dispatch tensor module.
51pub mod tensor;
52
53/// The backend dispatch behind `burn::server::RemoteServer`. Not a user path.
54#[doc(hidden)]
55#[cfg(feature = "remote-server")]
56pub mod remote_server;
57
58pub use backend::*;
59pub use device::*;
60pub use tensor::*;
61
62extern crate alloc;
63
64// Keep backend-free dispatch types opaque to downstream crates. An actually empty
65// enum makes every API accepting/returning a tensor appear unreachable there.
66// The private field prevents construction without leaking uninhabitedness.
67#[cfg(not(backend_enabled))]
68#[doc(hidden)]
69#[derive(Clone, Debug, PartialEq, Eq)]
70pub struct NoBackend {
71    never: core::convert::Infallible,
72}
73
74#[cfg(not(backend_enabled))]
75impl NoBackend {
76    pub(crate) fn unreachable(&self) -> ! {
77        match self.never {}
78    }
79}
80
81/// Backends and devices used.
82pub mod backends {
83    #[cfg(feature = "autodiff")]
84    pub use burn_autodiff as autodiff;
85    #[cfg(feature = "autodiff")]
86    pub use burn_autodiff::Autodiff; // re-export for extensions
87
88    /// The cubecl backend: CUDA, ROCm, Metal, Vulkan, WebGPU, wgpu and the CPU
89    /// runtime are all this one type, and a tensor's device says which of them
90    /// it runs on. The features still decide which runtimes are compiled in.
91    #[cfg(cube_backend)]
92    pub use burn_cubecl::Cube;
93
94    #[cfg(feature = "flex")]
95    pub use burn_flex as flex;
96    #[cfg(feature = "flex")]
97    pub use burn_flex::Flex;
98
99    #[cfg(feature = "remote")]
100    pub use burn_remote::RemoteBackend as Remote;
101
102    /// Public graph-capture API types.
103    #[cfg(feature = "capture")]
104    pub mod capture {
105        pub use burn_capture::{
106            CaptureBackend, CaptureError, CaptureScope, CaptureTensor, CapturedGraph,
107            CompletedCaptureScope, IntoCaptureOutputs, TensorId,
108        };
109    }
110    #[cfg(feature = "capture")]
111    pub use burn_capture::CaptureBackend as Capture;
112
113    #[cfg(feature = "group")]
114    pub use burn_group::GroupBackend as Group;
115    #[cfg(feature = "group")]
116    pub use burn_group::{GroupError, Sharding, ShardingError};
117}
118
119// Re-export devices
120
121/// Backend devices.
122pub mod devices {
123    /// Wgpu initialization and interoperability types used by the device facade.
124    #[cfg(feature = "wgpu")]
125    pub mod wgpu {
126        #[cfg(not(target_family = "wasm"))]
127        pub use burn_cubecl::cubecl::wgpu::try_init_setup;
128        pub use burn_cubecl::cubecl::wgpu::{
129            AutoGraphicsApi, MemoryConfiguration, RuntimeOptions, WgpuBackend, WgpuDevice,
130            WgpuInitError, WgpuSetup, try_init_device, try_init_setup_async, wgpu,
131        };
132    }
133
134    #[cfg(feature = "cpu")]
135    pub use burn_cubecl::cubecl::cpu::CpuDevice;
136    #[cfg(feature = "cuda")]
137    pub use burn_cubecl::cubecl::cuda::CudaDevice;
138    #[cfg(feature = "rocm")]
139    pub use burn_cubecl::cubecl::hip::AmdDevice as RocmDevice;
140    #[cfg(feature = "wgpu")]
141    pub use burn_cubecl::cubecl::wgpu::{
142        AutoCompiler, AutoGraphicsApi, WgpuBackend, WgpuDevice, WgpuDeviceKind, init_setup_async,
143    };
144
145    /// The device every cubecl runtime shares; which runtime it names is a
146    /// property of the value, not of its type, and [`RuntimeId`] is how that
147    /// property is named.
148    #[cfg(cube_backend)]
149    pub use burn_cubecl::CubeDevice;
150    #[cfg(cube_backend)]
151    pub use burn_cubecl::cubecl::RuntimeId;
152    #[cfg(feature = "flex")]
153    pub use burn_flex::FlexDevice;
154
155    #[cfg(feature = "remote")]
156    pub use burn_remote::RemoteDevice;
157
158    #[cfg(feature = "group")]
159    pub use burn_group::GroupDevice;
160}
161
162/// The remote backend's crate, for `burn::remote` and `burn::server` to build on. Not a user path.
163#[doc(hidden)]
164#[cfg(feature = "remote")]
165pub use burn_remote as __remote;
burn_dispatch/lib.rs

burn_dispatch/
lib.rs
