burn_dispatch/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2#![warn(missing_docs)]
3#![cfg_attr(docsrs, feature(doc_cfg))]
4#![recursion_limit = "138"]
5// Wiring up the deprecated `NdArray` backend is this crate's job, and the backend registry macros
6// expand it into every dispatch impl, so the warnings land on `macros.rs` rather than on any site
7// we could annotate individually. `allow(deprecated)` is a lint level scoped to this crate, and
8// lint levels never propagate to dependents, so downstream code naming `NdArray` (directly or via
9// our re-export) still gets the warning. The `cfg_attr` keeps this confined to the ndarray-enabled
10// build: `ndarray` is not a default feature, so the default build that CI lints with
11// `--deny warnings` retains full deprecation signal for every other dependency.
12#![cfg_attr(feature = "ndarray", allow(deprecated))]
13
14//! Burn multi-backend dispatch.
15//!
16//! # Available Backends
17//!
18//! The dispatch backend supports the following variants, each enabled via cargo features:
19//!
20//! | Backend    | Feature    | Description |
21//! |------------|------------|-------------|
22//! | `Cpu`      | `cpu`      | Rust CPU backend (MLIR + LLVM) |
23//! | `Cuda`     | `cuda`     | NVIDIA CUDA backend |
24//! | `Metal`    | `metal`    | Apple Metal backend via `wgpu` (MSL) |
25//! | `Rocm`     | `rocm`     | AMD ROCm backend |
26//! | `Vulkan`   | `vulkan`   | Vulkan backend via `wgpu` (SPIR-V) |
27//! | `Wgpu`     | `webgpu`   | WebGPU backend via `wgpu` (WGSL) |
28//! | `Flex`     | `flex`     | Pure Rust CPU backend using `burn-flex` |
29//! | `NdArray`  | `ndarray`  | Pure Rust CPU backend using `ndarray` (deprecated - use `flex`) |
30//! | `LibTorch` | `tch`      | Libtorch backend via `tch` |
31//! | `Autodiff` | `autodiff` | Autodiff-enabled backend (used in combination with any of the backends above) |
32//!
33//! **Note:** All backends, including the WGPU-based ones (`wgpu`, `metal`, `vulkan`, `webgpu`),
34//! can be combined freely. Each enabled wgpu backend appears as its own
35//! [`DispatchDevice`] variant.
36
37#[macro_use]
38mod macros;
39
40/// Dispatch backend module.
41pub mod backend;
42/// Dispatch device module.
43pub mod device;
44mod ops;
45/// Dispatch tensor module.
46pub mod tensor;
47
48/// Entry points for hosting a remote-execution server.
49#[cfg(feature = "remote-server")]
50pub mod remote_server;
51
52pub use backend::*;
53pub use device::*;
54pub use tensor::*;
55
56extern crate alloc;
57
58/// Backends and devices used.
59pub mod backends {
60    #[cfg(feature = "autodiff")]
61    pub use burn_autodiff as autodiff;
62    #[cfg(feature = "autodiff")]
63    pub use burn_autodiff::Autodiff; // re-export for extensions
64
65    #[cfg(feature = "cpu")]
66    pub use burn_cpu as cpu;
67    #[cfg(feature = "cpu")]
68    pub use burn_cpu::Cpu;
69    #[cfg(feature = "cuda")]
70    pub use burn_cuda as cuda;
71    #[cfg(feature = "cuda")]
72    pub use burn_cuda::Cuda;
73    #[cfg(feature = "rocm")]
74    pub use burn_rocm as rocm;
75    #[cfg(feature = "rocm")]
76    pub use burn_rocm::Rocm;
77    #[cfg(feature = "wgpu")]
78    pub use burn_wgpu as wgpu;
79    #[cfg(feature = "metal")]
80    pub use burn_wgpu::Metal;
81    #[cfg(feature = "vulkan")]
82    pub use burn_wgpu::Vulkan;
83    #[cfg(feature = "webgpu")]
84    pub use burn_wgpu::WebGpu;
85    #[cfg(feature = "wgpu")]
86    pub use burn_wgpu::Wgpu;
87
88    #[cfg(any(feature = "flex", default_backend))]
89    pub use burn_flex as flex;
90    #[cfg(any(feature = "flex", default_backend))]
91    pub use burn_flex::Flex;
92    #[cfg(feature = "ndarray")]
93    pub use burn_ndarray as ndarray;
94    #[cfg(feature = "ndarray")]
95    pub use burn_ndarray::NdArray;
96    #[cfg(feature = "tch")]
97    pub use burn_tch as libtorch;
98    #[cfg(feature = "tch")]
99    pub use burn_tch::LibTorch;
100
101    #[cfg(feature = "remote")]
102    pub use burn_remote as remote;
103    #[cfg(feature = "remote")]
104    pub use burn_remote::RemoteBackend as Remote;
105
106    /// Public graph-capture API types.
107    #[cfg(feature = "capture")]
108    pub mod capture {
109        pub use burn_capture::{
110            CaptureBackend, CaptureError, CaptureScope, CapturedGraph, CompletedCaptureScope,
111            TensorId,
112        };
113    }
114    #[cfg(feature = "capture")]
115    pub use burn_capture::CaptureBackend as Capture;
116
117    pub use super::devices::*;
118}
119
120// Re-export devices
121
122/// Backend devices.
123pub mod devices {
124    #[cfg(feature = "cpu")]
125    pub use burn_cpu::CpuDevice;
126    #[cfg(feature = "cuda")]
127    pub use burn_cuda::CudaDevice;
128    #[cfg(feature = "rocm")]
129    pub use burn_rocm::RocmDevice;
130    #[cfg(feature = "wgpu")]
131    pub use burn_wgpu::WgpuDevice;
132
133    #[cfg(any(feature = "flex", default_backend))]
134    pub use burn_flex::FlexDevice;
135    #[cfg(feature = "ndarray")]
136    pub use burn_ndarray::NdArrayDevice;
137    #[cfg(feature = "tch")]
138    pub use burn_tch::LibTorchDevice;
139
140    #[cfg(feature = "remote")]
141    pub use burn_remote::RemoteDevice;
142
143    #[cfg(feature = "remote")]
144    pub use burn_remote::BURN_REMOTE_ALPN;
145}
burn_dispatch/lib.rs

burn_dispatch/
lib.rs
