hanzo_ml/
storage.rs

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

hanzo_ml/
storage.rs
