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

cubecl_ir/
scope.rs
