hanzo_ml/
custom_op.rs

1use crate::layout::LayoutRelation;
2use crate::op::{BackpropOp, Op};
3use crate::tensor::from_storage;
4#[cfg(feature = "rocm")]
5use crate::RocmStorage;
6#[cfg(feature = "vulkan")]
7use crate::VulkanStorage;
8#[cfg(feature = "wgpu")]
9use crate::WgpuStorage;
10use crate::{bail, CpuStorage, CudaStorage, Layout, MetalStorage, Result, Shape, Storage, Tensor};
11use std::sync::Arc;
12
13/// Name a custom/inplace op that has no native Vulkan path when `VK_PROFILE` is set, before
14/// it bails. The size-only readback profiler can't attribute a missing op to a name; this surfaces
15/// the exact op (+ its shape) so a GPU re-run knows which `vulkan_fwd` override to add next. The
16/// env read is on the cold bail path only (the op errors out right after), so it's effectively
17/// zero-cost for ops that DO have a native path. Vulkan-only (the default impls it guards are
18/// `#[cfg(feature = "vulkan")]`), so it never touches other backends.
19#[cfg(feature = "vulkan")]
20fn log_vulkan_custom_op_bail(name: &str, l: &Layout) {
21    if std::env::var("VK_PROFILE")
22        .map(|v| v != "0")
23        .unwrap_or(false)
24    {
25        eprintln!(
26            "[VK_PROFILE] custom-op bail op={name} shape={:?} (no vulkan_fwd; would round-trip/err)",
27            l.shape().dims()
28        );
29    }
30}
31
32/// Unary ops that can be defined in user-land.
33pub trait CustomOp1 {
34    // Box<dyn> does not support const yet, so use a function to get the name.
35    fn name(&self) -> &'static str;
36
37    /// The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides,
38    /// offsets etc so the associated layout should be used to access it.
39    fn cpu_fwd(&self, storage: &CpuStorage, layout: &Layout) -> Result<(CpuStorage, Shape)>;
40
41    /// The forward pass, as run on a gpu device. Note that the storage can use arbitrary strides,
42    /// offsets etc so the associated layout should be used to access it.
43    fn cuda_fwd(&self, _storage: &CudaStorage, _layout: &Layout) -> Result<(CudaStorage, Shape)> {
44        Err(crate::Error::Cuda(
45            format!("no cuda implementation for {}", self.name()).into(),
46        ))
47    }
48
49    #[cfg(feature = "rocm")]
50    fn rocm_fwd(&self, _storage: &RocmStorage, _layout: &Layout) -> Result<(RocmStorage, Shape)> {
51        Err(crate::Error::Msg(format!(
52            "no rocm implementation for {}",
53            self.name()
54        )))
55    }
56    #[cfg(feature = "vulkan")]
57    fn vulkan_fwd(
58        &self,
59        _storage: &VulkanStorage,
60        _layout: &Layout,
61    ) -> Result<(VulkanStorage, Shape)> {
62        log_vulkan_custom_op_bail(self.name(), _layout);
63        Err(crate::Error::Msg(format!(
64            "no vulkan implementation for {}",
65            self.name()
66        )))
67    }
68    #[cfg(feature = "wgpu")]
69    fn wgpu_fwd(&self, _storage: &WgpuStorage, _layout: &Layout) -> Result<(WgpuStorage, Shape)> {
70        Err(crate::Error::Msg(format!(
71            "no wgpu implementation for {}",
72            self.name()
73        )))
74    }
75
76    /// The forward pass, as run on a metal gpu device. Note that the storage can use arbitrary strides,
77    /// offsets etc so the associated layout should be used to access it.
78    fn metal_fwd(
79        &self,
80        _storage: &MetalStorage,
81        _layout: &Layout,
82    ) -> Result<(MetalStorage, Shape)> {
83        Err(crate::Error::Metal(
84            format!("no metal implementation for {}", self.name()).into(),
85        ))
86    }
87
88    /// This function takes as argument the argument `arg` used in the forward pass, the result
89    /// produced by the forward operation `res` and the gradient of the result `grad_res`.
90    /// The function should return the gradient of the argument.
91    fn bwd(&self, _arg: &Tensor, _res: &Tensor, _grad_res: &Tensor) -> Result<Option<Tensor>> {
92        Err(crate::Error::BackwardNotSupported { op: self.name() })
93    }
94}
95
96pub trait CustomOp2 {
97    fn name(&self) -> &'static str;
98
99    /// The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides,
100    /// offsets etc so the associated layout should be used to access it.
101    fn cpu_fwd(
102        &self,
103        s1: &CpuStorage,
104        l1: &Layout,
105        s2: &CpuStorage,
106        l2: &Layout,
107    ) -> Result<(CpuStorage, Shape)>;
108
109    /// The forward pass, as run on a gpu device. Note that the storage can use arbitrary strides,
110    /// offsets etc so the associated layout should be used to access it.
111    fn cuda_fwd(
112        &self,
113        _: &CudaStorage,
114        _: &Layout,
115        _: &CudaStorage,
116        _: &Layout,
117    ) -> Result<(CudaStorage, Shape)> {
118        Err(crate::Error::Cuda(
119            format!("no cuda implementation for {}", self.name()).into(),
120        ))
121    }
122
123    #[cfg(feature = "rocm")]
124    fn rocm_fwd(
125        &self,
126        _: &RocmStorage,
127        _: &Layout,
128        _: &RocmStorage,
129        _: &Layout,
130    ) -> Result<(RocmStorage, Shape)> {
131        Err(crate::Error::Msg(format!(
132            "no rocm implementation for {}",
133            self.name()
134        )))
135    }
136    #[cfg(feature = "vulkan")]
137    fn vulkan_fwd(
138        &self,
139        _: &VulkanStorage,
140        l1: &Layout,
141        _: &VulkanStorage,
142        _: &Layout,
143    ) -> Result<(VulkanStorage, Shape)> {
144        log_vulkan_custom_op_bail(self.name(), l1);
145        Err(crate::Error::Msg(format!(
146            "no vulkan implementation for {}",
147            self.name()
148        )))
149    }
150    #[cfg(feature = "wgpu")]
151    fn wgpu_fwd(
152        &self,
153        _: &WgpuStorage,
154        _l1: &Layout,
155        _: &WgpuStorage,
156        _: &Layout,
157    ) -> Result<(WgpuStorage, Shape)> {
158        Err(crate::Error::Msg(format!(
159            "no wgpu implementation for {}",
160            self.name()
161        )))
162    }
163
164    /// The forward pass, as run on a metal gpu device. Note that the storage can use arbitrary strides,
165    /// offsets etc so the associated layout should be used to access it.
166    fn metal_fwd(
167        &self,
168        _: &MetalStorage,
169        _: &Layout,
170        _: &MetalStorage,
171        _: &Layout,
172    ) -> Result<(MetalStorage, Shape)> {
173        Err(crate::Error::Metal(
174            format!("no metal implementation for {}", self.name()).into(),
175        ))
176    }
177
178    fn bwd(
179        &self,
180        _arg1: &Tensor,
181        _arg2: &Tensor,
182        _res: &Tensor,
183        _grad_res: &Tensor,
184    ) -> Result<(Option<Tensor>, Option<Tensor>)> {
185        Err(crate::Error::BackwardNotSupported { op: self.name() })
186    }
187}
188
189pub trait CustomOp3 {
190    fn name(&self) -> &'static str;
191
192    /// The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides,
193    /// offsets etc so the associated layout should be used to access it.
194    fn cpu_fwd(
195        &self,
196        s1: &CpuStorage,
197        l1: &Layout,
198        s2: &CpuStorage,
199        l2: &Layout,
200        s3: &CpuStorage,
201        l3: &Layout,
202    ) -> Result<(CpuStorage, Shape)>;
203
204    /// The forward pass, as run on a gpu device. Note that the storage can use arbitrary strides,
205    /// offsets etc so the associated layout should be used to access it.
206    fn cuda_fwd(
207        &self,
208        _: &CudaStorage,
209        _: &Layout,
210        _: &CudaStorage,
211        _: &Layout,
212        _: &CudaStorage,
213        _: &Layout,
214    ) -> Result<(CudaStorage, Shape)> {
215        Err(crate::Error::Cuda(
216            format!("no cuda implementation for {}", self.name()).into(),
217        ))
218    }
219
220    #[cfg(feature = "rocm")]
221    fn rocm_fwd(
222        &self,
223        _: &RocmStorage,
224        _: &Layout,
225        _: &RocmStorage,
226        _: &Layout,
227        _: &RocmStorage,
228        _: &Layout,
229    ) -> Result<(RocmStorage, Shape)> {
230        Err(crate::Error::Msg(format!(
231            "no rocm implementation for {}",
232            self.name()
233        )))
234    }
235    #[cfg(feature = "vulkan")]
236    fn vulkan_fwd(
237        &self,
238        _: &VulkanStorage,
239        l1: &Layout,
240        _: &VulkanStorage,
241        _: &Layout,
242        _: &VulkanStorage,
243        _: &Layout,
244    ) -> Result<(VulkanStorage, Shape)> {
245        log_vulkan_custom_op_bail(self.name(), l1);
246        Err(crate::Error::Msg(format!(
247            "no vulkan implementation for {}",
248            self.name()
249        )))
250    }
251    #[cfg(feature = "wgpu")]
252    fn wgpu_fwd(
253        &self,
254        _: &WgpuStorage,
255        _l1: &Layout,
256        _: &WgpuStorage,
257        _: &Layout,
258        _: &WgpuStorage,
259        _: &Layout,
260    ) -> Result<(WgpuStorage, Shape)> {
261        Err(crate::Error::Msg(format!(
262            "no wgpu implementation for {}",
263            self.name()
264        )))
265    }
266
267    /// The forward pass, as run on a metal gpu device. Note that the storage can use arbitrary strides,
268    /// offsets etc so the associated layout should be used to access it.
269    fn metal_fwd(
270        &self,
271        _: &MetalStorage,
272        _: &Layout,
273        _: &MetalStorage,
274        _: &Layout,
275        _: &MetalStorage,
276        _: &Layout,
277    ) -> Result<(MetalStorage, Shape)> {
278        Err(crate::Error::Metal(
279            format!("no metal implementation for {}", self.name()).into(),
280        ))
281    }
282
283    fn bwd(
284        &self,
285        _arg1: &Tensor,
286        _arg2: &Tensor,
287        _arg3: &Tensor,
288        _res: &Tensor,
289        _grad_res: &Tensor,
290    ) -> Result<(Option<Tensor>, Option<Tensor>, Option<Tensor>)> {
291        Err(crate::Error::BackwardNotSupported { op: self.name() })
292    }
293}
294
295impl Tensor {
296    /// Applies a unary custom op without backward support
297    pub fn apply_op1_no_bwd<C: CustomOp1>(&self, c: &C) -> Result<Self> {
298        let (storage, shape) = self.storage().apply_op1(self.layout(), c)?;
299        Ok(from_storage(storage, shape, BackpropOp::none(), false))
300    }
301
302    /// Applies a binary custom op without backward support
303    pub fn apply_op2_no_bwd<C: CustomOp2>(&self, rhs: &Self, c: &C) -> Result<Self> {
304        let (storage, shape) =
305            self.storage()
306                .apply_op2(self.layout(), &rhs.storage(), rhs.layout(), c)?;
307        Ok(from_storage(storage, shape, BackpropOp::none(), false))
308    }
309
310    /// Applies a ternary custom op without backward support
311    pub fn apply_op3_no_bwd<C: CustomOp3>(&self, t2: &Self, t3: &Self, c: &C) -> Result<Self> {
312        let (storage, shape) = self.storage().apply_op3(
313            self.layout(),
314            &t2.storage(),
315            t2.layout(),
316            &t3.storage(),
317            t3.layout(),
318            c,
319        )?;
320        Ok(from_storage(storage, shape, BackpropOp::none(), false))
321    }
322
323    /// Applies a unary custom op.
324    pub fn apply_op1_arc(&self, c: Arc<Box<dyn CustomOp1 + Send + Sync>>) -> Result<Self> {
325        let (storage, shape) = self
326            .storage()
327            .apply_op1(self.layout(), c.as_ref().as_ref())?;
328        let op = BackpropOp::new1(self, |s| Op::CustomOp1(s, c.clone()));
329        Ok(from_storage(storage, shape, op, false))
330    }
331
332    pub fn apply_op1<C: 'static + CustomOp1 + Send + Sync>(&self, c: C) -> Result<Self> {
333        self.apply_op1_arc(Arc::new(Box::new(c)))
334    }
335
336    /// Applies a binary custom op.
337    pub fn apply_op2_arc(
338        &self,
339        rhs: &Self,
340        c: Arc<Box<dyn CustomOp2 + Send + Sync>>,
341    ) -> Result<Self> {
342        let (storage, shape) = self.storage().apply_op2(
343            self.layout(),
344            &rhs.storage(),
345            rhs.layout(),
346            c.as_ref().as_ref(),
347        )?;
348        let op = BackpropOp::new2(self, rhs, |t1, t2| Op::CustomOp2(t1, t2, c.clone()));
349        Ok(from_storage(storage, shape, op, false))
350    }
351
352    pub fn apply_op2<C: 'static + CustomOp2 + Send + Sync>(&self, r: &Self, c: C) -> Result<Self> {
353        self.apply_op2_arc(r, Arc::new(Box::new(c)))
354    }
355
356    /// Applies a ternary custom op.
357    pub fn apply_op3_arc(
358        &self,
359        t2: &Self,
360        t3: &Self,
361        c: Arc<Box<dyn CustomOp3 + Send + Sync>>,
362    ) -> Result<Self> {
363        let (storage, shape) = self.storage().apply_op3(
364            self.layout(),
365            &t2.storage(),
366            t2.layout(),
367            &t3.storage(),
368            t3.layout(),
369            c.as_ref().as_ref(),
370        )?;
371        let op = BackpropOp::new3(self, t2, t3, |t1, t2, t3| {
372            Op::CustomOp3(t1, t2, t3, c.clone())
373        });
374        Ok(from_storage(storage, shape, op, false))
375    }
376
377    pub fn apply_op3<C: 'static + CustomOp3 + Send + Sync>(
378        &self,
379        t2: &Self,
380        t3: &Self,
381        c: C,
382    ) -> Result<Self> {
383        self.apply_op3_arc(t2, t3, Arc::new(Box::new(c)))
384    }
385}
386
387// In place ops.
388
389/// Unary ops that can be defined in user-land.
390/// These ops work in place and as such back-prop is unsupported.
391pub trait InplaceOp1 {
392    // Box<dyn> does not support const yet, so use a function to get the name.
393    fn name(&self) -> &'static str;
394
395    /// The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides,
396    /// offsets etc so the associated layout should be used to access it.
397    fn cpu_fwd(&self, storage: &mut CpuStorage, layout: &Layout) -> Result<()>;
398
399    /// The forward pass, as run on a gpu device. Note that the storage can use arbitrary strides,
400    /// offsets etc so the associated layout should be used to access it.
401    fn cuda_fwd(&self, _storage: &mut CudaStorage, _layout: &Layout) -> Result<()> {
402        Err(crate::Error::Cuda(
403            format!("no cuda implementation for {}", self.name()).into(),
404        ))
405    }
406
407    #[cfg(feature = "rocm")]
408    fn rocm_fwd(&self, _storage: &mut RocmStorage, _layout: &Layout) -> Result<()> {
409        Err(crate::Error::Msg(format!(
410            "no rocm implementation for {}",
411            self.name()
412        )))
413    }
414    #[cfg(feature = "vulkan")]
415    fn vulkan_fwd(&self, _storage: &mut VulkanStorage, _layout: &Layout) -> Result<()> {
416        log_vulkan_custom_op_bail(self.name(), _layout);
417        Err(crate::Error::Msg(format!(
418            "no vulkan implementation for {}",
419            self.name()
420        )))
421    }
422    #[cfg(feature = "wgpu")]
423    fn wgpu_fwd(&self, _storage: &mut WgpuStorage, _layout: &Layout) -> Result<()> {
424        Err(crate::Error::Msg(format!(
425            "no wgpu implementation for {}",
426            self.name()
427        )))
428    }
429
430    /// The forward pass, as run on a metal gpu device. Note that the storage can use arbitrary strides,
431    /// offsets etc so the associated layout should be used to access it.
432    fn metal_fwd(&self, _storage: &mut MetalStorage, _layout: &Layout) -> Result<()> {
433        Err(crate::Error::Metal(
434            format!("no metal implementation for {}", self.name()).into(),
435        ))
436    }
437}
438
439/// In-place ops that can be defined in user-land.
440/// These ops work in-place and as such back-propagation is unsupported.
441pub trait InplaceOpN<const N: usize> {
442    fn name(&self) -> &'static str;
443
444    /// Defines the source access pattern of this in-place op.
445    /// Defaults to `None`, which rejects all in-place source aliasing.
446    fn src_access_pattern(&self) -> Option<AccessPattern> {
447        None
448    }
449
450    fn cpu_fwd(
451        &self,
452        dst: &mut CpuStorage,
453        dst_l: &Layout,
454        srcs: [(&CpuStorage, &Layout); N],
455    ) -> Result<()> {
456        let _ = (dst, dst_l, srcs);
457        bail!("no cpu implementation for {}", self.name())
458    }
459
460    fn cpu_fwd_aliased(
461        &self,
462        dst: &mut CpuStorage,
463        dst_l: &Layout,
464        srcs: [(Src<'_, CpuStorage>, &Layout); N],
465    ) -> Result<()> {
466        let _ = (dst, dst_l, srcs);
467        bail!("no aliased cpu implementation for {}", self.name())
468    }
469
470    fn cuda_fwd(
471        &self,
472        dst: &mut CudaStorage,
473        dst_l: &Layout,
474        srcs: [(&CudaStorage, &Layout); N],
475    ) -> Result<()> {
476        let _ = (dst, dst_l, srcs);
477        bail!("no cuda implementation for {}", self.name())
478    }
479
480    fn cuda_fwd_aliased(
481        &self,
482        dst: &mut CudaStorage,
483        dst_l: &Layout,
484        srcs: [(Src<'_, CudaStorage>, &Layout); N],
485    ) -> Result<()> {
486        let _ = (dst, dst_l, srcs);
487        bail!("no aliased cpu implementation for {}", self.name())
488    }
489
490    fn metal_fwd(
491        &self,
492        dst: &mut MetalStorage,
493        dst_l: &Layout,
494        srcs: [(&MetalStorage, &Layout); N],
495    ) -> Result<()> {
496        let _ = (dst, dst_l, srcs);
497        bail!("no metal implementation for {}", self.name())
498    }
499
500    fn metal_fwd_aliased(
501        &self,
502        dst: &mut MetalStorage,
503        dst_l: &Layout,
504        srcs: [(Src<'_, MetalStorage>, &Layout); N],
505    ) -> Result<()> {
506        let _ = (dst, dst_l, srcs);
507        bail!("no aliased metal implementation for {}", self.name())
508    }
509
510    #[cfg(feature = "rocm")]
511    fn rocm_fwd(
512        &self,
513        dst: &mut RocmStorage,
514        dst_l: &Layout,
515        srcs: [(&RocmStorage, &Layout); N],
516    ) -> Result<()> {
517        let _ = (dst, dst_l, srcs);
518        bail!("no rocm implementation for {}", self.name())
519    }
520
521    #[cfg(feature = "rocm")]
522    fn rocm_fwd_aliased(
523        &self,
524        dst: &mut RocmStorage,
525        dst_l: &Layout,
526        srcs: [(Src<'_, RocmStorage>, &Layout); N],
527    ) -> Result<()> {
528        let _ = (dst, dst_l, srcs);
529        bail!("no aliased rocm implementation for {}", self.name())
530    }
531
532    #[cfg(feature = "vulkan")]
533    fn vulkan_fwd(
534        &self,
535        dst: &mut VulkanStorage,
536        dst_l: &Layout,
537        srcs: [(&VulkanStorage, &Layout); N],
538    ) -> Result<()> {
539        log_vulkan_custom_op_bail(self.name(), dst_l);
540        let _ = (dst, dst_l, srcs);
541        bail!("no vulkan implementation for {}", self.name())
542    }
543
544    #[cfg(feature = "vulkan")]
545    fn vulkan_fwd_aliased(
546        &self,
547        dst: &mut VulkanStorage,
548        dst_l: &Layout,
549        srcs: [(Src<'_, VulkanStorage>, &Layout); N],
550    ) -> Result<()> {
551        log_vulkan_custom_op_bail(self.name(), dst_l);
552        let _ = (dst, dst_l, srcs);
553        bail!("no aliased vulkan implementation for {}", self.name())
554    }
555
556    #[cfg(feature = "wgpu")]
557    fn wgpu_fwd(
558        &self,
559        dst: &mut WgpuStorage,
560        dst_l: &Layout,
561        srcs: [(&WgpuStorage, &Layout); N],
562    ) -> Result<()> {
563        let _ = (dst, dst_l, srcs);
564        bail!("no wgpu implementation for {}", self.name())
565    }
566
567    #[cfg(feature = "wgpu")]
568    fn wgpu_fwd_aliased(
569        &self,
570        dst: &mut WgpuStorage,
571        dst_l: &Layout,
572        srcs: [(Src<'_, WgpuStorage>, &Layout); N],
573    ) -> Result<()> {
574        let _ = (dst, dst_l, srcs);
575        bail!("no aliased wgpu implementation for {}", self.name())
576    }
577}
578
579#[derive(Debug)]
580pub enum Src<'a, S> {
581    Distinct(&'a S),
582    Aliased(LayoutRelation),
583}
584
585impl<S> Copy for Src<'_, S> {}
586
587impl<S> Clone for Src<'_, S> {
588    fn clone(&self) -> Self {
589        *self
590    }
591}
592
593// If all sources are distinct we return the entire array of tensors without the `Src` wrapper.
594pub(crate) fn all_distinct<'a, B, const N: usize>(
595    srcs: &[(Src<'a, B>, &'a Layout); N],
596) -> Option<[(&'a B, &'a Layout); N]> {
597    srcs.iter()
598        .all(|(s, _)| matches!(s, Src::Distinct(_)))
599        .then(|| {
600            std::array::from_fn(|i| match srcs[i] {
601                (Src::Distinct(s), l) => (s, l),
602                _ => unreachable!("checked immediately above"),
603            })
604        })
605}
606
607/// Indicates which indices a kernel reads, relative to the destination indices it writes.
608///
609/// When used with [`crate::LayoutRelation`] we can describe safe access patterns.
610/// For example `Elementwise` is safe to use with both `LayoutRelation::Identical` and `LayoutRelation::Disjoint`,
611/// while `Arbitrary` is only guaranteed to be safe with `LayoutRelation::Disjoint`.
612#[derive(Copy, Clone, PartialEq, Eq, Debug)]
613pub enum AccessPattern {
614    /// Reads only source index `i` when writing destination index `i`.
615    Elementwise,
616    /// Reads arbitrary source indices.
617    Arbitrary,
618}
619
620impl AccessPattern {
621    fn supports(self, rel: LayoutRelation) -> bool {
622        matches!(
623            (self, rel),
624            (
625                AccessPattern::Elementwise,
626                LayoutRelation::Identical | LayoutRelation::Disjoint
627            ) | (AccessPattern::Arbitrary, LayoutRelation::Disjoint)
628        )
629    }
630}
631
632macro_rules! forward_op1 {
633    ($fwd:ident, $fwd_aliased:ident, $storage:ty) => {
634        fn $fwd(
635            &self,
636            dst: &mut $storage,
637            dl: &Layout,
638            _: [(&$storage, &Layout); 0],
639        ) -> Result<()> {
640            InplaceOp1::$fwd(self, dst, dl)
641        }
642
643        fn $fwd_aliased(
644            &self,
645            dst: &mut $storage,
646            dl: &Layout,
647            _: [(Src<'_, $storage>, &Layout); 0],
648        ) -> Result<()> {
649            InplaceOp1::$fwd(self, dst, dl)
650        }
651    };
652}
653
654impl<C: InplaceOp1> InplaceOpN<0> for C {
655    fn name(&self) -> &'static str {
656        InplaceOp1::name(self)
657    }
658
659    forward_op1!(cpu_fwd, cpu_fwd_aliased, CpuStorage);
660    forward_op1!(cuda_fwd, cuda_fwd_aliased, CudaStorage);
661    forward_op1!(metal_fwd, metal_fwd_aliased, MetalStorage);
662    #[cfg(feature = "rocm")]
663    forward_op1!(rocm_fwd, rocm_fwd_aliased, RocmStorage);
664    #[cfg(feature = "vulkan")]
665    forward_op1!(vulkan_fwd, vulkan_fwd_aliased, VulkanStorage);
666    #[cfg(feature = "wgpu")]
667    forward_op1!(wgpu_fwd, wgpu_fwd_aliased, WgpuStorage);
668}
669
670macro_rules! forward_op2 {
671    ($fwd:ident, $fwd_aliased:ident, $storage:ty) => {
672        fn $fwd(
673            &self,
674            dst: &mut $storage,
675            dl: &Layout,
676            srcs: [(&$storage, &Layout); 1],
677        ) -> Result<()> {
678            let [(s, sl)] = srcs;
679            InplaceOp2::$fwd(self, dst, dl, s, sl)
680        }
681
682        fn $fwd_aliased(
683            &self,
684            dst: &mut $storage,
685            dl: &Layout,
686            srcs: [(Src<'_, $storage>, &Layout); 1],
687        ) -> Result<()> {
688            match srcs {
689                [(Src::Distinct(s), sl)] => InplaceOp2::$fwd(self, dst, dl, s, sl),
690                _ => bail!(
691                    "{}: aliased input requires migrating to InplaceOpN",
692                    self.name()
693                ),
694            }
695        }
696    };
697}
698
699impl<C: InplaceOp2> InplaceOpN<1> for C {
700    fn name(&self) -> &'static str {
701        InplaceOp2::name(self)
702    }
703
704    forward_op2!(cpu_fwd, cpu_fwd_aliased, CpuStorage);
705    forward_op2!(cuda_fwd, cuda_fwd_aliased, CudaStorage);
706    forward_op2!(metal_fwd, metal_fwd_aliased, MetalStorage);
707    #[cfg(feature = "rocm")]
708    forward_op2!(rocm_fwd, rocm_fwd_aliased, RocmStorage);
709    #[cfg(feature = "vulkan")]
710    forward_op2!(vulkan_fwd, vulkan_fwd_aliased, VulkanStorage);
711    #[cfg(feature = "wgpu")]
712    forward_op2!(wgpu_fwd, wgpu_fwd_aliased, WgpuStorage);
713}
714
715pub trait InplaceOp2 {
716    fn name(&self) -> &'static str;
717
718    /// The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides,
719    /// offsets etc so the associated layout should be used to access it.
720    fn cpu_fwd(&self, s1: &mut CpuStorage, l1: &Layout, s2: &CpuStorage, l2: &Layout)
721        -> Result<()>;
722
723    /// The forward pass, as run on a gpu device. Note that the storage can use arbitrary strides,
724    /// offsets etc so the associated layout should be used to access it.
725    fn cuda_fwd(
726        &self,
727        s1: &mut CudaStorage,
728        l1: &Layout,
729        s2: &CudaStorage,
730        l2: &Layout,
731    ) -> Result<()> {
732        _ = (s1, l1, s2, l2);
733        Err(crate::Error::Cuda(
734            format!("no cuda implementation for {}", self.name()).into(),
735        ))
736    }
737
738    #[cfg(feature = "rocm")]
739    fn rocm_fwd(&self, _: &mut RocmStorage, _: &Layout, _: &RocmStorage, _: &Layout) -> Result<()> {
740        Err(crate::Error::Msg(format!(
741            "no rocm implementation for {}",
742            self.name()
743        )))
744    }
745    #[cfg(feature = "vulkan")]
746    fn vulkan_fwd(
747        &self,
748        _: &mut VulkanStorage,
749        l1: &Layout,
750        _: &VulkanStorage,
751        _: &Layout,
752    ) -> Result<()> {
753        log_vulkan_custom_op_bail(self.name(), l1);
754        Err(crate::Error::Msg(format!(
755            "no vulkan implementation for {}",
756            self.name()
757        )))
758    }
759    #[cfg(feature = "wgpu")]
760    fn wgpu_fwd(
761        &self,
762        _: &mut WgpuStorage,
763        _l1: &Layout,
764        _: &WgpuStorage,
765        _: &Layout,
766    ) -> Result<()> {
767        Err(crate::Error::Msg(format!(
768            "no wgpu implementation for {}",
769            self.name()
770        )))
771    }
772
773    /// The forward pass, as run on a metal gpu device. Note that the storage can use arbitrary strides,
774    /// offsets etc so the associated layout should be used to access it.
775    fn metal_fwd(
776        &self,
777        s1: &mut MetalStorage,
778        l1: &Layout,
779        s2: &MetalStorage,
780        l2: &Layout,
781    ) -> Result<()> {
782        _ = (s1, l1, s2, l2);
783        Err(crate::Error::Metal(
784            format!("no metal implementation for {}", self.name()).into(),
785        ))
786    }
787}
788
789macro_rules! forward_op3 {
790    ($fwd:ident, $fwd_aliased:ident, $storage:ty) => {
791        fn $fwd(
792            &self,
793            dst: &mut $storage,
794            dl: &Layout,
795            srcs: [(&$storage, &Layout); 2],
796        ) -> Result<()> {
797            let [(s1, l1), (s2, l2)] = srcs;
798            InplaceOp3::$fwd(self, dst, dl, s1, l1, s2, l2)
799        }
800
801        fn $fwd_aliased(
802            &self,
803            dst: &mut $storage,
804            dl: &Layout,
805            srcs: [(Src<'_, $storage>, &Layout); 2],
806        ) -> Result<()> {
807            match srcs {
808                [(Src::Distinct(s1), l1), (Src::Distinct(s2), l2)] => {
809                    InplaceOp3::$fwd(self, dst, dl, s1, l1, s2, l2)
810                }
811                _ => bail!(
812                    "{}: aliased input requires migrating to InplaceOpN",
813                    self.name()
814                ),
815            }
816        }
817    };
818}
819
820impl<C: InplaceOp3> InplaceOpN<2> for C {
821    fn name(&self) -> &'static str {
822        InplaceOp3::name(self)
823    }
824
825    forward_op3!(cpu_fwd, cpu_fwd_aliased, CpuStorage);
826    forward_op3!(cuda_fwd, cuda_fwd_aliased, CudaStorage);
827    forward_op3!(metal_fwd, metal_fwd_aliased, MetalStorage);
828    #[cfg(feature = "rocm")]
829    forward_op3!(rocm_fwd, rocm_fwd_aliased, RocmStorage);
830    #[cfg(feature = "vulkan")]
831    forward_op3!(vulkan_fwd, vulkan_fwd_aliased, VulkanStorage);
832    #[cfg(feature = "wgpu")]
833    forward_op3!(wgpu_fwd, wgpu_fwd_aliased, WgpuStorage);
834}
835
836pub trait InplaceOp3 {
837    fn name(&self) -> &'static str;
838
839    /// The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides,
840    /// offsets etc so the associated layout should be used to access it.
841    fn cpu_fwd(
842        &self,
843        s1: &mut CpuStorage,
844        l1: &Layout,
845        s2: &CpuStorage,
846        l2: &Layout,
847        s3: &CpuStorage,
848        l3: &Layout,
849    ) -> Result<()>;
850
851    /// The forward pass, as run on a gpu device. Note that the storage can use arbitrary strides,
852    /// offsets etc so the associated layout should be used to access it.
853    fn cuda_fwd(
854        &self,
855        _: &mut CudaStorage,
856        _: &Layout,
857        _: &CudaStorage,
858        _: &Layout,
859        _: &CudaStorage,
860        _: &Layout,
861    ) -> Result<()> {
862        Err(crate::Error::Cuda(
863            format!("no cuda implementation for {}", self.name()).into(),
864        ))
865    }
866
867    #[cfg(feature = "rocm")]
868    fn rocm_fwd(
869        &self,
870        _: &mut RocmStorage,
871        _: &Layout,
872        _: &RocmStorage,
873        _: &Layout,
874        _: &RocmStorage,
875        _: &Layout,
876    ) -> Result<()> {
877        Err(crate::Error::Msg(format!(
878            "no rocm implementation for {}",
879            self.name()
880        )))
881    }
882    #[cfg(feature = "vulkan")]
883    fn vulkan_fwd(
884        &self,
885        _: &mut VulkanStorage,
886        l1: &Layout,
887        _: &VulkanStorage,
888        _: &Layout,
889        _: &VulkanStorage,
890        _: &Layout,
891    ) -> Result<()> {
892        log_vulkan_custom_op_bail(self.name(), l1);
893        Err(crate::Error::Msg(format!(
894            "no vulkan implementation for {}",
895            self.name()
896        )))
897    }
898    #[cfg(feature = "wgpu")]
899    fn wgpu_fwd(
900        &self,
901        _: &mut WgpuStorage,
902        _l1: &Layout,
903        _: &WgpuStorage,
904        _: &Layout,
905        _: &WgpuStorage,
906        _: &Layout,
907    ) -> Result<()> {
908        Err(crate::Error::Msg(format!(
909            "no wgpu implementation for {}",
910            self.name()
911        )))
912    }
913
914    /// The forward pass, as run on a metal gpu device. Note that the storage can use arbitrary strides,
915    /// offsets etc so the associated layout should be used to access it.
916    fn metal_fwd(
917        &self,
918        _: &mut MetalStorage,
919        _: &Layout,
920        _: &MetalStorage,
921        _: &Layout,
922        _: &MetalStorage,
923        _: &Layout,
924    ) -> Result<()> {
925        Err(crate::Error::Metal(
926            format!("no metal implementation for {}", self.name()).into(),
927        ))
928    }
929}
930
931impl Tensor {
932    /// Applies a custom op in-place for the `self` tensor.
933    ///
934    /// Tensors sharing underlying storage with `self` are classified and passed as [`Src::Aliased`].
935    /// Separate tensors are locked and passed as [`Src::Distinct`].
936    fn inplace_op<const N: usize, C: InplaceOpN<N>>(&self, srcs: [&Self; N], c: &C) -> Result<()> {
937        let name = c.name();
938
939        // Ensure writes cannot collide with themselves
940        if self.layout().has_internal_overlap() {
941            bail!("{name}: dst has repeated elements (zero-stride). Can not write in-place")
942        }
943
944        // Classify srcs wrt dst
945        let access = c.src_access_pattern();
946        let mut rels: [Option<LayoutRelation>; N] = [None; N];
947        for i in 0..N {
948            if !self.same_storage(srcs[i]) {
949                continue;
950            }
951            let rel = Layout::relation(self.layout(), srcs[i].layout());
952            match access {
953                Some(a) if a.supports(rel) => rels[i] = Some(rel),
954                Some(a) => bail!(
955                    "src {i} shares storage with dst ({rel:?}), which is not supported for the access pattern of `{name}` ({a:?})."
956                ),
957                None => bail!(
958                    "src {i} shares storage with dst, and `{name}` does not support aliased operands."
959                ),
960            }
961        }
962
963        // Acquire locks in order sorted by `Tensor::storage_key`.
964        // Avoids deadlock from cycle of waiting locks.
965        let dst_key = self.storage_key();
966
967        let mut order: [usize; N] = std::array::from_fn(|i| i);
968        order.sort_unstable_by_key(|&i| srcs[i].storage_key());
969
970        let mut guards: [Option<_>; N] = std::array::from_fn(|_| None);
971        let mut dst: Option<_> = None;
972
973        for &i in order.iter() {
974            if rels[i].is_some() {
975                continue; // Aliased. Read through `dst`
976            }
977            let key = srcs[i].storage_key();
978            if key > dst_key && dst.is_none() {
979                dst = Some(self.storage_mut());
980            }
981            // Two sources sharing the same allocation, but distinct from dst.
982            // `read_recursive` allows for shared read access
983            guards[i] = Some(srcs[i].storage());
984        }
985        // If `dst` is not yet set we acquire it from `self` now.
986        let mut dst = match dst {
987            Some(g) => g,
988            None => self.storage_mut(),
989        };
990
991        let operands: [(Src<'_, Storage>, &Layout); N] = std::array::from_fn(|i| {
992            let s = match (&guards[i], rels[i]) {
993                (Some(g), None) => Src::Distinct(&**g),
994                (None, Some(rel)) => Src::Aliased(rel),
995                _ => unreachable!(
996                    "Source is either distinct or aliased. Other match patterns should be impossible"
997                ),
998            };
999            (s, srcs[i].layout())
1000        });
1001
1002        dst.inplace_op(self.layout(), operands, c)
1003    }
1004
1005    /// Applies a unary custom op in place.
1006    pub fn inplace_op1<C: InplaceOp1>(&self, c: &C) -> Result<()> {
1007        self.inplace_op([], c)
1008    }
1009
1010    /// Applies a binary custom op in place (for the first tensor).
1011    pub fn inplace_op2<C: InplaceOpN<1>>(&self, rhs: &Self, c: &C) -> Result<()> {
1012        self.inplace_op([rhs], c)
1013    }
1014
1015    /// Applies a ternary custom op in place (for the first tensor).
1016    pub fn inplace_op3<C: InplaceOpN<2>>(&self, t2: &Self, t3: &Self, c: &C) -> Result<()> {
1017        self.inplace_op([t2, t3], c)
1018    }
1019}
hanzo_ml/custom_op.rs

hanzo_ml/
custom_op.rs
