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

cubecl_ir/dialect/
memory.rs
