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//! | `Autodiff` | `autodiff` | Autodiff-enabled backend (used in combination with any of the backends above) |
42//!
43//! **Note:** The features can be combined freely. The cubecl-backed ones all
44//! select the same backend, so they share the one `DispatchDevice::Cube`
45//! variant — enabling several compiles several runtimes in, and the device a
46//! tensor carries is what picks between them.
47
48#[macro_use]
49mod macros;
50
51/// Dispatch backend module.
52pub mod backend;
53/// Dispatch device module.
54pub mod device;
55mod ops;
56/// Dispatch tensor module.
57pub mod tensor;
58
59/// Entry points for hosting a remote-execution server.
60#[cfg(feature = "remote-server")]
61pub mod remote_server;
62
63pub use backend::*;
64pub use device::*;
65pub use tensor::*;
66
67extern crate alloc;
68
69// Keep backend-free dispatch types opaque to downstream crates. An actually empty
70// enum makes every API accepting/returning a tensor appear unreachable there.
71// The private field prevents construction without leaking uninhabitedness.
72#[cfg(not(backend_enabled))]
73#[doc(hidden)]
74#[derive(Clone, Debug, PartialEq, Eq)]
75pub struct NoBackend {
76 never: core::convert::Infallible,
77}
78
79#[cfg(not(backend_enabled))]
80impl NoBackend {
81 pub(crate) fn unreachable(&self) -> ! {
82 match self.never {}
83 }
84}
85
86/// Backends and devices used.
87pub mod backends {
88 #[cfg(feature = "autodiff")]
89 pub use burn_autodiff as autodiff;
90 #[cfg(feature = "autodiff")]
91 pub use burn_autodiff::Autodiff; // re-export for extensions
92
93 /// The cubecl backend: CUDA, ROCm, Metal, Vulkan, WebGPU, wgpu and the CPU
94 /// runtime are all this one type, and a tensor's device says which of them
95 /// it runs on. The features still decide which runtimes are compiled in.
96 #[cfg(cube_backend)]
97 pub use burn_cubecl::Cube;
98
99 #[cfg(feature = "flex")]
100 pub use burn_flex as flex;
101 #[cfg(feature = "flex")]
102 pub use burn_flex::Flex;
103 #[cfg(feature = "ndarray")]
104 pub use burn_ndarray as ndarray;
105 #[cfg(feature = "ndarray")]
106 pub use burn_ndarray::NdArray;
107 #[cfg(feature = "tch")]
108 pub use burn_tch as libtorch;
109 #[cfg(feature = "tch")]
110 pub use burn_tch::LibTorch;
111
112 #[cfg(feature = "remote")]
113 pub use burn_remote as remote;
114 #[cfg(feature = "remote")]
115 pub use burn_remote::RemoteBackend as Remote;
116
117 /// Public graph-capture API types.
118 #[cfg(feature = "capture")]
119 pub mod capture {
120 pub use burn_capture::{
121 CaptureBackend, CaptureError, CaptureScope, CapturedGraph, CompletedCaptureScope,
122 TensorId,
123 };
124 }
125 #[cfg(feature = "capture")]
126 pub use burn_capture::CaptureBackend as Capture;
127}
128
129// Re-export devices
130
131/// Backend devices.
132pub mod devices {
133 #[cfg(feature = "cpu")]
134 pub use burn_cubecl::cubecl::cpu::CpuDevice;
135 #[cfg(feature = "cuda")]
136 pub use burn_cubecl::cubecl::cuda::CudaDevice;
137 #[cfg(feature = "rocm")]
138 pub use burn_cubecl::cubecl::hip::AmdDevice as RocmDevice;
139 #[cfg(feature = "wgpu")]
140 pub use burn_cubecl::cubecl::wgpu::{
141 AutoCompiler, AutoGraphicsApi, WgpuBackend, WgpuDevice, WgpuDeviceKind, init_setup_async,
142 };
143
144 /// The device every cubecl runtime shares; which runtime it names is a
145 /// property of the value, not of its type, and [`RuntimeId`] is how that
146 /// property is named.
147 #[cfg(cube_backend)]
148 pub use burn_cubecl::CubeDevice;
149 #[cfg(cube_backend)]
150 pub use burn_cubecl::cubecl::RuntimeId;
151 #[cfg(feature = "flex")]
152 pub use burn_flex::FlexDevice;
153 #[cfg(feature = "ndarray")]
154 pub use burn_ndarray::NdArrayDevice;
155 #[cfg(feature = "tch")]
156 pub use burn_tch::LibTorchDevice;
157
158 #[cfg(feature = "remote")]
159 pub use burn_remote::RemoteDevice;
160
161 #[cfg(feature = "remote")]
162 pub use burn_remote::BURN_REMOTE_ALPN;
163}