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
6//! Burn multi-backend dispatch.
7//!
8//! # Available Backends
9//!
10//! The dispatch backend supports the following variants, each enabled via cargo features:
11//!
12//! | Backend    | Feature    | Description |
13//! |------------|------------|-------------|
14//! | `Cpu`      | `cpu`      | Rust CPU backend (MLIR + LLVM) |
15//! | `Cuda`     | `cuda`     | NVIDIA CUDA backend |
16//! | `Metal`    | `metal`    | Apple Metal backend via `wgpu` (MSL) |
17//! | `Rocm`     | `rocm`     | AMD ROCm backend |
18//! | `Vulkan`   | `vulkan`   | Vulkan backend via `wgpu` (SPIR-V) |
19//! | `Wgpu`     | `webgpu`   | WebGPU backend via `wgpu` (WGSL) |
20//! | `Flex`     | `flex`     | Pure Rust CPU backend using `burn-flex` |
21//! | `NdArray`  | `ndarray`  | Pure Rust CPU backend using `ndarray` (legacy - prefer `flex`) |
22//! | `LibTorch` | `tch`      | Libtorch backend via `tch` |
23//! | `Autodiff` | `autodiff` | Autodiff-enabled backend (used in combination with any of the backends above) |
24//!
25//! **Note:** All backends, including the WGPU-based ones (`wgpu`, `metal`, `vulkan`, `webgpu`),
26//! can be combined freely. Each enabled wgpu backend appears as its own
27//! [`DispatchDevice`] variant.
28
29#[macro_use]
30mod macros;
31
32/// Dispatch backend module.
33pub mod backend;
34/// Dispatch device module.
35pub mod device;
36mod ops;
37/// Dispatch tensor module.
38pub mod tensor;
39
40/// Entry points for hosting a remote-execution server.
41#[cfg(feature = "remote-server")]
42pub mod remote_server;
43
44pub use backend::*;
45pub use device::*;
46pub use tensor::*;
47
48extern crate alloc;
49
50/// Backends and devices used.
51pub mod backends {
52    #[cfg(feature = "autodiff")]
53    pub use burn_autodiff as autodiff;
54    #[cfg(feature = "autodiff")]
55    pub use burn_autodiff::Autodiff; // re-export for extensions
56
57    #[cfg(feature = "cpu")]
58    pub use burn_cpu as cpu;
59    #[cfg(feature = "cpu")]
60    pub use burn_cpu::Cpu;
61    #[cfg(feature = "cuda")]
62    pub use burn_cuda as cuda;
63    #[cfg(feature = "cuda")]
64    pub use burn_cuda::Cuda;
65    #[cfg(feature = "rocm")]
66    pub use burn_rocm as rocm;
67    #[cfg(feature = "rocm")]
68    pub use burn_rocm::Rocm;
69    #[cfg(feature = "wgpu")]
70    pub use burn_wgpu as wgpu;
71    #[cfg(feature = "metal")]
72    pub use burn_wgpu::Metal;
73    #[cfg(feature = "vulkan")]
74    pub use burn_wgpu::Vulkan;
75    #[cfg(feature = "webgpu")]
76    pub use burn_wgpu::WebGpu;
77    #[cfg(feature = "wgpu")]
78    pub use burn_wgpu::Wgpu;
79
80    #[cfg(any(feature = "flex", default_backend))]
81    pub use burn_flex as flex;
82    #[cfg(any(feature = "flex", default_backend))]
83    pub use burn_flex::Flex;
84    #[cfg(feature = "ndarray")]
85    pub use burn_ndarray as ndarray;
86    #[cfg(feature = "ndarray")]
87    pub use burn_ndarray::NdArray;
88    #[cfg(feature = "tch")]
89    pub use burn_tch as libtorch;
90    #[cfg(feature = "tch")]
91    pub use burn_tch::LibTorch;
92
93    #[cfg(feature = "remote")]
94    pub use burn_remote as remote;
95    #[cfg(feature = "remote")]
96    pub use burn_remote::RemoteBackend as Remote;
97
98    pub use super::devices::*;
99}
100
101// Re-export devices
102
103/// Backend devices.
104pub mod devices {
105    #[cfg(feature = "cpu")]
106    pub use burn_cpu::CpuDevice;
107    #[cfg(feature = "cuda")]
108    pub use burn_cuda::CudaDevice;
109    #[cfg(feature = "rocm")]
110    pub use burn_rocm::RocmDevice;
111    #[cfg(feature = "wgpu")]
112    pub use burn_wgpu::WgpuDevice;
113
114    #[cfg(any(feature = "flex", default_backend))]
115    pub use burn_flex::FlexDevice;
116    #[cfg(feature = "ndarray")]
117    pub use burn_ndarray::NdArrayDevice;
118    #[cfg(feature = "tch")]
119    pub use burn_tch::LibTorchDevice;
120
121    #[cfg(feature = "remote")]
122    pub use burn_remote::RemoteDevice;
123
124    #[cfg(feature = "remote")]
125    pub use burn_remote::BURN_REMOTE_ALPN;
126}
burn_dispatch/lib.rs

burn_dispatch/
lib.rs
