cubecl_std/tensor/layout/
chain.rs

1use cubecl::prelude::*;
2use cubecl_core as cubecl;
3
4use crate::tensor::layout::{Layout, LayoutExpand};
5
6/// Chain of layouts, can be used to launch with multiple layouts
7#[derive(CubeType)]
8pub struct Chain<L0: Layout, L1: Layout<SourceCoordinates = L0::Coordinates>> {
9    l0: L0,
10    l1: L1,
11}
12
13#[cube]
14impl<L0: Layout, L1: Layout<SourceCoordinates = L0::Coordinates>> Chain<L0, L1> {
15    pub fn new(l0: L0, l1: L1) -> Self {
16        Chain::<L0, L1> { l0, l1 }
17    }
18}
19
20#[cube]
21impl<L0: Layout, L1: Layout<SourceCoordinates = L0::Coordinates>> Layout for Chain<L0, L1> {
22    type Coordinates = L1::Coordinates;
23    type SourceCoordinates = L0::SourceCoordinates;
24
25    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
26        let pos = self.l1.to_source_pos(pos);
27        self.l0.to_source_pos(pos)
28    }
29
30    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
31        let (pos, l1_in_bounds) = self.l1.to_source_pos_checked(pos);
32        self.l0.is_in_bounds(pos) && l1_in_bounds
33    }
34
35    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
36        let (pos, l1_in_bounds) = self.l1.to_source_pos_checked(pos);
37        let (pos, l0_in_bounds) = self.l0.to_source_pos_checked(pos);
38        (pos, l0_in_bounds && l1_in_bounds)
39    }
40
41    fn shape(&self) -> Self::Coordinates {
42        self.l1.shape()
43    }
44}
45
46pub use launch::*;
47mod launch {
48
49    use crate::tensor::launch::{MemoryArg, ViewLayoutLaunchArg};
50
51    use super::*;
52
53    pub struct ChainLaunch<
54        L0: Layout + ViewLayoutLaunchArg,
55        L1: Layout<SourceCoordinates = L0::Coordinates> + ViewLayoutLaunchArg,
56    > {
57        l0: L0::RuntimeArg,
58        l1: L1::RuntimeArg,
59    }
60    impl<
61        L0: Layout + ViewLayoutLaunchArg,
62        L1: Layout<SourceCoordinates = L0::Coordinates> + ViewLayoutLaunchArg,
63    > ChainLaunch<L0, L1>
64    {
65        pub fn new(l0: L0::RuntimeArg, l1: L1::RuntimeArg) -> Self {
66            Self { l0, l1 }
67        }
68    }
69
70    pub struct ChainCompilationArg<
71        L0: Layout + ViewLayoutLaunchArg,
72        L1: Layout<SourceCoordinates = L0::Coordinates> + ViewLayoutLaunchArg,
73    > {
74        l0: L0::CompilationArg,
75        l1: L1::CompilationArg,
76    }
77    impl<
78        L0: Layout + ViewLayoutLaunchArg,
79        L1: Layout<SourceCoordinates = L0::Coordinates> + ViewLayoutLaunchArg,
80    > Clone for ChainCompilationArg<L0, L1>
81    {
82        fn clone(&self) -> Self {
83            Self {
84                l0: self.l0.clone(),
85                l1: self.l1.clone(),
86            }
87        }
88    }
89
90    impl<
91        L0: Layout + ViewLayoutLaunchArg,
92        L1: Layout<SourceCoordinates = L0::Coordinates> + ViewLayoutLaunchArg,
93    > core::hash::Hash for ChainCompilationArg<L0, L1>
94    {
95        fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
96            self.l0.hash(state);
97            self.l1.hash(state);
98        }
99    }
100    impl<
101        L0: Layout + ViewLayoutLaunchArg,
102        L1: Layout<SourceCoordinates = L0::Coordinates> + ViewLayoutLaunchArg,
103    > core::cmp::PartialEq for ChainCompilationArg<L0, L1>
104    {
105        fn eq(&self, other: &Self) -> bool {
106            self.l0.eq(&other.l0) && self.l1.eq(&other.l1)
107        }
108    }
109    impl<
110        L0: Layout + ViewLayoutLaunchArg,
111        L1: Layout<SourceCoordinates = L0::Coordinates> + ViewLayoutLaunchArg,
112    > core::fmt::Debug for ChainCompilationArg<L0, L1>
113    {
114        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
115            f.debug_struct(stringify!(Chain))
116                .field(stringify!(l0), &self.l0)
117                .field(stringify!(l1), &self.l1)
118                .finish()
119        }
120    }
121    impl<
122        L0: Layout + ViewLayoutLaunchArg,
123        L1: Layout<SourceCoordinates = L0::Coordinates> + ViewLayoutLaunchArg,
124    > core::cmp::Eq for ChainCompilationArg<L0, L1>
125    {
126    }
127
128    impl<
129        L0: Layout + ViewLayoutLaunchArg,
130        L1: Layout<SourceCoordinates = L0::Coordinates> + ViewLayoutLaunchArg,
131    > ViewLayoutLaunchArg for Chain<L0, L1>
132    {
133        type RuntimeArg = ChainLaunch<L0, L1>;
134        type CompilationArg = ChainCompilationArg<L0, L1>;
135
136        fn register<B: MemoryArg>(
137            arg: Self::RuntimeArg,
138            buffer: &B,
139            ty: Type,
140            launcher: &mut KernelLauncher,
141        ) -> Self::CompilationArg {
142            ChainCompilationArg {
143                l0: L0::register(arg.l0, buffer, ty, launcher),
144                l1: L1::register(arg.l1, buffer, ty, launcher),
145            }
146        }
147        fn expand(
148            arg: &Self::CompilationArg,
149            ty: Type,
150            builder: &mut KernelBuilder,
151        ) -> <Self as CubeType>::ExpandType {
152            ChainExpand {
153                l0: L0::expand(&arg.l0, ty, builder),
154                l1: L1::expand(&arg.l1, ty, builder),
155            }
156        }
157        fn expand_output(
158            arg: &Self::CompilationArg,
159            ty: Type,
160            builder: &mut KernelBuilder,
161        ) -> <Self as CubeType>::ExpandType {
162            ChainExpand {
163                l0: L0::expand_output(&arg.l0, ty, builder),
164                l1: L1::expand_output(&arg.l1, ty, builder),
165            }
166        }
167    }
168}
cubecl_std/tensor/layout/chain.rs

cubecl_std/tensor/layout/
chain.rs
