ruda_tensor_dispatch/
tensor.rs

1use ruda_autodiff::{Autodiff, AutodiffValue};
2use ruda_tensor::{Backend, AutodiffContext, TensorMetadata, QTensorPrimitive, DeviceOps, DType, Shape};
3use ruda_tensor::ops::*;
4use ruda_tensor::tensor::{FloatTensor, IntTensor, BoolTensor, QuantizedTensor};
5use ruda_core::tensor::{QuantScheme, QuantAcc, QuantPropagation};
6use crate::{DispatchDevice, backends};
7
8#[derive(Clone, Debug)]
9pub enum BackendTensor<B: Backend> {
10Float(FloatTensor<B>),
11AdFloat(FloatTensor<Autodiff<B>>),
12Int(IntTensor<B>),
13AdInt(IntTensor<Autodiff<B>>),
14Bool(BoolTensor<B>),
15AdBool(BoolTensor<Autodiff<B>>),
16Quantized(QuantizedTensor<B>),
17AdQuantized(QuantizedTensor<Autodiff<B>>),
18}
19
20impl<B: Backend> BackendTensor<B> {
21pub(crate) fn into_native_float(self, _: AutodiffContext) -> FloatTensor<B> {
22    match self {
23        Self::Float(tensor) => tensor,
24        Self::AdFloat(tensor) => tensor.primitive,
25        _ => panic!("Expected a float tensor"),
26    }
27}
28pub(crate) fn into_ad_float(self, context: AutodiffContext) -> FloatTensor<Autodiff<B>> {
29    let tensor = match self {
30        Self::AdFloat(tensor) => tensor,
31        Self::Float(tensor) => <FloatTensor<Autodiff<B>>>::new(tensor),
32        _ => panic!("Expected a float tensor"),
33    };
34    Autodiff::<B>::float_with_autodiff_context(tensor, context)
35}
36pub(crate) fn into_native_int(self, _: AutodiffContext) -> IntTensor<B> {
37    match self {
38        Self::Int(tensor) => tensor,
39        Self::AdInt(tensor) => tensor.primitive,
40        _ => panic!("Expected a int tensor"),
41    }
42}
43pub(crate) fn into_ad_int(self, context: AutodiffContext) -> IntTensor<Autodiff<B>> {
44    let tensor = match self {
45        Self::AdInt(tensor) => tensor,
46        Self::Int(tensor) => AutodiffValue::with_context(tensor, context),
47        _ => panic!("Expected a int tensor"),
48    };
49    Autodiff::<B>::int_with_autodiff_context(tensor, context)
50}
51pub(crate) fn into_native_bool(self, _: AutodiffContext) -> BoolTensor<B> {
52    match self {
53        Self::Bool(tensor) => tensor,
54        Self::AdBool(tensor) => tensor.primitive,
55        _ => panic!("Expected a bool tensor"),
56    }
57}
58pub(crate) fn into_ad_bool(self, context: AutodiffContext) -> BoolTensor<Autodiff<B>> {
59    let tensor = match self {
60        Self::AdBool(tensor) => tensor,
61        Self::Bool(tensor) => AutodiffValue::with_context(tensor, context),
62        _ => panic!("Expected a bool tensor"),
63    };
64    Autodiff::<B>::bool_with_autodiff_context(tensor, context)
65}
66pub(crate) fn into_native_quantized(self, _: AutodiffContext) -> QuantizedTensor<B> {
67    match self {
68        Self::Quantized(tensor) => tensor,
69        Self::AdQuantized(tensor) => tensor.primitive,
70        _ => panic!("Expected a quantized tensor"),
71    }
72}
73pub(crate) fn into_ad_quantized(self, context: AutodiffContext) -> QuantizedTensor<Autodiff<B>> {
74    let tensor = match self {
75        Self::AdQuantized(tensor) => tensor,
76        Self::Quantized(tensor) => ruda_autodiff::AutodiffQTensor::untracked(tensor),
77        _ => panic!("Expected a quantized tensor"),
78    };
79    Autodiff::<B>::q_with_autodiff_context(tensor, context)
80}
81
82    pub fn context(&self) -> AutodiffContext {
83        match self {
84Self::AdFloat(tensor) => Autodiff::<B>::float_autodiff_context(tensor),
85Self::AdInt(tensor) => Autodiff::<B>::int_autodiff_context(tensor),
86Self::AdBool(tensor) => Autodiff::<B>::bool_autodiff_context(tensor),
87Self::AdQuantized(tensor) => Autodiff::<B>::q_autodiff_context(tensor),
88            _ => AutodiffContext::Disabled,
89        }
90    }
91
92    pub fn with_context(self, context: AutodiffContext) -> Self {
93        match context {
94            AutodiffContext::Disabled => match self {
95value @ (Self::Float(_) | Self::AdFloat(_)) => Self::Float(value.into_native_float(context)),
96value @ (Self::Int(_) | Self::AdInt(_)) => Self::Int(value.into_native_int(context)),
97value @ (Self::Bool(_) | Self::AdBool(_)) => Self::Bool(value.into_native_bool(context)),
98value @ (Self::Quantized(_) | Self::AdQuantized(_)) => Self::Quantized(value.into_native_quantized(context)),
99            },
100            AutodiffContext::Enabled(_) => match self {
101value @ (Self::Float(_) | Self::AdFloat(_)) => Self::AdFloat(value.into_ad_float(context)),
102value @ (Self::Int(_) | Self::AdInt(_)) => Self::AdInt(value.into_ad_int(context)),
103value @ (Self::Bool(_) | Self::AdBool(_)) => Self::AdBool(value.into_ad_bool(context)),
104value @ (Self::Quantized(_) | Self::AdQuantized(_)) => Self::AdQuantized(value.into_ad_quantized(context)),
105            },
106        }
107    }
108
109    pub fn device(&self) -> B::Device {
110        match self {
111Self::Float(tensor) => B::float_device(tensor),
112Self::AdFloat(tensor) => B::float_device(&tensor.primitive),
113Self::Int(tensor) => B::int_device(tensor),
114Self::AdInt(tensor) => B::int_device(&tensor.primitive),
115Self::Bool(tensor) => B::bool_device(tensor),
116Self::AdBool(tensor) => B::bool_device(&tensor.primitive),
117Self::Quantized(tensor) => B::q_device(tensor),
118Self::AdQuantized(tensor) => B::q_device(&tensor.primitive),
119        }
120    }
121}
122
123impl<B: Backend> TensorMetadata for BackendTensor<B> {
124fn dtype(&self) -> DType {
125    match self {
126Self::Float(tensor) => tensor.dtype(),
127Self::AdFloat(tensor) => tensor.dtype(),
128Self::Int(tensor) => tensor.dtype(),
129Self::AdInt(tensor) => tensor.dtype(),
130Self::Bool(tensor) => tensor.dtype(),
131Self::AdBool(tensor) => tensor.dtype(),
132Self::Quantized(tensor) => tensor.dtype(),
133Self::AdQuantized(tensor) => tensor.dtype(),
134    }
135}
136fn shape(&self) -> Shape {
137    match self {
138Self::Float(tensor) => tensor.shape(),
139Self::AdFloat(tensor) => tensor.shape(),
140Self::Int(tensor) => tensor.shape(),
141Self::AdInt(tensor) => tensor.shape(),
142Self::Bool(tensor) => tensor.shape(),
143Self::AdBool(tensor) => tensor.shape(),
144Self::Quantized(tensor) => tensor.shape(),
145Self::AdQuantized(tensor) => tensor.shape(),
146    }
147}
148fn rank(&self) -> usize {
149    match self {
150Self::Float(tensor) => tensor.rank(),
151Self::AdFloat(tensor) => tensor.rank(),
152Self::Int(tensor) => tensor.rank(),
153Self::AdInt(tensor) => tensor.rank(),
154Self::Bool(tensor) => tensor.rank(),
155Self::AdBool(tensor) => tensor.rank(),
156Self::Quantized(tensor) => tensor.rank(),
157Self::AdQuantized(tensor) => tensor.rank(),
158    }
159}
160fn can_mut(&self) -> bool {
161    match self {
162Self::Float(tensor) => tensor.can_mut(),
163Self::AdFloat(tensor) => tensor.can_mut(),
164Self::Int(tensor) => tensor.can_mut(),
165Self::AdInt(tensor) => tensor.can_mut(),
166Self::Bool(tensor) => tensor.can_mut(),
167Self::AdBool(tensor) => tensor.can_mut(),
168Self::Quantized(tensor) => tensor.can_mut(),
169Self::AdQuantized(tensor) => tensor.can_mut(),
170    }
171}
172}
173
174impl<B: Backend> QTensorPrimitive for BackendTensor<B> {
175fn scheme(&self) -> &QuantScheme {
176    match self {
177        Self::Quantized(tensor) => tensor.scheme(),
178        Self::AdQuantized(tensor) => tensor.primitive.scheme(),
179        _ => panic!("Expected a quantized tensor"),
180    }
181}
182fn acc_precision(&self) -> QuantAcc {
183    match self {
184        Self::Quantized(tensor) => tensor.acc_precision(),
185        Self::AdQuantized(tensor) => tensor.primitive.acc_precision(),
186        _ => panic!("Expected a quantized tensor"),
187    }
188}
189fn propagation(&self) -> QuantPropagation {
190    match self {
191        Self::Quantized(tensor) => tensor.propagation(),
192        Self::AdQuantized(tensor) => tensor.primitive.propagation(),
193        _ => panic!("Expected a quantized tensor"),
194    }
195}
196}
197
198#[derive(Clone, Debug)]
199pub enum DispatchTensor {
200Host(BackendTensor<backends::Host>),
201#[cfg(feature = "cuda")]
202Cuda(BackendTensor<backends::Cuda>),
203#[cfg(feature = "cpu")]
204Cpu(BackendTensor<backends::Cpu>),
205#[cfg(feature = "wgpu")]
206Wgpu(BackendTensor<backends::Wgpu>),
207#[cfg(feature = "rocm")]
208Rocm(BackendTensor<backends::Rocm>),
209#[cfg(feature = "tch")]
210LibTorch(BackendTensor<backends::LibTorch>),
211#[cfg(feature = "remote")]
212Remote(BackendTensor<backends::Remote>),
213#[cfg(feature = "capture")]
214Capture(BackendTensor<backends::Capture>),
215}
216
217impl DispatchTensor {
218    /// Identifier to declare a captured tensor as a graph boundary.
219    #[cfg(feature = "capture")]
220    pub fn capture_tensor_id(&self) -> Option<ruda_tensor_capture::TensorId> {
221        match self {
222            Self::Capture(BackendTensor::Float(tensor) | BackendTensor::Int(tensor)
223                | BackendTensor::Bool(tensor) | BackendTensor::Quantized(tensor)) => Some(tensor.id()),
224            _ => None,
225        }
226    }
227
228    pub fn context(&self) -> AutodiffContext {
229        match self {
230Self::Host(tensor) => tensor.context(),
231#[cfg(feature = "cuda")]
232Self::Cuda(tensor) => tensor.context(),
233#[cfg(feature = "cpu")]
234Self::Cpu(tensor) => tensor.context(),
235#[cfg(feature = "wgpu")]
236Self::Wgpu(tensor) => tensor.context(),
237#[cfg(feature = "rocm")]
238Self::Rocm(tensor) => tensor.context(),
239#[cfg(feature = "tch")]
240Self::LibTorch(tensor) => tensor.context(),
241#[cfg(feature = "remote")]
242Self::Remote(tensor) => tensor.context(),
243#[cfg(feature = "capture")]
244Self::Capture(tensor) => tensor.context(),
245        }
246    }
247
248    pub fn with_context(self, context: AutodiffContext) -> Self {
249        match self {
250Self::Host(tensor) => Self::Host(tensor.with_context(context)),
251#[cfg(feature = "cuda")]
252Self::Cuda(tensor) => Self::Cuda(tensor.with_context(context)),
253#[cfg(feature = "cpu")]
254Self::Cpu(tensor) => Self::Cpu(tensor.with_context(context)),
255#[cfg(feature = "wgpu")]
256Self::Wgpu(tensor) => Self::Wgpu(tensor.with_context(context)),
257#[cfg(feature = "rocm")]
258Self::Rocm(tensor) => Self::Rocm(tensor.with_context(context)),
259#[cfg(feature = "tch")]
260Self::LibTorch(tensor) => Self::LibTorch(tensor.with_context(context)),
261#[cfg(feature = "remote")]
262Self::Remote(tensor) => Self::Remote(tensor.with_context(context)),
263#[cfg(feature = "capture")]
264Self::Capture(tensor) => {
265    assert_eq!(context, AutodiffContext::Disabled, "Capture does not support autodiff");
266    Self::Capture(tensor.with_context(context))
267},
268        }
269    }
270
271    pub fn device(&self) -> DispatchDevice {
272        let device = match self {
273Self::Host(tensor) => DispatchDevice::Host(tensor.device()),
274#[cfg(feature = "cuda")]
275Self::Cuda(tensor) => DispatchDevice::Cuda(tensor.device()),
276#[cfg(feature = "cpu")]
277Self::Cpu(tensor) => DispatchDevice::Cpu(tensor.device()),
278#[cfg(feature = "wgpu")]
279Self::Wgpu(tensor) => DispatchDevice::Wgpu(tensor.device()),
280#[cfg(feature = "rocm")]
281Self::Rocm(tensor) => DispatchDevice::Rocm(tensor.device()),
282#[cfg(feature = "tch")]
283Self::LibTorch(tensor) => DispatchDevice::LibTorch(tensor.device()),
284#[cfg(feature = "remote")]
285Self::Remote(tensor) => DispatchDevice::Remote(tensor.device()),
286#[cfg(feature = "capture")]
287Self::Capture(tensor) => DispatchDevice::Capture(tensor.device()),
288        };
289        device.with_context(self.context())
290    }
291}
292
293impl TensorMetadata for DispatchTensor {
294fn dtype(&self) -> DType {
295    match self {
296Self::Host(tensor) => tensor.dtype(),
297#[cfg(feature = "cuda")]
298Self::Cuda(tensor) => tensor.dtype(),
299#[cfg(feature = "cpu")]
300Self::Cpu(tensor) => tensor.dtype(),
301#[cfg(feature = "wgpu")]
302Self::Wgpu(tensor) => tensor.dtype(),
303#[cfg(feature = "rocm")]
304Self::Rocm(tensor) => tensor.dtype(),
305#[cfg(feature = "tch")]
306Self::LibTorch(tensor) => tensor.dtype(),
307#[cfg(feature = "remote")]
308Self::Remote(tensor) => tensor.dtype(),
309#[cfg(feature = "capture")]
310Self::Capture(tensor) => tensor.dtype(),
311    }
312}
313fn shape(&self) -> Shape {
314    match self {
315Self::Host(tensor) => tensor.shape(),
316#[cfg(feature = "cuda")]
317Self::Cuda(tensor) => tensor.shape(),
318#[cfg(feature = "cpu")]
319Self::Cpu(tensor) => tensor.shape(),
320#[cfg(feature = "wgpu")]
321Self::Wgpu(tensor) => tensor.shape(),
322#[cfg(feature = "rocm")]
323Self::Rocm(tensor) => tensor.shape(),
324#[cfg(feature = "tch")]
325Self::LibTorch(tensor) => tensor.shape(),
326#[cfg(feature = "remote")]
327Self::Remote(tensor) => tensor.shape(),
328#[cfg(feature = "capture")]
329Self::Capture(tensor) => tensor.shape(),
330    }
331}
332fn rank(&self) -> usize {
333    match self {
334Self::Host(tensor) => tensor.rank(),
335#[cfg(feature = "cuda")]
336Self::Cuda(tensor) => tensor.rank(),
337#[cfg(feature = "cpu")]
338Self::Cpu(tensor) => tensor.rank(),
339#[cfg(feature = "wgpu")]
340Self::Wgpu(tensor) => tensor.rank(),
341#[cfg(feature = "rocm")]
342Self::Rocm(tensor) => tensor.rank(),
343#[cfg(feature = "tch")]
344Self::LibTorch(tensor) => tensor.rank(),
345#[cfg(feature = "remote")]
346Self::Remote(tensor) => tensor.rank(),
347#[cfg(feature = "capture")]
348Self::Capture(tensor) => tensor.rank(),
349    }
350}
351fn can_mut(&self) -> bool {
352    match self {
353Self::Host(tensor) => tensor.can_mut(),
354#[cfg(feature = "cuda")]
355Self::Cuda(tensor) => tensor.can_mut(),
356#[cfg(feature = "cpu")]
357Self::Cpu(tensor) => tensor.can_mut(),
358#[cfg(feature = "wgpu")]
359Self::Wgpu(tensor) => tensor.can_mut(),
360#[cfg(feature = "rocm")]
361Self::Rocm(tensor) => tensor.can_mut(),
362#[cfg(feature = "tch")]
363Self::LibTorch(tensor) => tensor.can_mut(),
364#[cfg(feature = "remote")]
365Self::Remote(tensor) => tensor.can_mut(),
366#[cfg(feature = "capture")]
367Self::Capture(tensor) => tensor.can_mut(),
368    }
369}
370}
371
372impl QTensorPrimitive for DispatchTensor {
373fn scheme(&self) -> &QuantScheme {
374    match self {
375Self::Host(tensor) => tensor.scheme(),
376#[cfg(feature = "cuda")]
377Self::Cuda(tensor) => tensor.scheme(),
378#[cfg(feature = "cpu")]
379Self::Cpu(tensor) => tensor.scheme(),
380#[cfg(feature = "wgpu")]
381Self::Wgpu(tensor) => tensor.scheme(),
382#[cfg(feature = "rocm")]
383Self::Rocm(tensor) => tensor.scheme(),
384#[cfg(feature = "tch")]
385Self::LibTorch(tensor) => tensor.scheme(),
386#[cfg(feature = "remote")]
387Self::Remote(tensor) => tensor.scheme(),
388#[cfg(feature = "capture")]
389Self::Capture(tensor) => tensor.scheme(),
390    }
391}
392fn acc_precision(&self) -> QuantAcc {
393    match self {
394Self::Host(tensor) => tensor.acc_precision(),
395#[cfg(feature = "cuda")]
396Self::Cuda(tensor) => tensor.acc_precision(),
397#[cfg(feature = "cpu")]
398Self::Cpu(tensor) => tensor.acc_precision(),
399#[cfg(feature = "wgpu")]
400Self::Wgpu(tensor) => tensor.acc_precision(),
401#[cfg(feature = "rocm")]
402Self::Rocm(tensor) => tensor.acc_precision(),
403#[cfg(feature = "tch")]
404Self::LibTorch(tensor) => tensor.acc_precision(),
405#[cfg(feature = "remote")]
406Self::Remote(tensor) => tensor.acc_precision(),
407#[cfg(feature = "capture")]
408Self::Capture(tensor) => tensor.acc_precision(),
409    }
410}
411fn propagation(&self) -> QuantPropagation {
412    match self {
413Self::Host(tensor) => tensor.propagation(),
414#[cfg(feature = "cuda")]
415Self::Cuda(tensor) => tensor.propagation(),
416#[cfg(feature = "cpu")]
417Self::Cpu(tensor) => tensor.propagation(),
418#[cfg(feature = "wgpu")]
419Self::Wgpu(tensor) => tensor.propagation(),
420#[cfg(feature = "rocm")]
421Self::Rocm(tensor) => tensor.propagation(),
422#[cfg(feature = "tch")]
423Self::LibTorch(tensor) => tensor.propagation(),
424#[cfg(feature = "remote")]
425Self::Remote(tensor) => tensor.propagation(),
426#[cfg(feature = "capture")]
427Self::Capture(tensor) => tensor.propagation(),
428    }
429}
430}
ruda_tensor_dispatch/tensor.rs

ruda_tensor_dispatch/
tensor.rs
