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    /// Run `f` with device specific context.
342    ///
343    /// On CPU this installs hanzo's private rayon thread pool for the
344    /// duration of `f`, keeping worker threads warm across the many short
345    /// parallel sections in a model forward pass. Currently noop for other backends.
346    pub fn with_context<F, R>(&self, f: F) -> R
347    where
348        F: FnOnce() -> R + Send,
349        R: Send,
350    {
351        match self {
352            Self::Cpu => crate::utils::with_threadpool(f),
353            _ => f(),
354        }
355    }
356
357    pub fn set_seed(&self, seed: u64) -> Result<()> {
358        match self {
359            Self::Cpu => CpuDevice.set_seed(seed),
360            Self::Cuda(c) => c.set_seed(seed),
361            Self::Metal(m) => m.set_seed(seed),
362            #[cfg(feature = "rocm")]
363            Self::Rocm(r) => r.set_seed(seed),
364            #[cfg(feature = "vulkan")]
365            Self::Vulkan(r) => r.set_seed(seed),
366            #[cfg(feature = "wgpu")]
367            Self::Wgpu(r) => r.set_seed(seed),
368        }
369    }
370
371    pub fn get_current_seed(&self) -> Result<u64> {
372        match self {
373            Self::Cpu => CpuDevice.get_current_seed(),
374            Self::Cuda(c) => c.get_current_seed(),
375            Self::Metal(m) => m.get_current_seed(),
376            #[cfg(feature = "rocm")]
377            Self::Rocm(r) => r.get_current_seed(),
378            #[cfg(feature = "vulkan")]
379            Self::Vulkan(r) => r.get_current_seed(),
380            #[cfg(feature = "wgpu")]
381            Self::Wgpu(r) => r.get_current_seed(),
382        }
383    }
384
385    pub fn same_device(&self, rhs: &Self) -> bool {
386        match (self, rhs) {
387            (Self::Cpu, Self::Cpu) => true,
388            (Self::Cuda(lhs), Self::Cuda(rhs)) => lhs.same_device(rhs),
389            (Self::Metal(lhs), Self::Metal(rhs)) => lhs.same_device(rhs),
390            #[cfg(feature = "rocm")]
391            (Self::Rocm(lhs), Self::Rocm(rhs)) => lhs.same_device(rhs),
392            #[cfg(feature = "vulkan")]
393            (Self::Vulkan(lhs), Self::Vulkan(rhs)) => lhs.same_device(rhs),
394            #[cfg(feature = "wgpu")]
395            (Self::Wgpu(lhs), Self::Wgpu(rhs)) => lhs.same_device(rhs),
396            _ => false,
397        }
398    }
399
400    pub fn location(&self) -> DeviceLocation {
401        match self {
402            Self::Cpu => DeviceLocation::Cpu,
403            Self::Cuda(device) => device.location(),
404            Device::Metal(device) => device.location(),
405            #[cfg(feature = "rocm")]
406            Self::Rocm(device) => device.location(),
407            #[cfg(feature = "vulkan")]
408            Self::Vulkan(device) => device.location(),
409            #[cfg(feature = "wgpu")]
410            Self::Wgpu(device) => device.location(),
411        }
412    }
413
414    pub fn is_cpu(&self) -> bool {
415        matches!(self, Self::Cpu)
416    }
417
418    pub fn is_cuda(&self) -> bool {
419        matches!(self, Self::Cuda(_))
420    }
421
422    pub fn is_metal(&self) -> bool {
423        matches!(self, Self::Metal(_))
424    }
425
426    pub fn is_rocm(&self) -> bool {
427        #[cfg(feature = "rocm")]
428        {
429            matches!(self, Self::Rocm(_))
430        }
431        #[cfg(not(feature = "rocm"))]
432        {
433            false
434        }
435    }
436
437    pub fn is_vulkan(&self) -> bool {
438        #[cfg(feature = "vulkan")]
439        {
440            matches!(self, Self::Vulkan(_))
441        }
442        #[cfg(not(feature = "vulkan"))]
443        {
444            false
445        }
446    }
447
448    pub fn is_wgpu(&self) -> bool {
449        #[cfg(feature = "wgpu")]
450        {
451            matches!(self, Self::Wgpu(_))
452        }
453        #[cfg(not(feature = "wgpu"))]
454        {
455            false
456        }
457    }
458
459    pub fn supports_bf16(&self) -> bool {
460        match self {
461            Self::Cuda(_) | Self::Metal(_) => true,
462            Self::Cpu => false,
463            #[cfg(feature = "rocm")]
464            Self::Rocm(_) => true,
465            // Dozen/D3D12 Vulkan path on the 8060S has no native bf16; the
466            // backend is f32/u32-only, so default away from bf16.
467            #[cfg(feature = "vulkan")]
468            Self::Vulkan(_) => false,
469            // wgpu/WGSL path on the GB10 computes in f32/u32; no native bf16.
470            #[cfg(feature = "wgpu")]
471            Self::Wgpu(_) => false,
472        }
473    }
474
475    /// Return `BF16` for devices that support it, otherwise default to `F32`.
476    pub fn bf16_default_to_f32(&self) -> DType {
477        if self.supports_bf16() {
478            DType::BF16
479        } else {
480            DType::F32
481        }
482    }
483
484    pub fn cuda_if_available(ordinal: usize) -> Result<Self> {
485        if crate::utils::cuda_is_available() {
486            Self::new_cuda(ordinal)
487        } else {
488            Ok(Self::Cpu)
489        }
490    }
491
492    pub fn metal_if_available(ordinal: usize) -> Result<Self> {
493        if crate::utils::metal_is_available() {
494            Self::new_metal(ordinal)
495        } else {
496            Ok(Self::Cpu)
497        }
498    }
499
500    pub(crate) fn rand_uniform_f64(
501        &self,
502        lo: f64,
503        up: f64,
504        shape: &Shape,
505        dtype: DType,
506    ) -> Result<Storage> {
507        match self {
508            Device::Cpu => {
509                let storage = CpuDevice.rand_uniform(shape, dtype, lo, up)?;
510                Ok(Storage::Cpu(storage))
511            }
512            Device::Cuda(device) => {
513                // TODO: Remove the special case if we start supporting generating f16/bf16 directly.
514                if dtype == DType::F16 || dtype == DType::BF16 {
515                    let storage = device.rand_uniform(shape, DType::F32, lo, up)?;
516                    Storage::Cuda(storage).to_dtype(&crate::Layout::contiguous(shape), dtype)
517                } else {
518                    let storage = device.rand_uniform(shape, dtype, lo, up)?;
519                    Ok(Storage::Cuda(storage))
520                }
521            }
522            Device::Metal(device) => {
523                let storage = device.rand_uniform(shape, dtype, lo, up)?;
524                Ok(Storage::Metal(storage))
525            }
526            #[cfg(feature = "rocm")]
527            Device::Rocm(device) => {
528                let storage = device.rand_uniform(shape, dtype, lo, up)?;
529                Ok(Storage::Rocm(storage))
530            }
531            #[cfg(feature = "vulkan")]
532            Device::Vulkan(device) => {
533                let storage = device.rand_uniform(shape, dtype, lo, up)?;
534                Ok(Storage::Vulkan(storage))
535            }
536            #[cfg(feature = "wgpu")]
537            Device::Wgpu(device) => {
538                let storage = device.rand_uniform(shape, dtype, lo, up)?;
539                Ok(Storage::Wgpu(storage))
540            }
541        }
542    }
543
544    pub(crate) fn rand_uniform<T: crate::FloatDType>(
545        &self,
546        lo: T,
547        up: T,
548        shape: &Shape,
549    ) -> Result<Storage> {
550        self.rand_uniform_f64(lo.to_f64(), up.to_f64(), shape, T::DTYPE)
551    }
552
553    pub(crate) fn rand_normal_f64(
554        &self,
555        mean: f64,
556        std: f64,
557        shape: &Shape,
558        dtype: DType,
559    ) -> Result<Storage> {
560        match self {
561            Device::Cpu => {
562                let storage = CpuDevice.rand_normal(shape, dtype, mean, std)?;
563                Ok(Storage::Cpu(storage))
564            }
565            Device::Cuda(device) => {
566                // TODO: Remove the special case if we start supporting generating f16/bf16 directly.
567                if dtype == DType::F16 || dtype == DType::BF16 {
568                    let storage = device.rand_normal(shape, DType::F32, mean, std)?;
569                    Storage::Cuda(storage).to_dtype(&crate::Layout::contiguous(shape), dtype)
570                } else {
571                    let storage = device.rand_normal(shape, dtype, mean, std)?;
572                    Ok(Storage::Cuda(storage))
573                }
574            }
575            Device::Metal(device) => {
576                let storage = device.rand_normal(shape, dtype, mean, std)?;
577                Ok(Storage::Metal(storage))
578            }
579            #[cfg(feature = "rocm")]
580            Device::Rocm(device) => {
581                let storage = device.rand_normal(shape, dtype, mean, std)?;
582                Ok(Storage::Rocm(storage))
583            }
584            #[cfg(feature = "vulkan")]
585            Device::Vulkan(device) => {
586                let storage = device.rand_normal(shape, dtype, mean, std)?;
587                Ok(Storage::Vulkan(storage))
588            }
589            #[cfg(feature = "wgpu")]
590            Device::Wgpu(device) => {
591                let storage = device.rand_normal(shape, dtype, mean, std)?;
592                Ok(Storage::Wgpu(storage))
593            }
594        }
595    }
596
597    pub(crate) fn rand_normal<T: crate::FloatDType>(
598        &self,
599        mean: T,
600        std: T,
601        shape: &Shape,
602    ) -> Result<Storage> {
603        self.rand_normal_f64(mean.to_f64(), std.to_f64(), shape, T::DTYPE)
604    }
605
606    pub(crate) fn zeros(&self, shape: &Shape, dtype: DType) -> Result<Storage> {
607        match self {
608            Device::Cpu => {
609                let storage = CpuDevice.zeros_impl(shape, dtype)?;
610                Ok(Storage::Cpu(storage))
611            }
612            Device::Cuda(device) => {
613                let storage = device.zeros_impl(shape, dtype)?;
614                Ok(Storage::Cuda(storage))
615            }
616            Device::Metal(device) => {
617                let storage = device.zeros_impl(shape, dtype)?;
618                Ok(Storage::Metal(storage))
619            }
620            #[cfg(feature = "rocm")]
621            Device::Rocm(device) => {
622                let storage = device.zeros_impl(shape, dtype)?;
623                Ok(Storage::Rocm(storage))
624            }
625            #[cfg(feature = "vulkan")]
626            Device::Vulkan(device) => {
627                let storage = device.zeros_impl(shape, dtype)?;
628                Ok(Storage::Vulkan(storage))
629            }
630            #[cfg(feature = "wgpu")]
631            Device::Wgpu(device) => {
632                let storage = device.zeros_impl(shape, dtype)?;
633                Ok(Storage::Wgpu(storage))
634            }
635        }
636    }
637
638    pub(crate) unsafe fn alloc_uninit(&self, shape: &Shape, dtype: DType) -> Result<Storage> {
639        match self {
640            Device::Cpu => {
641                let storage = CpuDevice.alloc_uninit(shape, dtype)?;
642                Ok(Storage::Cpu(storage))
643            }
644            Device::Cuda(device) => {
645                let storage = device.alloc_uninit(shape, dtype)?;
646                Ok(Storage::Cuda(storage))
647            }
648            Device::Metal(device) => {
649                let storage = device.alloc_uninit(shape, dtype)?;
650                Ok(Storage::Metal(storage))
651            }
652            #[cfg(feature = "rocm")]
653            Device::Rocm(device) => {
654                let storage = device.alloc_uninit(shape, dtype)?;
655                Ok(Storage::Rocm(storage))
656            }
657            #[cfg(feature = "vulkan")]
658            Device::Vulkan(device) => {
659                let storage = device.alloc_uninit(shape, dtype)?;
660                Ok(Storage::Vulkan(storage))
661            }
662            #[cfg(feature = "wgpu")]
663            Device::Wgpu(device) => {
664                let storage = device.alloc_uninit(shape, dtype)?;
665                Ok(Storage::Wgpu(storage))
666            }
667        }
668    }
669
670    pub(crate) fn storage_from_slice<D: WithDType>(&self, data: &[D]) -> Result<Storage> {
671        match self {
672            Device::Cpu => Ok(Storage::Cpu(data.to_cpu_storage())),
673            Device::Cuda(device) => {
674                let storage = device.storage_from_slice(data)?;
675                Ok(Storage::Cuda(storage))
676            }
677            Device::Metal(device) => {
678                let storage = device.storage_from_slice(data)?;
679                Ok(Storage::Metal(storage))
680            }
681            #[cfg(feature = "rocm")]
682            Device::Rocm(device) => {
683                let storage = device.storage_from_slice(data)?;
684                Ok(Storage::Rocm(storage))
685            }
686            #[cfg(feature = "vulkan")]
687            Device::Vulkan(device) => {
688                let storage = device.storage_from_slice(data)?;
689                Ok(Storage::Vulkan(storage))
690            }
691            #[cfg(feature = "wgpu")]
692            Device::Wgpu(device) => {
693                let storage = device.storage_from_slice(data)?;
694                Ok(Storage::Wgpu(storage))
695            }
696        }
697    }
698
699    pub(crate) fn storage<A: NdArray>(&self, array: A) -> Result<Storage> {
700        match self {
701            Device::Cpu => Ok(Storage::Cpu(array.to_cpu_storage())),
702            Device::Cuda(device) => {
703                let storage = array.to_cpu_storage();
704                let storage = device.storage_from_cpu_storage_owned(storage)?;
705                Ok(Storage::Cuda(storage))
706            }
707            Device::Metal(device) => {
708                let storage = array.to_cpu_storage();
709                let storage = device.storage_from_cpu_storage_owned(storage)?;
710                Ok(Storage::Metal(storage))
711            }
712            #[cfg(feature = "rocm")]
713            Device::Rocm(device) => {
714                let storage = array.to_cpu_storage();
715                let storage = device.storage_from_cpu_storage_owned(storage)?;
716                Ok(Storage::Rocm(storage))
717            }
718            #[cfg(feature = "vulkan")]
719            Device::Vulkan(device) => {
720                let storage = array.to_cpu_storage();
721                let storage = device.storage_from_cpu_storage_owned(storage)?;
722                Ok(Storage::Vulkan(storage))
723            }
724            #[cfg(feature = "wgpu")]
725            Device::Wgpu(device) => {
726                let storage = array.to_cpu_storage();
727                let storage = device.storage_from_cpu_storage_owned(storage)?;
728                Ok(Storage::Wgpu(storage))
729            }
730        }
731    }
732
733    pub(crate) fn storage_owned<S: WithDType>(&self, data: Vec<S>) -> Result<Storage> {
734        match self {
735            Device::Cpu => Ok(Storage::Cpu(S::to_cpu_storage_owned(data))),
736            Device::Cuda(device) => {
737                let storage = S::to_cpu_storage_owned(data);
738                let storage = device.storage_from_cpu_storage_owned(storage)?;
739                Ok(Storage::Cuda(storage))
740            }
741            Device::Metal(device) => {
742                let storage = S::to_cpu_storage_owned(data);
743                let storage = device.storage_from_cpu_storage_owned(storage)?;
744                Ok(Storage::Metal(storage))
745            }
746            #[cfg(feature = "rocm")]
747            Device::Rocm(device) => {
748                let storage = S::to_cpu_storage_owned(data);
749                let storage = device.storage_from_cpu_storage_owned(storage)?;
750                Ok(Storage::Rocm(storage))
751            }
752            #[cfg(feature = "vulkan")]
753            Device::Vulkan(device) => {
754                let storage = S::to_cpu_storage_owned(data);
755                let storage = device.storage_from_cpu_storage_owned(storage)?;
756                Ok(Storage::Vulkan(storage))
757            }
758            #[cfg(feature = "wgpu")]
759            Device::Wgpu(device) => {
760                let storage = S::to_cpu_storage_owned(data);
761                let storage = device.storage_from_cpu_storage_owned(storage)?;
762                Ok(Storage::Wgpu(storage))
763            }
764        }
765    }
766
767    /// Return memory the backend holds idle for reuse. A burst of work with large temporaries
768    /// (a prefill) calls it when done; it may synchronize the device, so not inside a forward.
769    pub fn trim(&self) {
770        #[cfg(feature = "rocm")]
771        if let Self::Rocm(d) = self {
772            d.trim();
773        }
774    }
775
776    pub fn synchronize(&self) -> Result<()> {
777        match self {
778            Self::Cpu => Ok(()),
779            Self::Cuda(d) => d.synchronize(),
780            Self::Metal(d) => d.synchronize(),
781            #[cfg(feature = "rocm")]
782            Self::Rocm(d) => d.synchronize(),
783            #[cfg(feature = "vulkan")]
784            Self::Vulkan(d) => d.synchronize(),
785            #[cfg(feature = "wgpu")]
786            Self::Wgpu(d) => d.synchronize(),
787        }
788    }
789}
hanzo_ml/device.rs

hanzo_ml/
device.rs
