1use alloc::string::String;
2use ruda_tensor::{Backend, BackendTypes, AutodiffBackend, AutodiffContext, GradientCheckpointingStrategy,
3 DType, DTypeUsageSet, DeviceId, DeviceOps, ExecutionError, TensorData};
4use ruda_tensor::backend::{ProfileOptions, ProfileToken, ProfileDuration, BackendGraph};
5use crate::{DispatchDevice, DispatchTensor, BackendTensor, context};
6
7#[derive(Clone, Copy, Debug, Default)]
8pub struct Dispatch;
9
10impl BackendTypes for Dispatch {
11 type Device = DispatchDevice;
12 type FloatTensorPrimitive = DispatchTensor;
13 type IntTensorPrimitive = DispatchTensor;
14 type BoolTensorPrimitive = DispatchTensor;
15 type QuantizedTensorPrimitive = DispatchTensor;
16 type FloatElem = f32;
17 type IntElem = i32;
18 type BoolElem = bool;
19}
20
21impl Backend for Dispatch {
22 fn ad_enabled(device: &DispatchDevice) -> bool {
23 matches!(context::current(device), AutodiffContext::Enabled(_))
24 }
25 fn autodiff_creation_context(device: &DispatchDevice) -> AutodiffContext { context::current(device) }
26 fn with_autodiff_creation_context<R>(device: &DispatchDevice, context: AutodiffContext, func: impl FnOnce() -> R) -> R {
27 context::run(device, context, func)
28 }
29 fn name(device: &DispatchDevice) -> String {
30 macro_rules! call {
31 ($B:ty, $variant:ident, $device:expr) => { <$B>::name($device) };
32 }
33 native_dispatch!(device, call)
34 }
35 fn seed(device: &DispatchDevice, seed: u64) {
36 macro_rules! call {
37 ($B:ty, $variant:ident, $device:expr) => { <$B>::seed($device, seed) };
38 }
39 native_dispatch!(device, call)
40 }
41 fn rng_state(device: &DispatchDevice) -> Result<alloc::vec::Vec<u8>, ExecutionError> {
42 macro_rules! call {
43 ($B:ty, $variant:ident, $device:expr) => { <$B>::rng_state($device) };
44 }
45 native_dispatch!(device, call)
46 }
47 fn restore_rng_state(device: &DispatchDevice, state: &[u8]) -> Result<(), ExecutionError> {
48 macro_rules! call {
49 ($B:ty, $variant:ident, $device:expr) => { <$B>::restore_rng_state($device, state) };
50 }
51 native_dispatch!(device, call)
52 }
53 fn sync(device: &DispatchDevice) -> Result<(), ExecutionError> {
54 macro_rules! call {
55 ($B:ty, $variant:ident, $device:expr) => { <$B>::sync($device) };
56 }
57 native_dispatch!(device, call)
58 }
59 fn flush(device: &DispatchDevice) -> Result<(), ExecutionError> {
60 macro_rules! call {
61 ($B:ty, $variant:ident, $device:expr) => { <$B>::flush($device) };
62 }
63 native_dispatch!(device, call)
64 }
65 fn memory_cleanup(device: &DispatchDevice) {
66 macro_rules! call {
67 ($B:ty, $variant:ident, $device:expr) => { <$B>::memory_cleanup($device) };
68 }
69 native_dispatch!(device, call)
70 }
71 fn memory_persistent_allocations<Output: Send, Input: Send, Func: Fn(Input) -> Output + Send>(device: &DispatchDevice, input: Input, func: Func) -> Output {
72 macro_rules! call {
73 ($B:ty, $variant:ident, $device:expr) => { <$B>::memory_persistent_allocations($device, input, func) };
74 }
75 native_dispatch!(device, call)
76 }
77 fn profile<O: Send + 'static>(device: &DispatchDevice, options: ProfileOptions, func: impl FnOnce() -> O + Send) -> Result<(O, ProfileDuration), ExecutionError> {
78 macro_rules! call {
79 ($B:ty, $variant:ident, $device:expr) => { <$B>::profile($device, options, func) };
80 }
81 native_dispatch!(device, call)
82 }
83 fn profile_start(device: &DispatchDevice) -> Result<Option<ProfileToken>, ExecutionError> {
84 macro_rules! call {
85 ($B:ty, $variant:ident, $device:expr) => { <$B>::profile_start($device) };
86 }
87 native_dispatch!(device, call)
88 }
89 fn profile_end(device: &DispatchDevice, token: ProfileToken, options: ProfileOptions) -> Result<ProfileDuration, ExecutionError> {
90 macro_rules! call {
91 ($B:ty, $variant:ident, $device:expr) => { <$B>::profile_end($device, token, options) };
92 }
93 native_dispatch!(device, call)
94 }
95 fn profile_abandon(device: &DispatchDevice, token: ProfileToken) {
96 macro_rules! call {
97 ($B:ty, $variant:ident, $device:expr) => { <$B>::profile_abandon($device, token) };
98 }
99 native_dispatch!(device, call)
100 }
101 fn graph_prepare(device: &DispatchDevice) -> Result<(), ExecutionError> {
102 macro_rules! call {
103 ($B:ty, $variant:ident, $device:expr) => { <$B>::graph_prepare($device) };
104 }
105 native_dispatch!(device, call)
106 }
107 fn graph_start_capture(device: &DispatchDevice) -> Result<(), ExecutionError> {
108 macro_rules! call {
109 ($B:ty, $variant:ident, $device:expr) => { <$B>::graph_start_capture($device) };
110 }
111 native_dispatch!(device, call)
112 }
113 fn graph_stop_capture(device: &DispatchDevice) -> Result<BackendGraph, ExecutionError> {
114 macro_rules! call {
115 ($B:ty, $variant:ident, $device:expr) => { <$B>::graph_stop_capture($device) };
116 }
117 native_dispatch!(device, call)
118 }
119 fn graph_abort_capture(device: &DispatchDevice) {
120 macro_rules! call {
121 ($B:ty, $variant:ident, $device:expr) => { <$B>::graph_abort_capture($device) };
122 }
123 native_dispatch!(device, call)
124 }
125 fn dtype_usage(device: &DispatchDevice, dtype: DType) -> DTypeUsageSet {
126 macro_rules! call {
127 ($B:ty, $variant:ident, $device:expr) => { <$B>::dtype_usage($device, dtype) };
128 }
129 native_dispatch!(device, call)
130 }
131 fn staging<'a, Iter: Iterator<Item = &'a mut TensorData>>(data: Iter, device: &DispatchDevice) {
132 macro_rules! call {
133 ($B:ty, $variant:ident, $device:expr) => { <$B>::staging(data, $device) };
134 }
135 native_dispatch!(device, call)
136 }
137 fn device_count(type_id: u16) -> usize {
138 let native = type_id & 255;
139 match (type_id & 0x7fff) >> 8 {
1401 => <crate::backends::Host as Backend>::device_count(native),
141#[cfg(feature = "cuda")]
1422 => <crate::backends::Cuda as Backend>::device_count(native),
143#[cfg(feature = "cpu")]
1443 => <crate::backends::Cpu as Backend>::device_count(native),
145#[cfg(feature = "wgpu")]
1464 => <crate::backends::Wgpu as Backend>::device_count(native),
147#[cfg(feature = "rocm")]
1485 => <crate::backends::Rocm as Backend>::device_count(native),
149#[cfg(feature = "tch")]
1506 => <crate::backends::LibTorch as Backend>::device_count(native),
151#[cfg(feature = "remote")]
1527 => <crate::backends::Remote as Backend>::device_count(native),
153#[cfg(feature = "capture")]
1548 => 0,
155 _ => 0,
156 }
157 }
158}
159
160impl AutodiffBackend for Dispatch {
161 type InnerBackend = Self;
162 type Gradients = ruda_autodiff::grads::Gradients;
163
164 #[cfg(feature = "distributed")]
165 fn set_distributed_params(tensor: DispatchTensor, param_id: ruda_tensor::distributed::DistributedParamId) -> DispatchTensor {
166 crate::distributed::mark(tensor, param_id)
167 }
168
169 #[cfg(feature = "distributed")]
170 fn distributed_params(tensor: &DispatchTensor) -> Option<ruda_tensor::distributed::DistributedParams> {
171 crate::distributed::params(tensor)
172 }
173
174 #[cfg(feature = "distributed")]
175 fn is_distributed(tensor: &DispatchTensor) -> bool { crate::distributed::params(tensor).is_some() }
176
177 fn backward(tensor: DispatchTensor) -> Self::Gradients {
178 match tensor {
179DispatchTensor::Host(BackendTensor::AdFloat(tensor)) => tensor.backward_local(),
180#[cfg(feature = "cuda")]
181DispatchTensor::Cuda(BackendTensor::AdFloat(tensor)) => {
182 #[cfg(feature = "distributed")]
183 { tensor.backward() }
184 #[cfg(not(feature = "distributed"))]
185 { tensor.backward_local() }
186},
187#[cfg(feature = "cpu")]
188DispatchTensor::Cpu(BackendTensor::AdFloat(tensor)) => tensor.backward_local(),
189#[cfg(feature = "wgpu")]
190DispatchTensor::Wgpu(BackendTensor::AdFloat(tensor)) => tensor.backward_local(),
191#[cfg(feature = "rocm")]
192DispatchTensor::Rocm(BackendTensor::AdFloat(tensor)) => tensor.backward_local(),
193#[cfg(feature = "tch")]
194DispatchTensor::LibTorch(BackendTensor::AdFloat(tensor)) => {
195 #[cfg(feature = "distributed")]
196 { tensor.backward() }
197 #[cfg(not(feature = "distributed"))]
198 { tensor.backward_local() }
199},
200#[cfg(feature = "remote")]
201DispatchTensor::Remote(BackendTensor::AdFloat(tensor)) => {
202 #[cfg(feature = "distributed")]
203 { tensor.backward() }
204 #[cfg(not(feature = "distributed"))]
205 { tensor.backward_local() }
206},
207 _ => panic!("Backward requires an autodiff float tensor"),
208 }
209 }
210 fn grad(tensor: &DispatchTensor, grads: &Self::Gradients) -> Option<DispatchTensor> {
211 match tensor {
212DispatchTensor::Host(BackendTensor::AdFloat(tensor)) => tensor.grad(grads).map(|value| DispatchTensor::Host(BackendTensor::Float(value))),
213#[cfg(feature = "cuda")]
214DispatchTensor::Cuda(BackendTensor::AdFloat(tensor)) => tensor.grad(grads).map(|value| DispatchTensor::Cuda(BackendTensor::Float(value))),
215#[cfg(feature = "cpu")]
216DispatchTensor::Cpu(BackendTensor::AdFloat(tensor)) => tensor.grad(grads).map(|value| DispatchTensor::Cpu(BackendTensor::Float(value))),
217#[cfg(feature = "wgpu")]
218DispatchTensor::Wgpu(BackendTensor::AdFloat(tensor)) => tensor.grad(grads).map(|value| DispatchTensor::Wgpu(BackendTensor::Float(value))),
219#[cfg(feature = "rocm")]
220DispatchTensor::Rocm(BackendTensor::AdFloat(tensor)) => tensor.grad(grads).map(|value| DispatchTensor::Rocm(BackendTensor::Float(value))),
221#[cfg(feature = "tch")]
222DispatchTensor::LibTorch(BackendTensor::AdFloat(tensor)) => tensor.grad(grads).map(|value| DispatchTensor::LibTorch(BackendTensor::Float(value))),
223#[cfg(feature = "remote")]
224DispatchTensor::Remote(BackendTensor::AdFloat(tensor)) => tensor.grad(grads).map(|value| DispatchTensor::Remote(BackendTensor::Float(value))),
225 _ => None,
226 }
227 }
228 fn grad_remove(tensor: &DispatchTensor, grads: &mut Self::Gradients) -> Option<DispatchTensor> {
229 match tensor {
230DispatchTensor::Host(BackendTensor::AdFloat(tensor)) => tensor.grad_remove(grads).map(|value| DispatchTensor::Host(BackendTensor::Float(value))),
231#[cfg(feature = "cuda")]
232DispatchTensor::Cuda(BackendTensor::AdFloat(tensor)) => tensor.grad_remove(grads).map(|value| DispatchTensor::Cuda(BackendTensor::Float(value))),
233#[cfg(feature = "cpu")]
234DispatchTensor::Cpu(BackendTensor::AdFloat(tensor)) => tensor.grad_remove(grads).map(|value| DispatchTensor::Cpu(BackendTensor::Float(value))),
235#[cfg(feature = "wgpu")]
236DispatchTensor::Wgpu(BackendTensor::AdFloat(tensor)) => tensor.grad_remove(grads).map(|value| DispatchTensor::Wgpu(BackendTensor::Float(value))),
237#[cfg(feature = "rocm")]
238DispatchTensor::Rocm(BackendTensor::AdFloat(tensor)) => tensor.grad_remove(grads).map(|value| DispatchTensor::Rocm(BackendTensor::Float(value))),
239#[cfg(feature = "tch")]
240DispatchTensor::LibTorch(BackendTensor::AdFloat(tensor)) => tensor.grad_remove(grads).map(|value| DispatchTensor::LibTorch(BackendTensor::Float(value))),
241#[cfg(feature = "remote")]
242DispatchTensor::Remote(BackendTensor::AdFloat(tensor)) => tensor.grad_remove(grads).map(|value| DispatchTensor::Remote(BackendTensor::Float(value))),
243 _ => None,
244 }
245 }
246 fn q_grad(tensor: &DispatchTensor, grads: &Self::Gradients) -> Option<DispatchTensor> {
247 match tensor {
248DispatchTensor::Host(BackendTensor::AdQuantized(tensor)) => tensor.grad(grads).map(|value| DispatchTensor::Host(BackendTensor::Float(value))),
249#[cfg(feature = "cuda")]
250DispatchTensor::Cuda(BackendTensor::AdQuantized(tensor)) => tensor.grad(grads).map(|value| DispatchTensor::Cuda(BackendTensor::Float(value))),
251#[cfg(feature = "cpu")]
252DispatchTensor::Cpu(BackendTensor::AdQuantized(tensor)) => tensor.grad(grads).map(|value| DispatchTensor::Cpu(BackendTensor::Float(value))),
253#[cfg(feature = "wgpu")]
254DispatchTensor::Wgpu(BackendTensor::AdQuantized(tensor)) => tensor.grad(grads).map(|value| DispatchTensor::Wgpu(BackendTensor::Float(value))),
255#[cfg(feature = "rocm")]
256DispatchTensor::Rocm(BackendTensor::AdQuantized(tensor)) => tensor.grad(grads).map(|value| DispatchTensor::Rocm(BackendTensor::Float(value))),
257#[cfg(feature = "tch")]
258DispatchTensor::LibTorch(BackendTensor::AdQuantized(tensor)) => tensor.grad(grads).map(|value| DispatchTensor::LibTorch(BackendTensor::Float(value))),
259#[cfg(feature = "remote")]
260DispatchTensor::Remote(BackendTensor::AdQuantized(tensor)) => tensor.grad(grads).map(|value| DispatchTensor::Remote(BackendTensor::Float(value))),
261 _ => None,
262 }
263 }
264 fn q_grad_remove(tensor: &DispatchTensor, grads: &mut Self::Gradients) -> Option<DispatchTensor> {
265 match tensor {
266DispatchTensor::Host(BackendTensor::AdQuantized(tensor)) => tensor.grad_remove(grads).map(|value| DispatchTensor::Host(BackendTensor::Float(value))),
267#[cfg(feature = "cuda")]
268DispatchTensor::Cuda(BackendTensor::AdQuantized(tensor)) => tensor.grad_remove(grads).map(|value| DispatchTensor::Cuda(BackendTensor::Float(value))),
269#[cfg(feature = "cpu")]
270DispatchTensor::Cpu(BackendTensor::AdQuantized(tensor)) => tensor.grad_remove(grads).map(|value| DispatchTensor::Cpu(BackendTensor::Float(value))),
271#[cfg(feature = "wgpu")]
272DispatchTensor::Wgpu(BackendTensor::AdQuantized(tensor)) => tensor.grad_remove(grads).map(|value| DispatchTensor::Wgpu(BackendTensor::Float(value))),
273#[cfg(feature = "rocm")]
274DispatchTensor::Rocm(BackendTensor::AdQuantized(tensor)) => tensor.grad_remove(grads).map(|value| DispatchTensor::Rocm(BackendTensor::Float(value))),
275#[cfg(feature = "tch")]
276DispatchTensor::LibTorch(BackendTensor::AdQuantized(tensor)) => tensor.grad_remove(grads).map(|value| DispatchTensor::LibTorch(BackendTensor::Float(value))),
277#[cfg(feature = "remote")]
278DispatchTensor::Remote(BackendTensor::AdQuantized(tensor)) => tensor.grad_remove(grads).map(|value| DispatchTensor::Remote(BackendTensor::Float(value))),
279 _ => None,
280 }
281 }
282 fn grad_replace(tensor: &DispatchTensor, grads: &mut Self::Gradients, grad: DispatchTensor) {
283 match (tensor, grad) {
284(DispatchTensor::Host(BackendTensor::AdFloat(tensor)), DispatchTensor::Host(value)) =>
285 tensor.grad_replace(grads, value.into_native_float(AutodiffContext::Disabled)),
286#[cfg(feature = "cuda")]
287(DispatchTensor::Cuda(BackendTensor::AdFloat(tensor)), DispatchTensor::Cuda(value)) =>
288 tensor.grad_replace(grads, value.into_native_float(AutodiffContext::Disabled)),
289#[cfg(feature = "cpu")]
290(DispatchTensor::Cpu(BackendTensor::AdFloat(tensor)), DispatchTensor::Cpu(value)) =>
291 tensor.grad_replace(grads, value.into_native_float(AutodiffContext::Disabled)),
292#[cfg(feature = "wgpu")]
293(DispatchTensor::Wgpu(BackendTensor::AdFloat(tensor)), DispatchTensor::Wgpu(value)) =>
294 tensor.grad_replace(grads, value.into_native_float(AutodiffContext::Disabled)),
295#[cfg(feature = "rocm")]
296(DispatchTensor::Rocm(BackendTensor::AdFloat(tensor)), DispatchTensor::Rocm(value)) =>
297 tensor.grad_replace(grads, value.into_native_float(AutodiffContext::Disabled)),
298#[cfg(feature = "tch")]
299(DispatchTensor::LibTorch(BackendTensor::AdFloat(tensor)), DispatchTensor::LibTorch(value)) =>
300 tensor.grad_replace(grads, value.into_native_float(AutodiffContext::Disabled)),
301#[cfg(feature = "remote")]
302(DispatchTensor::Remote(BackendTensor::AdFloat(tensor)), DispatchTensor::Remote(value)) =>
303 tensor.grad_replace(grads, value.into_native_float(AutodiffContext::Disabled)),
304 _ => panic!("Gradient and autodiff tensor must use the same backend"),
305 }
306 }
307 fn q_grad_replace(tensor: &DispatchTensor, grads: &mut Self::Gradients, grad: DispatchTensor) {
308 match (tensor, grad) {
309(DispatchTensor::Host(BackendTensor::AdQuantized(tensor)), DispatchTensor::Host(value)) =>
310 tensor.grad_replace(grads, value.into_native_float(AutodiffContext::Disabled)),
311#[cfg(feature = "cuda")]
312(DispatchTensor::Cuda(BackendTensor::AdQuantized(tensor)), DispatchTensor::Cuda(value)) =>
313 tensor.grad_replace(grads, value.into_native_float(AutodiffContext::Disabled)),
314#[cfg(feature = "cpu")]
315(DispatchTensor::Cpu(BackendTensor::AdQuantized(tensor)), DispatchTensor::Cpu(value)) =>
316 tensor.grad_replace(grads, value.into_native_float(AutodiffContext::Disabled)),
317#[cfg(feature = "wgpu")]
318(DispatchTensor::Wgpu(BackendTensor::AdQuantized(tensor)), DispatchTensor::Wgpu(value)) =>
319 tensor.grad_replace(grads, value.into_native_float(AutodiffContext::Disabled)),
320#[cfg(feature = "rocm")]
321(DispatchTensor::Rocm(BackendTensor::AdQuantized(tensor)), DispatchTensor::Rocm(value)) =>
322 tensor.grad_replace(grads, value.into_native_float(AutodiffContext::Disabled)),
323#[cfg(feature = "tch")]
324(DispatchTensor::LibTorch(BackendTensor::AdQuantized(tensor)), DispatchTensor::LibTorch(value)) =>
325 tensor.grad_replace(grads, value.into_native_float(AutodiffContext::Disabled)),
326#[cfg(feature = "remote")]
327(DispatchTensor::Remote(BackendTensor::AdQuantized(tensor)), DispatchTensor::Remote(value)) =>
328 tensor.grad_replace(grads, value.into_native_float(AutodiffContext::Disabled)),
329 _ => panic!("Gradient and autodiff tensor must use the same backend"),
330 }
331 }
332 fn inner(tensor: DispatchTensor) -> DispatchTensor { tensor.with_context(AutodiffContext::Disabled) }
333 fn from_inner(tensor: DispatchTensor) -> DispatchTensor {
334 let context = match context::current(&tensor.device()) {
335 AutodiffContext::Disabled => AutodiffContext::Enabled(GradientCheckpointingStrategy::Disabled),
336 context => context,
337 };
338 tensor.with_context(context)
339 }
340 fn int_inner(tensor: DispatchTensor) -> DispatchTensor { tensor.with_context(AutodiffContext::Disabled) }
341 fn int_from_inner(tensor: DispatchTensor) -> DispatchTensor {
342 let context = match context::current(&tensor.device()) {
343 AutodiffContext::Disabled => AutodiffContext::Enabled(GradientCheckpointingStrategy::Disabled),
344 context => context,
345 };
346 tensor.with_context(context)
347 }
348 fn bool_inner(tensor: DispatchTensor) -> DispatchTensor { tensor.with_context(AutodiffContext::Disabled) }
349 fn bool_from_inner(tensor: DispatchTensor) -> DispatchTensor {
350 let context = match context::current(&tensor.device()) {
351 AutodiffContext::Disabled => AutodiffContext::Enabled(GradientCheckpointingStrategy::Disabled),
352 context => context,
353 };
354 tensor.with_context(context)
355 }
356 fn q_inner(tensor: DispatchTensor) -> DispatchTensor { tensor.with_context(AutodiffContext::Disabled) }
357 fn q_from_inner(tensor: DispatchTensor) -> DispatchTensor {
358 let context = match context::current(&tensor.device()) {
359 AutodiffContext::Disabled => AutodiffContext::Enabled(GradientCheckpointingStrategy::Disabled),
360 context => context,
361 };
362 tensor.with_context(context)
363 }
364}