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

hanzo_ml/
custom_op.rs
