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// Wiring up the deprecated `NdArray` and `LibTorch` backends is this crate's job, and the backend
20// registry macros expand them into every dispatch impl, so the warnings land on `macros.rs` rather
21// than on any site we could annotate individually. `allow(deprecated)` is a lint level scoped to
22// this crate, and lint levels never propagate to dependents, so downstream code naming `NdArray` or
23// `LibTorch` (directly or via our re-export) still gets the warning. The `cfg_attr` keeps this
24// confined to the builds that enable them: neither `ndarray` nor `tch` is a default feature, so the
25// default build that CI lints with `--deny warnings` retains full deprecation signal for every
26// other dependency.
27#![cfg_attr(any(feature = "ndarray", feature = "tch"), allow(deprecated))]
28
29//! Burn multi-backend dispatch.
30//!
31//! # Available Backends
32//!
33//! The dispatch backend supports the following variants, each enabled via cargo features:
34//!
35//! | Backend    | Feature    | Description |
36//! |------------|------------|-------------|
37//! | `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 |
38//! | `Flex`     | `flex`     | Pure Rust CPU backend using `burn-flex` |
39//! | `NdArray`  | `ndarray`  | Pure Rust CPU backend using `ndarray` (deprecated - use `flex`) |
40//! | `LibTorch` | `tch`      | Libtorch backend via `tch` (deprecated - use a CubeCL backend) |
41//! | `Remote`   | `remote`   | Devices hosted by another process or machine, through `burn-remote` |
42//! | `Capture`  | `capture`  | Records operation graphs instead of executing them, through `burn-capture` |
43//! | `Autodiff` | `autodiff` | Autodiff-enabled backend (used in combination with any of the backends above) |
44//!
45//! **Note:** The features can be combined freely. The cubecl-backed ones all
46//! select the same backend, so they share the one `DispatchDevice::Cube`
47//! variant — enabling several compiles several runtimes in, and the device a
48//! tensor carries is what picks between them.
49
50#[macro_use]
51mod macros;
52
53/// Dispatch backend module.
54pub mod backend;
55/// Dispatch device module.
56pub mod device;
57mod ops;
58/// Dispatch tensor module.
59pub mod tensor;
60
61/// The backend dispatch behind `burn::server::RemoteServer`. Not a user path.
62#[doc(hidden)]
63#[cfg(feature = "remote-server")]
64pub mod remote_server;
65
66pub use backend::*;
67pub use device::*;
68pub use tensor::*;
69
70extern crate alloc;
71
72// Keep backend-free dispatch types opaque to downstream crates. An actually empty
73// enum makes every API accepting/returning a tensor appear unreachable there.
74// The private field prevents construction without leaking uninhabitedness.
75#[cfg(not(backend_enabled))]
76#[doc(hidden)]
77#[derive(Clone, Debug, PartialEq, Eq)]
78pub struct NoBackend {
79    never: core::convert::Infallible,
80}
81
82#[cfg(not(backend_enabled))]
83impl NoBackend {
84    pub(crate) fn unreachable(&self) -> ! {
85        match self.never {}
86    }
87}
88
89/// Backends and devices used.
90pub mod backends {
91    #[cfg(feature = "autodiff")]
92    pub use burn_autodiff as autodiff;
93    #[cfg(feature = "autodiff")]
94    pub use burn_autodiff::Autodiff; // re-export for extensions
95
96    /// The cubecl backend: CUDA, ROCm, Metal, Vulkan, WebGPU, wgpu and the CPU
97    /// runtime are all this one type, and a tensor's device says which of them
98    /// it runs on. The features still decide which runtimes are compiled in.
99    #[cfg(cube_backend)]
100    pub use burn_cubecl::Cube;
101
102    #[cfg(feature = "flex")]
103    pub use burn_flex as flex;
104    #[cfg(feature = "flex")]
105    pub use burn_flex::Flex;
106    #[cfg(feature = "ndarray")]
107    pub use burn_ndarray as ndarray;
108    #[cfg(feature = "ndarray")]
109    pub use burn_ndarray::NdArray;
110    #[cfg(feature = "tch")]
111    pub use burn_tch as libtorch;
112    #[cfg(feature = "tch")]
113    pub use burn_tch::LibTorch;
114
115    #[cfg(feature = "remote")]
116    pub use burn_remote::RemoteBackend as Remote;
117
118    /// Public graph-capture API types.
119    #[cfg(feature = "capture")]
120    pub mod capture {
121        pub use burn_capture::{
122            CaptureBackend, CaptureError, CaptureScope, CapturedGraph, CompletedCaptureScope,
123            TensorId,
124        };
125    }
126    #[cfg(feature = "capture")]
127    pub use burn_capture::CaptureBackend as Capture;
128}
129
130// Re-export devices
131
132/// Backend devices.
133pub mod devices {
134    /// Wgpu initialization and interoperability types used by the device facade.
135    #[cfg(feature = "wgpu")]
136    pub mod wgpu {
137        #[cfg(not(target_family = "wasm"))]
138        pub use burn_cubecl::cubecl::wgpu::try_init_setup;
139        pub use burn_cubecl::cubecl::wgpu::{
140            AutoGraphicsApi, MemoryConfiguration, RuntimeOptions, WgpuBackend, WgpuDevice,
141            WgpuInitError, WgpuSetup, try_init_device, try_init_setup_async, wgpu,
142        };
143    }
144
145    #[cfg(feature = "cpu")]
146    pub use burn_cubecl::cubecl::cpu::CpuDevice;
147    #[cfg(feature = "cuda")]
148    pub use burn_cubecl::cubecl::cuda::CudaDevice;
149    #[cfg(feature = "rocm")]
150    pub use burn_cubecl::cubecl::hip::AmdDevice as RocmDevice;
151    #[cfg(feature = "wgpu")]
152    pub use burn_cubecl::cubecl::wgpu::{
153        AutoCompiler, AutoGraphicsApi, WgpuBackend, WgpuDevice, WgpuDeviceKind, init_setup_async,
154    };
155
156    /// The device every cubecl runtime shares; which runtime it names is a
157    /// property of the value, not of its type, and [`RuntimeId`] is how that
158    /// property is named.
159    #[cfg(cube_backend)]
160    pub use burn_cubecl::CubeDevice;
161    #[cfg(cube_backend)]
162    pub use burn_cubecl::cubecl::RuntimeId;
163    #[cfg(feature = "flex")]
164    pub use burn_flex::FlexDevice;
165    #[cfg(feature = "ndarray")]
166    pub use burn_ndarray::NdArrayDevice;
167    #[cfg(feature = "tch")]
168    pub use burn_tch::LibTorchDevice;
169
170    #[cfg(feature = "remote")]
171    pub use burn_remote::RemoteDevice;
172}
173
174/// The remote backend's crate, for `burn::remote` and `burn::server` to build on. Not a user path.
175#[doc(hidden)]
176#[cfg(feature = "remote")]
177pub use burn_remote as __remote;
burn_dispatch/lib.rs

burn_dispatch/
lib.rs
