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

cubecl_ir/dialect/
tma.rs
