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;
8#[cfg(feature = "wgpu")]
9use crate::WgpuStorage;
10use crate::{CpuStorage, CudaStorage, DType, Device, Error, Layout, MetalStorage, Result, Shape};
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)) => c.wgpu_fwd(s1, l1, s2, l2, s3, l3),
516            _ => unreachable!(),
517        }
518    }
519
520    pub(crate) fn unary_impl<B: op::UnaryOpT>(&self, layout: &Layout) -> Result<Self> {
521        match self {
522            Storage::Cpu(storage) => {
523                let storage = storage.unary_impl::<B>(layout)?;
524                Ok(Self::Cpu(storage))
525            }
526            Self::Cuda(storage) => {
527                let storage = storage.unary_impl::<B>(layout)?;
528                Ok(Self::Cuda(storage))
529            }
530            Self::Metal(storage) => {
531                let storage = storage.unary_impl::<B>(layout)?;
532                Ok(Self::Metal(storage))
533            }
534            #[cfg(feature = "rocm")]
535            Self::Rocm(storage) => {
536                let storage = storage.unary_impl::<B>(layout)?;
537                Ok(Self::Rocm(storage))
538            }
539            #[cfg(feature = "vulkan")]
540            Self::Vulkan(storage) => {
541                let storage = storage.unary_impl::<B>(layout)?;
542                Ok(Self::Vulkan(storage))
543            }
544            #[cfg(feature = "wgpu")]
545            Self::Wgpu(storage) => {
546                let storage = storage.unary_impl::<B>(layout)?;
547                Ok(Self::Wgpu(storage))
548            }
549        }
550    }
551
552    pub(crate) fn binary_impl<B: op::BinaryOpT>(
553        &self,
554        rhs: &Self,
555        lhs_layout: &Layout,
556        rhs_layout: &Layout,
557    ) -> Result<Self> {
558        self.same_device(rhs, B::NAME)?;
559        self.same_dtype(rhs, B::NAME)?;
560        match (self, rhs) {
561            (Storage::Cpu(lhs), Storage::Cpu(rhs)) => {
562                let storage = lhs.binary_impl::<B>(rhs, lhs_layout, rhs_layout)?;
563                Ok(Self::Cpu(storage))
564            }
565            (Self::Cuda(lhs), Self::Cuda(rhs)) => {
566                let storage = lhs.binary_impl::<B>(rhs, lhs_layout, rhs_layout)?;
567                Ok(Self::Cuda(storage))
568            }
569            (Self::Metal(lhs), Self::Metal(rhs)) => {
570                let storage = lhs.binary_impl::<B>(rhs, lhs_layout, rhs_layout)?;
571                Ok(Self::Metal(storage))
572            }
573            #[cfg(feature = "rocm")]
574            (Self::Rocm(lhs), Self::Rocm(rhs)) => {
575                let storage = lhs.binary_impl::<B>(rhs, lhs_layout, rhs_layout)?;
576                Ok(Self::Rocm(storage))
577            }
578            #[cfg(feature = "vulkan")]
579            (Self::Vulkan(lhs), Self::Vulkan(rhs)) => {
580                let storage = lhs.binary_impl::<B>(rhs, lhs_layout, rhs_layout)?;
581                Ok(Self::Vulkan(storage))
582            }
583            #[cfg(feature = "wgpu")]
584            (Self::Wgpu(lhs), Self::Wgpu(rhs)) => {
585                let storage = lhs.binary_impl::<B>(rhs, lhs_layout, rhs_layout)?;
586                Ok(Self::Wgpu(storage))
587            }
588            (lhs, rhs) => {
589                // Should not happen because of the same device check above but we're defensive
590                // anyway.
591                Err(Error::DeviceMismatchBinaryOp {
592                    lhs: lhs.device().location(),
593                    rhs: rhs.device().location(),
594                    op: B::NAME,
595                }
596                .bt())
597            }
598        }
599    }
600
601    pub(crate) fn conv1d(
602        &self,
603        l: &Layout,
604        kernel: &Self,
605        kernel_l: &Layout,
606        params: &crate::conv::ParamsConv1D,
607    ) -> Result<Self> {
608        self.same_device(kernel, "conv1d")?;
609        self.same_dtype(kernel, "conv1d")?;
610        match (self, &kernel) {
611            (Storage::Cpu(inp), Storage::Cpu(kernel)) => {
612                let s = inp.conv1d(l, kernel, kernel_l, params)?;
613                Ok(Self::Cpu(s))
614            }
615            (Storage::Cuda(inp), Storage::Cuda(kernel)) => {
616                let s = inp.conv1d(l, kernel, kernel_l, params)?;
617                Ok(Self::Cuda(s))
618            }
619            (Storage::Metal(inp), Storage::Metal(kernel)) => {
620                let s = inp.conv1d(l, kernel, kernel_l, params)?;
621                Ok(Self::Metal(s))
622            }
623            #[cfg(feature = "rocm")]
624            (Storage::Rocm(inp), Storage::Rocm(kernel)) => {
625                let s = inp.conv1d(l, kernel, kernel_l, params)?;
626                Ok(Self::Rocm(s))
627            }
628            #[cfg(feature = "vulkan")]
629            (Storage::Vulkan(inp), Storage::Vulkan(kernel)) => {
630                let s = inp.conv1d(l, kernel, kernel_l, params)?;
631                Ok(Self::Vulkan(s))
632            }
633            #[cfg(feature = "wgpu")]
634            (Storage::Wgpu(inp), Storage::Wgpu(kernel)) => {
635                let s = inp.conv1d(l, kernel, kernel_l, params)?;
636                Ok(Self::Wgpu(s))
637            }
638            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
639                lhs: lhs.device().location(),
640                rhs: rhs.device().location(),
641                op: "conv1d",
642            }
643            .bt()),
644        }
645    }
646
647    pub(crate) fn conv_transpose1d(
648        &self,
649        l: &Layout,
650        kernel: &Self,
651        kernel_l: &Layout,
652        params: &crate::conv::ParamsConvTranspose1D,
653    ) -> Result<Self> {
654        self.same_device(kernel, "conv-transpose1d")?;
655        self.same_dtype(kernel, "conv-transpose1d")?;
656        match (self, &kernel) {
657            (Storage::Cpu(inp), Storage::Cpu(kernel)) => {
658                let s = inp.conv_transpose1d(l, kernel, kernel_l, params)?;
659                Ok(Self::Cpu(s))
660            }
661            (Storage::Cuda(inp), Storage::Cuda(kernel)) => {
662                let s = inp.conv_transpose1d(l, kernel, kernel_l, params)?;
663                Ok(Self::Cuda(s))
664            }
665            (Storage::Metal(inp), Storage::Metal(kernel)) => {
666                let s = inp.conv_transpose1d(l, kernel, kernel_l, params)?;
667                Ok(Self::Metal(s))
668            }
669            #[cfg(feature = "rocm")]
670            (Storage::Rocm(inp), Storage::Rocm(kernel)) => {
671                let s = inp.conv_transpose1d(l, kernel, kernel_l, params)?;
672                Ok(Self::Rocm(s))
673            }
674            #[cfg(feature = "vulkan")]
675            (Storage::Vulkan(inp), Storage::Vulkan(kernel)) => {
676                let s = inp.conv_transpose1d(l, kernel, kernel_l, params)?;
677                Ok(Self::Vulkan(s))
678            }
679            #[cfg(feature = "wgpu")]
680            (Storage::Wgpu(inp), Storage::Wgpu(kernel)) => {
681                let s = inp.conv_transpose1d(l, kernel, kernel_l, params)?;
682                Ok(Self::Wgpu(s))
683            }
684            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
685                lhs: lhs.device().location(),
686                rhs: rhs.device().location(),
687                op: "conv-transpose1d",
688            }
689            .bt()),
690        }
691    }
692
693    pub(crate) fn conv2d(
694        &self,
695        l: &Layout,
696        kernel: &Self,
697        kernel_l: &Layout,
698        params: &crate::conv::ParamsConv2D,
699    ) -> Result<Self> {
700        self.same_device(kernel, "conv2d")?;
701        self.same_dtype(kernel, "conv2d")?;
702        match (self, &kernel) {
703            (Storage::Cpu(inp), Storage::Cpu(kernel)) => {
704                let s = inp.conv2d(l, kernel, kernel_l, params)?;
705                Ok(Self::Cpu(s))
706            }
707            (Storage::Cuda(inp), Storage::Cuda(kernel)) => {
708                let s = inp.conv2d(l, kernel, kernel_l, params)?;
709                Ok(Self::Cuda(s))
710            }
711            (Storage::Metal(inp), Storage::Metal(kernel)) => {
712                let s = inp.conv2d(l, kernel, kernel_l, params)?;
713                Ok(Self::Metal(s))
714            }
715            #[cfg(feature = "rocm")]
716            (Storage::Rocm(inp), Storage::Rocm(kernel)) => {
717                let s = inp.conv2d(l, kernel, kernel_l, params)?;
718                Ok(Self::Rocm(s))
719            }
720            #[cfg(feature = "vulkan")]
721            (Storage::Vulkan(inp), Storage::Vulkan(kernel)) => {
722                let s = inp.conv2d(l, kernel, kernel_l, params)?;
723                Ok(Self::Vulkan(s))
724            }
725            #[cfg(feature = "wgpu")]
726            (Storage::Wgpu(inp), Storage::Wgpu(kernel)) => {
727                let s = inp.conv2d(l, kernel, kernel_l, params)?;
728                Ok(Self::Wgpu(s))
729            }
730            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
731                lhs: lhs.device().location(),
732                rhs: rhs.device().location(),
733                op: "conv2d",
734            }
735            .bt()),
736        }
737    }
738
739    pub(crate) fn conv_transpose2d(
740        &self,
741        l: &Layout,
742        kernel: &Self,
743        kernel_l: &Layout,
744        params: &crate::conv::ParamsConvTranspose2D,
745    ) -> Result<Self> {
746        self.same_device(kernel, "conv_transpose2d")?;
747        self.same_dtype(kernel, "conv_transpose2d")?;
748        match (self, &kernel) {
749            (Storage::Cpu(inp), Storage::Cpu(kernel)) => {
750                let s = inp.conv_transpose2d(l, kernel, kernel_l, params)?;
751                Ok(Self::Cpu(s))
752            }
753            (Storage::Cuda(inp), Storage::Cuda(kernel)) => {
754                let s = inp.conv_transpose2d(l, kernel, kernel_l, params)?;
755                Ok(Self::Cuda(s))
756            }
757            (Storage::Metal(inp), Storage::Metal(kernel)) => {
758                let s = inp.conv_transpose2d(l, kernel, kernel_l, params)?;
759                Ok(Self::Metal(s))
760            }
761            #[cfg(feature = "rocm")]
762            (Storage::Rocm(inp), Storage::Rocm(kernel)) => {
763                let s = inp.conv_transpose2d(l, kernel, kernel_l, params)?;
764                Ok(Self::Rocm(s))
765            }
766            #[cfg(feature = "vulkan")]
767            (Storage::Vulkan(inp), Storage::Vulkan(kernel)) => {
768                let s = inp.conv_transpose2d(l, kernel, kernel_l, params)?;
769                Ok(Self::Vulkan(s))
770            }
771            #[cfg(feature = "wgpu")]
772            (Storage::Wgpu(inp), Storage::Wgpu(kernel)) => {
773                let s = inp.conv_transpose2d(l, kernel, kernel_l, params)?;
774                Ok(Self::Wgpu(s))
775            }
776            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
777                lhs: lhs.device().location(),
778                rhs: rhs.device().location(),
779                op: "conv_transpose2d",
780            }
781            .bt()),
782        }
783    }
784
785    pub(crate) fn avg_pool2d(
786        &self,
787        layout: &Layout,
788        kernel_size: (usize, usize),
789        stride: (usize, usize),
790    ) -> Result<Self> {
791        match self {
792            Storage::Cpu(storage) => {
793                let storage = storage.avg_pool2d(layout, kernel_size, stride)?;
794                Ok(Self::Cpu(storage))
795            }
796            Self::Cuda(storage) => {
797                let storage = storage.avg_pool2d(layout, kernel_size, stride)?;
798                Ok(Self::Cuda(storage))
799            }
800            Self::Metal(storage) => {
801                let storage = storage.avg_pool2d(layout, kernel_size, stride)?;
802                Ok(Self::Metal(storage))
803            }
804            #[cfg(feature = "rocm")]
805            Self::Rocm(storage) => {
806                let storage = storage.avg_pool2d(layout, kernel_size, stride)?;
807                Ok(Self::Rocm(storage))
808            }
809            #[cfg(feature = "vulkan")]
810            Self::Vulkan(storage) => {
811                let storage = storage.avg_pool2d(layout, kernel_size, stride)?;
812                Ok(Self::Vulkan(storage))
813            }
814            #[cfg(feature = "wgpu")]
815            Self::Wgpu(storage) => {
816                let storage = storage.avg_pool2d(layout, kernel_size, stride)?;
817                Ok(Self::Wgpu(storage))
818            }
819        }
820    }
821
822    pub(crate) fn max_pool2d(
823        &self,
824        layout: &Layout,
825        kernel_size: (usize, usize),
826        stride: (usize, usize),
827    ) -> Result<Self> {
828        match self {
829            Storage::Cpu(storage) => {
830                let storage = storage.max_pool2d(layout, kernel_size, stride)?;
831                Ok(Self::Cpu(storage))
832            }
833            Self::Cuda(storage) => {
834                let storage = storage.max_pool2d(layout, kernel_size, stride)?;
835                Ok(Self::Cuda(storage))
836            }
837            Self::Metal(storage) => {
838                let storage = storage.max_pool2d(layout, kernel_size, stride)?;
839                Ok(Self::Metal(storage))
840            }
841            #[cfg(feature = "rocm")]
842            Self::Rocm(storage) => {
843                let storage = storage.max_pool2d(layout, kernel_size, stride)?;
844                Ok(Self::Rocm(storage))
845            }
846            #[cfg(feature = "vulkan")]
847            Self::Vulkan(storage) => {
848                let storage = storage.max_pool2d(layout, kernel_size, stride)?;
849                Ok(Self::Vulkan(storage))
850            }
851            #[cfg(feature = "wgpu")]
852            Self::Wgpu(storage) => {
853                let storage = storage.max_pool2d(layout, kernel_size, stride)?;
854                Ok(Self::Wgpu(storage))
855            }
856        }
857    }
858
859    pub(crate) fn upsample_nearest1d(&self, layout: &Layout, sz: usize) -> Result<Self> {
860        match self {
861            Storage::Cpu(storage) => {
862                let storage = storage.upsample_nearest1d(layout, sz)?;
863                Ok(Self::Cpu(storage))
864            }
865            Self::Cuda(storage) => {
866                let storage = storage.upsample_nearest1d(layout, sz)?;
867                Ok(Self::Cuda(storage))
868            }
869            Self::Metal(storage) => {
870                let storage = storage.upsample_nearest1d(layout, sz)?;
871                Ok(Self::Metal(storage))
872            }
873            #[cfg(feature = "rocm")]
874            Self::Rocm(storage) => {
875                let storage = storage.upsample_nearest1d(layout, sz)?;
876                Ok(Self::Rocm(storage))
877            }
878            #[cfg(feature = "vulkan")]
879            Self::Vulkan(storage) => {
880                let storage = storage.upsample_nearest1d(layout, sz)?;
881                Ok(Self::Vulkan(storage))
882            }
883            #[cfg(feature = "wgpu")]
884            Self::Wgpu(storage) => {
885                let storage = storage.upsample_nearest1d(layout, sz)?;
886                Ok(Self::Wgpu(storage))
887            }
888        }
889    }
890
891    pub(crate) fn upsample_nearest2d(&self, layout: &Layout, h: usize, w: usize) -> Result<Self> {
892        match self {
893            Storage::Cpu(storage) => {
894                let storage = storage.upsample_nearest2d(layout, h, w)?;
895                Ok(Self::Cpu(storage))
896            }
897            Self::Cuda(storage) => {
898                let storage = storage.upsample_nearest2d(layout, h, w)?;
899                Ok(Self::Cuda(storage))
900            }
901            Self::Metal(storage) => {
902                let storage = storage.upsample_nearest2d(layout, h, w)?;
903                Ok(Self::Metal(storage))
904            }
905            #[cfg(feature = "rocm")]
906            Self::Rocm(storage) => {
907                let storage = storage.upsample_nearest2d(layout, h, w)?;
908                Ok(Self::Rocm(storage))
909            }
910            #[cfg(feature = "vulkan")]
911            Self::Vulkan(storage) => {
912                let storage = storage.upsample_nearest2d(layout, h, w)?;
913                Ok(Self::Vulkan(storage))
914            }
915            #[cfg(feature = "wgpu")]
916            Self::Wgpu(storage) => {
917                let storage = storage.upsample_nearest2d(layout, h, w)?;
918                Ok(Self::Wgpu(storage))
919            }
920        }
921    }
922
923    pub(crate) fn upsample_bilinear2d(
924        &self,
925        layout: &Layout,
926        h: usize,
927        w: usize,
928        align_corners: bool,
929        scale_h: Option<f64>,
930        scale_w: Option<f64>,
931    ) -> Result<Self> {
932        match self {
933            Storage::Cpu(storage) => {
934                let storage =
935                    storage.upsample_bilinear2d(layout, h, w, align_corners, scale_h, scale_w)?;
936                Ok(Self::Cpu(storage))
937            }
938            Self::Cuda(storage) => {
939                let storage =
940                    storage.upsample_bilinear2d(layout, h, w, align_corners, scale_h, scale_w)?;
941                Ok(Self::Cuda(storage))
942            }
943            Self::Metal(storage) => {
944                let storage =
945                    storage.upsample_bilinear2d(layout, h, w, align_corners, scale_h, scale_w)?;
946                Ok(Self::Metal(storage))
947            }
948            #[cfg(feature = "rocm")]
949            Self::Rocm(storage) => {
950                let storage =
951                    storage.upsample_bilinear2d(layout, h, w, align_corners, scale_h, scale_w)?;
952                Ok(Self::Rocm(storage))
953            }
954            #[cfg(feature = "vulkan")]
955            Self::Vulkan(storage) => {
956                let storage =
957                    storage.upsample_bilinear2d(layout, h, w, align_corners, scale_h, scale_w)?;
958                Ok(Self::Vulkan(storage))
959            }
960            #[cfg(feature = "wgpu")]
961            Self::Wgpu(storage) => {
962                let storage =
963                    storage.upsample_bilinear2d(layout, h, w, align_corners, scale_h, scale_w)?;
964                Ok(Self::Wgpu(storage))
965            }
966        }
967    }
968
969    pub(crate) fn where_cond(
970        &self,
971        layout: &Layout,
972        t: &Self,
973        layout_t: &Layout,
974        f: &Self,
975        layout_f: &Layout,
976    ) -> Result<Self> {
977        self.same_device(t, "where")?;
978        self.same_device(f, "where")?;
979        t.same_dtype(f, "where")?;
980        match (self, t, f) {
981            (Storage::Cpu(cond), Storage::Cpu(t), Storage::Cpu(f)) => {
982                let storage = cond.where_cond(layout, t, layout_t, f, layout_f)?;
983                Ok(Self::Cpu(storage))
984            }
985            (Self::Cuda(cond), Self::Cuda(t), Self::Cuda(f)) => {
986                let storage = cond.where_cond(layout, t, layout_t, f, layout_f)?;
987                Ok(Self::Cuda(storage))
988            }
989            (Self::Metal(cond), Self::Metal(t), Self::Metal(f)) => {
990                let storage = cond.where_cond(layout, t, layout_t, f, layout_f)?;
991                Ok(Self::Metal(storage))
992            }
993            #[cfg(feature = "rocm")]
994            (Self::Rocm(cond), Self::Rocm(t), Self::Rocm(f)) => {
995                let storage = cond.where_cond(layout, t, layout_t, f, layout_f)?;
996                Ok(Self::Rocm(storage))
997            }
998            #[cfg(feature = "vulkan")]
999            (Self::Vulkan(cond), Self::Vulkan(t), Self::Vulkan(f)) => {
1000                let storage = cond.where_cond(layout, t, layout_t, f, layout_f)?;
1001                Ok(Self::Vulkan(storage))
1002            }
1003            #[cfg(feature = "wgpu")]
1004            (Self::Wgpu(cond), Self::Wgpu(t), Self::Wgpu(f)) => {
1005                let storage = cond.where_cond(layout, t, layout_t, f, layout_f)?;
1006                Ok(Self::Wgpu(storage))
1007            }
1008            (_, lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
1009                lhs: lhs.device().location(),
1010                rhs: rhs.device().location(),
1011                op: "where",
1012            }
1013            .bt()),
1014        }
1015    }
1016
1017    pub(crate) fn gather(
1018        &self,
1019        l: &Layout,
1020        indexes: &Self,
1021        indexes_l: &Layout,
1022        d: usize,
1023    ) -> Result<Self> {
1024        self.same_device(indexes, "index-add")?;
1025        match (self, indexes) {
1026            (Self::Cpu(s), Self::Cpu(indexes)) => {
1027                let storage = s.gather(l, indexes, indexes_l, d)?;
1028                Ok(Self::Cpu(storage))
1029            }
1030            (Self::Cuda(s), Self::Cuda(indexes)) => {
1031                let storage = s.gather(l, indexes, indexes_l, d)?;
1032                Ok(Self::Cuda(storage))
1033            }
1034            (Self::Metal(s), Self::Metal(indexes)) => {
1035                let storage = s.gather(l, indexes, indexes_l, d)?;
1036                Ok(Self::Metal(storage))
1037            }
1038            #[cfg(feature = "rocm")]
1039            (Self::Rocm(s), Self::Rocm(indexes)) => {
1040                let storage = s.gather(l, indexes, indexes_l, d)?;
1041                Ok(Self::Rocm(storage))
1042            }
1043            #[cfg(feature = "vulkan")]
1044            (Self::Vulkan(s), Self::Vulkan(indexes)) => {
1045                let storage = s.gather(l, indexes, indexes_l, d)?;
1046                Ok(Self::Vulkan(storage))
1047            }
1048            #[cfg(feature = "wgpu")]
1049            (Self::Wgpu(s), Self::Wgpu(indexes)) => {
1050                let storage = s.gather(l, indexes, indexes_l, d)?;
1051                Ok(Self::Wgpu(storage))
1052            }
1053            _ => unreachable!(),
1054        }
1055    }
1056
1057    pub(crate) fn scatter_set(
1058        &mut self,
1059        l: &Layout,
1060        indexes: &Self,
1061        indexes_l: &Layout,
1062        source: &Self,
1063        source_l: &Layout,
1064        d: usize,
1065    ) -> Result<()> {
1066        self.same_device(indexes, "scatter-set")?;
1067        self.same_device(source, "scatter-set")?;
1068        match (self, indexes, source) {
1069            (Self::Cpu(s), Self::Cpu(indexes), Self::Cpu(source)) => {
1070                s.scatter_set(l, indexes, indexes_l, source, source_l, d)?;
1071            }
1072            (Self::Cuda(s), Self::Cuda(indexes), Self::Cuda(source)) => {
1073                s.scatter_set(l, indexes, indexes_l, source, source_l, d)?;
1074            }
1075            (Self::Metal(s), Self::Metal(indexes), Self::Metal(source)) => {
1076                s.scatter_set(l, indexes, indexes_l, source, source_l, d)?;
1077            }
1078            #[cfg(feature = "rocm")]
1079            (Self::Rocm(s), Self::Rocm(indexes), Self::Rocm(source)) => {
1080                s.scatter_set(l, indexes, indexes_l, source, source_l, d)?;
1081            }
1082            #[cfg(feature = "vulkan")]
1083            (Self::Vulkan(s), Self::Vulkan(indexes), Self::Vulkan(source)) => {
1084                s.scatter_set(l, indexes, indexes_l, source, source_l, d)?;
1085            }
1086            #[cfg(feature = "wgpu")]
1087            (Self::Wgpu(s), Self::Wgpu(indexes), Self::Wgpu(source)) => {
1088                s.scatter_set(l, indexes, indexes_l, source, source_l, d)?;
1089            }
1090            _ => unreachable!(),
1091        }
1092        Ok(())
1093    }
1094
1095    pub(crate) fn scatter_add(
1096        &mut self,
1097        l: &Layout,
1098        indexes: &Self,
1099        indexes_l: &Layout,
1100        source: &Self,
1101        source_l: &Layout,
1102        d: usize,
1103    ) -> Result<()> {
1104        self.same_device(indexes, "scatter-add")?;
1105        self.same_device(source, "scatter-add")?;
1106        match (self, indexes, source) {
1107            (Self::Cpu(s), Self::Cpu(indexes), Self::Cpu(source)) => {
1108                s.scatter_add_set(l, indexes, indexes_l, source, source_l, d)?;
1109            }
1110            (Self::Cuda(s), Self::Cuda(indexes), Self::Cuda(source)) => {
1111                s.scatter_add_set(l, indexes, indexes_l, source, source_l, d)?;
1112            }
1113            (Self::Metal(s), Self::Metal(indexes), Self::Metal(source)) => {
1114                s.scatter_add_set(l, indexes, indexes_l, source, source_l, d)?;
1115            }
1116            #[cfg(feature = "rocm")]
1117            (Self::Rocm(s), Self::Rocm(indexes), Self::Rocm(source)) => {
1118                s.scatter_add_set(l, indexes, indexes_l, source, source_l, d)?;
1119            }
1120            #[cfg(feature = "vulkan")]
1121            (Self::Vulkan(s), Self::Vulkan(indexes), Self::Vulkan(source)) => {
1122                s.scatter_add_set(l, indexes, indexes_l, source, source_l, d)?;
1123            }
1124            #[cfg(feature = "wgpu")]
1125            (Self::Wgpu(s), Self::Wgpu(indexes), Self::Wgpu(source)) => {
1126                s.scatter_add_set(l, indexes, indexes_l, source, source_l, d)?;
1127            }
1128            _ => unreachable!(),
1129        }
1130        Ok(())
1131    }
1132
1133    pub(crate) fn index_add(
1134        &self,
1135        l: &Layout,
1136        indexes: &Self,
1137        indexes_l: &Layout,
1138        source: &Self,
1139        source_l: &Layout,
1140        d: usize,
1141    ) -> Result<Self> {
1142        self.same_device(indexes, "index-add")?;
1143        self.same_device(source, "index-add")?;
1144        match (self, indexes, source) {
1145            (Self::Cpu(s), Self::Cpu(indexes), Self::Cpu(source)) => {
1146                let storage = s.index_add(l, indexes, indexes_l, source, source_l, d)?;
1147                Ok(Self::Cpu(storage))
1148            }
1149            (Self::Cuda(s), Self::Cuda(indexes), Self::Cuda(source)) => {
1150                let storage = s.index_add(l, indexes, indexes_l, source, source_l, d)?;
1151                Ok(Self::Cuda(storage))
1152            }
1153            (Self::Metal(s), Self::Metal(indexes), Self::Metal(source)) => {
1154                let storage = s.index_add(l, indexes, indexes_l, source, source_l, d)?;
1155                Ok(Self::Metal(storage))
1156            }
1157            #[cfg(feature = "rocm")]
1158            (Self::Rocm(s), Self::Rocm(indexes), Self::Rocm(source)) => {
1159                let storage = s.index_add(l, indexes, indexes_l, source, source_l, d)?;
1160                Ok(Self::Rocm(storage))
1161            }
1162            #[cfg(feature = "vulkan")]
1163            (Self::Vulkan(s), Self::Vulkan(indexes), Self::Vulkan(source)) => {
1164                let storage = s.index_add(l, indexes, indexes_l, source, source_l, d)?;
1165                Ok(Self::Vulkan(storage))
1166            }
1167            #[cfg(feature = "wgpu")]
1168            (Self::Wgpu(s), Self::Wgpu(indexes), Self::Wgpu(source)) => {
1169                let storage = s.index_add(l, indexes, indexes_l, source, source_l, d)?;
1170                Ok(Self::Wgpu(storage))
1171            }
1172            _ => unreachable!(),
1173        }
1174    }
1175
1176    pub(crate) fn index_select(
1177        &self,
1178        rhs: &Self,
1179        lhs_l: &Layout,
1180        rhs_l: &Layout,
1181        d: usize,
1182    ) -> Result<Self> {
1183        self.same_device(rhs, "index-select")?;
1184        match (self, rhs) {
1185            (Self::Cpu(lhs), Self::Cpu(rhs)) => {
1186                let storage = lhs.index_select(rhs, lhs_l, rhs_l, d)?;
1187                Ok(Self::Cpu(storage))
1188            }
1189            (Self::Cuda(lhs), Self::Cuda(rhs)) => {
1190                let storage = lhs.index_select(rhs, lhs_l, rhs_l, d)?;
1191                Ok(Self::Cuda(storage))
1192            }
1193            (Self::Metal(lhs), Self::Metal(rhs)) => {
1194                let storage = lhs.index_select(rhs, lhs_l, rhs_l, d)?;
1195                Ok(Self::Metal(storage))
1196            }
1197            #[cfg(feature = "rocm")]
1198            (Self::Rocm(lhs), Self::Rocm(rhs)) => {
1199                let storage = lhs.index_select(rhs, lhs_l, rhs_l, d)?;
1200                Ok(Self::Rocm(storage))
1201            }
1202            #[cfg(feature = "vulkan")]
1203            (Self::Vulkan(lhs), Self::Vulkan(rhs)) => {
1204                let storage = lhs.index_select(rhs, lhs_l, rhs_l, d)?;
1205                Ok(Self::Vulkan(storage))
1206            }
1207            #[cfg(feature = "wgpu")]
1208            (Self::Wgpu(lhs), Self::Wgpu(rhs)) => {
1209                let storage = lhs.index_select(rhs, lhs_l, rhs_l, d)?;
1210                Ok(Self::Wgpu(storage))
1211            }
1212            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
1213                lhs: lhs.device().location(),
1214                rhs: rhs.device().location(),
1215                op: "index-select",
1216            }
1217            .bt()),
1218        }
1219    }
1220
1221    pub(crate) fn matmul(
1222        &self,
1223        rhs: &Self,
1224        bmnk: (usize, usize, usize, usize),
1225        lhs_layout: &Layout,
1226        rhs_layout: &Layout,
1227    ) -> Result<Self> {
1228        self.same_device(rhs, "matmul")?;
1229        self.same_dtype(rhs, "matmul")?;
1230        match (self, rhs) {
1231            (Self::Cpu(lhs), Self::Cpu(rhs)) => {
1232                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
1233                Ok(Self::Cpu(storage))
1234            }
1235            (Self::Cuda(lhs), Self::Cuda(rhs)) => {
1236                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
1237                Ok(Self::Cuda(storage))
1238            }
1239            (Self::Metal(lhs), Self::Metal(rhs)) => {
1240                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
1241                Ok(Self::Metal(storage))
1242            }
1243            #[cfg(feature = "rocm")]
1244            (Self::Rocm(lhs), Self::Rocm(rhs)) => {
1245                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
1246                Ok(Self::Rocm(storage))
1247            }
1248            #[cfg(feature = "vulkan")]
1249            (Self::Vulkan(lhs), Self::Vulkan(rhs)) => {
1250                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
1251                Ok(Self::Vulkan(storage))
1252            }
1253            #[cfg(feature = "wgpu")]
1254            (Self::Wgpu(lhs), Self::Wgpu(rhs)) => {
1255                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
1256                Ok(Self::Wgpu(storage))
1257            }
1258            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
1259                lhs: lhs.device().location(),
1260                rhs: rhs.device().location(),
1261                op: "matmul",
1262            }
1263            .bt()),
1264        }
1265    }
1266
1267    // self, the source can be strided whereas dst is contiguous.
1268    pub(crate) fn copy_strided_src(
1269        &self,
1270        dst: &mut Self,
1271        dst_offset: usize,
1272        src_l: &Layout,
1273    ) -> Result<()> {
1274        match (self, dst) {
1275            (Self::Cpu(src), Self::Cpu(dst)) => src.copy_strided_src(dst, dst_offset, src_l),
1276            (Self::Cuda(src), Self::Cuda(dst)) => Ok(src.copy_strided_src(dst, dst_offset, src_l)?),
1277            (Self::Metal(src), Self::Metal(dst)) => {
1278                Ok(src.copy_strided_src(dst, dst_offset, src_l)?)
1279            }
1280            #[cfg(feature = "rocm")]
1281            (Self::Rocm(src), Self::Rocm(dst)) => Ok(src.copy_strided_src(dst, dst_offset, src_l)?),
1282            #[cfg(feature = "vulkan")]
1283            (Self::Vulkan(src), Self::Vulkan(dst)) => {
1284                Ok(src.copy_strided_src(dst, dst_offset, src_l)?)
1285            }
1286            #[cfg(feature = "wgpu")]
1287            (Self::Wgpu(src), Self::Wgpu(dst)) => Ok(src.copy_strided_src(dst, dst_offset, src_l)?),
1288            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
1289                lhs: lhs.device().location(),
1290                rhs: rhs.device().location(),
1291                op: "copy",
1292            }
1293            .bt()),
1294        }
1295    }
1296
1297    #[allow(clippy::too_many_arguments)]
1298    pub(crate) fn copy2d(
1299        &self,
1300        dst: &mut Self,
1301        d1: usize,
1302        d2: usize,
1303        src_s: usize,
1304        dst_s: usize,
1305        src_o: usize,
1306        dst_o: usize,
1307    ) -> Result<()> {
1308        match (self, dst) {
1309            (Self::Cpu(src), Self::Cpu(dst)) => src.copy2d(dst, d1, d2, src_s, dst_s, src_o, dst_o),
1310            (Self::Cuda(src), Self::Cuda(dst)) => {
1311                Ok(src.copy2d(dst, d1, d2, src_s, dst_s, src_o, dst_o)?)
1312            }
1313            (Self::Metal(src), Self::Metal(dst)) => {
1314                Ok(src.copy2d(dst, d1, d2, src_s, dst_s, src_o, dst_o)?)
1315            }
1316            #[cfg(feature = "rocm")]
1317            (Self::Rocm(src), Self::Rocm(dst)) => {
1318                Ok(src.copy2d(dst, d1, d2, src_s, dst_s, src_o, dst_o)?)
1319            }
1320            #[cfg(feature = "vulkan")]
1321            (Self::Vulkan(src), Self::Vulkan(dst)) => {
1322                Ok(src.copy2d(dst, d1, d2, src_s, dst_s, src_o, dst_o)?)
1323            }
1324            #[cfg(feature = "wgpu")]
1325            (Self::Wgpu(src), Self::Wgpu(dst)) => {
1326                Ok(src.copy2d(dst, d1, d2, src_s, dst_s, src_o, dst_o)?)
1327            }
1328            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
1329                lhs: lhs.device().location(),
1330                rhs: rhs.device().location(),
1331                op: "copy2d",
1332            }
1333            .bt()),
1334        }
1335    }
1336}
hanzo_ml/storage.rs

hanzo_ml/
storage.rs
