hanzo_ml/
custom_op.rs

1use crate::op::{BackpropOp, Op};
2use crate::tensor::from_storage;
3#[cfg(feature = "rocm")]
4use crate::RocmStorage;
5#[cfg(feature = "vulkan")]
6use crate::VulkanStorage;
7#[cfg(feature = "wgpu")]
8use crate::WgpuStorage;
9use crate::{CpuStorage, CudaStorage, Layout, MetalStorage, Result, Shape, Tensor};
10use std::sync::Arc;
11
12/// Name a custom/inplace op that has no native Vulkan path when `VK_PROFILE` is set, before
13/// it bails. The size-only readback profiler can't attribute a missing op to a name; this surfaces
14/// the exact op (+ its shape) so a GPU re-run knows which `vulkan_fwd` override to add next. The
15/// env read is on the cold bail path only (the op errors out right after), so it's effectively
16/// zero-cost for ops that DO have a native path. Vulkan-only (the default impls it guards are
17/// `#[cfg(feature = "vulkan")]`), so it never touches other backends.
18#[cfg(feature = "vulkan")]
19fn log_vulkan_custom_op_bail(name: &str, l: &Layout) {
20    if std::env::var("VK_PROFILE")
21        .map(|v| v != "0")
22        .unwrap_or(false)
23    {
24        eprintln!(
25            "[VK_PROFILE] custom-op bail op={name} shape={:?} (no vulkan_fwd; would round-trip/err)",
26            l.shape().dims()
27        );
28    }
29}
30
31/// Unary ops that can be defined in user-land.
32pub trait CustomOp1 {
33    // Box<dyn> does not support const yet, so use a function to get the name.
34    fn name(&self) -> &'static str;
35
36    /// The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides,
37    /// offsets etc so the associated layout should be used to access it.
38    fn cpu_fwd(&self, storage: &CpuStorage, layout: &Layout) -> Result<(CpuStorage, Shape)>;
39
40    /// The forward pass, as run on a gpu device. Note that the storage can use arbitrary strides,
41    /// offsets etc so the associated layout should be used to access it.
42    fn cuda_fwd(&self, _storage: &CudaStorage, _layout: &Layout) -> Result<(CudaStorage, Shape)> {
43        Err(crate::Error::Cuda(
44            format!("no cuda implementation for {}", self.name()).into(),
45        ))
46    }
47
48    #[cfg(feature = "rocm")]
49    fn rocm_fwd(&self, _storage: &RocmStorage, _layout: &Layout) -> Result<(RocmStorage, Shape)> {
50        Err(crate::Error::Msg(format!(
51            "no rocm implementation for {}",
52            self.name()
53        )))
54    }
55    #[cfg(feature = "vulkan")]
56    fn vulkan_fwd(
57        &self,
58        _storage: &VulkanStorage,
59        _layout: &Layout,
60    ) -> Result<(VulkanStorage, Shape)> {
61        log_vulkan_custom_op_bail(self.name(), _layout);
62        Err(crate::Error::Msg(format!(
63            "no vulkan implementation for {}",
64            self.name()
65        )))
66    }
67    #[cfg(feature = "wgpu")]
68    fn wgpu_fwd(&self, _storage: &WgpuStorage, _layout: &Layout) -> Result<(WgpuStorage, Shape)> {
69        Err(crate::Error::Msg(format!(
70            "no wgpu implementation for {}",
71            self.name()
72        )))
73    }
74
75    /// The forward pass, as run on a metal gpu device. Note that the storage can use arbitrary strides,
76    /// offsets etc so the associated layout should be used to access it.
77    fn metal_fwd(
78        &self,
79        _storage: &MetalStorage,
80        _layout: &Layout,
81    ) -> Result<(MetalStorage, Shape)> {
82        Err(crate::Error::Metal(
83            format!("no metal implementation for {}", self.name()).into(),
84        ))
85    }
86
87    /// This function takes as argument the argument `arg` used in the forward pass, the result
88    /// produced by the forward operation `res` and the gradient of the result `grad_res`.
89    /// The function should return the gradient of the argument.
90    fn bwd(&self, _arg: &Tensor, _res: &Tensor, _grad_res: &Tensor) -> Result<Option<Tensor>> {
91        Err(crate::Error::BackwardNotSupported { op: self.name() })
92    }
93}
94
95pub trait CustomOp2 {
96    fn name(&self) -> &'static str;
97
98    /// The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides,
99    /// offsets etc so the associated layout should be used to access it.
100    fn cpu_fwd(
101        &self,
102        s1: &CpuStorage,
103        l1: &Layout,
104        s2: &CpuStorage,
105        l2: &Layout,
106    ) -> Result<(CpuStorage, Shape)>;
107
108    /// The forward pass, as run on a gpu device. Note that the storage can use arbitrary strides,
109    /// offsets etc so the associated layout should be used to access it.
110    fn cuda_fwd(
111        &self,
112        _: &CudaStorage,
113        _: &Layout,
114        _: &CudaStorage,
115        _: &Layout,
116    ) -> Result<(CudaStorage, Shape)> {
117        Err(crate::Error::Cuda(
118            format!("no cuda implementation for {}", self.name()).into(),
119        ))
120    }
121
122    #[cfg(feature = "rocm")]
123    fn rocm_fwd(
124        &self,
125        _: &RocmStorage,
126        _: &Layout,
127        _: &RocmStorage,
128        _: &Layout,
129    ) -> Result<(RocmStorage, Shape)> {
130        Err(crate::Error::Msg(format!(
131            "no rocm implementation for {}",
132            self.name()
133        )))
134    }
135    #[cfg(feature = "vulkan")]
136    fn vulkan_fwd(
137        &self,
138        _: &VulkanStorage,
139        l1: &Layout,
140        _: &VulkanStorage,
141        _: &Layout,
142    ) -> Result<(VulkanStorage, Shape)> {
143        log_vulkan_custom_op_bail(self.name(), l1);
144        Err(crate::Error::Msg(format!(
145            "no vulkan implementation for {}",
146            self.name()
147        )))
148    }
149    #[cfg(feature = "wgpu")]
150    fn wgpu_fwd(
151        &self,
152        _: &WgpuStorage,
153        _l1: &Layout,
154        _: &WgpuStorage,
155        _: &Layout,
156    ) -> Result<(WgpuStorage, Shape)> {
157        Err(crate::Error::Msg(format!(
158            "no wgpu implementation for {}",
159            self.name()
160        )))
161    }
162
163    /// The forward pass, as run on a metal gpu device. Note that the storage can use arbitrary strides,
164    /// offsets etc so the associated layout should be used to access it.
165    fn metal_fwd(
166        &self,
167        _: &MetalStorage,
168        _: &Layout,
169        _: &MetalStorage,
170        _: &Layout,
171    ) -> Result<(MetalStorage, Shape)> {
172        Err(crate::Error::Metal(
173            format!("no metal implementation for {}", self.name()).into(),
174        ))
175    }
176
177    fn bwd(
178        &self,
179        _arg1: &Tensor,
180        _arg2: &Tensor,
181        _res: &Tensor,
182        _grad_res: &Tensor,
183    ) -> Result<(Option<Tensor>, Option<Tensor>)> {
184        Err(crate::Error::BackwardNotSupported { op: self.name() })
185    }
186}
187
188pub trait CustomOp3 {
189    fn name(&self) -> &'static str;
190
191    /// The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides,
192    /// offsets etc so the associated layout should be used to access it.
193    fn cpu_fwd(
194        &self,
195        s1: &CpuStorage,
196        l1: &Layout,
197        s2: &CpuStorage,
198        l2: &Layout,
199        s3: &CpuStorage,
200        l3: &Layout,
201    ) -> Result<(CpuStorage, Shape)>;
202
203    /// The forward pass, as run on a gpu device. Note that the storage can use arbitrary strides,
204    /// offsets etc so the associated layout should be used to access it.
205    fn cuda_fwd(
206        &self,
207        _: &CudaStorage,
208        _: &Layout,
209        _: &CudaStorage,
210        _: &Layout,
211        _: &CudaStorage,
212        _: &Layout,
213    ) -> Result<(CudaStorage, Shape)> {
214        Err(crate::Error::Cuda(
215            format!("no cuda implementation for {}", self.name()).into(),
216        ))
217    }
218
219    #[cfg(feature = "rocm")]
220    fn rocm_fwd(
221        &self,
222        _: &RocmStorage,
223        _: &Layout,
224        _: &RocmStorage,
225        _: &Layout,
226        _: &RocmStorage,
227        _: &Layout,
228    ) -> Result<(RocmStorage, Shape)> {
229        Err(crate::Error::Msg(format!(
230            "no rocm implementation for {}",
231            self.name()
232        )))
233    }
234    #[cfg(feature = "vulkan")]
235    fn vulkan_fwd(
236        &self,
237        _: &VulkanStorage,
238        l1: &Layout,
239        _: &VulkanStorage,
240        _: &Layout,
241        _: &VulkanStorage,
242        _: &Layout,
243    ) -> Result<(VulkanStorage, Shape)> {
244        log_vulkan_custom_op_bail(self.name(), l1);
245        Err(crate::Error::Msg(format!(
246            "no vulkan implementation for {}",
247            self.name()
248        )))
249    }
250    #[cfg(feature = "wgpu")]
251    fn wgpu_fwd(
252        &self,
253        _: &WgpuStorage,
254        _l1: &Layout,
255        _: &WgpuStorage,
256        _: &Layout,
257        _: &WgpuStorage,
258        _: &Layout,
259    ) -> Result<(WgpuStorage, Shape)> {
260        Err(crate::Error::Msg(format!(
261            "no wgpu implementation for {}",
262            self.name()
263        )))
264    }
265
266    /// The forward pass, as run on a metal gpu device. Note that the storage can use arbitrary strides,
267    /// offsets etc so the associated layout should be used to access it.
268    fn metal_fwd(
269        &self,
270        _: &MetalStorage,
271        _: &Layout,
272        _: &MetalStorage,
273        _: &Layout,
274        _: &MetalStorage,
275        _: &Layout,
276    ) -> Result<(MetalStorage, Shape)> {
277        Err(crate::Error::Metal(
278            format!("no metal implementation for {}", self.name()).into(),
279        ))
280    }
281
282    fn bwd(
283        &self,
284        _arg1: &Tensor,
285        _arg2: &Tensor,
286        _arg3: &Tensor,
287        _res: &Tensor,
288        _grad_res: &Tensor,
289    ) -> Result<(Option<Tensor>, Option<Tensor>, Option<Tensor>)> {
290        Err(crate::Error::BackwardNotSupported { op: self.name() })
291    }
292}
293
294impl Tensor {
295    /// Applies a unary custom op without backward support
296    pub fn apply_op1_no_bwd<C: CustomOp1>(&self, c: &C) -> Result<Self> {
297        let (storage, shape) = self.storage().apply_op1(self.layout(), c)?;
298        Ok(from_storage(storage, shape, BackpropOp::none(), false))
299    }
300
301    /// Applies a binary custom op without backward support
302    pub fn apply_op2_no_bwd<C: CustomOp2>(&self, rhs: &Self, c: &C) -> Result<Self> {
303        let (storage, shape) =
304            self.storage()
305                .apply_op2(self.layout(), &rhs.storage(), rhs.layout(), c)?;
306        Ok(from_storage(storage, shape, BackpropOp::none(), false))
307    }
308
309    /// Applies a ternary custom op without backward support
310    pub fn apply_op3_no_bwd<C: CustomOp3>(&self, t2: &Self, t3: &Self, c: &C) -> Result<Self> {
311        let (storage, shape) = self.storage().apply_op3(
312            self.layout(),
313            &t2.storage(),
314            t2.layout(),
315            &t3.storage(),
316            t3.layout(),
317            c,
318        )?;
319        Ok(from_storage(storage, shape, BackpropOp::none(), false))
320    }
321
322    /// Applies a unary custom op.
323    pub fn apply_op1_arc(&self, c: Arc<Box<dyn CustomOp1 + Send + Sync>>) -> Result<Self> {
324        let (storage, shape) = self
325            .storage()
326            .apply_op1(self.layout(), c.as_ref().as_ref())?;
327        let op = BackpropOp::new1(self, |s| Op::CustomOp1(s, c.clone()));
328        Ok(from_storage(storage, shape, op, false))
329    }
330
331    pub fn apply_op1<C: 'static + CustomOp1 + Send + Sync>(&self, c: C) -> Result<Self> {
332        self.apply_op1_arc(Arc::new(Box::new(c)))
333    }
334
335    /// Applies a binary custom op.
336    pub fn apply_op2_arc(
337        &self,
338        rhs: &Self,
339        c: Arc<Box<dyn CustomOp2 + Send + Sync>>,
340    ) -> Result<Self> {
341        let (storage, shape) = self.storage().apply_op2(
342            self.layout(),
343            &rhs.storage(),
344            rhs.layout(),
345            c.as_ref().as_ref(),
346        )?;
347        let op = BackpropOp::new2(self, rhs, |t1, t2| Op::CustomOp2(t1, t2, c.clone()));
348        Ok(from_storage(storage, shape, op, false))
349    }
350
351    pub fn apply_op2<C: 'static + CustomOp2 + Send + Sync>(&self, r: &Self, c: C) -> Result<Self> {
352        self.apply_op2_arc(r, Arc::new(Box::new(c)))
353    }
354
355    /// Applies a ternary custom op.
356    pub fn apply_op3_arc(
357        &self,
358        t2: &Self,
359        t3: &Self,
360        c: Arc<Box<dyn CustomOp3 + Send + Sync>>,
361    ) -> Result<Self> {
362        let (storage, shape) = self.storage().apply_op3(
363            self.layout(),
364            &t2.storage(),
365            t2.layout(),
366            &t3.storage(),
367            t3.layout(),
368            c.as_ref().as_ref(),
369        )?;
370        let op = BackpropOp::new3(self, t2, t3, |t1, t2, t3| {
371            Op::CustomOp3(t1, t2, t3, c.clone())
372        });
373        Ok(from_storage(storage, shape, op, false))
374    }
375
376    pub fn apply_op3<C: 'static + CustomOp3 + Send + Sync>(
377        &self,
378        t2: &Self,
379        t3: &Self,
380        c: C,
381    ) -> Result<Self> {
382        self.apply_op3_arc(t2, t3, Arc::new(Box::new(c)))
383    }
384}
385
386// In place ops.
387
388/// Unary ops that can be defined in user-land.
389/// These ops work in place and as such back-prop is unsupported.
390pub trait InplaceOp1 {
391    // Box<dyn> does not support const yet, so use a function to get the name.
392    fn name(&self) -> &'static str;
393
394    /// The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides,
395    /// offsets etc so the associated layout should be used to access it.
396    fn cpu_fwd(&self, storage: &mut CpuStorage, layout: &Layout) -> Result<()>;
397
398    /// The forward pass, as run on a gpu device. Note that the storage can use arbitrary strides,
399    /// offsets etc so the associated layout should be used to access it.
400    fn cuda_fwd(&self, _storage: &mut CudaStorage, _layout: &Layout) -> Result<()> {
401        Err(crate::Error::Cuda(
402            format!("no cuda implementation for {}", self.name()).into(),
403        ))
404    }
405
406    #[cfg(feature = "rocm")]
407    fn rocm_fwd(&self, _storage: &mut RocmStorage, _layout: &Layout) -> Result<()> {
408        Err(crate::Error::Msg(format!(
409            "no rocm implementation for {}",
410            self.name()
411        )))
412    }
413    #[cfg(feature = "vulkan")]
414    fn vulkan_fwd(&self, _storage: &mut VulkanStorage, _layout: &Layout) -> Result<()> {
415        log_vulkan_custom_op_bail(self.name(), _layout);
416        Err(crate::Error::Msg(format!(
417            "no vulkan implementation for {}",
418            self.name()
419        )))
420    }
421    #[cfg(feature = "wgpu")]
422    fn wgpu_fwd(&self, _storage: &mut WgpuStorage, _layout: &Layout) -> Result<()> {
423        Err(crate::Error::Msg(format!(
424            "no wgpu implementation for {}",
425            self.name()
426        )))
427    }
428
429    /// The forward pass, as run on a metal gpu device. Note that the storage can use arbitrary strides,
430    /// offsets etc so the associated layout should be used to access it.
431    fn metal_fwd(&self, _storage: &mut MetalStorage, _layout: &Layout) -> Result<()> {
432        Err(crate::Error::Metal(
433            format!("no metal implementation for {}", self.name()).into(),
434        ))
435    }
436}
437
438pub trait InplaceOp2 {
439    fn name(&self) -> &'static str;
440
441    /// The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides,
442    /// offsets etc so the associated layout should be used to access it.
443    fn cpu_fwd(&self, s1: &mut CpuStorage, l1: &Layout, s2: &CpuStorage, l2: &Layout)
444        -> Result<()>;
445
446    /// The forward pass, as run on a gpu device. Note that the storage can use arbitrary strides,
447    /// offsets etc so the associated layout should be used to access it.
448    fn cuda_fwd(&self, _: &mut CudaStorage, _: &Layout, _: &CudaStorage, _: &Layout) -> Result<()> {
449        Err(crate::Error::Cuda(
450            format!("no cuda implementation for {}", self.name()).into(),
451        ))
452    }
453
454    #[cfg(feature = "rocm")]
455    fn rocm_fwd(&self, _: &mut RocmStorage, _: &Layout, _: &RocmStorage, _: &Layout) -> Result<()> {
456        Err(crate::Error::Msg(format!(
457            "no rocm implementation for {}",
458            self.name()
459        )))
460    }
461    #[cfg(feature = "vulkan")]
462    fn vulkan_fwd(
463        &self,
464        _: &mut VulkanStorage,
465        l1: &Layout,
466        _: &VulkanStorage,
467        _: &Layout,
468    ) -> Result<()> {
469        log_vulkan_custom_op_bail(self.name(), l1);
470        Err(crate::Error::Msg(format!(
471            "no vulkan implementation for {}",
472            self.name()
473        )))
474    }
475    #[cfg(feature = "wgpu")]
476    fn wgpu_fwd(
477        &self,
478        _: &mut WgpuStorage,
479        _l1: &Layout,
480        _: &WgpuStorage,
481        _: &Layout,
482    ) -> Result<()> {
483        Err(crate::Error::Msg(format!(
484            "no wgpu implementation for {}",
485            self.name()
486        )))
487    }
488
489    /// The forward pass, as run on a metal gpu device. Note that the storage can use arbitrary strides,
490    /// offsets etc so the associated layout should be used to access it.
491    fn metal_fwd(
492        &self,
493        _: &mut MetalStorage,
494        _: &Layout,
495        _: &MetalStorage,
496        _: &Layout,
497    ) -> Result<()> {
498        Err(crate::Error::Metal(
499            format!("no metal implementation for {}", self.name()).into(),
500        ))
501    }
502}
503
504pub trait InplaceOp3 {
505    fn name(&self) -> &'static str;
506
507    /// The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides,
508    /// offsets etc so the associated layout should be used to access it.
509    fn cpu_fwd(
510        &self,
511        s1: &mut CpuStorage,
512        l1: &Layout,
513        s2: &CpuStorage,
514        l2: &Layout,
515        s3: &CpuStorage,
516        l3: &Layout,
517    ) -> Result<()>;
518
519    /// The forward pass, as run on a gpu device. Note that the storage can use arbitrary strides,
520    /// offsets etc so the associated layout should be used to access it.
521    fn cuda_fwd(
522        &self,
523        _: &mut CudaStorage,
524        _: &Layout,
525        _: &CudaStorage,
526        _: &Layout,
527        _: &CudaStorage,
528        _: &Layout,
529    ) -> Result<()> {
530        Err(crate::Error::Cuda(
531            format!("no cuda implementation for {}", self.name()).into(),
532        ))
533    }
534
535    #[cfg(feature = "rocm")]
536    fn rocm_fwd(
537        &self,
538        _: &mut RocmStorage,
539        _: &Layout,
540        _: &RocmStorage,
541        _: &Layout,
542        _: &RocmStorage,
543        _: &Layout,
544    ) -> Result<()> {
545        Err(crate::Error::Msg(format!(
546            "no rocm implementation for {}",
547            self.name()
548        )))
549    }
550    #[cfg(feature = "vulkan")]
551    fn vulkan_fwd(
552        &self,
553        _: &mut VulkanStorage,
554        l1: &Layout,
555        _: &VulkanStorage,
556        _: &Layout,
557        _: &VulkanStorage,
558        _: &Layout,
559    ) -> Result<()> {
560        log_vulkan_custom_op_bail(self.name(), l1);
561        Err(crate::Error::Msg(format!(
562            "no vulkan implementation for {}",
563            self.name()
564        )))
565    }
566    #[cfg(feature = "wgpu")]
567    fn wgpu_fwd(
568        &self,
569        _: &mut WgpuStorage,
570        _l1: &Layout,
571        _: &WgpuStorage,
572        _: &Layout,
573        _: &WgpuStorage,
574        _: &Layout,
575    ) -> Result<()> {
576        Err(crate::Error::Msg(format!(
577            "no wgpu implementation for {}",
578            self.name()
579        )))
580    }
581
582    /// The forward pass, as run on a metal gpu device. Note that the storage can use arbitrary strides,
583    /// offsets etc so the associated layout should be used to access it.
584    fn metal_fwd(
585        &self,
586        _: &mut MetalStorage,
587        _: &Layout,
588        _: &MetalStorage,
589        _: &Layout,
590        _: &MetalStorage,
591        _: &Layout,
592    ) -> Result<()> {
593        Err(crate::Error::Metal(
594            format!("no metal implementation for {}", self.name()).into(),
595        ))
596    }
597}
598
599impl Tensor {
600    /// Applies a unary custom op in place.
601    pub fn inplace_op1<C: InplaceOp1>(&self, c: &C) -> Result<()> {
602        self.storage_mut().inplace_op1(self.layout(), c)
603    }
604
605    /// Applies a unary custom op in place (for the first tensor).
606    pub fn inplace_op2<C: InplaceOp2>(&self, rhs: &Self, c: &C) -> Result<()> {
607        self.storage_mut()
608            .inplace_op2(self.layout(), &rhs.storage(), rhs.layout(), c)
609    }
610
611    /// Applies a ternary custom op in place (for the first tensor).
612    pub fn inplace_op3<C: InplaceOp3>(&self, t2: &Self, t3: &Self, c: &C) -> Result<()> {
613        self.storage_mut().inplace_op3(
614            self.layout(),
615            &t2.storage(),
616            t2.layout(),
617            &t3.storage(),
618            t3.layout(),
619            c,
620        )
621    }
622}
623
624#[cfg(feature = "ug")]
625pub struct UgIOp1 {
626    name: &'static str,
627    #[cfg(feature = "cuda")]
628    func: cudarc::driver::CudaFunction,
629    #[cfg(feature = "metal")]
630    func: hanzo_metal_kernels::metal::ComputePipeline,
631}
632
633#[cfg(feature = "ug")]
634impl UgIOp1 {
635    #[allow(unused)]
636    #[cfg(all(not(target_arch = "wasm32"), not(target_os = "ios")))]
637    pub fn new(
638        name: &'static str,
639        kernel: hanzo_ug::lang::ssa::Kernel,
640        device: &crate::Device,
641    ) -> Result<Self> {
642        #[cfg(feature = "cuda")]
643        {
644            let device = device.as_cuda_device()?;
645            let func = device.compile(name, kernel)?;
646            Ok(Self {
647                name,
648                func: func.into_cuda_function(),
649            })
650        }
651        #[cfg(feature = "metal")]
652        {
653            let device = device.as_metal_device()?;
654            let func = device.compile(name, kernel)?;
655            Ok(Self { name, func })
656        }
657        #[cfg(not(any(feature = "cuda", feature = "metal")))]
658        {
659            Ok(Self { name })
660        }
661    }
662}
663
664#[cfg(feature = "ug")]
665impl InplaceOp1 for UgIOp1 {
666    fn name(&self) -> &'static str {
667        self.name
668    }
669
670    fn cpu_fwd(&self, _: &mut CpuStorage, _: &Layout) -> Result<()> {
671        crate::bail!("ug ops are only supported on metal/cuda at the moment")
672    }
673
674    #[cfg(feature = "metal")]
675    fn metal_fwd(&self, sto: &mut MetalStorage, layout: &Layout) -> Result<()> {
676        use crate::backend::BackendStorage;
677        use objc2_metal;
678
679        let elem_count = layout.shape().elem_count();
680        if sto.dtype() != crate::DType::F32 {
681            // TODO: support more dtypes.
682            crate::bail!("input is not a f32 tensor")
683        }
684        let device = sto.device();
685        let encoder = device.command_encoder()?;
686        encoder.set_compute_pipeline_state(&self.func);
687        let (g, b) = if elem_count.is_multiple_of(32) {
688            (elem_count / 32, 32)
689        } else {
690            (elem_count, 1)
691        };
692        let grid_dims = objc2_metal::MTLSize {
693            width: g,
694            height: 1,
695            depth: 1,
696        };
697        let group_dims = hanzo_metal_kernels::utils::get_block_dims(b, 1, 1);
698        encoder.set_output_buffer(0, Some(sto.buffer()), 0);
699        encoder.dispatch_threads(grid_dims, group_dims);
700
701        Ok(())
702    }
703
704    #[cfg(feature = "cuda")]
705    fn cuda_fwd(&self, sto: &mut CudaStorage, layout: &Layout) -> Result<()> {
706        use crate::cuda_backend::WrapErr;
707        use cudarc::driver::PushKernelArg;
708
709        let elem_count = layout.shape().elem_count();
710        let stream = sto.device.cuda_stream();
711        // TODO: support more dtypes.
712        let sto = sto.as_cuda_slice::<f32>()?;
713        let sto = match layout.contiguous_offsets() {
714            None => crate::bail!("input has to be contiguous"),
715            Some((o1, o2)) => sto.slice(o1..o2),
716        };
717        let (g, b) = if elem_count % 32 == 0 {
718            (elem_count / 32, 32)
719        } else {
720            (elem_count, 1)
721        };
722        let cfg = cudarc::driver::LaunchConfig {
723            grid_dim: (g as u32, 1, 1),
724            block_dim: (b as u32, 1, 1),
725            shared_mem_bytes: 0,
726        };
727        let mut builder = stream.launch_builder(&self.func);
728        builder.arg(&sto);
729        unsafe { builder.launch(cfg) }.w()?;
730        Ok(())
731    }
732}
hanzo_ml/custom_op.rs

hanzo_ml/
custom_op.rs
