cubecl_ir/interfaces/
memory_slot.rs

1use core::{
2    fmt::{self, Display},
3    hash::Hash,
4};
5
6use derive_more::{Eq, PartialEq};
7use derive_new::new;
8use pliron::{
9    attribute::AttrObj,
10    basic_block::BasicBlock,
11    graph::HasLabel,
12    opts::mem2reg::AllocInfo,
13    printable::{self, Printable},
14    region::Region,
15    utils::table::{HMap, SmallMap, SmallSet},
16    value::DefiningEntity,
17};
18
19use crate::prelude::*;
20
21pub type LogicalResult = core::result::Result<(), ()>;
22
23#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
24pub enum DeletionKind {
25    Keep,
26    Delete,
27}
28
29/// The defining entity of a [`MemoryValue`]:
30/// Either an [`Operation`], a [`BasicBlock`] or the special live-on-entry state.
31#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
32pub enum MemoryDefiningEntity {
33    Op(Ptr<Operation>),
34    Block(Ptr<BasicBlock>),
35    LiveOnEntry,
36}
37
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub struct MemoryValue {
40    val_uid: u64,
41    #[eq(skip)]
42    defining_entity: MemoryDefiningEntity,
43}
44
45impl Hash for MemoryValue {
46    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
47        self.val_uid.hash(state);
48    }
49}
50
51impl Display for MemoryValue {
52    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53        if self == &MemoryValue::LIVE_ON_ENTRY {
54            write!(f, "LiveOnEntry")
55        } else {
56            write!(f, "M{}", self.val_uid)
57        }
58    }
59}
60
61impl MemoryValue {
62    pub const LIVE_ON_ENTRY: MemoryValue = MemoryValue {
63        val_uid: 0,
64        defining_entity: MemoryDefiningEntity::LiveOnEntry,
65    };
66
67    pub fn defining_entity(&self) -> Option<DefiningEntity> {
68        match self.defining_entity {
69            MemoryDefiningEntity::Op(op) => Some(DefiningEntity::Op(op)),
70            MemoryDefiningEntity::Block(block) => Some(DefiningEntity::Block(block)),
71            MemoryDefiningEntity::LiveOnEntry => None,
72        }
73    }
74}
75
76#[derive(Default)]
77pub struct MemorySSAContext {
78    last_idx: u64,
79}
80
81impl MemorySSAContext {
82    pub fn new_value_in_block(&mut self, block: Ptr<BasicBlock>) -> MemoryValue {
83        self.last_idx += 1;
84        MemoryValue {
85            val_uid: self.last_idx,
86            defining_entity: MemoryDefiningEntity::Block(block),
87        }
88    }
89
90    pub fn new_value_at_op(&mut self, op: Ptr<Operation>) -> MemoryValue {
91        self.last_idx += 1;
92        MemoryValue {
93            val_uid: self.last_idx,
94            defining_entity: MemoryDefiningEntity::Op(op),
95        }
96    }
97}
98
99/// Different from `RegionPredecessor` because terminators don't actually matter for memory SSA
100#[derive(Clone, Copy, PartialEq, Eq, Hash)]
101pub enum MemoryRegionPredecessor {
102    Parent,
103    Block(Ptr<BasicBlock>),
104}
105
106impl Printable for MemoryRegionPredecessor {
107    fn fmt(&self, ctx: &Context, _: &printable::State, f: &mut fmt::Formatter<'_>) -> fmt::Result {
108        match self {
109            MemoryRegionPredecessor::Parent => f.write_str("Parent"),
110            MemoryRegionPredecessor::Block(block) => {
111                write!(f, "{}", block.label(ctx))
112            }
113        }
114    }
115}
116
117pub type RegionMemoryPhiInputs = SmallMap<MemoryRegionPredecessor, MemoryValue, 2>;
118pub enum RegionMemoryValue {
119    Forward(MemoryValue),
120    RegionPhi(RegionMemoryPhiInputs),
121}
122
123#[op_interface]
124pub trait PromotableRegionOpInterface {
125    verify_op_succ!();
126
127    /// Returns true if `region` (a child of this op) can be analysed for
128    /// promotion with respect to `alloc`.
129    /// `has_value_stores` is a hint: true when the region contains stores to alloc.
130    fn is_region_promotable(
131        &self,
132        ctx: &Context,
133        alloc: &AllocInfo,
134        region: Ptr<Region>,
135        has_value_stores: bool,
136    ) -> bool;
137
138    /// Called before descending into nested regions.
139    /// `reaching_def` is the value in `slot` on entry to this op.
140    /// Populate `regions_to_process` with the reaching def each region starts with.
141    /// You may mutate the op in place, but do NOT delete ops or touch terminators.
142    fn setup_promotion(
143        &self,
144        ctx: &mut Context,
145        alloc: &AllocInfo,
146        reaching_def: Value,
147        has_value_stores: bool,
148        regions_to_process: &mut SmallMap<Ptr<Region>, Value, 2>,
149    );
150
151    /// Called after reaching defs are computed for all regions, but before
152    /// blocking uses are removed. Returns the new reaching def at the op's exit.
153    /// Mutation is allowed, but you must not change control flow or add ops that
154    /// interact with the slot's value.
155    fn finalize_promotion(
156        &self,
157        ctx: &mut Context,
158        alloc: &AllocInfo,
159        entry_reaching_def: Value,
160        has_value_stores: bool,
161        reaching_at_block_end: &HMap<Ptr<BasicBlock>, Value>,
162    ) -> Value;
163}
164
165#[op_interface]
166pub trait MemorySSARegionOpInterface {
167    verify_op_succ!();
168
169    /// Called before descending into nested regions.
170    /// `reaching_def` is the state of the memory on entry to this op.
171    /// Populate `regions_to_process` with the reaching def each region starts with.
172    fn setup_memory_ssa(
173        &self,
174        ctx: &Context,
175        state: &mut MemorySSAContext,
176        reaching_def: MemoryValue,
177        has_memory_defs: bool,
178        regions_to_process: &mut SmallMap<Ptr<Region>, MemoryValue, 2>,
179    );
180
181    /// Called after reaching defs are computed for all regions.
182    /// Returns the new reaching def at the op's exit.
183    /// If values need to be merged at the start of a nested region, the input values must be added
184    /// to `region_phis`. This will insert a memory phi at the beginning of the region's entry block,
185    /// with the result value being the one added to `regions_to_process` in `setup_memory_ssa`.
186    #[allow(
187        clippy::too_many_arguments,
188        reason = "packing them into structs would be more complex"
189    )]
190    fn finalize_memory_ssa(
191        &self,
192        ctx: &Context,
193        state: &mut MemorySSAContext,
194        entry_reaching_def: MemoryValue,
195        has_memory_defs: bool,
196        reaching_at_region_entry: &HMap<Ptr<Region>, MemoryValue>,
197        reaching_at_block_end: &HMap<Ptr<BasicBlock>, MemoryValue>,
198        region_phis: &mut SmallMap<Ptr<Region>, RegionMemoryPhiInputs, 2>,
199    ) -> RegionMemoryValue;
200}
201
202/// Describes a type that can be broken down into indexable sub-element types.
203#[type_interface]
204pub trait DestructurableTypeInterface {
205    verify_ty_succ!();
206
207    /// Destructures the type into subelements into a map of indices to
208    /// types of subelements. Returns nothing if the type cannot be destructured.
209    fn subelement_index_map(&self, ctx: &Context) -> Option<HMap<AttrObj, TypeHandle>>;
210
211    /// Indicates which type is held at the provided index, returning None
212    /// if no type could be computed. While this can return information
213    /// even when the type cannot be completely destructured, it must be coherent
214    /// with the types returned by `subelement_index_map` when they exist.
215    fn type_at_index(&self, ctx: &Context, index: &AttrObj) -> TypeHandle;
216}
217
218/// Describes operations creating values of aggregates that can be
219/// destructured into multiple smaller values.
220#[op_interface]
221pub trait DestructurableConstructorOpInterface {
222    verify_op_succ!();
223
224    /// Returns the list of value for which destructuring should be attempted,
225    /// specifying in which way the value should be destructured into subvalues.
226    /// This computes the type of the value for each subvalue to be generated. The type of the value
227    /// must implement [`DestructurableTypeInterface`].
228    ///
229    /// No IR mutation is allowed in this method.
230    fn destructurable_values(&self, ctx: &Context) -> Vec<DestructurableValueSlot>;
231
232    /// Destructures this value into multiple subvalues. The original value must still exist
233    /// at the end of this call. Only generates subvalues for the indices found in
234    /// `used_indices` since all other subvalues are unused.
235    ///
236    /// The rewriter is located before this op.
237    fn destructure(
238        &self,
239        ctx: &mut Context,
240        value: &DestructurableValueSlot,
241        used_indices: &SmallSet<AttrObj, 8>,
242        rewriter: &mut PassRewriter,
243        new_constructors: &mut Vec<TraitOp<dyn DestructurableConstructorOpInterface>>,
244    ) -> HMap<AttrObj, ValueSlot>;
245
246    /// Hook triggered once the destructuring of a value is complete, meaning the
247    /// original value is no longer being referred to and could be deleted.
248    /// This will only be called for values declared by this operation.
249    ///
250    /// Must return a new destructurable constructor op if this hook creates
251    /// a new destructurable op, `None` otherwise.
252    fn handle_destructuring_complete(
253        &self,
254        ctx: &mut Context,
255        value: &DestructurableValueSlot,
256        rewriter: &mut PassRewriter,
257    ) -> Option<TraitOp<dyn DestructurableConstructorOpInterface>>;
258}
259
260/// Describes operations that can access a sub-element of a destructurable value.
261#[op_interface]
262pub trait DestructurableAccessorOpInterface {
263    verify_op_succ!();
264
265    /// For a given destructurable value, returns whether this operation can
266    /// rewire its uses of the value to use the values generated after
267    /// destructuring. This may involve creating new operations.
268    ///
269    /// This method must also register the indices it will access within the
270    /// `used_indices` set. If the accessor generates new values mapping to
271    /// subelements, they must be registered in `must_be_safely_used` to ensure
272    /// they are used in a safe manner.
273    ///
274    /// No IR mutation is allowed in this method.
275    fn can_rewire(
276        &self,
277        ctx: &Context,
278        value: &DestructurableValueSlot,
279        used_indices: &mut SmallSet<AttrObj, 8>,
280        must_be_safely_used: &mut Vec<ValueSlot>,
281    ) -> bool;
282
283    /// Rewires the use of a slot to the generated subvalues, without deleting
284    /// any operation. Returns whether the accessor should be deleted.
285    ///
286    /// Deletion of operations is not allowed, only the accessor can be
287    /// scheduled for deletion by returning the appropriate value.
288    fn rewire(
289        &self,
290        ctx: &mut Context,
291        value: &DestructurableValueSlot,
292        subvalues: &HMap<AttrObj, ValueSlot>,
293        rewriter: &mut PassRewriter,
294    ) -> DeletionKind;
295}
296
297#[op_interface]
298pub trait SafeMemorySlotAccessOpInterface {
299    verify_op_succ!();
300
301    #[allow(clippy::result_unit_err)]
302    /// Returns whether all accesses in this operation to the provided value are
303    /// done in a safe manner. To be safe, the access must only access the value
304    /// inside the bounds that its type implies.
305    ///
306    /// If the safety of the accesses depends on the safety of the accesses to
307    /// further value, the result of this method will be conditioned to
308    /// the safety of the accesses to the value added by this method to
309    /// `must_be_safely_used`.
310    ///
311    /// No IR mutation is allowed in this method.
312    fn ensure_only_safe_accesses(
313        &self,
314        ctx: &Context,
315        value: &ValueSlot,
316        must_be_safely_used: &mut Vec<ValueSlot>,
317    ) -> LogicalResult;
318}
319
320#[derive(new, Debug)]
321pub struct ValueSlot {
322    pub value: Value,
323    pub elem_ty: TypeHandle,
324}
325
326#[derive(Debug)]
327pub struct DestructurableValueSlot {
328    pub slot: ValueSlot,
329    pub subelement_types: HMap<AttrObj, TypeHandle>,
330}
cubecl_ir/interfaces/memory_slot.rs

cubecl_ir/interfaces/
memory_slot.rs
