cubecl_ir/
scope.rs

1use crate::{
2    EnumSet,
3    interfaces::control_flow::{SymbolOpInterface, SymbolVisibility},
4};
5use alloc::{boxed::Box, format, rc::Rc, string::String, vec, vec::Vec};
6use core::{
7    any::{TypeId, type_name},
8    cell::{Ref, RefCell, RefMut, UnsafeCell},
9    fmt::{Debug, Display},
10    sync::atomic::Ordering,
11};
12use cubecl_common::format::type_name_sanitized;
13use cubecl_environment::{collections::HashMap, sync::Mutex};
14use derive_more::{Eq, PartialEq};
15use pliron::{
16    attribute::AttrObj,
17    basic_block::BasicBlock,
18    builtin::{
19        attributes::{TypeAttr, VecAttr},
20        given_names::set_operation_result_name,
21        op_interfaces::{OneResultInterface, SingleBlockRegionInterface},
22        ops::{ConstantOp, FuncOp, ModuleOp},
23        type_interfaces::FunctionTypeInterface,
24        types::{FunctionType, UnitType},
25    },
26    context::{AuxDataIndex, Context},
27    dict_key,
28    identifier::Identifier,
29    irbuild::{
30        inserter::{IRInserter, Inserter},
31        listener::DummyListener,
32    },
33    op::Op,
34    operation::Operation,
35    printable::Printable,
36    r#type::{TypeHandle, Typed, type_cast},
37    value::Value,
38};
39use portable_atomic::AtomicUsize;
40use spin::LazyLock;
41
42use crate::{
43    AddressSpace, AddressType, DeviceProperties, ElemType, FastMath, TargetProperties, TypeHash,
44    arena::DropBump,
45    attributes::{
46        ATTR_BUFFER_BINDING, ATTR_KEY_ARG_ATTRS, ATTR_TENSOR_MAP_BINDING, BoolAttr,
47        BufferBindingAttr, EntrypointAbiAttr, EntrypointInterface, FuncInterface, IndexAttr,
48    },
49    dialect::{
50        OperationPtrExt,
51        branch::{IfOp, ReturnOp, YieldOp},
52        memory::DeclareVariableOp,
53        vector::CompositeExtractOp,
54    },
55    interfaces::{ScalarType, TypedExt},
56    read_value,
57    settings::KernelSettings,
58    types::{PointerType, RuntimeArrayType, cuda::TensorMapType, scalar::BoolType},
59};
60
61pub type Types = HashMap<TypeId, ElemType>;
62pub type Sizes = HashMap<TypeId, usize>;
63
64pub type OpInserter = IRInserter<DummyListener>;
65
66/// SAFETY: This should be fine for parsing the AST, hopefully. There's just no good way to
67/// have both owned and borrowed contexts in scopes.
68#[derive(Clone)]
69enum CtxHandle {
70    Rc(Rc<UnsafeCell<Context>>),
71    Ref(*mut Context),
72}
73
74impl CtxHandle {
75    pub fn borrow(&self) -> &Context {
76        match self {
77            CtxHandle::Rc(cell) => unsafe { &*cell.get() },
78            CtxHandle::Ref(ptr) => unsafe { &**ptr },
79        }
80    }
81
82    #[allow(clippy::mut_from_ref)]
83    pub fn borrow_mut(&self) -> &mut Context {
84        match self {
85            CtxHandle::Rc(cell) => unsafe { &mut *cell.get() },
86            CtxHandle::Ref(ptr) => unsafe { &mut **ptr },
87        }
88    }
89}
90
91/// SAFETY: This should be fine for parsing the AST, hopefully. There's just no good way to
92/// have both owned and borrowed inserters in scopes.
93enum InserterHandle {
94    Owned(Box<UnsafeCell<dyn Inserter>>),
95    Ref(*mut dyn Inserter),
96}
97
98impl InserterHandle {
99    pub fn owned(inserter: impl Inserter + 'static) -> Self {
100        Self::Owned(Box::new(UnsafeCell::new(inserter)))
101    }
102
103    #[allow(clippy::mut_from_ref)]
104    pub fn borrow_mut(&self) -> &mut dyn Inserter {
105        match self {
106            InserterHandle::Owned(cell) => unsafe { &mut *cell.get() },
107            InserterHandle::Ref(ptr) => unsafe { &mut **ptr },
108        }
109    }
110}
111
112/// The scope represents the region currently being parsed, as well as the current insertion point.
113/// It is created from scratch for the initial codegen phase, or rebuilt from an existing scope and
114/// a rewriter for passes/conversions.
115///
116/// All state apart from the insertion point is shared. Global state like the type map and errors
117/// are stored in the context's auxiliary storage, so they can be reconstructed from the rewriter
118/// state.
119#[allow(missing_docs)]
120pub struct Scope {
121    ctx: CtxHandle,
122    inserter: InserterHandle,
123    expand_state: RefCell<ExpandState>,
124}
125
126#[derive(Clone, Copy, Default)]
127pub struct ExpandState {
128    pub may_return: bool,
129    pub may_break: bool,
130    // Whether the kernel has *not* returned. Inverted to save a not on the loop condition.
131    pub inv_return_flag: Option<Value>,
132    /// Whether the loop is *not* broken. Inverted to save a not on the loop condition.
133    pub inv_break_flag: Option<Value>,
134}
135
136impl Debug for Scope {
137    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
138        f.debug_struct("Scope").finish()
139    }
140}
141
142#[track_caller]
143pub fn ident(name: impl Into<String>) -> Identifier {
144    Identifier::try_new(name.into()).unwrap()
145}
146
147pub struct GlobalState {
148    pub reference_arena: DropBump,
149    pub errors: Vec<String>,
150    pub warnings: Vec<String>,
151
152    pub module: ModuleOp,
153    pub module_inserter: OpInserter,
154    pub entry_func: FuncOp,
155    pub ident_unique_id: AtomicUsize,
156    pub typemap: Types,
157    pub sizemap: Sizes,
158    pub modes: InstructionModes,
159    pub target_properties: TargetProperties,
160    pub device_properties: Option<Rc<DeviceProperties>>,
161}
162
163unsafe impl Send for GlobalState {}
164
165impl GlobalState {
166    /// Register the element type for the given generic type.
167    pub fn register_type<T: 'static>(&mut self, elem: ElemType) {
168        self.typemap.insert(TypeId::of::<T>(), elem);
169    }
170}
171
172dict_key!(ADDRESS_TYPE_KEY, "kernel_address_type");
173
174static TY_IDENTS: LazyLock<Mutex<HashMap<TypeId, Identifier>>> = LazyLock::new(Default::default);
175
176fn ty_ident<T: 'static>() -> Identifier {
177    let mut idents = TY_IDENTS.lock();
178    let ident = idents
179        .entry(TypeId::of::<T>())
180        .or_insert_with(|| Identifier::try_from(type_name_sanitized::<T>()).unwrap());
181    ident.clone()
182}
183
184fn ty_key<T: 'static>(ctx: &Context) -> Option<AuxDataIndex> {
185    let mut idents = TY_IDENTS.lock();
186    let ident = idents
187        .entry(TypeId::of::<T>())
188        .or_insert_with(|| Identifier::try_from(type_name_sanitized::<T>()).unwrap());
189    ctx.aux_data_map.get(ident).copied()
190}
191
192pub trait ContextExt {
193    fn aux_ty<T: Send + 'static>(&self) -> &T;
194    fn aux_ty_mut<T: Send + 'static>(&mut self) -> &mut T;
195    fn set_aux_ty<T: Send + 'static>(&mut self, value: T);
196    fn set_address_type(&mut self, addr: AddressType);
197    fn address_type(&self) -> AddressType;
198}
199
200impl ContextExt for Context {
201    #[track_caller]
202    fn aux_ty<T: Send + 'static>(&self) -> &T {
203        let key = ty_key::<T>(self)
204            .ok_or_else(|| format!("Key for {} should exist", type_name::<T>()))
205            .unwrap();
206        self.aux_data[key].downcast_ref().unwrap()
207    }
208
209    #[track_caller]
210    fn aux_ty_mut<T: Send + 'static>(&mut self) -> &mut T {
211        let key = ty_key::<T>(self)
212            .ok_or_else(|| format!("Key for {} should exist", type_name::<T>()))
213            .unwrap();
214        self.aux_data[key].downcast_mut().unwrap()
215    }
216
217    fn set_aux_ty<T: Send + 'static>(&mut self, value: T) {
218        if let Some(key) = ty_key::<T>(self) {
219            *self.aux_data.get_mut(key).unwrap() = Box::new(value);
220        } else {
221            let ident = ty_ident::<T>();
222            let key = self.aux_data.insert(Box::new(value));
223            self.aux_data_map.insert(ident, key);
224        }
225    }
226
227    fn set_address_type(&mut self, addr: AddressType) {
228        if let Some(key) = self.aux_data_map.get(&ADDRESS_TYPE_KEY).copied() {
229            *self.aux_data.get_mut(key).unwrap() = Box::new(addr);
230        } else {
231            let key = self.aux_data.insert(Box::new(addr));
232            self.aux_data_map.insert(ADDRESS_TYPE_KEY.clone(), key);
233        }
234    }
235
236    fn address_type(&self) -> AddressType {
237        let key = self.aux_data_map[&ADDRESS_TYPE_KEY];
238        *self.aux_data[key].downcast_ref::<AddressType>().unwrap()
239    }
240}
241
242pub trait FuncOpExt {
243    fn push_argument(&self, ctx: &Context, ty: TypeHandle) -> usize;
244    fn pop_argument(&self, ctx: &Context);
245    fn remove_argument(&self, ctx: &Context, arg_idx: usize);
246    fn return_type(&self, ctx: &Context) -> TypeHandle;
247}
248
249impl FuncOpExt for FuncOp {
250    fn push_argument(&self, ctx: &Context, ty: TypeHandle) -> usize {
251        let id = BasicBlock::push_argument(self.get_entry_block(ctx), ctx, ty);
252
253        let (mut arg_types, res_types) = {
254            let current_func_ty = self.get_type(ctx).deref(ctx);
255            let current_func_ty = current_func_ty.downcast_ref::<FunctionType>().unwrap();
256            (current_func_ty.arg_types(), current_func_ty.res_types())
257        };
258
259        arg_types.insert(id, ty);
260        let new_func_ty = FunctionType::get(ctx, arg_types, res_types).to_handle();
261        self.set_attr_builtin_func_type(ctx, new_func_ty.into());
262        id
263    }
264
265    fn pop_argument(&self, ctx: &Context) {
266        let last_idx = self.get_entry_block(ctx).deref(ctx).get_num_arguments() - 1;
267        BasicBlock::pop_argument(self.get_entry_block(ctx), ctx);
268
269        let (mut arg_types, res_types) = {
270            let current_func_ty = self.get_type(ctx).deref(ctx);
271            let current_func_ty = current_func_ty.downcast_ref::<FunctionType>().unwrap();
272            (current_func_ty.arg_types(), current_func_ty.res_types())
273        };
274
275        arg_types.pop();
276        let new_func_ty = FunctionType::get(ctx, arg_types, res_types).to_handle();
277        self.set_attr_builtin_func_type(ctx, new_func_ty.into());
278        let mut op = self.get_operation().deref_mut(ctx);
279        let arg_attrs = op.attributes.0.get_mut(&ATTR_KEY_ARG_ATTRS);
280        if let Some(arg_attrs) = arg_attrs.and_then(|attr| attr.downcast_mut::<VecAttr>()) {
281            arg_attrs.0.truncate(last_idx);
282        }
283    }
284
285    fn remove_argument(&self, ctx: &Context, arg_idx: usize) {
286        BasicBlock::remove_argument(self.get_entry_block(ctx), ctx, arg_idx);
287
288        let (mut arg_types, res_types) = {
289            let current_func_ty = self.get_type(ctx).deref(ctx);
290            let current_func_ty = current_func_ty.downcast_ref::<FunctionType>().unwrap();
291            (current_func_ty.arg_types(), current_func_ty.res_types())
292        };
293
294        arg_types.remove(arg_idx);
295        let new_func_ty = FunctionType::get(ctx, arg_types, res_types).to_handle();
296        self.set_attr_builtin_func_type(ctx, new_func_ty.into());
297
298        let mut op = self.get_operation().deref_mut(ctx);
299        let arg_attrs = op.attributes.0.get_mut(&ATTR_KEY_ARG_ATTRS);
300        if let Some(arg_attrs) = arg_attrs.and_then(|attr| attr.downcast_mut::<VecAttr>()) {
301            arg_attrs.0.remove(arg_idx);
302        }
303    }
304
305    fn return_type(&self, ctx: &Context) -> TypeHandle {
306        let ty = self.get_type(ctx).deref(ctx);
307        ty.downcast_ref::<FunctionType>().unwrap().res_types()[0]
308    }
309}
310
311fn new_context(settings: KernelSettings) -> Rc<UnsafeCell<Context>> {
312    let mut ctx = Context::default();
313    ctx.set_address_type(settings.address_type);
314
315    let module = ModuleOp::new(&mut ctx, ident("kernel"));
316    let module_block = module.get_body(&ctx, 0);
317    let mut module_inserter = OpInserter::new_at_block_end(module_block);
318
319    // Start out empty and fill in once args register themselves
320    let entry_func_ty = FunctionType::get(&ctx, vec![], vec![UnitType::get(&ctx).into()]);
321    let entry_name = Identifier::try_new(settings.kernel_name).unwrap_or(ident("kernel_entry"));
322    let abi = EntrypointAbiAttr::new(settings.cube_dim, settings.cluster_dim);
323    let entry_func = FuncOp::new(&mut ctx, entry_name, entry_func_ty);
324    entry_func.set_visibility(&mut ctx, SymbolVisibility::Public);
325    entry_func.set_entrypoint_abi(&mut ctx, abi);
326    module_inserter.append_op(&ctx, &entry_func);
327
328    let mut state = GlobalState {
329        reference_arena: Default::default(),
330        module,
331        module_inserter,
332        entry_func,
333        ident_unique_id: Default::default(),
334        typemap: Default::default(),
335        sizemap: Default::default(),
336        modes: Default::default(),
337        target_properties: Default::default(),
338        device_properties: Default::default(),
339        errors: Default::default(),
340        warnings: Default::default(),
341    };
342    settings.address_type.register(&mut state);
343
344    ctx.set_aux_ty(state);
345    Rc::new(UnsafeCell::new(ctx))
346}
347
348/// Create a dummy context that can't be used for actual codegen. Useful for registering and
349/// resolving types without making the interface for that overly complex.
350/// Maybe we can replace the `&Scope` with a type registration trait on that function only at some point.
351fn dummy_context() -> Rc<UnsafeCell<Context>> {
352    let mut ctx = Context::default();
353
354    let module = ModuleOp::new(&mut ctx, ident("dummy_module"));
355    let module_block = module.get_body(&ctx, 0);
356    let mut module_inserter = OpInserter::new_at_block_end(module_block);
357
358    let entry_func_ty = FunctionType::get(&ctx, vec![], vec![UnitType::get(&ctx).into()]);
359    let entry_name = ident("dummy_entry");
360    let entry_func = FuncOp::new(&mut ctx, entry_name, entry_func_ty);
361    module_inserter.append_op(&ctx, &entry_func);
362
363    let state = GlobalState {
364        reference_arena: Default::default(),
365        module,
366        module_inserter,
367        entry_func,
368        ident_unique_id: Default::default(),
369        typemap: Default::default(),
370        sizemap: Default::default(),
371        modes: Default::default(),
372        target_properties: Default::default(),
373        device_properties: Default::default(),
374        errors: Default::default(),
375        warnings: Default::default(),
376    };
377
378    ctx.set_aux_ty(state);
379    Rc::new(UnsafeCell::new(ctx))
380}
381
382impl Debug for GlobalState {
383    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
384        f.debug_struct("GlobalStateInner")
385            .field("reference_arena", &self.reference_arena)
386            .field("typemap", &self.typemap)
387            .field("sizemap", &self.sizemap)
388            .field("modes", &self.modes)
389            .field("target_properties", &self.target_properties)
390            .field("device_properties", &self.device_properties)
391            .finish()
392    }
393}
394
395/// Modes set and reset during expansion
396#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
397#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, TypeHash)]
398pub struct InstructionModes {
399    pub fp_math_mode: EnumSet<FastMath>,
400}
401
402impl Scope {
403    /// Set the device properties.
404    pub fn device_properties(&self, properties: &DeviceProperties) {
405        self.state_mut().device_properties = Some(Rc::new(properties.clone()));
406    }
407
408    #[track_caller]
409    pub fn state(&self) -> &GlobalState {
410        self.ctx().aux_ty()
411    }
412
413    #[track_caller]
414    pub fn state_mut(&self) -> &mut GlobalState {
415        self.ctx_mut().aux_ty_mut()
416    }
417
418    fn ident_id(&self) -> usize {
419        self.state().ident_unique_id.fetch_add(1, Ordering::SeqCst)
420    }
421
422    #[track_caller]
423    pub fn expand_state(&self) -> Ref<'_, ExpandState> {
424        self.expand_state.borrow()
425    }
426
427    #[track_caller]
428    pub fn expand_state_mut(&self) -> RefMut<'_, ExpandState> {
429        self.expand_state.borrow_mut()
430    }
431
432    pub fn ctx(&self) -> &Context {
433        self.ctx.borrow()
434    }
435
436    // #[allow(clippy::mut_from_ref)]
437    pub fn ctx_mut(&self) -> &mut Context {
438        self.ctx.borrow_mut()
439    }
440
441    pub fn inserter(&self) -> &mut dyn Inserter {
442        self.inserter.borrow_mut()
443    }
444
445    /// Create a parse scope that is at the root of a kernel definition.
446    ///
447    /// A local scope can be created with the [child](Self::child) method.
448    pub fn root(settings: KernelSettings) -> Self {
449        let ctx = new_context(settings);
450        let mut inserter = {
451            let ctx = unsafe { &*ctx.get() };
452            let state = ctx.aux_ty::<GlobalState>();
453            let entry_block = state.entry_func.get_entry_block(ctx);
454            OpInserter::new_at_block_end(entry_block)
455        };
456        let return_flag =
457            init_bool_flag(unsafe { &mut *ctx.get() }, &mut inserter, "inv_return_flag");
458        Self {
459            ctx: CtxHandle::Rc(ctx),
460            inserter: InserterHandle::owned(inserter),
461            expand_state: RefCell::new(ExpandState {
462                may_break: false,
463                may_return: false,
464                inv_return_flag: Some(return_flag),
465                inv_break_flag: None,
466            }),
467        }
468    }
469
470    /// Create a parse scope that only exists for type registration in the kernel builder.
471    /// The state will be incomplete and should never be used for actual codegen.
472    pub fn dummy() -> Self {
473        let ctx = dummy_context();
474        let inserter = {
475            let ctx = unsafe { &*ctx.get() };
476            let state = ctx.aux_ty::<GlobalState>();
477            let entry_block = state.entry_func.get_entry_block(ctx);
478            OpInserter::new_at_block_end(entry_block)
479        };
480        Self {
481            ctx: CtxHandle::Rc(ctx),
482            inserter: InserterHandle::owned(inserter),
483            expand_state: Default::default(),
484        }
485    }
486
487    /// Create a rewrite scope from an existing context and inserter/rewriter
488    pub fn from_context_and_inserter(
489        ctx: &mut Context,
490        inserter: &mut (impl Inserter + 'static),
491    ) -> Self {
492        let inserter: *mut dyn Inserter = inserter;
493        Self {
494            ctx: CtxHandle::Ref(ctx),
495            inserter: InserterHandle::Ref(inserter),
496            expand_state: RefCell::new(ExpandState {
497                may_return: false,
498                may_break: false,
499                inv_return_flag: None,
500                inv_break_flag: None,
501            }),
502        }
503    }
504
505    /// Create a new mutable local variable of type specified by `value_ty`.
506    /// `initializer` is a constant attribute and has the same rules as `OpConstant`. This is because
507    /// SPIR-V does not allow non-constant (technically non-global, but constants are the only
508    /// non-pointer globals) initializers.
509    pub fn create_local_mut(
510        &self,
511        value_ty: impl Into<TypeHandle>,
512        init: Option<AttrObj>,
513    ) -> Value {
514        let value_ty = value_ty.into();
515        let ctx = self.ctx_mut();
516        let align = value_ty.align(ctx);
517        let op = DeclareVariableOp::new(ctx, value_ty, AddressSpace::Local, align, init);
518        let out = op.get_result(ctx);
519        self.inserter().append_op(ctx, &op);
520        out
521    }
522
523    /// Create a shared variable of the given item type.
524    pub fn create_shared(
525        &self,
526        value_ty: impl Into<TypeHandle>,
527        alignment: Option<usize>,
528    ) -> Value {
529        let value_ty = value_ty.into();
530        let ctx = self.ctx_mut();
531        let align = alignment.unwrap_or_else(|| value_ty.align(ctx));
532        let op = DeclareVariableOp::new(ctx, value_ty, AddressSpace::Shared, align, None);
533        let out = op.get_result(ctx);
534        self.inserter().append_op(ctx, &op);
535        out
536    }
537
538    pub fn func_ident(&self, label: Option<&str>) -> Identifier {
539        let unique_id = self.ident_id();
540        match label {
541            Some(label) => ident(format!("{label}_{unique_id}")),
542            None => ident(format!("func_{unique_id}")),
543        }
544    }
545
546    /// Create a new function.
547    pub fn register_func(&self, func: FuncOp) {
548        let ctx = self.ctx_mut();
549        let state = self.state_mut();
550        func.set_visibility(ctx, SymbolVisibility::Private);
551        state.module_inserter.append_op(ctx, &func);
552    }
553
554    /// Register an [`Instruction`] into the scope.
555    pub fn register(&self, op: &dyn Op) {
556        let ctx = self.ctx();
557        self.inserter().append_op(ctx, op);
558    }
559
560    /// Register an [`Instruction`] into the scope and return its result.
561    pub fn register_with_result(&self, op: &dyn OneResultInterface) -> Value {
562        self.register(op);
563        op.get_result(self.ctx())
564    }
565
566    /// Terminate block with a `cube.yield` if not already terminated
567    pub fn terminate_yield(&self) {
568        let block = self.inserter().get_insertion_block(self.ctx());
569        let block = block.expect("Should have insertion block");
570        if block.deref(self.ctx()).get_terminator(self.ctx()).is_none() {
571            self.register(&YieldOp::new(self.ctx_mut()));
572        }
573    }
574
575    pub fn set_break_return(&self, children: &[Scope]) {
576        self.set_may_break(children);
577        self.set_may_return(children);
578    }
579
580    pub fn set_may_return(&self, children: &[Scope]) {
581        let child_may_return = children.iter().any(|scope| scope.expand_state().may_return);
582        if child_may_return {
583            self.expand_state_mut().may_return = true;
584            let flag = self.expand_state().inv_return_flag;
585            self.predicate_on_flag(flag.expect("Can't return in rewrite context"));
586        }
587    }
588
589    pub fn set_may_break(&self, children: &[Scope]) {
590        let child_may_return = children.iter().any(|scope| scope.expand_state().may_break);
591        if child_may_return {
592            self.expand_state_mut().may_break = true;
593            let flag = self.expand_state().inv_break_flag;
594            self.predicate_on_flag(flag.expect("Should have break flag"));
595        }
596    }
597
598    fn predicate_on_flag(&self, flag: Value) {
599        let ctx = self.ctx_mut();
600        let cond = read_value(self, flag);
601        let predication = IfOp::new(ctx, cond);
602        let then_block = predication.then_block(ctx);
603        let else_block = predication.else_block(ctx);
604        let yield_ = YieldOp::new(ctx).get_operation();
605        yield_.insert_at_back(then_block, ctx);
606        let yield_ = YieldOp::new(ctx).get_operation();
607        yield_.insert_at_back(else_block, ctx);
608        self.register(&predication);
609        // The insertion cursor can sit inside a block that already ends with a
610        // terminator — e.g. the then-block of an earlier predication, whose
611        // trailing yield was placed when that predication was built. Appending
612        // another yield there would leave the block with two terminators, which
613        // fails module verification. Only terminate the block if it isn't yet.
614        self.terminate_yield();
615        self.inserter()
616            .set_insertion_point_to_block_start(then_block);
617    }
618
619    /// Add a value to the global arena so we can create a kernel-wide reference to it.
620    /// The reference is the same as the type for simplicity, but is only valid for the duration of
621    /// the root scope. Ensure the reference lifetime is shortened to the lifetime of the underlying
622    /// variable being referenced.
623    pub fn create_kernel_ref<'a, T>(&self, value: T) -> &'a mut T
624    where
625        T: 'a,
626    {
627        let state = self.state_mut();
628        let reference = state.reference_arena.alloc(value);
629        unsafe { core::mem::transmute(reference) }
630    }
631
632    /// Resolve the element type of the given generic type.
633    pub fn resolve_type<T: 'static>(&self) -> Option<ElemType> {
634        let state = self.state();
635        let result = state.typemap.get(&TypeId::of::<T>());
636
637        result.cloned()
638    }
639
640    /// Resolve the comptime size of the given generic size.
641    pub fn resolve_size<T: 'static>(&self) -> Option<usize> {
642        let state = self.state();
643        let result = state.sizemap.get(&TypeId::of::<T>());
644
645        result.cloned()
646    }
647
648    /// Register the element type for the given generic type.
649    pub fn register_type<T: 'static>(&self, elem: ElemType) {
650        self.state_mut().register_type::<T>(elem);
651    }
652
653    /// Register the comptime size for the given generic size.
654    pub fn register_size<T: 'static>(&self, size: usize) {
655        let state = self.state_mut();
656
657        state.sizemap.insert(TypeId::of::<T>(), size);
658    }
659
660    /// Register the type and size of a scalarizable type
661    pub fn register_value_type<T: 'static, N: 'static>(&self, value: impl Typed) {
662        let ty = value.get_type(self.ctx());
663        let scalar_ty = ty.scalar_ty(self.ctx());
664        let vector_size = ty.vector_size(self.ctx());
665        let storage_ty = {
666            let ctx = self.ctx();
667            let scalar_ty = scalar_ty.deref(ctx);
668            let scalar = type_cast::<dyn ScalarType>(&*scalar_ty).unwrap();
669            scalar.elem_type(ctx)
670        };
671        self.register_type::<T>(storage_ty);
672        self.register_size::<N>(vector_size);
673    }
674
675    /// Create an empty child scope.
676    pub fn child(&self, inserter: impl Inserter + 'static) -> Self {
677        Self {
678            ctx: self.ctx.clone(),
679            inserter: InserterHandle::owned(inserter),
680            expand_state: RefCell::new(ExpandState {
681                may_break: false,
682                may_return: false,
683                inv_return_flag: self.expand_state().inv_return_flag,
684                inv_break_flag: self.expand_state().inv_break_flag,
685            }),
686        }
687    }
688
689    /// Create a child scope with a new break condition.
690    pub fn loop_child(&self, inserter: impl Inserter + 'static) -> Self {
691        let break_flag = init_bool_flag(self.ctx_mut(), self.inserter(), "inv_break_flag");
692        Self {
693            ctx: self.ctx.clone(),
694            inserter: InserterHandle::owned(inserter),
695            expand_state: RefCell::new(ExpandState {
696                may_return: false,
697                may_break: false,
698                inv_return_flag: self.expand_state().inv_return_flag,
699                inv_break_flag: Some(break_flag),
700            }),
701        }
702    }
703
704    /// Create a child that's at the root of a new function
705    pub fn func_child(&self, mut inserter: impl Inserter + 'static) -> Self {
706        let return_flag = init_bool_flag(self.ctx_mut(), &mut inserter, "inv_return_flag");
707        Self {
708            ctx: self.ctx.clone(),
709            inserter: InserterHandle::owned(inserter),
710            expand_state: RefCell::new(ExpandState {
711                may_break: false,
712                may_return: false,
713                inv_return_flag: Some(return_flag),
714                inv_break_flag: None,
715            }),
716        }
717    }
718
719    // Adds a validation error.
720    pub fn push_error(&self, msg: impl Into<String>) {
721        self.state_mut().errors.push(msg.into());
722    }
723
724    /// Returns all validation errors.
725    pub fn pop_errors(&self) -> Vec<String> {
726        core::mem::take(&mut self.state_mut().errors)
727    }
728
729    /// Adds a non-fatal validation warning. Logged when the kernel is built.
730    pub fn push_warning(&self, msg: impl Into<String>) {
731        self.state_mut().warnings.push(msg.into());
732    }
733
734    /// Returns all validation warnings.
735    pub fn pop_warnings(&self) -> Vec<String> {
736        core::mem::take(&mut self.state_mut().warnings)
737    }
738
739    /// Obtain the index-th buffer
740    pub fn global(
741        &self,
742        buffer_pos: usize,
743        ext_meta_pos: Option<usize>,
744        value_ty: TypeHandle,
745    ) -> Value {
746        let entry_func = self.state().entry_func;
747        let ctx = self.ctx_mut();
748
749        let ty_arr = RuntimeArrayType::get(ctx, value_ty);
750        let ty = PointerType::get(ctx, ty_arr.into(), AddressSpace::Global(buffer_pos));
751
752        let id = entry_func.push_argument(ctx, ty.to_handle());
753        entry_func.set_arg_attr(
754            ctx,
755            id,
756            &ATTR_BUFFER_BINDING,
757            Box::new(BufferBindingAttr::new(buffer_pos, ext_meta_pos)),
758        );
759
760        entry_func.get_entry_block(ctx).deref(ctx).get_argument(id)
761    }
762
763    /// Obtain the index-th tensor map
764    pub fn tensor_map(&self, buffer_pos: usize, ext_meta_pos: usize) -> Value {
765        let entry_func = self.state().entry_func;
766        let ctx = self.ctx();
767        let ty = TensorMapType::get(ctx);
768
769        let id = entry_func.push_argument(ctx, ty.to_handle());
770        entry_func.set_arg_attr_unit(ctx, id, &ATTR_TENSOR_MAP_BINDING);
771        entry_func.set_arg_attr(
772            ctx,
773            id,
774            &ATTR_BUFFER_BINDING,
775            Box::new(BufferBindingAttr::new(buffer_pos, Some(ext_meta_pos))),
776        );
777        entry_func.get_entry_block(ctx).deref(ctx).get_argument(id)
778    }
779
780    pub fn kernel_arg(&self, idx: usize) -> Value {
781        let entry_block = self.state().entry_func.get_entry_block(self.ctx());
782        entry_block.deref(self.ctx()).get_argument(idx)
783    }
784
785    pub fn extract_field(&self, aggregate: Value, field: usize) -> Value {
786        let ctx = self.ctx_mut();
787        let op = CompositeExtractOp::new(ctx, aggregate, field);
788        self.register_with_result(&op)
789    }
790
791    pub fn const_usize(&self, value: usize) -> Value {
792        let op = ConstantOp::new(self.ctx_mut(), Box::new(IndexAttr::new(value)));
793        self.register_with_result(&op)
794    }
795
796    pub fn const_bool(&self, value: bool) -> Value {
797        let op = ConstantOp::new(self.ctx_mut(), Box::new(BoolAttr::new(value)));
798        self.register_with_result(&op)
799    }
800
801    pub fn into_context(self) -> Option<Context> {
802        let entry = self.state().entry_func.get_entry_block(self.ctx());
803        let term = entry.deref(self.ctx()).get_terminator(self.ctx());
804        let is_yield = term.is_some_and(|term| term.is_op::<YieldOp>(self.ctx()));
805        if let Some(term) = term
806            && is_yield
807        {
808            self.inserter().set_insertion_point_to_block_end(entry);
809            self.register(&ReturnOp::new(self.ctx_mut()));
810            Operation::erase(term, self.ctx_mut());
811        } else if term.is_none() {
812            self.inserter().set_insertion_point_to_block_end(entry);
813            self.register(&ReturnOp::new(self.ctx_mut()));
814        }
815        match self.ctx {
816            CtxHandle::Rc(ctx) => Some(Rc::into_inner(ctx)?.into_inner()),
817            CtxHandle::Ref(_) => None,
818        }
819    }
820}
821
822impl Display for Scope {
823    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
824        let ctx = self.ctx();
825        let state = self.state();
826        write!(f, "{}", state.module.disp(ctx))
827    }
828}
829
830fn init_bool_flag(ctx: &mut Context, inserter: &mut dyn Inserter, name: &str) -> Value {
831    let bool = TypeAttr::new(BoolType::get(ctx).to_handle());
832    let r#true = BoolAttr::new(true).into();
833    let flag = DeclareVariableOp::new(ctx, bool, AddressSpace::Local, 1, Some(r#true));
834    inserter.append_op(ctx, &flag);
835    set_operation_result_name(ctx, flag.get_operation(), 0, Some(ident(name)));
836    flag.get_result(ctx)
837}
cubecl_ir/scope.rs

cubecl_ir/
scope.rs
