cubecl_ir/
scope.rs

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

cubecl_ir/
scope.rs
