hanzo_ml/
storage.rs

1use crate::backend::BackendStorage;
2use crate::custom_op::{all_distinct, InplaceOpN, Src};
3use crate::op::{self, CmpOp, ReduceOp};
4use crate::scalar::Scalar;
5#[cfg(feature = "rocm")]
6use crate::RocmStorage;
7#[cfg(feature = "vulkan")]
8use crate::VulkanStorage;
9#[cfg(feature = "wgpu")]
10use crate::WgpuStorage;
11use crate::{CpuStorage, CudaStorage, DType, Device, Error, Layout, MetalStorage, Result, Shape};
12use crate::{CustomOp1, CustomOp2, CustomOp3};
13use parking_lot::{RwLockReadGuard, RwLockWriteGuard};
14
15// We do not want to implement Clone on Storage as cloning may fail because of
16// out of memory. Instead try_clone should be used.
17#[derive(Debug)]
18pub enum Storage {
19    Cpu(CpuStorage),
20    Cuda(CudaStorage),
21    Metal(MetalStorage),
22    #[cfg(feature = "rocm")]
23    Rocm(RocmStorage),
24    #[cfg(feature = "vulkan")]
25    Vulkan(VulkanStorage),
26    #[cfg(feature = "wgpu")]
27    Wgpu(WgpuStorage),
28}
29
30pub type StorageRef<'a> = RwLockReadGuard<'a, Storage>;
31pub type StorageMutRef<'a> = RwLockWriteGuard<'a, Storage>;
32
33impl Storage {
34    pub fn try_clone(&self, layout: &Layout) -> Result<Self> {
35        match self {
36            Self::Cpu(storage) => Ok(Self::Cpu(storage.clone())),
37            Self::Cuda(storage) => {
38                let storage = storage.try_clone(layout)?;
39                Ok(Self::Cuda(storage))
40            }
41            Self::Metal(storage) => {
42                let storage = storage.try_clone(layout)?;
43                Ok(Self::Metal(storage))
44            }
45            #[cfg(feature = "rocm")]
46            Self::Rocm(storage) => {
47                let storage = storage.try_clone(layout)?;
48                Ok(Self::Rocm(storage))
49            }
50            #[cfg(feature = "vulkan")]
51            Self::Vulkan(storage) => {
52                let storage = storage.try_clone(layout)?;
53                Ok(Self::Vulkan(storage))
54            }
55            #[cfg(feature = "wgpu")]
56            Self::Wgpu(storage) => {
57                let storage = storage.try_clone(layout)?;
58                Ok(Self::Wgpu(storage))
59            }
60        }
61    }
62
63    pub fn device(&self) -> Device {
64        match self {
65            Self::Cpu(_) => Device::Cpu,
66            Self::Cuda(storage) => Device::Cuda(storage.device().clone()),
67            Self::Metal(storage) => Device::Metal(storage.device().clone()),
68            #[cfg(feature = "rocm")]
69            Self::Rocm(storage) => Device::Rocm(storage.device().clone()),
70            #[cfg(feature = "vulkan")]
71            Self::Vulkan(storage) => Device::Vulkan(storage.device().clone()),
72            #[cfg(feature = "wgpu")]
73            Self::Wgpu(storage) => Device::Wgpu(storage.device().clone()),
74        }
75    }
76
77    pub fn dtype(&self) -> DType {
78        match self {
79            Self::Cpu(storage) => storage.dtype(),
80            Self::Cuda(storage) => storage.dtype(),
81            Self::Metal(storage) => storage.dtype(),
82            #[cfg(feature = "rocm")]
83            Self::Rocm(storage) => storage.dtype(),
84            #[cfg(feature = "vulkan")]
85            Self::Vulkan(storage) => storage.dtype(),
86            #[cfg(feature = "wgpu")]
87            Self::Wgpu(storage) => storage.dtype(),
88        }
89    }
90
91    pub(crate) fn same_device(&self, rhs: &Self, op: &'static str) -> Result<()> {
92        let lhs_device = self.device();
93        let rhs_device = rhs.device();
94        let lhs = lhs_device.location();
95        let rhs = rhs_device.location();
96        let same_device = if self.device().is_metal() {
97            // On metal, we require the device to be exactly the same rather than
98            // having the same location. In cuda this is not necessary as all CudaDevice on the
99            // same GPU will use the same cuda stream.
100            lhs_device.same_device(&rhs_device)
101        } else {
102            lhs == rhs
103        };
104        if !same_device {
105            Err(Error::DeviceMismatchBinaryOp { lhs, rhs, op }.bt())
106        } else {
107            Ok(())
108        }
109    }
110
111    pub(crate) fn same_dtype(&self, rhs: &Self, op: &'static str) -> Result<()> {
112        let lhs = self.dtype();
113        let rhs = rhs.dtype();
114        if lhs != rhs {
115            Err(Error::DTypeMismatchBinaryOp { lhs, rhs, op }.bt())
116        } else {
117            Ok(())
118        }
119    }
120
121    pub(crate) fn const_set(&mut self, v: Scalar, l: &Layout) -> Result<()> {
122        match self {
123            Storage::Cpu(storage) => storage.const_set(v, l),
124            Storage::Cuda(storage) => storage.const_set(v, l),
125            Storage::Metal(storage) => storage.const_set(v, l),
126            #[cfg(feature = "rocm")]
127            Storage::Rocm(storage) => storage.const_set(v, l),
128            #[cfg(feature = "vulkan")]
129            Storage::Vulkan(storage) => storage.const_set(v, l),
130            #[cfg(feature = "wgpu")]
131            Storage::Wgpu(storage) => storage.const_set(v, l),
132        }
133    }
134
135    pub(crate) fn affine(&self, layout: &Layout, mul: f64, add: f64) -> Result<Self> {
136        match self {
137            Storage::Cpu(storage) => {
138                let storage = storage.affine(layout, mul, add)?;
139                Ok(Self::Cpu(storage))
140            }
141            Self::Cuda(storage) => {
142                let storage = storage.affine(layout, mul, add)?;
143                Ok(Self::Cuda(storage))
144            }
145            Self::Metal(storage) => {
146                let storage = storage.affine(layout, mul, add)?;
147                Ok(Self::Metal(storage))
148            }
149            #[cfg(feature = "rocm")]
150            Self::Rocm(storage) => {
151                let storage = storage.affine(layout, mul, add)?;
152                Ok(Self::Rocm(storage))
153            }
154            #[cfg(feature = "vulkan")]
155            Self::Vulkan(storage) => {
156                let storage = storage.affine(layout, mul, add)?;
157                Ok(Self::Vulkan(storage))
158            }
159            #[cfg(feature = "wgpu")]
160            Self::Wgpu(storage) => {
161                let storage = storage.affine(layout, mul, add)?;
162                Ok(Self::Wgpu(storage))
163            }
164        }
165    }
166
167    pub(crate) fn powf(&self, layout: &Layout, alpha: f64) -> Result<Self> {
168        match self {
169            Storage::Cpu(storage) => {
170                let storage = storage.powf(layout, alpha)?;
171                Ok(Self::Cpu(storage))
172            }
173            Self::Cuda(storage) => {
174                let storage = storage.powf(layout, alpha)?;
175                Ok(Self::Cuda(storage))
176            }
177            Self::Metal(storage) => {
178                let storage = storage.powf(layout, alpha)?;
179                Ok(Self::Metal(storage))
180            }
181            #[cfg(feature = "rocm")]
182            Self::Rocm(storage) => {
183                let storage = storage.powf(layout, alpha)?;
184                Ok(Self::Rocm(storage))
185            }
186            #[cfg(feature = "vulkan")]
187            Self::Vulkan(storage) => {
188                let storage = storage.powf(layout, alpha)?;
189                Ok(Self::Vulkan(storage))
190            }
191            #[cfg(feature = "wgpu")]
192            Self::Wgpu(storage) => {
193                let storage = storage.powf(layout, alpha)?;
194                Ok(Self::Wgpu(storage))
195            }
196        }
197    }
198
199    pub(crate) fn elu(&self, layout: &Layout, alpha: f64) -> Result<Self> {
200        match self {
201            Storage::Cpu(storage) => {
202                let storage = storage.elu(layout, alpha)?;
203                Ok(Self::Cpu(storage))
204            }
205            Self::Cuda(storage) => {
206                let storage = storage.elu(layout, alpha)?;
207                Ok(Self::Cuda(storage))
208            }
209            Self::Metal(storage) => {
210                let storage = storage.elu(layout, alpha)?;
211                Ok(Self::Metal(storage))
212            }
213            #[cfg(feature = "rocm")]
214            Self::Rocm(storage) => {
215                let storage = storage.elu(layout, alpha)?;
216                Ok(Self::Rocm(storage))
217            }
218            #[cfg(feature = "vulkan")]
219            Self::Vulkan(storage) => {
220                let storage = storage.elu(layout, alpha)?;
221                Ok(Self::Vulkan(storage))
222            }
223            #[cfg(feature = "wgpu")]
224            Self::Wgpu(storage) => {
225                let storage = storage.elu(layout, alpha)?;
226                Ok(Self::Wgpu(storage))
227            }
228        }
229    }
230
231    pub(crate) fn cmp(
232        &self,
233        op: CmpOp,
234        rhs: &Self,
235        lhs_layout: &Layout,
236        rhs_layout: &Layout,
237    ) -> Result<Self> {
238        self.same_device(rhs, "cmp")?;
239        self.same_dtype(rhs, "cmp")?;
240        match (self, rhs) {
241            (Storage::Cpu(lhs), Storage::Cpu(rhs)) => {
242                let storage = lhs.cmp(op, rhs, lhs_layout, rhs_layout)?;
243                Ok(Self::Cpu(storage))
244            }
245            (Self::Cuda(lhs), Self::Cuda(rhs)) => {
246                let storage = lhs.cmp(op, rhs, lhs_layout, rhs_layout)?;
247                Ok(Self::Cuda(storage))
248            }
249            (Self::Metal(lhs), Self::Metal(rhs)) => {
250                let storage = lhs.cmp(op, rhs, lhs_layout, rhs_layout)?;
251                Ok(Self::Metal(storage))
252            }
253            #[cfg(feature = "rocm")]
254            (Self::Rocm(lhs), Self::Rocm(rhs)) => {
255                let storage = lhs.cmp(op, rhs, lhs_layout, rhs_layout)?;
256                Ok(Self::Rocm(storage))
257            }
258            #[cfg(feature = "vulkan")]
259            (Self::Vulkan(lhs), Self::Vulkan(rhs)) => {
260                let storage = lhs.cmp(op, rhs, lhs_layout, rhs_layout)?;
261                Ok(Self::Vulkan(storage))
262            }
263            #[cfg(feature = "wgpu")]
264            (Self::Wgpu(lhs), Self::Wgpu(rhs)) => {
265                let storage = lhs.cmp(op, rhs, lhs_layout, rhs_layout)?;
266                Ok(Self::Wgpu(storage))
267            }
268            (lhs, rhs) => {
269                // Should not happen because of the same device check above but we're defensive
270                // anyway.
271                Err(Error::DeviceMismatchBinaryOp {
272                    lhs: lhs.device().location(),
273                    rhs: rhs.device().location(),
274                    op: "cmp",
275                }
276                .bt())
277            }
278        }
279    }
280
281    pub(crate) fn reduce_op(&self, op: ReduceOp, layout: &Layout, s: &[usize]) -> Result<Self> {
282        match self {
283            Storage::Cpu(storage) => {
284                let storage = storage.reduce_op(op, layout, s)?;
285                Ok(Self::Cpu(storage))
286            }
287            Self::Cuda(storage) => {
288                let storage = storage.reduce_op(op, layout, s)?;
289                Ok(Self::Cuda(storage))
290            }
291            Self::Metal(storage) => {
292                let storage = storage.reduce_op(op, layout, s)?;
293                Ok(Self::Metal(storage))
294            }
295            #[cfg(feature = "rocm")]
296            Self::Rocm(storage) => {
297                let storage = storage.reduce_op(op, layout, s)?;
298                Ok(Self::Rocm(storage))
299            }
300            #[cfg(feature = "vulkan")]
301            Self::Vulkan(storage) => {
302                let storage = storage.reduce_op(op, layout, s)?;
303                Ok(Self::Vulkan(storage))
304            }
305            #[cfg(feature = "wgpu")]
306            Self::Wgpu(storage) => {
307                let storage = storage.reduce_op(op, layout, s)?;
308                Ok(Self::Wgpu(storage))
309            }
310        }
311    }
312
313    pub(crate) fn to_dtype(&self, layout: &Layout, dtype: DType) -> Result<Self> {
314        match self {
315            Storage::Cpu(storage) => {
316                let storage = storage.to_dtype(layout, dtype)?;
317                Ok(Self::Cpu(storage))
318            }
319            Self::Cuda(storage) => {
320                let storage = storage.to_dtype(layout, dtype)?;
321                Ok(Self::Cuda(storage))
322            }
323            Self::Metal(storage) => {
324                let storage = storage.to_dtype(layout, dtype)?;
325                Ok(Self::Metal(storage))
326            }
327            #[cfg(feature = "rocm")]
328            Self::Rocm(storage) => {
329                let storage = storage.to_dtype(layout, dtype)?;
330                Ok(Self::Rocm(storage))
331            }
332            #[cfg(feature = "vulkan")]
333            Self::Vulkan(storage) => {
334                let storage = storage.to_dtype(layout, dtype)?;
335                Ok(Self::Vulkan(storage))
336            }
337            #[cfg(feature = "wgpu")]
338            Self::Wgpu(storage) => {
339                let storage = storage.to_dtype(layout, dtype)?;
340                Ok(Self::Wgpu(storage))
341            }
342        }
343    }
344
345    pub(crate) fn apply_op1(&self, l: &Layout, c: &dyn CustomOp1) -> Result<(Self, Shape)> {
346        match self {
347            Self::Cpu(storage) => {
348                let (storage, shape) = c.cpu_fwd(storage, l)?;
349                Ok((Self::Cpu(storage), shape))
350            }
351            Self::Cuda(storage) => {
352                let (storage, shape) = c.cuda_fwd(storage, l)?;
353                Ok((Self::Cuda(storage), shape))
354            }
355            Self::Metal(storage) => {
356                let (storage, shape) = c.metal_fwd(storage, l)?;
357                Ok((Self::Metal(storage), shape))
358            }
359            #[cfg(feature = "rocm")]
360            Self::Rocm(storage) => {
361                let (storage, shape) = c.rocm_fwd(storage, l)?;
362                Ok((Self::Rocm(storage), shape))
363            }
364            #[cfg(feature = "vulkan")]
365            Self::Vulkan(storage) => {
366                let (storage, shape) = c.vulkan_fwd(storage, l)?;
367                Ok((Self::Vulkan(storage), shape))
368            }
369            #[cfg(feature = "wgpu")]
370            Self::Wgpu(storage) => {
371                let (storage, shape) = c.wgpu_fwd(storage, l)?;
372                Ok((Self::Wgpu(storage), shape))
373            }
374        }
375    }
376
377    pub(crate) fn apply_op2(
378        &self,
379        l1: &Layout,
380        t2: &Self,
381        l2: &Layout,
382        c: &dyn CustomOp2,
383    ) -> Result<(Self, Shape)> {
384        self.same_device(t2, c.name())?;
385        match (self, t2) {
386            (Self::Cpu(s1), Self::Cpu(s2)) => {
387                let (s, shape) = c.cpu_fwd(s1, l1, s2, l2)?;
388                Ok((Self::Cpu(s), shape))
389            }
390            (Self::Cuda(s1), Self::Cuda(s2)) => {
391                let (s, shape) = c.cuda_fwd(s1, l1, s2, l2)?;
392                Ok((Self::Cuda(s), shape))
393            }
394            (Self::Metal(s1), Self::Metal(s2)) => {
395                let (s, shape) = c.metal_fwd(s1, l1, s2, l2)?;
396                Ok((Self::Metal(s), shape))
397            }
398            #[cfg(feature = "rocm")]
399            (Self::Rocm(s1), Self::Rocm(s2)) => {
400                let (s, shape) = c.rocm_fwd(s1, l1, s2, l2)?;
401                Ok((Self::Rocm(s), shape))
402            }
403            #[cfg(feature = "vulkan")]
404            (Self::Vulkan(s1), Self::Vulkan(s2)) => {
405                let (s, shape) = c.vulkan_fwd(s1, l1, s2, l2)?;
406                Ok((Self::Vulkan(s), shape))
407            }
408            #[cfg(feature = "wgpu")]
409            (Self::Wgpu(s1), Self::Wgpu(s2)) => {
410                let (s, shape) = c.wgpu_fwd(s1, l1, s2, l2)?;
411                Ok((Self::Wgpu(s), shape))
412            }
413            _ => unreachable!(),
414        }
415    }
416
417    pub(crate) fn apply_op3(
418        &self,
419        l1: &Layout,
420        t2: &Self,
421        l2: &Layout,
422        t3: &Self,
423        l3: &Layout,
424        c: &dyn CustomOp3,
425    ) -> Result<(Self, Shape)> {
426        self.same_device(t2, c.name())?;
427        self.same_device(t3, c.name())?;
428        match (self, t2, t3) {
429            (Self::Cpu(s1), Self::Cpu(s2), Self::Cpu(s3)) => {
430                let (s, shape) = c.cpu_fwd(s1, l1, s2, l2, s3, l3)?;
431                Ok((Self::Cpu(s), shape))
432            }
433            (Self::Cuda(s1), Self::Cuda(s2), Self::Cuda(s3)) => {
434                let (s, shape) = c.cuda_fwd(s1, l1, s2, l2, s3, l3)?;
435                Ok((Self::Cuda(s), shape))
436            }
437            (Self::Metal(s1), Self::Metal(s2), Self::Metal(s3)) => {
438                let (s, shape) = c.metal_fwd(s1, l1, s2, l2, s3, l3)?;
439                Ok((Self::Metal(s), shape))
440            }
441            #[cfg(feature = "rocm")]
442            (Self::Rocm(s1), Self::Rocm(s2), Self::Rocm(s3)) => {
443                let (s, shape) = c.rocm_fwd(s1, l1, s2, l2, s3, l3)?;
444                Ok((Self::Rocm(s), shape))
445            }
446            #[cfg(feature = "vulkan")]
447            (Self::Vulkan(s1), Self::Vulkan(s2), Self::Vulkan(s3)) => {
448                let (s, shape) = c.vulkan_fwd(s1, l1, s2, l2, s3, l3)?;
449                Ok((Self::Vulkan(s), shape))
450            }
451            #[cfg(feature = "wgpu")]
452            (Self::Wgpu(s1), Self::Wgpu(s2), Self::Wgpu(s3)) => {
453                let (s, shape) = c.wgpu_fwd(s1, l1, s2, l2, s3, l3)?;
454                Ok((Self::Wgpu(s), shape))
455            }
456            _ => unreachable!(),
457        }
458    }
459
460    /// Applies an custom in-place op for the `self` tensor.
461    ///
462    /// [`Src::Aliased`] use the same underlying storage as `self`, while [`Src::Distinct`] does not.
463    /// If there are aliases present the aliased forward function on `InplaceOpN` is called. Otherwise
464    /// the normal forward function can be used.
465    pub(crate) fn inplace_op<const N: usize, C: InplaceOpN<N>>(
466        &mut self,
467        dst_l: &Layout,
468        srcs: [(Src<'_, Storage>, &Layout); N],
469        c: &C,
470    ) -> Result<()> {
471        // Aliased sources are `self`. Only distinct needs checking.
472        for (s, _) in srcs.iter() {
473            if let Src::Distinct(s) = s {
474                self.same_device(s, c.name())?;
475            }
476        }
477        macro_rules! inplace_dispatch {
478            ($dst:expr, $variant:ident, $fwd:ident, $fwd_aliased:ident) => {{
479                // Extract underlying storage variant on same device
480                let operands: [(Src<'_, _>, &Layout); N] = std::array::from_fn(|i| {
481                    let (s, l) = srcs[i];
482                    let s = match s {
483                        Src::Distinct(Storage::$variant(s)) => Src::Distinct(s),
484                        Src::Aliased(rel) => Src::Aliased(rel),
485                        Src::Distinct(_) => {
486                            unreachable!("same_device above rejects mismatched backends")
487                        }
488                    };
489                    (s, l)
490                });
491
492                match all_distinct(&operands) {
493                    Some(distinct) => c.$fwd($dst, dst_l, distinct),
494                    None => c.$fwd_aliased($dst, dst_l, operands),
495                }
496            }};
497        }
498
499        match self {
500            Storage::Cpu(dst) => inplace_dispatch!(dst, Cpu, cpu_fwd, cpu_fwd_aliased),
501            Storage::Cuda(dst) => inplace_dispatch!(dst, Cuda, cuda_fwd, cuda_fwd_aliased),
502            Storage::Metal(dst) => inplace_dispatch!(dst, Metal, metal_fwd, metal_fwd_aliased),
503            #[cfg(feature = "rocm")]
504            Storage::Rocm(dst) => inplace_dispatch!(dst, Rocm, rocm_fwd, rocm_fwd_aliased),
505            #[cfg(feature = "vulkan")]
506            Storage::Vulkan(dst) => inplace_dispatch!(dst, Vulkan, vulkan_fwd, vulkan_fwd_aliased),
507            #[cfg(feature = "wgpu")]
508            Storage::Wgpu(dst) => inplace_dispatch!(dst, Wgpu, wgpu_fwd, wgpu_fwd_aliased),
509        }
510    }
511
512    pub(crate) fn unary_impl<B: op::UnaryOpT>(&self, layout: &Layout) -> Result<Self> {
513        match self {
514            Storage::Cpu(storage) => {
515                let storage = storage.unary_impl::<B>(layout)?;
516                Ok(Self::Cpu(storage))
517            }
518            Self::Cuda(storage) => {
519                let storage = storage.unary_impl::<B>(layout)?;
520                Ok(Self::Cuda(storage))
521            }
522            Self::Metal(storage) => {
523                let storage = storage.unary_impl::<B>(layout)?;
524                Ok(Self::Metal(storage))
525            }
526            #[cfg(feature = "rocm")]
527            Self::Rocm(storage) => {
528                let storage = storage.unary_impl::<B>(layout)?;
529                Ok(Self::Rocm(storage))
530            }
531            #[cfg(feature = "vulkan")]
532            Self::Vulkan(storage) => {
533                let storage = storage.unary_impl::<B>(layout)?;
534                Ok(Self::Vulkan(storage))
535            }
536            #[cfg(feature = "wgpu")]
537            Self::Wgpu(storage) => {
538                let storage = storage.unary_impl::<B>(layout)?;
539                Ok(Self::Wgpu(storage))
540            }
541        }
542    }
543
544    pub(crate) fn binary_impl<B: op::BinaryOpT>(
545        &self,
546        rhs: &Self,
547        lhs_layout: &Layout,
548        rhs_layout: &Layout,
549    ) -> Result<Self> {
550        self.same_device(rhs, B::NAME)?;
551        self.same_dtype(rhs, B::NAME)?;
552        match (self, rhs) {
553            (Storage::Cpu(lhs), Storage::Cpu(rhs)) => {
554                let storage = lhs.binary_impl::<B>(rhs, lhs_layout, rhs_layout)?;
555                Ok(Self::Cpu(storage))
556            }
557            (Self::Cuda(lhs), Self::Cuda(rhs)) => {
558                let storage = lhs.binary_impl::<B>(rhs, lhs_layout, rhs_layout)?;
559                Ok(Self::Cuda(storage))
560            }
561            (Self::Metal(lhs), Self::Metal(rhs)) => {
562                let storage = lhs.binary_impl::<B>(rhs, lhs_layout, rhs_layout)?;
563                Ok(Self::Metal(storage))
564            }
565            #[cfg(feature = "rocm")]
566            (Self::Rocm(lhs), Self::Rocm(rhs)) => {
567                let storage = lhs.binary_impl::<B>(rhs, lhs_layout, rhs_layout)?;
568                Ok(Self::Rocm(storage))
569            }
570            #[cfg(feature = "vulkan")]
571            (Self::Vulkan(lhs), Self::Vulkan(rhs)) => {
572                let storage = lhs.binary_impl::<B>(rhs, lhs_layout, rhs_layout)?;
573                Ok(Self::Vulkan(storage))
574            }
575            #[cfg(feature = "wgpu")]
576            (Self::Wgpu(lhs), Self::Wgpu(rhs)) => {
577                let storage = lhs.binary_impl::<B>(rhs, lhs_layout, rhs_layout)?;
578                Ok(Self::Wgpu(storage))
579            }
580            (lhs, rhs) => {
581                // Should not happen because of the same device check above but we're defensive
582                // anyway.
583                Err(Error::DeviceMismatchBinaryOp {
584                    lhs: lhs.device().location(),
585                    rhs: rhs.device().location(),
586                    op: B::NAME,
587                }
588                .bt())
589            }
590        }
591    }
592
593    pub(crate) fn conv1d(
594        &self,
595        l: &Layout,
596        kernel: &Self,
597        kernel_l: &Layout,
598        params: &crate::conv::ParamsConv1D,
599    ) -> Result<Self> {
600        self.same_device(kernel, "conv1d")?;
601        self.same_dtype(kernel, "conv1d")?;
602        match (self, &kernel) {
603            (Storage::Cpu(inp), Storage::Cpu(kernel)) => {
604                let s = inp.conv1d(l, kernel, kernel_l, params)?;
605                Ok(Self::Cpu(s))
606            }
607            (Storage::Cuda(inp), Storage::Cuda(kernel)) => {
608                let s = inp.conv1d(l, kernel, kernel_l, params)?;
609                Ok(Self::Cuda(s))
610            }
611            (Storage::Metal(inp), Storage::Metal(kernel)) => {
612                let s = inp.conv1d(l, kernel, kernel_l, params)?;
613                Ok(Self::Metal(s))
614            }
615            #[cfg(feature = "rocm")]
616            (Storage::Rocm(inp), Storage::Rocm(kernel)) => {
617                let s = inp.conv1d(l, kernel, kernel_l, params)?;
618                Ok(Self::Rocm(s))
619            }
620            #[cfg(feature = "vulkan")]
621            (Storage::Vulkan(inp), Storage::Vulkan(kernel)) => {
622                let s = inp.conv1d(l, kernel, kernel_l, params)?;
623                Ok(Self::Vulkan(s))
624            }
625            #[cfg(feature = "wgpu")]
626            (Storage::Wgpu(inp), Storage::Wgpu(kernel)) => {
627                let s = inp.conv1d(l, kernel, kernel_l, params)?;
628                Ok(Self::Wgpu(s))
629            }
630            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
631                lhs: lhs.device().location(),
632                rhs: rhs.device().location(),
633                op: "conv1d",
634            }
635            .bt()),
636        }
637    }
638
639    pub(crate) fn conv_transpose1d(
640        &self,
641        l: &Layout,
642        kernel: &Self,
643        kernel_l: &Layout,
644        params: &crate::conv::ParamsConvTranspose1D,
645    ) -> Result<Self> {
646        self.same_device(kernel, "conv-transpose1d")?;
647        self.same_dtype(kernel, "conv-transpose1d")?;
648        match (self, &kernel) {
649            (Storage::Cpu(inp), Storage::Cpu(kernel)) => {
650                let s = inp.conv_transpose1d(l, kernel, kernel_l, params)?;
651                Ok(Self::Cpu(s))
652            }
653            (Storage::Cuda(inp), Storage::Cuda(kernel)) => {
654                let s = inp.conv_transpose1d(l, kernel, kernel_l, params)?;
655                Ok(Self::Cuda(s))
656            }
657            (Storage::Metal(inp), Storage::Metal(kernel)) => {
658                let s = inp.conv_transpose1d(l, kernel, kernel_l, params)?;
659                Ok(Self::Metal(s))
660            }
661            #[cfg(feature = "rocm")]
662            (Storage::Rocm(inp), Storage::Rocm(kernel)) => {
663                let s = inp.conv_transpose1d(l, kernel, kernel_l, params)?;
664                Ok(Self::Rocm(s))
665            }
666            #[cfg(feature = "vulkan")]
667            (Storage::Vulkan(inp), Storage::Vulkan(kernel)) => {
668                let s = inp.conv_transpose1d(l, kernel, kernel_l, params)?;
669                Ok(Self::Vulkan(s))
670            }
671            #[cfg(feature = "wgpu")]
672            (Storage::Wgpu(inp), Storage::Wgpu(kernel)) => {
673                let s = inp.conv_transpose1d(l, kernel, kernel_l, params)?;
674                Ok(Self::Wgpu(s))
675            }
676            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
677                lhs: lhs.device().location(),
678                rhs: rhs.device().location(),
679                op: "conv-transpose1d",
680            }
681            .bt()),
682        }
683    }
684
685    pub(crate) fn conv2d(
686        &self,
687        l: &Layout,
688        kernel: &Self,
689        kernel_l: &Layout,
690        params: &crate::conv::ParamsConv2D,
691    ) -> Result<Self> {
692        self.same_device(kernel, "conv2d")?;
693        self.same_dtype(kernel, "conv2d")?;
694        match (self, &kernel) {
695            (Storage::Cpu(inp), Storage::Cpu(kernel)) => {
696                let s = inp.conv2d(l, kernel, kernel_l, params)?;
697                Ok(Self::Cpu(s))
698            }
699            (Storage::Cuda(inp), Storage::Cuda(kernel)) => {
700                let s = inp.conv2d(l, kernel, kernel_l, params)?;
701                Ok(Self::Cuda(s))
702            }
703            (Storage::Metal(inp), Storage::Metal(kernel)) => {
704                let s = inp.conv2d(l, kernel, kernel_l, params)?;
705                Ok(Self::Metal(s))
706            }
707            #[cfg(feature = "rocm")]
708            (Storage::Rocm(inp), Storage::Rocm(kernel)) => {
709                let s = inp.conv2d(l, kernel, kernel_l, params)?;
710                Ok(Self::Rocm(s))
711            }
712            #[cfg(feature = "vulkan")]
713            (Storage::Vulkan(inp), Storage::Vulkan(kernel)) => {
714                let s = inp.conv2d(l, kernel, kernel_l, params)?;
715                Ok(Self::Vulkan(s))
716            }
717            #[cfg(feature = "wgpu")]
718            (Storage::Wgpu(inp), Storage::Wgpu(kernel)) => {
719                let s = inp.conv2d(l, kernel, kernel_l, params)?;
720                Ok(Self::Wgpu(s))
721            }
722            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
723                lhs: lhs.device().location(),
724                rhs: rhs.device().location(),
725                op: "conv2d",
726            }
727            .bt()),
728        }
729    }
730
731    pub(crate) fn conv_transpose2d(
732        &self,
733        l: &Layout,
734        kernel: &Self,
735        kernel_l: &Layout,
736        params: &crate::conv::ParamsConvTranspose2D,
737    ) -> Result<Self> {
738        self.same_device(kernel, "conv_transpose2d")?;
739        self.same_dtype(kernel, "conv_transpose2d")?;
740        match (self, &kernel) {
741            (Storage::Cpu(inp), Storage::Cpu(kernel)) => {
742                let s = inp.conv_transpose2d(l, kernel, kernel_l, params)?;
743                Ok(Self::Cpu(s))
744            }
745            (Storage::Cuda(inp), Storage::Cuda(kernel)) => {
746                let s = inp.conv_transpose2d(l, kernel, kernel_l, params)?;
747                Ok(Self::Cuda(s))
748            }
749            (Storage::Metal(inp), Storage::Metal(kernel)) => {
750                let s = inp.conv_transpose2d(l, kernel, kernel_l, params)?;
751                Ok(Self::Metal(s))
752            }
753            #[cfg(feature = "rocm")]
754            (Storage::Rocm(inp), Storage::Rocm(kernel)) => {
755                let s = inp.conv_transpose2d(l, kernel, kernel_l, params)?;
756                Ok(Self::Rocm(s))
757            }
758            #[cfg(feature = "vulkan")]
759            (Storage::Vulkan(inp), Storage::Vulkan(kernel)) => {
760                let s = inp.conv_transpose2d(l, kernel, kernel_l, params)?;
761                Ok(Self::Vulkan(s))
762            }
763            #[cfg(feature = "wgpu")]
764            (Storage::Wgpu(inp), Storage::Wgpu(kernel)) => {
765                let s = inp.conv_transpose2d(l, kernel, kernel_l, params)?;
766                Ok(Self::Wgpu(s))
767            }
768            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
769                lhs: lhs.device().location(),
770                rhs: rhs.device().location(),
771                op: "conv_transpose2d",
772            }
773            .bt()),
774        }
775    }
776
777    pub(crate) fn avg_pool2d(
778        &self,
779        layout: &Layout,
780        kernel_size: (usize, usize),
781        stride: (usize, usize),
782    ) -> Result<Self> {
783        match self {
784            Storage::Cpu(storage) => {
785                let storage = storage.avg_pool2d(layout, kernel_size, stride)?;
786                Ok(Self::Cpu(storage))
787            }
788            Self::Cuda(storage) => {
789                let storage = storage.avg_pool2d(layout, kernel_size, stride)?;
790                Ok(Self::Cuda(storage))
791            }
792            Self::Metal(storage) => {
793                let storage = storage.avg_pool2d(layout, kernel_size, stride)?;
794                Ok(Self::Metal(storage))
795            }
796            #[cfg(feature = "rocm")]
797            Self::Rocm(storage) => {
798                let storage = storage.avg_pool2d(layout, kernel_size, stride)?;
799                Ok(Self::Rocm(storage))
800            }
801            #[cfg(feature = "vulkan")]
802            Self::Vulkan(storage) => {
803                let storage = storage.avg_pool2d(layout, kernel_size, stride)?;
804                Ok(Self::Vulkan(storage))
805            }
806            #[cfg(feature = "wgpu")]
807            Self::Wgpu(storage) => {
808                let storage = storage.avg_pool2d(layout, kernel_size, stride)?;
809                Ok(Self::Wgpu(storage))
810            }
811        }
812    }
813
814    pub(crate) fn max_pool2d(
815        &self,
816        layout: &Layout,
817        kernel_size: (usize, usize),
818        stride: (usize, usize),
819    ) -> Result<Self> {
820        match self {
821            Storage::Cpu(storage) => {
822                let storage = storage.max_pool2d(layout, kernel_size, stride)?;
823                Ok(Self::Cpu(storage))
824            }
825            Self::Cuda(storage) => {
826                let storage = storage.max_pool2d(layout, kernel_size, stride)?;
827                Ok(Self::Cuda(storage))
828            }
829            Self::Metal(storage) => {
830                let storage = storage.max_pool2d(layout, kernel_size, stride)?;
831                Ok(Self::Metal(storage))
832            }
833            #[cfg(feature = "rocm")]
834            Self::Rocm(storage) => {
835                let storage = storage.max_pool2d(layout, kernel_size, stride)?;
836                Ok(Self::Rocm(storage))
837            }
838            #[cfg(feature = "vulkan")]
839            Self::Vulkan(storage) => {
840                let storage = storage.max_pool2d(layout, kernel_size, stride)?;
841                Ok(Self::Vulkan(storage))
842            }
843            #[cfg(feature = "wgpu")]
844            Self::Wgpu(storage) => {
845                let storage = storage.max_pool2d(layout, kernel_size, stride)?;
846                Ok(Self::Wgpu(storage))
847            }
848        }
849    }
850
851    pub(crate) fn upsample_nearest1d(&self, layout: &Layout, sz: usize) -> Result<Self> {
852        match self {
853            Storage::Cpu(storage) => {
854                let storage = storage.upsample_nearest1d(layout, sz)?;
855                Ok(Self::Cpu(storage))
856            }
857            Self::Cuda(storage) => {
858                let storage = storage.upsample_nearest1d(layout, sz)?;
859                Ok(Self::Cuda(storage))
860            }
861            Self::Metal(storage) => {
862                let storage = storage.upsample_nearest1d(layout, sz)?;
863                Ok(Self::Metal(storage))
864            }
865            #[cfg(feature = "rocm")]
866            Self::Rocm(storage) => {
867                let storage = storage.upsample_nearest1d(layout, sz)?;
868                Ok(Self::Rocm(storage))
869            }
870            #[cfg(feature = "vulkan")]
871            Self::Vulkan(storage) => {
872                let storage = storage.upsample_nearest1d(layout, sz)?;
873                Ok(Self::Vulkan(storage))
874            }
875            #[cfg(feature = "wgpu")]
876            Self::Wgpu(storage) => {
877                let storage = storage.upsample_nearest1d(layout, sz)?;
878                Ok(Self::Wgpu(storage))
879            }
880        }
881    }
882
883    pub(crate) fn upsample_nearest2d(&self, layout: &Layout, h: usize, w: usize) -> Result<Self> {
884        match self {
885            Storage::Cpu(storage) => {
886                let storage = storage.upsample_nearest2d(layout, h, w)?;
887                Ok(Self::Cpu(storage))
888            }
889            Self::Cuda(storage) => {
890                let storage = storage.upsample_nearest2d(layout, h, w)?;
891                Ok(Self::Cuda(storage))
892            }
893            Self::Metal(storage) => {
894                let storage = storage.upsample_nearest2d(layout, h, w)?;
895                Ok(Self::Metal(storage))
896            }
897            #[cfg(feature = "rocm")]
898            Self::Rocm(storage) => {
899                let storage = storage.upsample_nearest2d(layout, h, w)?;
900                Ok(Self::Rocm(storage))
901            }
902            #[cfg(feature = "vulkan")]
903            Self::Vulkan(storage) => {
904                let storage = storage.upsample_nearest2d(layout, h, w)?;
905                Ok(Self::Vulkan(storage))
906            }
907            #[cfg(feature = "wgpu")]
908            Self::Wgpu(storage) => {
909                let storage = storage.upsample_nearest2d(layout, h, w)?;
910                Ok(Self::Wgpu(storage))
911            }
912        }
913    }
914
915    pub(crate) fn upsample_bilinear2d(
916        &self,
917        layout: &Layout,
918        h: usize,
919        w: usize,
920        align_corners: bool,
921        scale_h: Option<f64>,
922        scale_w: Option<f64>,
923    ) -> Result<Self> {
924        match self {
925            Storage::Cpu(storage) => {
926                let storage =
927                    storage.upsample_bilinear2d(layout, h, w, align_corners, scale_h, scale_w)?;
928                Ok(Self::Cpu(storage))
929            }
930            Self::Cuda(storage) => {
931                let storage =
932                    storage.upsample_bilinear2d(layout, h, w, align_corners, scale_h, scale_w)?;
933                Ok(Self::Cuda(storage))
934            }
935            Self::Metal(storage) => {
936                let storage =
937                    storage.upsample_bilinear2d(layout, h, w, align_corners, scale_h, scale_w)?;
938                Ok(Self::Metal(storage))
939            }
940            #[cfg(feature = "rocm")]
941            Self::Rocm(storage) => {
942                let storage =
943                    storage.upsample_bilinear2d(layout, h, w, align_corners, scale_h, scale_w)?;
944                Ok(Self::Rocm(storage))
945            }
946            #[cfg(feature = "vulkan")]
947            Self::Vulkan(storage) => {
948                let storage =
949                    storage.upsample_bilinear2d(layout, h, w, align_corners, scale_h, scale_w)?;
950                Ok(Self::Vulkan(storage))
951            }
952            #[cfg(feature = "wgpu")]
953            Self::Wgpu(storage) => {
954                let storage =
955                    storage.upsample_bilinear2d(layout, h, w, align_corners, scale_h, scale_w)?;
956                Ok(Self::Wgpu(storage))
957            }
958        }
959    }
960
961    pub(crate) fn where_cond(
962        &self,
963        layout: &Layout,
964        t: &Self,
965        layout_t: &Layout,
966        f: &Self,
967        layout_f: &Layout,
968    ) -> Result<Self> {
969        self.same_device(t, "where")?;
970        self.same_device(f, "where")?;
971        t.same_dtype(f, "where")?;
972        match (self, t, f) {
973            (Storage::Cpu(cond), Storage::Cpu(t), Storage::Cpu(f)) => {
974                let storage = cond.where_cond(layout, t, layout_t, f, layout_f)?;
975                Ok(Self::Cpu(storage))
976            }
977            (Self::Cuda(cond), Self::Cuda(t), Self::Cuda(f)) => {
978                let storage = cond.where_cond(layout, t, layout_t, f, layout_f)?;
979                Ok(Self::Cuda(storage))
980            }
981            (Self::Metal(cond), Self::Metal(t), Self::Metal(f)) => {
982                let storage = cond.where_cond(layout, t, layout_t, f, layout_f)?;
983                Ok(Self::Metal(storage))
984            }
985            #[cfg(feature = "rocm")]
986            (Self::Rocm(cond), Self::Rocm(t), Self::Rocm(f)) => {
987                let storage = cond.where_cond(layout, t, layout_t, f, layout_f)?;
988                Ok(Self::Rocm(storage))
989            }
990            #[cfg(feature = "vulkan")]
991            (Self::Vulkan(cond), Self::Vulkan(t), Self::Vulkan(f)) => {
992                let storage = cond.where_cond(layout, t, layout_t, f, layout_f)?;
993                Ok(Self::Vulkan(storage))
994            }
995            #[cfg(feature = "wgpu")]
996            (Self::Wgpu(cond), Self::Wgpu(t), Self::Wgpu(f)) => {
997                let storage = cond.where_cond(layout, t, layout_t, f, layout_f)?;
998                Ok(Self::Wgpu(storage))
999            }
1000            (_, lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
1001                lhs: lhs.device().location(),
1002                rhs: rhs.device().location(),
1003                op: "where",
1004            }
1005            .bt()),
1006        }
1007    }
1008
1009    pub(crate) fn gather(
1010        &self,
1011        l: &Layout,
1012        indexes: &Self,
1013        indexes_l: &Layout,
1014        d: usize,
1015    ) -> Result<Self> {
1016        self.same_device(indexes, "index-add")?;
1017        match (self, indexes) {
1018            (Self::Cpu(s), Self::Cpu(indexes)) => {
1019                let storage = s.gather(l, indexes, indexes_l, d)?;
1020                Ok(Self::Cpu(storage))
1021            }
1022            (Self::Cuda(s), Self::Cuda(indexes)) => {
1023                let storage = s.gather(l, indexes, indexes_l, d)?;
1024                Ok(Self::Cuda(storage))
1025            }
1026            (Self::Metal(s), Self::Metal(indexes)) => {
1027                let storage = s.gather(l, indexes, indexes_l, d)?;
1028                Ok(Self::Metal(storage))
1029            }
1030            #[cfg(feature = "rocm")]
1031            (Self::Rocm(s), Self::Rocm(indexes)) => {
1032                let storage = s.gather(l, indexes, indexes_l, d)?;
1033                Ok(Self::Rocm(storage))
1034            }
1035            #[cfg(feature = "vulkan")]
1036            (Self::Vulkan(s), Self::Vulkan(indexes)) => {
1037                let storage = s.gather(l, indexes, indexes_l, d)?;
1038                Ok(Self::Vulkan(storage))
1039            }
1040            #[cfg(feature = "wgpu")]
1041            (Self::Wgpu(s), Self::Wgpu(indexes)) => {
1042                let storage = s.gather(l, indexes, indexes_l, d)?;
1043                Ok(Self::Wgpu(storage))
1044            }
1045            _ => unreachable!(),
1046        }
1047    }
1048
1049    pub(crate) fn scatter_set(
1050        &mut self,
1051        l: &Layout,
1052        indexes: &Self,
1053        indexes_l: &Layout,
1054        source: &Self,
1055        source_l: &Layout,
1056        d: usize,
1057    ) -> Result<()> {
1058        self.same_device(indexes, "scatter-set")?;
1059        self.same_device(source, "scatter-set")?;
1060        match (self, indexes, source) {
1061            (Self::Cpu(s), Self::Cpu(indexes), Self::Cpu(source)) => {
1062                s.scatter_set(l, indexes, indexes_l, source, source_l, d)?;
1063            }
1064            (Self::Cuda(s), Self::Cuda(indexes), Self::Cuda(source)) => {
1065                s.scatter_set(l, indexes, indexes_l, source, source_l, d)?;
1066            }
1067            (Self::Metal(s), Self::Metal(indexes), Self::Metal(source)) => {
1068                s.scatter_set(l, indexes, indexes_l, source, source_l, d)?;
1069            }
1070            #[cfg(feature = "rocm")]
1071            (Self::Rocm(s), Self::Rocm(indexes), Self::Rocm(source)) => {
1072                s.scatter_set(l, indexes, indexes_l, source, source_l, d)?;
1073            }
1074            #[cfg(feature = "vulkan")]
1075            (Self::Vulkan(s), Self::Vulkan(indexes), Self::Vulkan(source)) => {
1076                s.scatter_set(l, indexes, indexes_l, source, source_l, d)?;
1077            }
1078            #[cfg(feature = "wgpu")]
1079            (Self::Wgpu(s), Self::Wgpu(indexes), Self::Wgpu(source)) => {
1080                s.scatter_set(l, indexes, indexes_l, source, source_l, d)?;
1081            }
1082            _ => unreachable!(),
1083        }
1084        Ok(())
1085    }
1086
1087    pub(crate) fn scatter_add(
1088        &mut self,
1089        l: &Layout,
1090        indexes: &Self,
1091        indexes_l: &Layout,
1092        source: &Self,
1093        source_l: &Layout,
1094        d: usize,
1095    ) -> Result<()> {
1096        self.same_device(indexes, "scatter-add")?;
1097        self.same_device(source, "scatter-add")?;
1098        match (self, indexes, source) {
1099            (Self::Cpu(s), Self::Cpu(indexes), Self::Cpu(source)) => {
1100                s.scatter_add_set(l, indexes, indexes_l, source, source_l, d)?;
1101            }
1102            (Self::Cuda(s), Self::Cuda(indexes), Self::Cuda(source)) => {
1103                s.scatter_add_set(l, indexes, indexes_l, source, source_l, d)?;
1104            }
1105            (Self::Metal(s), Self::Metal(indexes), Self::Metal(source)) => {
1106                s.scatter_add_set(l, indexes, indexes_l, source, source_l, d)?;
1107            }
1108            #[cfg(feature = "rocm")]
1109            (Self::Rocm(s), Self::Rocm(indexes), Self::Rocm(source)) => {
1110                s.scatter_add_set(l, indexes, indexes_l, source, source_l, d)?;
1111            }
1112            #[cfg(feature = "vulkan")]
1113            (Self::Vulkan(s), Self::Vulkan(indexes), Self::Vulkan(source)) => {
1114                s.scatter_add_set(l, indexes, indexes_l, source, source_l, d)?;
1115            }
1116            #[cfg(feature = "wgpu")]
1117            (Self::Wgpu(s), Self::Wgpu(indexes), Self::Wgpu(source)) => {
1118                s.scatter_add_set(l, indexes, indexes_l, source, source_l, d)?;
1119            }
1120            _ => unreachable!(),
1121        }
1122        Ok(())
1123    }
1124
1125    pub(crate) fn index_add(
1126        &self,
1127        l: &Layout,
1128        indexes: &Self,
1129        indexes_l: &Layout,
1130        source: &Self,
1131        source_l: &Layout,
1132        d: usize,
1133    ) -> Result<Self> {
1134        self.same_device(indexes, "index-add")?;
1135        self.same_device(source, "index-add")?;
1136        match (self, indexes, source) {
1137            (Self::Cpu(s), Self::Cpu(indexes), Self::Cpu(source)) => {
1138                let storage = s.index_add(l, indexes, indexes_l, source, source_l, d)?;
1139                Ok(Self::Cpu(storage))
1140            }
1141            (Self::Cuda(s), Self::Cuda(indexes), Self::Cuda(source)) => {
1142                let storage = s.index_add(l, indexes, indexes_l, source, source_l, d)?;
1143                Ok(Self::Cuda(storage))
1144            }
1145            (Self::Metal(s), Self::Metal(indexes), Self::Metal(source)) => {
1146                let storage = s.index_add(l, indexes, indexes_l, source, source_l, d)?;
1147                Ok(Self::Metal(storage))
1148            }
1149            #[cfg(feature = "rocm")]
1150            (Self::Rocm(s), Self::Rocm(indexes), Self::Rocm(source)) => {
1151                let storage = s.index_add(l, indexes, indexes_l, source, source_l, d)?;
1152                Ok(Self::Rocm(storage))
1153            }
1154            #[cfg(feature = "vulkan")]
1155            (Self::Vulkan(s), Self::Vulkan(indexes), Self::Vulkan(source)) => {
1156                let storage = s.index_add(l, indexes, indexes_l, source, source_l, d)?;
1157                Ok(Self::Vulkan(storage))
1158            }
1159            #[cfg(feature = "wgpu")]
1160            (Self::Wgpu(s), Self::Wgpu(indexes), Self::Wgpu(source)) => {
1161                let storage = s.index_add(l, indexes, indexes_l, source, source_l, d)?;
1162                Ok(Self::Wgpu(storage))
1163            }
1164            _ => unreachable!(),
1165        }
1166    }
1167
1168    pub(crate) fn index_select(
1169        &self,
1170        rhs: &Self,
1171        lhs_l: &Layout,
1172        rhs_l: &Layout,
1173        d: usize,
1174    ) -> Result<Self> {
1175        self.same_device(rhs, "index-select")?;
1176        match (self, rhs) {
1177            (Self::Cpu(lhs), Self::Cpu(rhs)) => {
1178                let storage = lhs.index_select(rhs, lhs_l, rhs_l, d)?;
1179                Ok(Self::Cpu(storage))
1180            }
1181            (Self::Cuda(lhs), Self::Cuda(rhs)) => {
1182                let storage = lhs.index_select(rhs, lhs_l, rhs_l, d)?;
1183                Ok(Self::Cuda(storage))
1184            }
1185            (Self::Metal(lhs), Self::Metal(rhs)) => {
1186                let storage = lhs.index_select(rhs, lhs_l, rhs_l, d)?;
1187                Ok(Self::Metal(storage))
1188            }
1189            #[cfg(feature = "rocm")]
1190            (Self::Rocm(lhs), Self::Rocm(rhs)) => {
1191                let storage = lhs.index_select(rhs, lhs_l, rhs_l, d)?;
1192                Ok(Self::Rocm(storage))
1193            }
1194            #[cfg(feature = "vulkan")]
1195            (Self::Vulkan(lhs), Self::Vulkan(rhs)) => {
1196                let storage = lhs.index_select(rhs, lhs_l, rhs_l, d)?;
1197                Ok(Self::Vulkan(storage))
1198            }
1199            #[cfg(feature = "wgpu")]
1200            (Self::Wgpu(lhs), Self::Wgpu(rhs)) => {
1201                let storage = lhs.index_select(rhs, lhs_l, rhs_l, d)?;
1202                Ok(Self::Wgpu(storage))
1203            }
1204            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
1205                lhs: lhs.device().location(),
1206                rhs: rhs.device().location(),
1207                op: "index-select",
1208            }
1209            .bt()),
1210        }
1211    }
1212
1213    pub(crate) fn matmul(
1214        &self,
1215        rhs: &Self,
1216        bmnk: (usize, usize, usize, usize),
1217        lhs_layout: &Layout,
1218        rhs_layout: &Layout,
1219    ) -> Result<Self> {
1220        self.same_device(rhs, "matmul")?;
1221        self.same_dtype(rhs, "matmul")?;
1222        match (self, rhs) {
1223            (Self::Cpu(lhs), Self::Cpu(rhs)) => {
1224                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
1225                Ok(Self::Cpu(storage))
1226            }
1227            (Self::Cuda(lhs), Self::Cuda(rhs)) => {
1228                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
1229                Ok(Self::Cuda(storage))
1230            }
1231            (Self::Metal(lhs), Self::Metal(rhs)) => {
1232                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
1233                Ok(Self::Metal(storage))
1234            }
1235            #[cfg(feature = "rocm")]
1236            (Self::Rocm(lhs), Self::Rocm(rhs)) => {
1237                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
1238                Ok(Self::Rocm(storage))
1239            }
1240            #[cfg(feature = "vulkan")]
1241            (Self::Vulkan(lhs), Self::Vulkan(rhs)) => {
1242                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
1243                Ok(Self::Vulkan(storage))
1244            }
1245            #[cfg(feature = "wgpu")]
1246            (Self::Wgpu(lhs), Self::Wgpu(rhs)) => {
1247                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
1248                Ok(Self::Wgpu(storage))
1249            }
1250            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
1251                lhs: lhs.device().location(),
1252                rhs: rhs.device().location(),
1253                op: "matmul",
1254            }
1255            .bt()),
1256        }
1257    }
1258
1259    // self, the source can be strided whereas dst is contiguous.
1260    pub(crate) fn copy_strided_src(
1261        &self,
1262        dst: &mut Self,
1263        dst_offset: usize,
1264        src_l: &Layout,
1265    ) -> Result<()> {
1266        match (self, dst) {
1267            (Self::Cpu(src), Self::Cpu(dst)) => src.copy_strided_src(dst, dst_offset, src_l),
1268            (Self::Cuda(src), Self::Cuda(dst)) => Ok(src.copy_strided_src(dst, dst_offset, src_l)?),
1269            (Self::Metal(src), Self::Metal(dst)) => {
1270                Ok(src.copy_strided_src(dst, dst_offset, src_l)?)
1271            }
1272            #[cfg(feature = "rocm")]
1273            (Self::Rocm(src), Self::Rocm(dst)) => Ok(src.copy_strided_src(dst, dst_offset, src_l)?),
1274            #[cfg(feature = "vulkan")]
1275            (Self::Vulkan(src), Self::Vulkan(dst)) => {
1276                Ok(src.copy_strided_src(dst, dst_offset, src_l)?)
1277            }
1278            #[cfg(feature = "wgpu")]
1279            (Self::Wgpu(src), Self::Wgpu(dst)) => Ok(src.copy_strided_src(dst, dst_offset, src_l)?),
1280            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
1281                lhs: lhs.device().location(),
1282                rhs: rhs.device().location(),
1283                op: "copy",
1284            }
1285            .bt()),
1286        }
1287    }
1288
1289    #[allow(clippy::too_many_arguments)]
1290    pub(crate) fn copy2d(
1291        &self,
1292        dst: &mut Self,
1293        d1: usize,
1294        d2: usize,
1295        src_s: usize,
1296        dst_s: usize,
1297        src_o: usize,
1298        dst_o: usize,
1299    ) -> Result<()> {
1300        match (self, dst) {
1301            (Self::Cpu(src), Self::Cpu(dst)) => src.copy2d(dst, d1, d2, src_s, dst_s, src_o, dst_o),
1302            (Self::Cuda(src), Self::Cuda(dst)) => {
1303                Ok(src.copy2d(dst, d1, d2, src_s, dst_s, src_o, dst_o)?)
1304            }
1305            (Self::Metal(src), Self::Metal(dst)) => {
1306                Ok(src.copy2d(dst, d1, d2, src_s, dst_s, src_o, dst_o)?)
1307            }
1308            #[cfg(feature = "rocm")]
1309            (Self::Rocm(src), Self::Rocm(dst)) => {
1310                Ok(src.copy2d(dst, d1, d2, src_s, dst_s, src_o, dst_o)?)
1311            }
1312            #[cfg(feature = "vulkan")]
1313            (Self::Vulkan(src), Self::Vulkan(dst)) => {
1314                Ok(src.copy2d(dst, d1, d2, src_s, dst_s, src_o, dst_o)?)
1315            }
1316            #[cfg(feature = "wgpu")]
1317            (Self::Wgpu(src), Self::Wgpu(dst)) => {
1318                Ok(src.copy2d(dst, d1, d2, src_s, dst_s, src_o, dst_o)?)
1319            }
1320            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
1321                lhs: lhs.device().location(),
1322                rhs: rhs.device().location(),
1323                op: "copy2d",
1324            }
1325            .bt()),
1326        }
1327    }
1328}
hanzo_ml/storage.rs

hanzo_ml/
storage.rs
