hanzo_ml/
device.rs

1use crate::backend::BackendDevice;
2use crate::cpu_backend::CpuDevice;
3use crate::{CpuStorage, DType, Result, Shape, Storage, WithDType};
4
5/// A `DeviceLocation` represents a physical device whereas multiple `Device`
6/// can live on the same location (typically for cuda devices).
7#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
8pub enum DeviceLocation {
9    Cpu,
10    Cuda {
11        gpu_id: usize,
12    },
13    Metal {
14        gpu_id: usize,
15    },
16    #[cfg(feature = "rocm")]
17    Rocm {
18        gpu_id: usize,
19    },
20    #[cfg(feature = "vulkan")]
21    Vulkan {
22        gpu_id: usize,
23    },
24    #[cfg(feature = "wgpu")]
25    Wgpu {
26        gpu_id: usize,
27    },
28}
29
30/// Cpu, Cuda, or Metal
31#[derive(Debug, Clone)]
32pub enum Device {
33    Cpu,
34    Cuda(crate::CudaDevice),
35    Metal(crate::MetalDevice),
36    #[cfg(feature = "rocm")]
37    Rocm(crate::RocmDevice),
38    #[cfg(feature = "vulkan")]
39    Vulkan(crate::VulkanDevice),
40    #[cfg(feature = "wgpu")]
41    Wgpu(crate::WgpuDevice),
42}
43
44pub trait NdArray {
45    fn shape(&self) -> Result<Shape>;
46
47    fn to_cpu_storage(&self) -> CpuStorage;
48}
49
50impl<S: WithDType> NdArray for S {
51    fn shape(&self) -> Result<Shape> {
52        Ok(Shape::from(()))
53    }
54
55    fn to_cpu_storage(&self) -> CpuStorage {
56        S::to_cpu_storage(&[*self])
57    }
58}
59
60impl<S: WithDType, const N: usize> NdArray for &[S; N] {
61    fn shape(&self) -> Result<Shape> {
62        Ok(Shape::from(self.len()))
63    }
64
65    fn to_cpu_storage(&self) -> CpuStorage {
66        S::to_cpu_storage(self.as_slice())
67    }
68}
69
70impl<S: WithDType> NdArray for &[S] {
71    fn shape(&self) -> Result<Shape> {
72        Ok(Shape::from(self.len()))
73    }
74
75    fn to_cpu_storage(&self) -> CpuStorage {
76        S::to_cpu_storage(self)
77    }
78}
79
80impl<S: WithDType, const N: usize, const M: usize> NdArray for &[[S; N]; M] {
81    fn shape(&self) -> Result<Shape> {
82        Ok(Shape::from((M, N)))
83    }
84
85    fn to_cpu_storage(&self) -> CpuStorage {
86        S::to_cpu_storage_owned(self.concat())
87    }
88}
89
90impl<S: WithDType, const N1: usize, const N2: usize, const N3: usize> NdArray
91    for &[[[S; N3]; N2]; N1]
92{
93    fn shape(&self) -> Result<Shape> {
94        Ok(Shape::from((N1, N2, N3)))
95    }
96
97    fn to_cpu_storage(&self) -> CpuStorage {
98        let mut vec = Vec::with_capacity(N1 * N2 * N3);
99        for i1 in 0..N1 {
100            for i2 in 0..N2 {
101                vec.extend(self[i1][i2])
102            }
103        }
104        S::to_cpu_storage_owned(vec)
105    }
106}
107
108impl<S: WithDType, const N1: usize, const N2: usize, const N3: usize, const N4: usize> NdArray
109    for &[[[[S; N4]; N3]; N2]; N1]
110{
111    fn shape(&self) -> Result<Shape> {
112        Ok(Shape::from((N1, N2, N3, N4)))
113    }
114
115    fn to_cpu_storage(&self) -> CpuStorage {
116        let mut vec = Vec::with_capacity(N1 * N2 * N3 * N4);
117        for i1 in 0..N1 {
118            for i2 in 0..N2 {
119                for i3 in 0..N3 {
120                    vec.extend(self[i1][i2][i3])
121                }
122            }
123        }
124        S::to_cpu_storage_owned(vec)
125    }
126}
127
128impl<S: WithDType> NdArray for Vec<S> {
129    fn shape(&self) -> Result<Shape> {
130        Ok(Shape::from(self.len()))
131    }
132
133    fn to_cpu_storage(&self) -> CpuStorage {
134        S::to_cpu_storage(self.as_slice())
135    }
136}
137
138impl<S: WithDType> NdArray for Vec<&[S]> {
139    fn shape(&self) -> Result<Shape> {
140        if self.is_empty() {
141            crate::bail!("empty array")
142        }
143        let n = self.len();
144        let m = self[0].len();
145        for v in self.iter() {
146            if v.len() != m {
147                crate::bail!("two elements have different len {m} {}", v.len())
148            }
149        }
150        Ok(Shape::from((n, m)))
151    }
152
153    fn to_cpu_storage(&self) -> CpuStorage {
154        let data = self.iter().copied().flatten().copied().collect::<Vec<_>>();
155        S::to_cpu_storage_owned(data)
156    }
157}
158
159impl<S: WithDType> NdArray for Vec<Vec<S>> {
160    fn shape(&self) -> Result<Shape> {
161        if self.is_empty() {
162            crate::bail!("empty array")
163        }
164        let n = self.len();
165        let m = self[0].len();
166        for v in self.iter() {
167            if v.len() != m {
168                crate::bail!("two elements have different len {m} {}", v.len())
169            }
170        }
171        Ok(Shape::from((n, m)))
172    }
173
174    fn to_cpu_storage(&self) -> CpuStorage {
175        let len: usize = self.iter().map(|v| v.len()).sum();
176        let mut dst = Vec::with_capacity(len);
177        for v in self.iter() {
178            dst.extend(v.iter().copied());
179        }
180        S::to_cpu_storage_owned(dst)
181    }
182}
183
184impl<S: WithDType> NdArray for Vec<Vec<Vec<S>>> {
185    fn shape(&self) -> Result<Shape> {
186        if self.is_empty() {
187            crate::bail!("empty array")
188        }
189        let shape0 = self[0].shape()?;
190        let n = self.len();
191        for v in self.iter() {
192            let shape = v.shape()?;
193            if shape != shape0 {
194                crate::bail!("two elements have different shapes {shape:?} {shape0:?}")
195            }
196        }
197        Ok(Shape::from([[n].as_slice(), shape0.dims()].concat()))
198    }
199
200    fn to_cpu_storage(&self) -> CpuStorage {
201        if self.is_empty() {
202            return S::to_cpu_storage_owned(vec![]);
203        }
204        let len: usize = self
205            .iter()
206            .map(|v| v.iter().map(|v| v.len()).sum::<usize>())
207            .sum();
208        let mut dst = Vec::with_capacity(len);
209        for v1 in self.iter() {
210            for v2 in v1.iter() {
211                dst.extend(v2.iter().copied());
212            }
213        }
214        S::to_cpu_storage_owned(dst)
215    }
216}
217
218impl<S: WithDType> NdArray for Vec<Vec<Vec<Vec<S>>>> {
219    fn shape(&self) -> Result<Shape> {
220        if self.is_empty() {
221            crate::bail!("empty array")
222        }
223        let shape0 = self[0].shape()?;
224        let n = self.len();
225        for v in self.iter() {
226            let shape = v.shape()?;
227            if shape != shape0 {
228                crate::bail!("two elements have different shapes {shape:?} {shape0:?}")
229            }
230        }
231        Ok(Shape::from([[n].as_slice(), shape0.dims()].concat()))
232    }
233
234    fn to_cpu_storage(&self) -> CpuStorage {
235        let len: usize = self
236            .iter()
237            .map(|v| {
238                v.iter()
239                    .map(|v| v.iter().map(|v| v.len()).sum::<usize>())
240                    .sum::<usize>()
241            })
242            .sum();
243        let mut dst = Vec::with_capacity(len);
244        for v1 in self.iter() {
245            for v2 in v1.iter() {
246                for v3 in v2.iter() {
247                    dst.extend(v3.iter().copied());
248                }
249            }
250        }
251        S::to_cpu_storage_owned(dst)
252    }
253}
254
255impl Device {
256    pub fn new_cuda(ordinal: usize) -> Result<Self> {
257        Ok(Self::Cuda(crate::CudaDevice::new(ordinal)?))
258    }
259
260    #[cfg(feature = "rocm")]
261    pub fn new_rocm(ordinal: usize) -> Result<Self> {
262        Ok(Self::Rocm(crate::RocmDevice::new(ordinal)?))
263    }
264    #[cfg(feature = "vulkan")]
265    pub fn new_vulkan(ordinal: usize) -> Result<Self> {
266        Ok(Self::Vulkan(crate::VulkanDevice::new(ordinal)?))
267    }
268    #[cfg(feature = "wgpu")]
269    pub fn new_wgpu(ordinal: usize) -> Result<Self> {
270        Ok(Self::Wgpu(crate::WgpuDevice::new(ordinal)?))
271    }
272
273    pub fn as_cuda_device(&self) -> Result<&crate::CudaDevice> {
274        match self {
275            Self::Cuda(d) => Ok(d),
276            Self::Cpu => crate::bail!("expected a cuda device, got cpu"),
277            Self::Metal(_) => crate::bail!("expected a cuda device, got Metal"),
278            #[cfg(feature = "rocm")]
279            Self::Rocm(_) => crate::bail!("expected a cuda device, got rocm"),
280            #[cfg(feature = "vulkan")]
281            Self::Vulkan(_) => crate::bail!("expected a cuda device, got vulkan"),
282            #[cfg(feature = "wgpu")]
283            Self::Wgpu(_) => crate::bail!("expected a cuda device, got wgpu"),
284        }
285    }
286
287    pub fn as_metal_device(&self) -> Result<&crate::MetalDevice> {
288        match self {
289            Self::Cuda(_) => crate::bail!("expected a metal device, got cuda"),
290            Self::Cpu => crate::bail!("expected a metal device, got cpu"),
291            Self::Metal(d) => Ok(d),
292            #[cfg(feature = "rocm")]
293            Self::Rocm(_) => crate::bail!("expected a metal device, got rocm"),
294            #[cfg(feature = "vulkan")]
295            Self::Vulkan(_) => crate::bail!("expected a metal device, got vulkan"),
296            #[cfg(feature = "wgpu")]
297            Self::Wgpu(_) => crate::bail!("expected a metal device, got wgpu"),
298        }
299    }
300
301    #[cfg(feature = "rocm")]
302    pub fn as_rocm_device(&self) -> Result<&crate::RocmDevice> {
303        match self {
304            Self::Cuda(_) => crate::bail!("expected a rocm device, got cuda"),
305            Self::Cpu => crate::bail!("expected a rocm device, got cpu"),
306            Self::Metal(_) => crate::bail!("expected a rocm device, got Metal"),
307            Self::Rocm(d) => Ok(d),
308        }
309    }
310    #[cfg(feature = "vulkan")]
311    pub fn as_vulkan_device(&self) -> Result<&crate::VulkanDevice> {
312        match self {
313            Self::Cuda(_) => crate::bail!("expected a vulkan device, got cuda"),
314            Self::Cpu => crate::bail!("expected a vulkan device, got cpu"),
315            Self::Metal(_) => crate::bail!("expected a vulkan device, got Metal"),
316            Self::Vulkan(d) => Ok(d),
317            #[cfg(feature = "wgpu")]
318            Self::Wgpu(_) => crate::bail!("expected a vulkan device, got wgpu"),
319        }
320    }
321    #[cfg(feature = "wgpu")]
322    pub fn as_wgpu_device(&self) -> Result<&crate::WgpuDevice> {
323        match self {
324            Self::Cuda(_) => crate::bail!("expected a wgpu device, got cuda"),
325            Self::Cpu => crate::bail!("expected a wgpu device, got cpu"),
326            Self::Metal(_) => crate::bail!("expected a wgpu device, got Metal"),
327            #[cfg(feature = "vulkan")]
328            Self::Vulkan(_) => crate::bail!("expected a wgpu device, got vulkan"),
329            Self::Wgpu(d) => Ok(d),
330        }
331    }
332
333    pub fn new_cuda_with_stream(ordinal: usize) -> Result<Self> {
334        Ok(Self::Cuda(crate::CudaDevice::new_with_stream(ordinal)?))
335    }
336
337    pub fn new_metal(ordinal: usize) -> Result<Self> {
338        Ok(Self::Metal(crate::MetalDevice::new(ordinal)?))
339    }
340
341    pub fn set_seed(&self, seed: u64) -> Result<()> {
342        match self {
343            Self::Cpu => CpuDevice.set_seed(seed),
344            Self::Cuda(c) => c.set_seed(seed),
345            Self::Metal(m) => m.set_seed(seed),
346            #[cfg(feature = "rocm")]
347            Self::Rocm(r) => r.set_seed(seed),
348            #[cfg(feature = "vulkan")]
349            Self::Vulkan(r) => r.set_seed(seed),
350            #[cfg(feature = "wgpu")]
351            Self::Wgpu(r) => r.set_seed(seed),
352        }
353    }
354
355    pub fn get_current_seed(&self) -> Result<u64> {
356        match self {
357            Self::Cpu => CpuDevice.get_current_seed(),
358            Self::Cuda(c) => c.get_current_seed(),
359            Self::Metal(m) => m.get_current_seed(),
360            #[cfg(feature = "rocm")]
361            Self::Rocm(r) => r.get_current_seed(),
362            #[cfg(feature = "vulkan")]
363            Self::Vulkan(r) => r.get_current_seed(),
364            #[cfg(feature = "wgpu")]
365            Self::Wgpu(r) => r.get_current_seed(),
366        }
367    }
368
369    pub fn same_device(&self, rhs: &Self) -> bool {
370        match (self, rhs) {
371            (Self::Cpu, Self::Cpu) => true,
372            (Self::Cuda(lhs), Self::Cuda(rhs)) => lhs.same_device(rhs),
373            (Self::Metal(lhs), Self::Metal(rhs)) => lhs.same_device(rhs),
374            #[cfg(feature = "rocm")]
375            (Self::Rocm(lhs), Self::Rocm(rhs)) => lhs.same_device(rhs),
376            #[cfg(feature = "vulkan")]
377            (Self::Vulkan(lhs), Self::Vulkan(rhs)) => lhs.same_device(rhs),
378            #[cfg(feature = "wgpu")]
379            (Self::Wgpu(lhs), Self::Wgpu(rhs)) => lhs.same_device(rhs),
380            _ => false,
381        }
382    }
383
384    pub fn location(&self) -> DeviceLocation {
385        match self {
386            Self::Cpu => DeviceLocation::Cpu,
387            Self::Cuda(device) => device.location(),
388            Device::Metal(device) => device.location(),
389            #[cfg(feature = "rocm")]
390            Self::Rocm(device) => device.location(),
391            #[cfg(feature = "vulkan")]
392            Self::Vulkan(device) => device.location(),
393            #[cfg(feature = "wgpu")]
394            Self::Wgpu(device) => device.location(),
395        }
396    }
397
398    pub fn is_cpu(&self) -> bool {
399        matches!(self, Self::Cpu)
400    }
401
402    pub fn is_cuda(&self) -> bool {
403        matches!(self, Self::Cuda(_))
404    }
405
406    pub fn is_metal(&self) -> bool {
407        matches!(self, Self::Metal(_))
408    }
409
410    pub fn is_rocm(&self) -> bool {
411        #[cfg(feature = "rocm")]
412        {
413            matches!(self, Self::Rocm(_))
414        }
415        #[cfg(not(feature = "rocm"))]
416        {
417            false
418        }
419    }
420
421    pub fn is_vulkan(&self) -> bool {
422        #[cfg(feature = "vulkan")]
423        {
424            matches!(self, Self::Vulkan(_))
425        }
426        #[cfg(not(feature = "vulkan"))]
427        {
428            false
429        }
430    }
431
432    pub fn is_wgpu(&self) -> bool {
433        #[cfg(feature = "wgpu")]
434        {
435            matches!(self, Self::Wgpu(_))
436        }
437        #[cfg(not(feature = "wgpu"))]
438        {
439            false
440        }
441    }
442
443    pub fn supports_bf16(&self) -> bool {
444        match self {
445            Self::Cuda(_) | Self::Metal(_) => true,
446            Self::Cpu => false,
447            #[cfg(feature = "rocm")]
448            Self::Rocm(_) => true,
449            // Dozen/D3D12 Vulkan path on the 8060S has no native bf16; the
450            // backend is f32/u32-only, so default away from bf16.
451            #[cfg(feature = "vulkan")]
452            Self::Vulkan(_) => false,
453            // wgpu/WGSL path on the GB10 computes in f32/u32; no native bf16.
454            #[cfg(feature = "wgpu")]
455            Self::Wgpu(_) => false,
456        }
457    }
458
459    /// Return `BF16` for devices that support it, otherwise default to `F32`.
460    pub fn bf16_default_to_f32(&self) -> DType {
461        if self.supports_bf16() {
462            DType::BF16
463        } else {
464            DType::F32
465        }
466    }
467
468    pub fn cuda_if_available(ordinal: usize) -> Result<Self> {
469        if crate::utils::cuda_is_available() {
470            Self::new_cuda(ordinal)
471        } else {
472            Ok(Self::Cpu)
473        }
474    }
475
476    pub fn metal_if_available(ordinal: usize) -> Result<Self> {
477        if crate::utils::metal_is_available() {
478            Self::new_metal(ordinal)
479        } else {
480            Ok(Self::Cpu)
481        }
482    }
483
484    pub(crate) fn rand_uniform_f64(
485        &self,
486        lo: f64,
487        up: f64,
488        shape: &Shape,
489        dtype: DType,
490    ) -> Result<Storage> {
491        match self {
492            Device::Cpu => {
493                let storage = CpuDevice.rand_uniform(shape, dtype, lo, up)?;
494                Ok(Storage::Cpu(storage))
495            }
496            Device::Cuda(device) => {
497                // TODO: Remove the special case if we start supporting generating f16/bf16 directly.
498                if dtype == DType::F16 || dtype == DType::BF16 {
499                    let storage = device.rand_uniform(shape, DType::F32, lo, up)?;
500                    Storage::Cuda(storage).to_dtype(&crate::Layout::contiguous(shape), dtype)
501                } else {
502                    let storage = device.rand_uniform(shape, dtype, lo, up)?;
503                    Ok(Storage::Cuda(storage))
504                }
505            }
506            Device::Metal(device) => {
507                let storage = device.rand_uniform(shape, dtype, lo, up)?;
508                Ok(Storage::Metal(storage))
509            }
510            #[cfg(feature = "rocm")]
511            Device::Rocm(device) => {
512                let storage = device.rand_uniform(shape, dtype, lo, up)?;
513                Ok(Storage::Rocm(storage))
514            }
515            #[cfg(feature = "vulkan")]
516            Device::Vulkan(device) => {
517                let storage = device.rand_uniform(shape, dtype, lo, up)?;
518                Ok(Storage::Vulkan(storage))
519            }
520            #[cfg(feature = "wgpu")]
521            Device::Wgpu(device) => {
522                let storage = device.rand_uniform(shape, dtype, lo, up)?;
523                Ok(Storage::Wgpu(storage))
524            }
525        }
526    }
527
528    pub(crate) fn rand_uniform<T: crate::FloatDType>(
529        &self,
530        lo: T,
531        up: T,
532        shape: &Shape,
533    ) -> Result<Storage> {
534        self.rand_uniform_f64(lo.to_f64(), up.to_f64(), shape, T::DTYPE)
535    }
536
537    pub(crate) fn rand_normal_f64(
538        &self,
539        mean: f64,
540        std: f64,
541        shape: &Shape,
542        dtype: DType,
543    ) -> Result<Storage> {
544        match self {
545            Device::Cpu => {
546                let storage = CpuDevice.rand_normal(shape, dtype, mean, std)?;
547                Ok(Storage::Cpu(storage))
548            }
549            Device::Cuda(device) => {
550                // TODO: Remove the special case if we start supporting generating f16/bf16 directly.
551                if dtype == DType::F16 || dtype == DType::BF16 {
552                    let storage = device.rand_normal(shape, DType::F32, mean, std)?;
553                    Storage::Cuda(storage).to_dtype(&crate::Layout::contiguous(shape), dtype)
554                } else {
555                    let storage = device.rand_normal(shape, dtype, mean, std)?;
556                    Ok(Storage::Cuda(storage))
557                }
558            }
559            Device::Metal(device) => {
560                let storage = device.rand_normal(shape, dtype, mean, std)?;
561                Ok(Storage::Metal(storage))
562            }
563            #[cfg(feature = "rocm")]
564            Device::Rocm(device) => {
565                let storage = device.rand_normal(shape, dtype, mean, std)?;
566                Ok(Storage::Rocm(storage))
567            }
568            #[cfg(feature = "vulkan")]
569            Device::Vulkan(device) => {
570                let storage = device.rand_normal(shape, dtype, mean, std)?;
571                Ok(Storage::Vulkan(storage))
572            }
573            #[cfg(feature = "wgpu")]
574            Device::Wgpu(device) => {
575                let storage = device.rand_normal(shape, dtype, mean, std)?;
576                Ok(Storage::Wgpu(storage))
577            }
578        }
579    }
580
581    pub(crate) fn rand_normal<T: crate::FloatDType>(
582        &self,
583        mean: T,
584        std: T,
585        shape: &Shape,
586    ) -> Result<Storage> {
587        self.rand_normal_f64(mean.to_f64(), std.to_f64(), shape, T::DTYPE)
588    }
589
590    pub(crate) fn zeros(&self, shape: &Shape, dtype: DType) -> Result<Storage> {
591        match self {
592            Device::Cpu => {
593                let storage = CpuDevice.zeros_impl(shape, dtype)?;
594                Ok(Storage::Cpu(storage))
595            }
596            Device::Cuda(device) => {
597                let storage = device.zeros_impl(shape, dtype)?;
598                Ok(Storage::Cuda(storage))
599            }
600            Device::Metal(device) => {
601                let storage = device.zeros_impl(shape, dtype)?;
602                Ok(Storage::Metal(storage))
603            }
604            #[cfg(feature = "rocm")]
605            Device::Rocm(device) => {
606                let storage = device.zeros_impl(shape, dtype)?;
607                Ok(Storage::Rocm(storage))
608            }
609            #[cfg(feature = "vulkan")]
610            Device::Vulkan(device) => {
611                let storage = device.zeros_impl(shape, dtype)?;
612                Ok(Storage::Vulkan(storage))
613            }
614            #[cfg(feature = "wgpu")]
615            Device::Wgpu(device) => {
616                let storage = device.zeros_impl(shape, dtype)?;
617                Ok(Storage::Wgpu(storage))
618            }
619        }
620    }
621
622    pub(crate) unsafe fn alloc_uninit(&self, shape: &Shape, dtype: DType) -> Result<Storage> {
623        match self {
624            Device::Cpu => {
625                let storage = CpuDevice.alloc_uninit(shape, dtype)?;
626                Ok(Storage::Cpu(storage))
627            }
628            Device::Cuda(device) => {
629                let storage = device.alloc_uninit(shape, dtype)?;
630                Ok(Storage::Cuda(storage))
631            }
632            Device::Metal(device) => {
633                let storage = device.alloc_uninit(shape, dtype)?;
634                Ok(Storage::Metal(storage))
635            }
636            #[cfg(feature = "rocm")]
637            Device::Rocm(device) => {
638                let storage = device.alloc_uninit(shape, dtype)?;
639                Ok(Storage::Rocm(storage))
640            }
641            #[cfg(feature = "vulkan")]
642            Device::Vulkan(device) => {
643                let storage = device.alloc_uninit(shape, dtype)?;
644                Ok(Storage::Vulkan(storage))
645            }
646            #[cfg(feature = "wgpu")]
647            Device::Wgpu(device) => {
648                let storage = device.alloc_uninit(shape, dtype)?;
649                Ok(Storage::Wgpu(storage))
650            }
651        }
652    }
653
654    pub(crate) fn storage_from_slice<D: WithDType>(&self, data: &[D]) -> Result<Storage> {
655        match self {
656            Device::Cpu => Ok(Storage::Cpu(data.to_cpu_storage())),
657            Device::Cuda(device) => {
658                let storage = device.storage_from_slice(data)?;
659                Ok(Storage::Cuda(storage))
660            }
661            Device::Metal(device) => {
662                let storage = device.storage_from_slice(data)?;
663                Ok(Storage::Metal(storage))
664            }
665            #[cfg(feature = "rocm")]
666            Device::Rocm(device) => {
667                let storage = device.storage_from_slice(data)?;
668                Ok(Storage::Rocm(storage))
669            }
670            #[cfg(feature = "vulkan")]
671            Device::Vulkan(device) => {
672                let storage = device.storage_from_slice(data)?;
673                Ok(Storage::Vulkan(storage))
674            }
675            #[cfg(feature = "wgpu")]
676            Device::Wgpu(device) => {
677                let storage = device.storage_from_slice(data)?;
678                Ok(Storage::Wgpu(storage))
679            }
680        }
681    }
682
683    pub(crate) fn storage<A: NdArray>(&self, array: A) -> Result<Storage> {
684        match self {
685            Device::Cpu => Ok(Storage::Cpu(array.to_cpu_storage())),
686            Device::Cuda(device) => {
687                let storage = array.to_cpu_storage();
688                let storage = device.storage_from_cpu_storage_owned(storage)?;
689                Ok(Storage::Cuda(storage))
690            }
691            Device::Metal(device) => {
692                let storage = array.to_cpu_storage();
693                let storage = device.storage_from_cpu_storage_owned(storage)?;
694                Ok(Storage::Metal(storage))
695            }
696            #[cfg(feature = "rocm")]
697            Device::Rocm(device) => {
698                let storage = array.to_cpu_storage();
699                let storage = device.storage_from_cpu_storage_owned(storage)?;
700                Ok(Storage::Rocm(storage))
701            }
702            #[cfg(feature = "vulkan")]
703            Device::Vulkan(device) => {
704                let storage = array.to_cpu_storage();
705                let storage = device.storage_from_cpu_storage_owned(storage)?;
706                Ok(Storage::Vulkan(storage))
707            }
708            #[cfg(feature = "wgpu")]
709            Device::Wgpu(device) => {
710                let storage = array.to_cpu_storage();
711                let storage = device.storage_from_cpu_storage_owned(storage)?;
712                Ok(Storage::Wgpu(storage))
713            }
714        }
715    }
716
717    pub(crate) fn storage_owned<S: WithDType>(&self, data: Vec<S>) -> Result<Storage> {
718        match self {
719            Device::Cpu => Ok(Storage::Cpu(S::to_cpu_storage_owned(data))),
720            Device::Cuda(device) => {
721                let storage = S::to_cpu_storage_owned(data);
722                let storage = device.storage_from_cpu_storage_owned(storage)?;
723                Ok(Storage::Cuda(storage))
724            }
725            Device::Metal(device) => {
726                let storage = S::to_cpu_storage_owned(data);
727                let storage = device.storage_from_cpu_storage_owned(storage)?;
728                Ok(Storage::Metal(storage))
729            }
730            #[cfg(feature = "rocm")]
731            Device::Rocm(device) => {
732                let storage = S::to_cpu_storage_owned(data);
733                let storage = device.storage_from_cpu_storage_owned(storage)?;
734                Ok(Storage::Rocm(storage))
735            }
736            #[cfg(feature = "vulkan")]
737            Device::Vulkan(device) => {
738                let storage = S::to_cpu_storage_owned(data);
739                let storage = device.storage_from_cpu_storage_owned(storage)?;
740                Ok(Storage::Vulkan(storage))
741            }
742            #[cfg(feature = "wgpu")]
743            Device::Wgpu(device) => {
744                let storage = S::to_cpu_storage_owned(data);
745                let storage = device.storage_from_cpu_storage_owned(storage)?;
746                Ok(Storage::Wgpu(storage))
747            }
748        }
749    }
750
751    pub fn synchronize(&self) -> Result<()> {
752        match self {
753            Self::Cpu => Ok(()),
754            Self::Cuda(d) => d.synchronize(),
755            Self::Metal(d) => d.synchronize(),
756            #[cfg(feature = "rocm")]
757            Self::Rocm(d) => d.synchronize(),
758            #[cfg(feature = "vulkan")]
759            Self::Vulkan(d) => d.synchronize(),
760            #[cfg(feature = "wgpu")]
761            Self::Wgpu(d) => d.synchronize(),
762        }
763    }
764}
hanzo_ml/device.rs

hanzo_ml/
device.rs
