cubecl_ir/dialect/
memory.rs

1use core::fmt::Debug;
2
3use ::pliron::parsable::ParseResult;
4use alloc::string::{String, ToString};
5use cubecl_macros_internal::cube_op;
6use derive_more::From;
7use derive_new::new;
8use pliron::{
9    arg_err,
10    attribute::AttrObj,
11    builtin::{
12        attributes::{TypeAttr, UnitAttr},
13        ops::ConstantOp,
14    },
15    combine::{
16        Parser, optional,
17        parser::char::{char, spaces, string},
18    },
19    derive::pliron_attr,
20    identifier::Identifier,
21    input_err,
22    irbuild::inserter::Inserter,
23    irfmt::parsers::{process_parsed_ssa_defs, spaced},
24    location::Location,
25    op::{OpBox, OpObj},
26    opts::mem2reg::{
27        AllocInfo, PromotableAllocationInterface, PromotableOpInterface, PromotableOpKind,
28    },
29    parsable::{IntoParseResult, Parsable},
30    printable::Printable,
31    r#type::{TypeHandle, type_cast},
32    utils::table::{HMap, SmallSet},
33    verify_err,
34};
35use thiserror::Error;
36
37use crate::{
38    AddressSpace, CanMaterialize, NoSideEffects, Pure,
39    attributes::{IndexAttr, ZeroAttr},
40    dialect::{general::PoisonOp, math::index_attr, ptr_value_ty},
41    interfaces::{
42        IndexableType, TriviallyUnrollable, TypedExt,
43        aliasing::AliasingOp,
44        memory_slot::{
45            DeletionKind, DestructurableAccessorOpInterface, DestructurableConstructorOpInterface,
46            DestructurableTypeInterface, DestructurableValueSlot, LogicalResult,
47            SafeMemorySlotAccessOpInterface, ValueSlot,
48        },
49    },
50    prelude::*,
51    try_cast_ty,
52    types::{PointerType, scalar::IndexType},
53};
54
55#[pliron_attr(name = "memory.address_space", format = "$0", verifier = "succ")]
56#[derive(new, From, PartialEq, Eq, Clone, Copy, Debug, Hash)]
57pub struct AddressSpaceAttr(pub AddressSpace);
58
59#[cube_op(name = "memory.declare_variable", format = "custom")]
60#[result_ty(from_inputs = variable_ptr_ty)]
61#[op_traits(NoSideEffects, CanMaterialize)]
62pub struct DeclareVariableOp {
63    pub value_ty: TypeAttr,
64    pub addr_space: AddressSpaceAttr,
65    pub alignment: IndexAttr,
66    #[attribute(optional, untyped)]
67    pub initializer: AttrObj,
68}
69
70impl Printable for DeclareVariableOp {
71    fn fmt(
72        &self,
73        ctx: &Context,
74        _state: &pliron::printable::State,
75        f: &mut core::fmt::Formatter<'_>,
76    ) -> core::fmt::Result {
77        write!(
78            f,
79            "{} = {} {} {}, align = {}",
80            self.get_result(ctx).disp(ctx),
81            self.get_opid(),
82            self.value_ty(ctx).disp(ctx),
83            self.addr_space(ctx).disp(ctx),
84            self.alignment(ctx).disp(ctx)
85        )?;
86        if let Some(init) = self.initializer(ctx) {
87            write!(f, ", init = {}", init.disp(ctx))?;
88        }
89
90        Ok(())
91    }
92}
93impl Parsable for DeclareVariableOp {
94    type Arg = Vec<(Identifier, Location)>;
95    type Parsed = OpObj;
96
97    fn parse<'a>(
98        input: &mut ::pliron::parsable::StateStream<'a>,
99        arg: Self::Arg,
100    ) -> ParseResult<'a, Self::Parsed> {
101        let cur_loc = input.loc();
102        let value_ty = TypeAttr::parse(input, ())?.0;
103        spaces().parse_stream(input).into_result()?;
104        let addr_space = AddressSpaceAttr::parse(input, ())?.0;
105        let mut label = (spaced(char(',')), string("align"), spaced(char('=')));
106        label.parse_stream(input).into_result()?;
107        let align = IndexAttr::parse(input, ())?.0;
108        let mut label = (spaced(char(',')), string("init"), spaced(char('=')));
109        label.parse_stream(input).into_result()?;
110        let mut init_parse = optional(AttrObj::parser(()));
111        let init = init_parse.parse_stream(input).into_result()?.0;
112
113        let ctx = &mut input.state.ctx;
114        if arg.len() != 1 {
115            input_err!(
116                cur_loc,
117                "Expected 1 result, got {} during parsing",
118                arg.len()
119            )?;
120        }
121        let op = DeclareVariableOp::new(ctx, value_ty, addr_space, align, init);
122        process_parsed_ssa_defs(input, &arg, op.get_operation())?;
123        Ok(OpBox::new(op)).into_parse_result()
124    }
125}
126
127#[op_interface_impl]
128impl PromotableAllocationInterface for DeclareVariableOp {
129    fn alloc_info(&self, ctx: &Context) -> Vec<AllocInfo> {
130        if self.addr_space(ctx).0 == AddressSpace::Local {
131            vec![AllocInfo {
132                ptr: self.get_result(ctx),
133                ty: self.value_ty(ctx).get_type(ctx),
134            }]
135        } else {
136            vec![]
137        }
138    }
139
140    fn default_value(
141        &self,
142        ctx: &mut Context,
143        inserter: &mut dyn Inserter,
144        alloc_info: &AllocInfo,
145    ) -> Result<Value> {
146        if alloc_info.ptr != self.get_result(ctx) {
147            return arg_err!(self.loc(ctx), UnrelatedAllocInfo);
148        }
149        if let Some(initializer) = self.initializer(ctx).map(|it| it.clone()) {
150            let constant = ConstantOp::new(ctx, initializer);
151            inserter.insert_op(ctx, &constant);
152            Ok(constant.get_result(ctx))
153        } else {
154            let poison = PoisonOp::new(ctx, alloc_info.ty);
155            inserter.insert_op(ctx, &poison);
156            Ok(poison.get_result(ctx))
157        }
158    }
159
160    fn promote(
161        &self,
162        ctx: &mut Context,
163        rewriter: &mut dyn Rewriter,
164        alloc_infos: &[AllocInfo],
165    ) -> Result<()> {
166        if alloc_infos.len() != 1 || alloc_infos[0].ptr != self.get_result(ctx) {
167            return arg_err!(self.loc(ctx), UnrelatedAllocInfo);
168        }
169        rewriter.erase_operation(ctx, self.get_operation());
170        Ok(())
171    }
172}
173
174#[op_interface_impl]
175impl DestructurableConstructorOpInterface for DeclareVariableOp {
176    fn destructurable_values(&self, ctx: &Context) -> Vec<DestructurableValueSlot> {
177        if self.addr_space(ctx).0 != AddressSpace::Local {
178            return vec![];
179        }
180        if let Some(init) = self.initializer(ctx)
181            && !init.is::<ZeroAttr>()
182        {
183            return vec![];
184        }
185        let value_ty = self.value_ty(ctx).get_type(ctx);
186        let ty = value_ty.deref(ctx);
187        let Some(destructurable) = type_cast::<dyn DestructurableTypeInterface>(&*ty) else {
188            return vec![];
189        };
190        let Some(destructured_type) = destructurable.subelement_index_map(ctx) else {
191            return vec![];
192        };
193
194        vec![DestructurableValueSlot {
195            slot: ValueSlot::new(self.get_result(ctx), value_ty),
196            subelement_types: destructured_type,
197        }]
198    }
199
200    fn destructure(
201        &self,
202        ctx: &mut Context,
203        _value: &DestructurableValueSlot,
204        used_indices: &SmallSet<AttrObj, 8>,
205        rewriter: &mut PassRewriter,
206        new_constructors: &mut Vec<TraitOp<dyn DestructurableConstructorOpInterface>>,
207    ) -> HMap<AttrObj, ValueSlot> {
208        let addr_space = self.addr_space(ctx).0;
209        let init = self.initializer(ctx).map(|it| {
210            assert!(it.is::<ZeroAttr>());
211        });
212        let destructured_type = {
213            let ty = self.value_ty(ctx).get_type(ctx).deref(ctx);
214            let destructurable = try_cast_ty!(ty, ctx, dyn DestructurableTypeInterface);
215            destructurable.subelement_index_map(ctx).unwrap()
216        };
217
218        let mut slot_map = HMap::new();
219        for used_index in used_indices {
220            let value_ty = destructured_type[used_index];
221            let init = init.map(|_| ZeroAttr::new(value_ty).into());
222            let align = value_ty.align(ctx);
223            let suballoc = DeclareVariableOp::new(ctx, value_ty, addr_space, align, init);
224            rewriter.append_op(ctx, &suballoc);
225
226            let slot = ValueSlot::new(suballoc.get_result(ctx), value_ty);
227            slot_map.insert(used_index.clone(), slot);
228            new_constructors.push(TraitOp::try_from_op(suballoc.get_operation(), ctx).unwrap());
229        }
230
231        slot_map
232    }
233
234    fn handle_destructuring_complete(
235        &self,
236        ctx: &mut Context,
237        value: &DestructurableValueSlot,
238        rewriter: &mut PassRewriter,
239    ) -> Option<TraitOp<dyn DestructurableConstructorOpInterface>> {
240        assert_eq!(value.slot.value, self.get_result(ctx));
241        rewriter.erase_operation(ctx, self.get_operation());
242        None
243    }
244}
245
246fn variable_ptr_ty(
247    ctx: &Context,
248    value_ty: &TypeAttr,
249    addr_space: &AddressSpaceAttr,
250    _align: &IndexAttr,
251) -> TypeHandle {
252    let value_ty = value_ty.get_type(ctx);
253    PointerType::get(ctx, value_ty, addr_space.0).into()
254}
255
256#[cube_op(
257    name = "memory.index",
258    format = "$0 `[` $1 `]` opt_attr($checked, $UnitAttr) ` : ` type($0)"
259)]
260#[result_ty(from_inputs = |ctx, base, _| indexed_ptr_ty(ctx, base))]
261#[op_interfaces(OperandNOfType<0, PointerType>, OperandNOfType<1, IndexType>)]
262#[op_traits(Pure, CanMaterialize)]
263pub struct IndexOp {
264    pub base: Value,
265    pub index: Value,
266    #[attribute(optional)]
267    pub checked: UnitAttr,
268}
269
270#[op_interface_impl]
271impl AliasingOp for IndexOp {
272    fn source_ptr(&self, ctx: &Context) -> Option<Value> {
273        Some(self.base(ctx))
274    }
275}
276
277fn const_index(ctx: &Context, value: Value) -> Option<usize> {
278    let def_op = value.defining_op()?;
279    let const_def = def_op.as_op::<ConstantOp>(ctx)?;
280    let attr = const_def.get_value(ctx);
281    let attr = attr.downcast_ref::<IndexAttr>()?;
282    Some(attr.0)
283}
284
285#[op_interface_impl]
286impl DestructurableAccessorOpInterface for IndexOp {
287    fn can_rewire(
288        &self,
289        ctx: &Context,
290        value: &DestructurableValueSlot,
291        used_indices: &mut SmallSet<AttrObj, 8>,
292        must_be_safely_used: &mut Vec<ValueSlot>,
293    ) -> bool {
294        if self.base(ctx) != value.slot.value {
295            return false;
296        }
297        let Some(index) = const_index(ctx, self.index(ctx)) else {
298            return false;
299        };
300        let attr = index_attr(index);
301        let elem_ty = value.subelement_types[&attr];
302        used_indices.insert(attr);
303
304        let used_slot = ValueSlot::new(self.get_result(ctx), elem_ty);
305        must_be_safely_used.push(used_slot);
306        true
307    }
308
309    fn rewire(
310        &self,
311        ctx: &mut Context,
312        _value: &DestructurableValueSlot,
313        subvalues: &HMap<AttrObj, ValueSlot>,
314        rewriter: &mut PassRewriter,
315    ) -> DeletionKind {
316        let index = const_index(ctx, self.index(ctx)).expect("checked before");
317        let index_attr = index_attr(index);
318        let new_slot = &subvalues[&index_attr];
319        rewriter.replace_value_uses_with(ctx, self.get_result(ctx), new_slot.value);
320        DeletionKind::Delete
321    }
322}
323
324impl IndexOp {
325    pub fn maybe_checked(ctx: &mut Context, base: Value, index: Value, checked: bool) -> Self {
326        let op = Self::new(ctx, base, index, checked.then_some(UnitAttr::new()));
327        if checked {
328            op.set_checked(ctx);
329        }
330        op
331    }
332}
333
334fn indexed_ptr_ty(ctx: &Context, base: &Value) -> TypeHandle {
335    let (value_ty, address_space) = {
336        let base_ty = base.get_type(ctx).deref(ctx);
337        let PointerType {
338            inner,
339            address_space,
340        } = base_ty.downcast_ref().expect("Should be pointer");
341        let list_ty = inner.deref(ctx);
342        let indexable = type_cast::<dyn IndexableType>(&*list_ty).expect("Should be indexable");
343        let value_ty = indexable.indexed_type(ctx);
344        (value_ty, *address_space)
345    };
346    PointerType::get(ctx, value_ty, address_space).into()
347}
348
349#[derive(Error, Debug)]
350#[error("Register Promotion: Allocation info provided is not related to this operation")]
351pub struct UnrelatedAllocInfo;
352
353#[cube_op(name = "memory.load")]
354#[result_ty(from_inputs = ptr_value_ty)]
355#[op_interfaces(OperandNOfType<0, PointerType>, TriviallyUnrollable)]
356#[op_traits(CanMaterialize, NoSideEffects)]
357pub struct LoadOp {
358    #[operand(ptr_read)]
359    pub ptr: Value,
360}
361
362#[op_interface_impl]
363impl PromotableOpInterface for LoadOp {
364    fn promotion_kind(&self, ctx: &Context, alloc_info: &AllocInfo) -> PromotableOpKind {
365        if self.ptr(ctx) == alloc_info.ptr {
366            PromotableOpKind::Load
367        } else {
368            PromotableOpKind::NonPromotableUse
369        }
370    }
371
372    fn promote(
373        &self,
374        ctx: &mut Context,
375        alloc_info_reaching_defs: &[(AllocInfo, Value)],
376        rewriter: &mut dyn Rewriter,
377    ) -> Result<()> {
378        if alloc_info_reaching_defs.len() != 1 {
379            return arg_err!(self.loc(ctx), UnrelatedAllocInfo);
380        }
381        let (alloc_info, reaching_def) = &alloc_info_reaching_defs[0];
382        if self.ptr(ctx) != alloc_info.ptr {
383            return arg_err!(self.loc(ctx), UnrelatedAllocInfo);
384        }
385        rewriter.replace_operation_with_values(ctx, self.get_operation(), vec![*reaching_def]);
386        Ok(())
387    }
388}
389
390#[op_interface_impl]
391impl SafeMemorySlotAccessOpInterface for LoadOp {
392    fn ensure_only_safe_accesses(
393        &self,
394        _: &Context,
395        _: &ValueSlot,
396        _: &mut Vec<ValueSlot>,
397    ) -> LogicalResult {
398        Ok(())
399    }
400}
401
402#[derive(Error)]
403pub enum StoreOpError {
404    #[error(
405        "[StoreOp]: Value type doesn't match the inner type of the pointer: expected {_0}, got {_1}"
406    )]
407    MismatchedValueType(String, String),
408}
409
410impl Debug for StoreOpError {
411    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
412        write!(f, "{self}")
413    }
414}
415
416#[cube_op(name = "memory.store", verifier = "custom")]
417#[result_ty(none)]
418#[op_interfaces(OperandNOfType<0, PointerType>, TriviallyUnrollable)]
419#[op_traits(CanMaterialize)]
420pub struct StoreOp {
421    #[operand(ptr_write)]
422    pub ptr: Value,
423    pub value: Value,
424}
425
426impl Verify for StoreOp {
427    fn verify(&self, ctx: &Context) -> Result<()> {
428        let loc = self.loc(ctx);
429        let ptr_value_ty = ptr_value_ty(ctx, &self.ptr(ctx));
430        let value_ty = self.value(ctx).get_type(ctx);
431        if ptr_value_ty != value_ty {
432            verify_err!(
433                loc,
434                StoreOpError::MismatchedValueType(
435                    ptr_value_ty.disp(ctx).to_string(),
436                    value_ty.disp(ctx).to_string()
437                )
438            )?;
439        }
440        Ok(())
441    }
442}
443
444#[op_interface_impl]
445impl PromotableOpInterface for StoreOp {
446    fn promotion_kind(&self, ctx: &Context, alloc_info: &AllocInfo) -> PromotableOpKind {
447        if self.ptr(ctx) == alloc_info.ptr {
448            PromotableOpKind::Store(self.value(ctx))
449        } else {
450            PromotableOpKind::NonPromotableUse
451        }
452    }
453
454    fn promote(
455        &self,
456        ctx: &mut Context,
457        alloc_info_reaching_defs: &[(AllocInfo, Value)],
458        rewriter: &mut dyn Rewriter,
459    ) -> Result<()> {
460        if alloc_info_reaching_defs.len() != 1 {
461            return arg_err!(self.loc(ctx), UnrelatedAllocInfo);
462        }
463        let (alloc_info, _reaching_def) = &alloc_info_reaching_defs[0];
464        if self.ptr(ctx) != alloc_info.ptr {
465            return arg_err!(self.loc(ctx), UnrelatedAllocInfo);
466        }
467        rewriter.erase_operation(ctx, self.get_operation());
468        Ok(())
469    }
470}
471
472#[op_interface_impl]
473impl SafeMemorySlotAccessOpInterface for StoreOp {
474    fn ensure_only_safe_accesses(
475        &self,
476        _: &Context,
477        _: &ValueSlot,
478        _: &mut Vec<ValueSlot>,
479    ) -> LogicalResult {
480        Ok(())
481    }
482}
483
484#[cube_op(name = "memory.copy")]
485#[result_ty(none)]
486#[op_interfaces(OperandNOfType<0, PointerType>, SameOperandsType)]
487#[op_traits(CanMaterialize)]
488pub struct CopyOp {
489    #[operand(ptr_read)]
490    pub source: Value,
491    #[operand(ptr_write)]
492    pub destination: Value,
493    pub len: IndexAttr,
494}
495
496#[op_interface_impl]
497impl SafeMemorySlotAccessOpInterface for CopyOp {
498    fn ensure_only_safe_accesses(
499        &self,
500        _: &Context,
501        _: &ValueSlot,
502        _: &mut Vec<ValueSlot>,
503    ) -> LogicalResult {
504        Ok(())
505    }
506}
cubecl_ir/dialect/memory.rs

cubecl_ir/dialect/
memory.rs
