cubecl_ir/dialect/
tma.rs

1use core::ops::RangeInclusive;
2
3use alloc::format;
4use cubecl_macros_internal::cube_op;
5use pliron::{derive::op_interface_impl, printable::Printable, verify_err};
6use thiserror::Error;
7
8use crate::{
9    AddressSpace,
10    attributes::IndexAttr,
11    interfaces::{
12        TypeExt,
13        side_effects::{MemoryEffect, MemoryEffectsOp},
14    },
15    prelude::*,
16    types::{
17        PointerType,
18        barrier::{BarrierLevel, BarrierType},
19        cuda::TensorMapType,
20    },
21};
22
23#[derive(Error, Debug)]
24pub enum TmaOpError {
25    #[error("[TmaOp]: Invalid address space for {_0}. Expected {_1}, got {_2}.")]
26    InvalidAddressSpace(&'static str, &'static str, AddressSpace),
27    #[error("[TmaOp]: Unsupported rank {_0}, expected rank to be in range {_1:?}.")]
28    UnsupportedRank(usize, RangeInclusive<usize>),
29}
30
31fn expected_barrier_ty(ctx: &Context) -> TypeHandle {
32    PointerType::get(
33        ctx,
34        BarrierType::get(ctx, BarrierLevel::Cube).into(),
35        AddressSpace::Shared,
36    )
37    .to_handle()
38}
39
40#[pliron_op(name = "tma.load", format, attributes = (tma_load_rank: IndexAttr))]
41#[op_interfaces(AtLeastNOpdsInterface<4>, OperandNOfType<0, PointerType>, OperandNOfType<1, TensorMapType>, OperandNOfType<2, PointerType>)]
42pub struct TmaLoadOp;
43
44impl TmaLoadOp {
45    pub fn new(
46        ctx: &mut Context,
47        barrier: Value,
48        tensor_map: Value,
49        destination: Value,
50        indices: Vec<Value>,
51    ) -> Self {
52        let rank = indices.len();
53        let mut operands = vec![barrier, tensor_map, destination];
54        operands.extend(indices);
55        let op = Self {
56            op: Operation::new(
57                ctx,
58                Self::get_concrete_op_info(),
59                vec![],
60                operands,
61                vec![],
62                0,
63            ),
64        };
65        op.set_attr_tma_load_rank(ctx, rank.into());
66        op
67    }
68
69    pub fn barrier(&self, ctx: &Context) -> Value {
70        self.get_operation().deref(ctx).get_operand(0)
71    }
72
73    pub fn tensor_map(&self, ctx: &Context) -> Value {
74        self.get_operation().deref(ctx).get_operand(1)
75    }
76
77    pub fn destination(&self, ctx: &Context) -> Value {
78        self.get_operation().deref(ctx).get_operand(2)
79    }
80
81    pub fn indices(&self, ctx: &Context) -> Vec<Value> {
82        self.get_operation().deref(ctx).operands().skip(3).collect()
83    }
84
85    pub fn rank(&self, ctx: &Context) -> usize {
86        self.get_attr_tma_load_rank(ctx).unwrap().0
87    }
88}
89
90#[op_interface_impl]
91impl MemoryEffectsOp for TmaLoadOp {
92    fn memory_effects(&self, ctx: &Context) -> Vec<MemoryEffect> {
93        vec![MemoryEffect::Write(self.destination(ctx))]
94    }
95}
96
97impl Verify for TmaLoadOp {
98    fn verify(&self, ctx: &Context) -> Result<()> {
99        let loc = self.loc(ctx);
100        let barrier_ty = self.barrier(ctx).get_type(ctx).as_ptr(ctx);
101        let dest_ty = self.barrier(ctx).get_type(ctx).as_ptr(ctx);
102
103        if !barrier_ty.inner.deref(ctx).is::<BarrierType>() {
104            let expected = expected_barrier_ty(ctx).deref(ctx);
105            return verify_err!(
106                loc,
107                OperandNOfTypeError::AllOperandsOfTypeVerifyErr(
108                    format!("{} {}", expected.get_type_id(), expected.disp(ctx)),
109                    format!("{} {}", barrier_ty.get_type_id(), barrier_ty.disp(ctx))
110                )
111            );
112        }
113
114        if dest_ty.address_space != AddressSpace::Shared {
115            return verify_err!(
116                loc,
117                TmaOpError::InvalidAddressSpace("destination", "Shared", dest_ty.address_space)
118            );
119        }
120
121        if !(1..=5).contains(&self.rank(ctx)) {
122            return verify_err!(loc, TmaOpError::UnsupportedRank(self.rank(ctx), 1..=5));
123        }
124
125        Ok(())
126    }
127}
128
129#[pliron_op(name = "tma.load_im2col", format, attributes = (tma_load_im2col_rank: IndexAttr))]
130#[op_interfaces(AtLeastNOpdsInterface<5>, OperandNOfType<0, PointerType>, OperandNOfType<1, TensorMapType>, OperandNOfType<2, PointerType>)]
131pub struct TmaLoadIm2colOp;
132
133impl TmaLoadIm2colOp {
134    pub fn new(
135        ctx: &mut Context,
136        barrier: Value,
137        tensor_map: Value,
138        destination: Value,
139        indices: Vec<Value>,
140        offsets: Vec<Value>,
141    ) -> Self {
142        let rank = indices.len();
143        let mut operands = vec![barrier, tensor_map, destination];
144        operands.extend(indices);
145        operands.extend(offsets);
146        let op = Self {
147            op: Operation::new(
148                ctx,
149                Self::get_concrete_op_info(),
150                vec![],
151                operands,
152                vec![],
153                0,
154            ),
155        };
156        op.set_attr_tma_load_im2col_rank(ctx, rank.into());
157        op
158    }
159
160    pub fn barrier(&self, ctx: &Context) -> Value {
161        self.get_operation().deref(ctx).get_operand(0)
162    }
163
164    pub fn tensor_map(&self, ctx: &Context) -> Value {
165        self.get_operation().deref(ctx).get_operand(1)
166    }
167
168    pub fn destination(&self, ctx: &Context) -> Value {
169        self.get_operation().deref(ctx).get_operand(2)
170    }
171
172    pub fn indices(&self, ctx: &Context) -> Vec<Value> {
173        let rank = self.get_attr_tma_load_im2col_rank(ctx).unwrap().0;
174        self.get_operation()
175            .deref(ctx)
176            .operands()
177            .skip(3)
178            .take(rank)
179            .collect()
180    }
181
182    pub fn offsets(&self, ctx: &Context) -> Vec<Value> {
183        let rank = self.get_attr_tma_load_im2col_rank(ctx).unwrap().0;
184        self.get_operation()
185            .deref(ctx)
186            .operands()
187            .skip(3 + rank)
188            .collect()
189    }
190
191    pub fn rank(&self, ctx: &Context) -> usize {
192        self.get_attr_tma_load_im2col_rank(ctx).unwrap().0
193    }
194}
195
196#[op_interface_impl]
197impl MemoryEffectsOp for TmaLoadIm2colOp {
198    fn memory_effects(&self, ctx: &Context) -> Vec<MemoryEffect> {
199        vec![MemoryEffect::Write(self.destination(ctx))]
200    }
201}
202
203impl Verify for TmaLoadIm2colOp {
204    fn verify(&self, ctx: &Context) -> Result<()> {
205        let loc = self.loc(ctx);
206        let barrier_ty = self.barrier(ctx).get_type(ctx).as_ptr(ctx);
207        let dest_ty = self.barrier(ctx).get_type(ctx).as_ptr(ctx);
208
209        if !barrier_ty.inner.deref(ctx).is::<BarrierType>() {
210            let expected = expected_barrier_ty(ctx).deref(ctx);
211            return verify_err!(
212                loc,
213                OperandNOfTypeError::AllOperandsOfTypeVerifyErr(
214                    format!("{} {}", expected.get_type_id(), expected.disp(ctx)),
215                    format!("{} {}", barrier_ty.get_type_id(), barrier_ty.disp(ctx))
216                )
217            );
218        }
219
220        if dest_ty.address_space != AddressSpace::Shared {
221            return verify_err!(
222                loc,
223                TmaOpError::InvalidAddressSpace("destination", "Shared", dest_ty.address_space)
224            );
225        }
226
227        if !(3..=5).contains(&self.rank(ctx)) {
228            return verify_err!(loc, TmaOpError::UnsupportedRank(self.rank(ctx), 3..=5));
229        }
230
231        Ok(())
232    }
233}
234
235#[pliron_op(name = "tma.store", format, attributes = (tma_store_rank: IndexAttr))]
236#[op_interfaces(AtLeastNOpdsInterface<4>, OperandNOfType<0, PointerType>, OperandNOfType<1, TensorMapType>)]
237pub struct TmaStoreOp;
238
239impl TmaStoreOp {
240    pub fn new(ctx: &mut Context, source: Value, tensor_map: Value, indices: Vec<Value>) -> Self {
241        let rank = indices.len();
242        let mut operands = vec![source, tensor_map];
243        operands.extend(indices);
244        let op = Self {
245            op: Operation::new(
246                ctx,
247                Self::get_concrete_op_info(),
248                vec![],
249                operands,
250                vec![],
251                0,
252            ),
253        };
254        op.set_attr_tma_store_rank(ctx, rank.into());
255        op
256    }
257
258    pub fn source(&self, ctx: &Context) -> Value {
259        self.get_operation().deref(ctx).get_operand(0)
260    }
261
262    pub fn tensor_map(&self, ctx: &Context) -> Value {
263        self.get_operation().deref(ctx).get_operand(1)
264    }
265
266    pub fn indices(&self, ctx: &Context) -> Vec<Value> {
267        self.get_operation().deref(ctx).operands().skip(2).collect()
268    }
269
270    pub fn rank(&self, ctx: &Context) -> usize {
271        self.get_attr_tma_store_rank(ctx).unwrap().0
272    }
273}
274
275#[op_interface_impl]
276impl MemoryEffectsOp for TmaStoreOp {
277    fn memory_effects(&self, ctx: &Context) -> Vec<MemoryEffect> {
278        vec![MemoryEffect::Read(self.source(ctx))]
279    }
280}
281
282impl Verify for TmaStoreOp {
283    fn verify(&self, ctx: &Context) -> Result<()> {
284        let loc = self.loc(ctx);
285        let src_ty = self.source(ctx).get_type(ctx).as_ptr(ctx);
286
287        if src_ty.address_space != AddressSpace::Shared {
288            return verify_err!(
289                loc,
290                TmaOpError::InvalidAddressSpace("source", "Shared", src_ty.address_space)
291            );
292        }
293
294        if !(1..=5).contains(&self.rank(ctx)) {
295            return verify_err!(loc, TmaOpError::UnsupportedRank(self.rank(ctx), 1..=5));
296        }
297
298        Ok(())
299    }
300}
301
302#[cube_op(name = "tma.commit_group")]
303#[result_ty(none)]
304pub struct CommitGroupOp {}
305
306#[cube_op(name = "tma.wait_group")]
307#[result_ty(none)]
308pub struct WaitGroupOp {
309    pub max_pending: IndexAttr,
310}
311
312#[cube_op(name = "tma.wait_group_read")]
313#[result_ty(none)]
314pub struct WaitGroupReadOp {
315    pub max_pending: IndexAttr,
316}
cubecl_ir/dialect/tma.rs

cubecl_ir/dialect/
tma.rs
