cubecl_ir/

scope.rs

use crate::{
    EnumSet,
    dialect::memory::StoreOp,
    interfaces::control_flow::{SymbolOpInterface, SymbolVisibility},
    rewrite::op_insertion_point_in_block,
};
use alloc::{boxed::Box, format, rc::Rc, string::String, sync::Arc, vec, vec::Vec};
use core::{
    any::{TypeId, type_name},
    cell::{Ref, RefCell, RefMut, UnsafeCell},
    fmt::{Debug, Display},
    sync::atomic::Ordering,
};
use cubecl_common::format::type_name_sanitized;
use cubecl_environment::{collections::HashMap, sync::Mutex};
use derive_more::{Eq, PartialEq};
use pliron::{
    attribute::AttrObj,
    basic_block::BasicBlock,
    builtin::{
        attributes::{TypeAttr, VecAttr},
        given_names::set_operation_result_name,
        op_interfaces::{OneResultInterface, SingleBlockRegionInterface},
        ops::{ConstantOp, FuncOp, ModuleOp},
        type_interfaces::FunctionTypeInterface,
        types::{FunctionType, UnitType},
    },
    context::{AuxDataIndex, Context},
    dict_key,
    identifier::Identifier,
    irbuild::inserter::{IRInserter, Inserter, OpInsertionPoint},
    op::Op,
    operation::Operation,
    printable::Printable,
    r#type::{TypeHandle, Typed, type_cast},
    value::Value,
};
use portable_atomic::AtomicUsize;
use spin::LazyLock;

use crate::{
    AddressSpace, AddressType, DeviceProperties, ElemType, FastMath, TargetProperties, TypeHash,
    arena::DropBump,
    attributes::{
        ATTR_BUFFER_BINDING, ATTR_KEY_ARG_ATTRS, ATTR_PEER_BINDING, ATTR_TENSOR_MAP_BINDING,
        BoolAttr, BufferBindingAttr, EntrypointAbiAttr, EntrypointInterface, FuncInterface,
        IndexAttr, peer_buffer_positions,
    },
    debug::{DebugState, LocationListener},
    dialect::{
        OperationPtrExt,
        branch::{IfOp, ReturnOp, YieldOp},
        memory::DeclareVariableOp,
        vector::CompositeExtractOp,
    },
    interfaces::{ScalarType, TypedExt},
    read_value,
    settings::{DebugInfo, KernelSettings},
    types::{PointerType, RuntimeArrayType, cuda::TensorMapType, scalar::BoolType},
};

pub type Types = HashMap<TypeId, ElemType>;
pub type Sizes = HashMap<TypeId, usize>;

pub type OpInserter = IRInserter<LocationListener>;

/// SAFETY: This should be fine for parsing the AST, hopefully. There's just no good way to
/// have both owned and borrowed contexts in scopes.
#[derive(Clone)]
enum CtxHandle {
    Rc(Rc<UnsafeCell<Context>>),
    Ref(*mut Context),
}

impl CtxHandle {
    pub fn borrow(&self) -> &Context {
        match self {
            CtxHandle::Rc(cell) => unsafe { &*cell.get() },
            CtxHandle::Ref(ptr) => unsafe { &**ptr },
        }
    }

    #[allow(clippy::mut_from_ref)]
    pub fn borrow_mut(&self) -> &mut Context {
        match self {
            CtxHandle::Rc(cell) => unsafe { &mut *cell.get() },
            CtxHandle::Ref(ptr) => unsafe { &mut **ptr },
        }
    }
}

/// SAFETY: This should be fine for parsing the AST, hopefully. There's just no good way to
/// have both owned and borrowed inserters in scopes.
enum InserterHandle {
    Owned(Box<UnsafeCell<dyn Inserter>>),
    Ref(*mut dyn Inserter),
}

impl InserterHandle {
    pub fn owned(inserter: impl Inserter + 'static) -> Self {
        Self::Owned(Box::new(UnsafeCell::new(inserter)))
    }

    #[allow(clippy::mut_from_ref)]
    pub fn borrow_mut(&self) -> &mut dyn Inserter {
        match self {
            InserterHandle::Owned(cell) => unsafe { &mut *cell.get() },
            InserterHandle::Ref(ptr) => unsafe { &mut **ptr },
        }
    }
}

/// The scope represents the region currently being parsed, as well as the current insertion point.
/// It is created from scratch for the initial codegen phase, or rebuilt from an existing scope and
/// a rewriter for passes/conversions.
///
/// All state apart from the insertion point is shared. Global state like the type map and errors
/// are stored in the context's auxiliary storage, so they can be reconstructed from the rewriter
/// state.
#[allow(missing_docs)]
pub struct Scope {
    ctx: CtxHandle,
    inserter: InserterHandle,
    expand_state: RefCell<ExpandState>,
}

#[derive(Clone, Copy, Default)]
pub struct ExpandState {
    pub may_terminate: bool,
    pub may_return: bool,
    pub may_break: bool,
    pub may_continue: bool,
    // Whether the kernel has *not* terminated. Inverted to save a not on the loop condition.
    pub inv_terminate_flag: Option<Value>,
    // Whether the current function has *not* returned. Inverted to save a not on the loop condition.
    pub inv_return_flag: Option<Value>,
    /// Whether the loop is *not* broken. Inverted to save a not on the loop condition.
    pub inv_break_flag: Option<Value>,
    /// Whether the loop is *not* continued. Inverted to save a not on the loop condition.
    pub inv_continue_flag: Option<Value>,
    /// The return value, if early return is used.
    pub return_value: Option<Value>,
}

impl Debug for Scope {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("Scope").finish()
    }
}

#[track_caller]
pub fn ident(name: impl Into<String>) -> Identifier {
    Identifier::try_new(name.into()).unwrap()
}

pub struct GlobalState {
    pub reference_arena: DropBump,
    pub errors: Vec<String>,
    pub warnings: Vec<String>,

    pub module: ModuleOp,
    pub module_inserter: OpInserter,
    pub entry_func: FuncOp,
    pub ident_unique_id: AtomicUsize,
    pub typemap: Types,
    pub sizemap: Sizes,
    pub modes: InstructionModes,
    pub target_properties: TargetProperties,
    pub device_properties: Option<Arc<DeviceProperties>>,
}

unsafe impl Send for GlobalState {}

impl GlobalState {
    /// Register the element type for the given generic type.
    pub fn register_type<T: 'static>(&mut self, elem: ElemType) {
        self.typemap.insert(TypeId::of::<T>(), elem);
    }
}

dict_key!(ADDRESS_TYPE_KEY, "kernel_address_type");

static TY_IDENTS: LazyLock<Mutex<HashMap<TypeId, Identifier>>> = LazyLock::new(Default::default);

fn ty_ident<T: 'static>() -> Identifier {
    let mut idents = TY_IDENTS.lock();
    let ident = idents
        .entry(TypeId::of::<T>())
        .or_insert_with(|| Identifier::try_from(type_name_sanitized::<T>()).unwrap());
    ident.clone()
}

fn ty_key<T: 'static>(ctx: &Context) -> Option<AuxDataIndex> {
    let mut idents = TY_IDENTS.lock();
    let ident = idents
        .entry(TypeId::of::<T>())
        .or_insert_with(|| Identifier::try_from(type_name_sanitized::<T>()).unwrap());
    ctx.aux_data_map.get(ident).copied()
}

pub trait ContextExt {
    fn aux_ty<T: Send + 'static>(&self) -> &T;
    fn try_aux_ty<T: Send + 'static>(&self) -> Option<&T>;
    fn aux_ty_mut<T: Send + 'static>(&mut self) -> &mut T;
    fn try_aux_ty_mut<T: Send + 'static>(&mut self) -> Option<&mut T>;
    fn set_aux_ty<T: Send + 'static>(&mut self, value: T);
    fn set_address_type(&mut self, addr: AddressType);
    fn address_type(&self) -> AddressType;
}

impl ContextExt for Context {
    #[track_caller]
    fn aux_ty<T: Send + 'static>(&self) -> &T {
        self.try_aux_ty()
            .ok_or_else(|| format!("Key for {} should exist", type_name::<T>()))
            .unwrap()
    }

    fn try_aux_ty<T: Send + 'static>(&self) -> Option<&T> {
        self.aux_data[ty_key::<T>(self)?].downcast_ref()
    }

    #[track_caller]
    fn aux_ty_mut<T: Send + 'static>(&mut self) -> &mut T {
        self.try_aux_ty_mut()
            .ok_or_else(|| format!("Key for {} should exist", type_name::<T>()))
            .unwrap()
    }

    fn try_aux_ty_mut<T: Send + 'static>(&mut self) -> Option<&mut T> {
        let key = ty_key::<T>(self)?;
        self.aux_data[key].downcast_mut()
    }

    fn set_aux_ty<T: Send + 'static>(&mut self, value: T) {
        if let Some(key) = ty_key::<T>(self) {
            *self.aux_data.get_mut(key).unwrap() = Box::new(value);
        } else {
            let ident = ty_ident::<T>();
            let key = self.aux_data.insert(Box::new(value));
            self.aux_data_map.insert(ident, key);
        }
    }

    fn set_address_type(&mut self, addr: AddressType) {
        if let Some(key) = self.aux_data_map.get(&ADDRESS_TYPE_KEY).copied() {
            *self.aux_data.get_mut(key).unwrap() = Box::new(addr);
        } else {
            let key = self.aux_data.insert(Box::new(addr));
            self.aux_data_map.insert(ADDRESS_TYPE_KEY.clone(), key);
        }
    }

    fn address_type(&self) -> AddressType {
        let key = self.aux_data_map[&ADDRESS_TYPE_KEY];
        *self.aux_data[key].downcast_ref::<AddressType>().unwrap()
    }
}

pub trait FuncOpExt {
    fn push_argument(&self, ctx: &Context, ty: TypeHandle) -> usize;
    fn pop_argument(&self, ctx: &Context);
    fn remove_argument(&self, ctx: &Context, arg_idx: usize);
    fn return_type(&self, ctx: &Context) -> TypeHandle;
}

impl FuncOpExt for FuncOp {
    fn push_argument(&self, ctx: &Context, ty: TypeHandle) -> usize {
        let id = BasicBlock::push_argument(self.get_entry_block(ctx), ctx, ty);

        let (mut arg_types, res_types) = {
            let current_func_ty = self.get_type(ctx).deref(ctx);
            let current_func_ty = current_func_ty.downcast_ref::<FunctionType>().unwrap();
            (current_func_ty.arg_types(), current_func_ty.res_types())
        };

        arg_types.insert(id, ty);
        let new_func_ty = FunctionType::get(ctx, arg_types, res_types).to_handle();
        self.set_attr_builtin_func_type(ctx, new_func_ty.into());
        id
    }

    fn pop_argument(&self, ctx: &Context) {
        let last_idx = self.get_entry_block(ctx).deref(ctx).get_num_arguments() - 1;
        BasicBlock::pop_argument(self.get_entry_block(ctx), ctx);

        let (mut arg_types, res_types) = {
            let current_func_ty = self.get_type(ctx).deref(ctx);
            let current_func_ty = current_func_ty.downcast_ref::<FunctionType>().unwrap();
            (current_func_ty.arg_types(), current_func_ty.res_types())
        };

        arg_types.pop();
        let new_func_ty = FunctionType::get(ctx, arg_types, res_types).to_handle();
        self.set_attr_builtin_func_type(ctx, new_func_ty.into());
        let mut op = self.get_operation().deref_mut(ctx);
        let arg_attrs = op.attributes.0.get_mut(&ATTR_KEY_ARG_ATTRS);
        if let Some(arg_attrs) = arg_attrs.and_then(|attr| attr.downcast_mut::<VecAttr>()) {
            arg_attrs.0.truncate(last_idx);
        }
    }

    fn remove_argument(&self, ctx: &Context, arg_idx: usize) {
        BasicBlock::remove_argument(self.get_entry_block(ctx), ctx, arg_idx);

        let (mut arg_types, res_types) = {
            let current_func_ty = self.get_type(ctx).deref(ctx);
            let current_func_ty = current_func_ty.downcast_ref::<FunctionType>().unwrap();
            (current_func_ty.arg_types(), current_func_ty.res_types())
        };

        arg_types.remove(arg_idx);
        let new_func_ty = FunctionType::get(ctx, arg_types, res_types).to_handle();
        self.set_attr_builtin_func_type(ctx, new_func_ty.into());

        let mut op = self.get_operation().deref_mut(ctx);
        let arg_attrs = op.attributes.0.get_mut(&ATTR_KEY_ARG_ATTRS);
        if let Some(arg_attrs) = arg_attrs.and_then(|attr| attr.downcast_mut::<VecAttr>()) {
            arg_attrs.0.remove(arg_idx);
        }
    }

    fn return_type(&self, ctx: &Context) -> TypeHandle {
        let ty = self.get_type(ctx).deref(ctx);
        ty.downcast_ref::<FunctionType>().unwrap().res_types()[0]
    }
}

fn new_context(settings: KernelSettings) -> Rc<UnsafeCell<Context>> {
    let mut ctx = Context::default();
    ctx.set_address_type(settings.address_type);

    let module = ModuleOp::new(&mut ctx, ident("kernel"));
    let module_block = module.get_body(&ctx, 0);
    let mut module_inserter = OpInserter::new_at_block_end(module_block);

    // Start out empty and fill in once args register themselves
    let entry_func_ty = FunctionType::get(&ctx, vec![], vec![UnitType::get(&ctx).into()]);
    let entry_name = Identifier::try_new(settings.kernel_name).unwrap_or(ident("kernel_entry"));
    let abi = EntrypointAbiAttr::new(settings.cube_dim, settings.cluster_dim);
    let entry_func = FuncOp::new(&mut ctx, entry_name, entry_func_ty);
    entry_func.set_visibility(&mut ctx, SymbolVisibility::Public);
    entry_func.set_entrypoint_abi(&mut ctx, abi);
    module_inserter.append_op(&ctx, &entry_func);

    let mut state = GlobalState {
        reference_arena: Default::default(),
        module,
        module_inserter,
        entry_func,
        ident_unique_id: Default::default(),
        typemap: Default::default(),
        sizemap: Default::default(),
        modes: Default::default(),
        target_properties: Default::default(),
        device_properties: Default::default(),
        errors: Default::default(),
        warnings: Default::default(),
    };
    settings.address_type.register(&mut state);

    ctx.set_aux_ty(state);
    ctx.set_aux_ty(DebugState::new(settings.debug_info != DebugInfo::None));
    Rc::new(UnsafeCell::new(ctx))
}

/// Create a dummy context that can't be used for actual codegen. Useful for registering and
/// resolving types without making the interface for that overly complex.
/// Maybe we can replace the `&Scope` with a type registration trait on that function only at some point.
fn dummy_context() -> Rc<UnsafeCell<Context>> {
    let mut ctx = Context::default();

    let module = ModuleOp::new(&mut ctx, ident("dummy_module"));
    let module_block = module.get_body(&ctx, 0);
    let mut module_inserter = OpInserter::new_at_block_end(module_block);

    let entry_func_ty = FunctionType::get(&ctx, vec![], vec![UnitType::get(&ctx).into()]);
    let entry_name = ident("dummy_entry");
    let entry_func = FuncOp::new(&mut ctx, entry_name, entry_func_ty);
    module_inserter.append_op(&ctx, &entry_func);

    let state = GlobalState {
        reference_arena: Default::default(),
        module,
        module_inserter,
        entry_func,
        ident_unique_id: Default::default(),
        typemap: Default::default(),
        sizemap: Default::default(),
        modes: Default::default(),
        target_properties: Default::default(),
        device_properties: Default::default(),
        errors: Default::default(),
        warnings: Default::default(),
    };

    ctx.set_aux_ty(state);
    Rc::new(UnsafeCell::new(ctx))
}

impl Debug for GlobalState {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("GlobalStateInner")
            .field("reference_arena", &self.reference_arena)
            .field("typemap", &self.typemap)
            .field("sizemap", &self.sizemap)
            .field("modes", &self.modes)
            .field("target_properties", &self.target_properties)
            .field("device_properties", &self.device_properties)
            .finish()
    }
}

/// Modes set and reset during expansion
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, TypeHash)]
pub struct InstructionModes {
    pub fp_math_mode: EnumSet<FastMath>,
}

impl Scope {
    /// Set the device properties, shared with the client that holds them: a launch sets them
    /// every time, so they are never copied.
    pub fn device_properties(&self, properties: Arc<DeviceProperties>) {
        self.state_mut().device_properties = Some(properties);
    }

    #[track_caller]
    pub fn state(&self) -> &GlobalState {
        self.ctx().aux_ty()
    }

    #[track_caller]
    pub fn state_mut(&self) -> &mut GlobalState {
        self.ctx_mut().aux_ty_mut()
    }

    /// The source location state, or `None` when this kernel records no locations.
    pub fn debug_state(&self) -> Option<&mut DebugState> {
        self.ctx_mut()
            .try_aux_ty_mut::<DebugState>()
            .filter(|debug| debug.is_enabled())
    }

    fn ident_id(&self) -> usize {
        self.state().ident_unique_id.fetch_add(1, Ordering::SeqCst)
    }

    #[track_caller]
    pub fn expand_state(&self) -> Ref<'_, ExpandState> {
        self.expand_state.borrow()
    }

    #[track_caller]
    pub fn expand_state_mut(&self) -> RefMut<'_, ExpandState> {
        self.expand_state.borrow_mut()
    }

    pub fn ctx(&self) -> &Context {
        self.ctx.borrow()
    }

    // #[allow(clippy::mut_from_ref)]
    pub fn ctx_mut(&self) -> &mut Context {
        self.ctx.borrow_mut()
    }

    pub fn inserter(&self) -> &mut dyn Inserter {
        self.inserter.borrow_mut()
    }

    /// Create a parse scope that is at the root of a kernel definition.
    ///
    /// A local scope can be created with the [child](Self::child) method.
    pub fn root(settings: KernelSettings) -> Self {
        let ctx = new_context(settings);
        let mut inserter = {
            let ctx = unsafe { &*ctx.get() };
            let state = ctx.aux_ty::<GlobalState>();
            let entry_block = state.entry_func.get_entry_block(ctx);
            OpInserter::new_at_block_end(entry_block)
        };
        let terminate_flag = init_bool_flag(
            unsafe { &mut *ctx.get() },
            &mut inserter,
            "inv_terminate_flag",
        );
        Self {
            ctx: CtxHandle::Rc(ctx),
            inserter: InserterHandle::owned(inserter),
            expand_state: RefCell::new(ExpandState {
                inv_terminate_flag: Some(terminate_flag),
                ..ExpandState::default()
            }),
        }
    }

    /// Create a parse scope that only exists for type registration in the kernel builder.
    /// The state will be incomplete and should never be used for actual codegen.
    pub fn dummy() -> Self {
        let ctx = dummy_context();
        let inserter = {
            let ctx = unsafe { &*ctx.get() };
            let state = ctx.aux_ty::<GlobalState>();
            let entry_block = state.entry_func.get_entry_block(ctx);
            OpInserter::new_at_block_end(entry_block)
        };
        Self {
            ctx: CtxHandle::Rc(ctx),
            inserter: InserterHandle::owned(inserter),
            expand_state: Default::default(),
        }
    }

    /// Create a rewrite scope from an existing context and inserter/rewriter
    pub fn from_context_and_inserter(
        ctx: &mut Context,
        inserter: &mut (impl Inserter + 'static),
    ) -> Self {
        let inserter: *mut dyn Inserter = inserter;
        Self {
            ctx: CtxHandle::Ref(ctx),
            inserter: InserterHandle::Ref(inserter),
            expand_state: RefCell::new(ExpandState::default()),
        }
    }

    /// Create a new mutable local variable of type specified by `value_ty`.
    /// `initializer` is a constant attribute and has the same rules as `OpConstant`. This is because
    /// SPIR-V does not allow non-constant (technically non-global, but constants are the only
    /// non-pointer globals) initializers.
    pub fn create_local_mut(
        &self,
        value_ty: impl Into<TypeHandle>,
        init: Option<AttrObj>,
    ) -> Value {
        let value_ty = value_ty.into();
        let ctx = self.ctx_mut();
        let align = value_ty.align(ctx);
        let op = DeclareVariableOp::new(ctx, value_ty, AddressSpace::Local, align, init);
        let out = op.get_result(ctx);
        self.inserter().append_op(ctx, &op);
        out
    }

    /// Create a shared variable of the given item type.
    pub fn create_shared(
        &self,
        value_ty: impl Into<TypeHandle>,
        alignment: Option<usize>,
    ) -> Value {
        let value_ty = value_ty.into();
        let ctx = self.ctx_mut();
        let align = alignment.unwrap_or_else(|| value_ty.align(ctx));
        let op = DeclareVariableOp::new(ctx, value_ty, AddressSpace::Shared, align, None);
        let out = op.get_result(ctx);
        self.inserter().append_op(ctx, &op);
        out
    }

    pub fn func_ident(&self, label: Option<&str>) -> Identifier {
        let unique_id = self.ident_id();
        match label {
            Some(label) => ident(format!("{label}_{unique_id}")),
            None => ident(format!("func_{unique_id}")),
        }
    }

    /// Create a new function.
    pub fn register_func(&self, func: FuncOp) {
        let ctx = self.ctx_mut();
        let state = self.state_mut();
        func.set_visibility(ctx, SymbolVisibility::Private);
        state.module_inserter.append_op(ctx, &func);
    }

    /// Register an [`Instruction`] into the scope.
    pub fn register(&self, op: &dyn Op) {
        let ctx = self.ctx();
        self.inserter().append_op(ctx, op);
    }

    /// Register an [`Instruction`] into the scope and return its result.
    pub fn register_with_result(&self, op: &dyn OneResultInterface) -> Value {
        self.register(op);
        op.get_result(self.ctx())
    }

    /// Terminate block with a `cube.yield` if not already terminated
    pub fn terminate_yield(&self) {
        let block = self.inserter().get_insertion_block(self.ctx());
        let block = block.expect("Should have insertion block");
        if block.deref(self.ctx()).get_terminator(self.ctx()).is_none() {
            self.register(&YieldOp::new(self.ctx_mut()));
        }
    }

    pub fn update_flags_after_branch(&self, children: &[Scope]) {
        self.set_may_continue(children);
        self.set_may_break(children);
        self.set_may_return(children);
        self.set_may_terminate(children);
    }

    pub fn update_flags_after_loop(&self, children: &[Scope]) {
        self.set_may_return(children);
        self.set_may_terminate(children);
    }

    pub fn update_flags_after_early_return(&self, children: &[Scope]) {
        self.set_may_terminate(children);
    }

    pub fn update_flags_before_unrolled_iteration(&self, body: &Scope) {
        body.update_flags_before_unrolled_iteration_at(self.unrolled_insertion_point(body));
    }

    fn update_flags_before_unrolled_iteration_at(&self, insertion_point: OpInsertionPoint) {
        self.inserter().set_insertion_point(insertion_point);

        // Reset continue to initial state
        if self.expand_state().may_continue {
            let flag = self.expand_state().inv_continue_flag.unwrap();
            self.register(&StoreOp::new(self.ctx_mut(), flag, self.const_bool(true)));
            self.expand_state_mut().may_continue = false;
        }

        self.set_may_break(core::slice::from_ref(self));
        self.set_may_return(core::slice::from_ref(self));
        self.set_may_terminate(core::slice::from_ref(self));
    }

    pub fn finalize_unrolled_loop(&self, body: &Scope) {
        self.inserter()
            .set_insertion_point(self.unrolled_insertion_point(body));

        self.set_may_return(core::slice::from_ref(body));
        self.set_may_terminate(core::slice::from_ref(body));
    }

    fn unrolled_insertion_point(&self, body: &Scope) -> OpInsertionPoint {
        let current_block = self.inserter().get_insertion_block(self.ctx()).unwrap();
        op_insertion_point_in_block(
            self.ctx(),
            body.inserter().get_insertion_point(),
            current_block,
        )
    }

    pub fn set_may_terminate(&self, children: &[Scope]) {
        let child_may_terminate = children
            .iter()
            .any(|scope| scope.expand_state().may_terminate);
        if child_may_terminate {
            self.expand_state_mut().may_terminate = true;
            let flag = self.expand_state().inv_terminate_flag;
            self.predicate_on_flag(flag.expect("Can't terminate in rewrite context"));
        }
    }

    pub fn set_may_return(&self, children: &[Scope]) {
        let child_may_return = children.iter().any(|scope| scope.expand_state().may_return);
        if child_may_return {
            self.expand_state_mut().may_return = true;
            let flag = self.expand_state().inv_return_flag;
            self.predicate_on_flag(flag.expect("Can't return in rewrite context"));
        }
    }

    pub fn set_may_break(&self, children: &[Scope]) {
        let child_may_break = children.iter().any(|scope| scope.expand_state().may_break);
        if child_may_break {
            self.expand_state_mut().may_break = true;
            let flag = self.expand_state().inv_break_flag;
            self.predicate_on_flag(flag.expect("Should have break flag"));
        }
    }

    pub fn set_may_continue(&self, children: &[Scope]) {
        let child_may_continue = children
            .iter()
            .any(|scope| scope.expand_state().may_continue);
        if child_may_continue {
            self.expand_state_mut().may_continue = true;
            let flag = self.expand_state().inv_continue_flag;
            self.predicate_on_flag(flag.expect("Should have continue flag"));
        }
    }

    fn predicate_on_flag(&self, flag: Value) {
        let ctx = self.ctx_mut();
        let cond = read_value(self, flag);
        let predication = IfOp::new(ctx, cond);
        let then_block = predication.then_block(ctx);
        let else_block = predication.else_block(ctx);
        let yield_ = YieldOp::new(ctx).get_operation();
        yield_.insert_at_back(then_block, ctx);
        let yield_ = YieldOp::new(ctx).get_operation();
        yield_.insert_at_back(else_block, ctx);
        self.register(&predication);
        // The insertion cursor can sit inside a block that already ends with a
        // terminator — e.g. the then-block of an earlier predication, whose
        // trailing yield was placed when that predication was built. Appending
        // another yield there would leave the block with two terminators, which
        // fails module verification. Only terminate the block if it isn't yet.
        self.terminate_yield();
        self.inserter()
            .set_insertion_point_to_block_start(then_block);
    }

    /// Add a value to the global arena so we can create a kernel-wide reference to it.
    /// The reference is the same as the type for simplicity, but is only valid for the duration of
    /// the root scope. Ensure the reference lifetime is shortened to the lifetime of the underlying
    /// variable being referenced.
    pub fn create_kernel_ref<'a, T>(&self, value: T) -> &'a mut T
    where
        T: 'a,
    {
        let state = self.state_mut();
        let reference = state.reference_arena.alloc(value);
        unsafe { core::mem::transmute(reference) }
    }

    /// Resolve the element type of the given generic type.
    pub fn resolve_type<T: 'static>(&self) -> Option<ElemType> {
        let state = self.state();
        let result = state.typemap.get(&TypeId::of::<T>());

        result.cloned()
    }

    /// Resolve the comptime size of the given generic size.
    pub fn resolve_size<T: 'static>(&self) -> Option<usize> {
        let state = self.state();
        let result = state.sizemap.get(&TypeId::of::<T>());

        result.cloned()
    }

    /// Register the element type for the given generic type.
    pub fn register_type<T: 'static>(&self, elem: ElemType) {
        self.state_mut().register_type::<T>(elem);
    }

    /// `T` names `elem` until the returned definition is dropped, which gives `T` back the
    /// element it named before, if any: what a `define!` in a function re-entered on another
    /// element leaves its outer call.
    pub fn define_type<T: 'static>(&self, elem: ElemType) -> TypeDefinition<'_> {
        let before = self.resolve_type::<T>();
        self.register_type::<T>(elem);
        TypeDefinition {
            scope: self,
            ty: TypeId::of::<T>(),
            before,
        }
    }

    /// `T` names `size` until the returned definition is dropped, which gives `T` back the size
    /// it named before, if any: what a `size!` in a function re-entered on another size leaves
    /// its outer call.
    pub fn define_size<T: 'static>(&self, size: usize) -> SizeDefinition<'_> {
        let before = self.resolve_size::<T>();
        self.register_size::<T>(size);
        SizeDefinition {
            scope: self,
            size: TypeId::of::<T>(),
            before,
        }
    }

    /// Register the comptime size for the given generic size.
    pub fn register_size<T: 'static>(&self, size: usize) {
        let state = self.state_mut();

        state.sizemap.insert(TypeId::of::<T>(), size);
    }

    /// Register the type and size of a scalarizable type
    pub fn register_value_type<T: 'static, N: 'static>(&self, value: impl Typed) {
        let ty = value.get_type(self.ctx());
        let scalar_ty = ty.scalar_ty(self.ctx());
        let vector_size = ty.vector_size(self.ctx());
        let storage_ty = {
            let ctx = self.ctx();
            let scalar_ty = scalar_ty.deref(ctx);
            let scalar = type_cast::<dyn ScalarType>(&*scalar_ty).unwrap();
            scalar.elem_type(ctx)
        };
        self.register_type::<T>(storage_ty);
        self.register_size::<N>(vector_size);
    }

    /// Create an empty child scope for branching ops that don't need new flags, i.e. if/else.
    pub fn branch_child(&self, inserter: impl Inserter + 'static) -> Self {
        Self {
            ctx: self.ctx.clone(),
            inserter: InserterHandle::owned(inserter),
            expand_state: RefCell::new(ExpandState {
                may_continue: false,
                may_break: false,
                may_return: false,
                may_terminate: false,
                inv_terminate_flag: self.expand_state().inv_terminate_flag,
                inv_return_flag: self.expand_state().inv_return_flag,
                inv_break_flag: self.expand_state().inv_break_flag,
                inv_continue_flag: self.expand_state().inv_continue_flag,
                return_value: self.expand_state().return_value,
            }),
        }
    }

    /// Create a child scope with a new break condition.
    pub fn loop_child(&self, mut inserter: impl Inserter + 'static) -> Self {
        let break_flag = init_bool_flag(self.ctx_mut(), self.inserter(), "inv_break_flag");
        let continue_flag = init_bool_flag(self.ctx_mut(), &mut inserter, "inv_continue_flag");
        Self {
            ctx: self.ctx.clone(),
            inserter: InserterHandle::owned(inserter),
            expand_state: RefCell::new(ExpandState {
                may_terminate: false,
                may_return: false,
                may_break: false,
                may_continue: false,
                inv_terminate_flag: self.expand_state().inv_terminate_flag,
                inv_return_flag: self.expand_state().inv_return_flag,
                inv_break_flag: Some(break_flag),
                inv_continue_flag: Some(continue_flag),
                return_value: self.expand_state().return_value,
            }),
        }
    }

    /// Create a child scope with a new break condition.
    pub fn unrolled_loop_child(&self) -> Self {
        let break_flag = init_bool_flag(self.ctx_mut(), self.inserter(), "inv_break_flag");
        let continue_flag = init_bool_flag(self.ctx_mut(), self.inserter(), "inv_continue_flag");
        let inserter = OpInserter::new(self.inserter().get_insertion_point());
        Self {
            ctx: self.ctx.clone(),
            inserter: InserterHandle::owned(inserter),
            expand_state: RefCell::new(ExpandState {
                may_terminate: false,
                may_return: false,
                may_break: false,
                may_continue: false,
                inv_terminate_flag: self.expand_state().inv_terminate_flag,
                inv_return_flag: self.expand_state().inv_return_flag,
                inv_break_flag: Some(break_flag),
                inv_continue_flag: Some(continue_flag),
                return_value: self.expand_state().return_value,
            }),
        }
    }

    /// Create a child that's at the root of a new inlined function and sets up a new early return flag.
    pub fn inlined_func_child(
        &self,
        mut inserter: impl Inserter + 'static,
        return_value: Option<Value>,
    ) -> Self {
        let return_flag = init_bool_flag(self.ctx_mut(), &mut inserter, "inv_return_flag");
        Self {
            ctx: self.ctx.clone(),
            inserter: InserterHandle::owned(inserter),
            expand_state: RefCell::new(ExpandState {
                may_continue: false,
                may_break: false,
                may_return: false,
                may_terminate: false,
                inv_terminate_flag: self.expand_state().inv_terminate_flag,
                inv_return_flag: Some(return_flag),
                inv_break_flag: None,
                inv_continue_flag: None,
                return_value: return_value,
            }),
        }
    }

    /// Create a child that's at the root of a new non-inlined function.
    /// `terminate!()` won't work for now so if non-inlined functions are ever made user accessible
    /// we need to deal with that.
    pub fn func_child(&self, mut inserter: impl Inserter + 'static) -> Self {
        let terminate_flag = init_bool_flag(self.ctx_mut(), &mut inserter, "inv_terminate_flag");
        Self {
            ctx: self.ctx.clone(),
            inserter: InserterHandle::owned(inserter),
            expand_state: RefCell::new(ExpandState {
                may_continue: false,
                may_break: false,
                may_return: false,
                may_terminate: false,
                inv_terminate_flag: Some(terminate_flag),
                inv_return_flag: None,
                inv_break_flag: None,
                inv_continue_flag: None,
                return_value: None,
            }),
        }
    }

    // Adds a validation error.
    pub fn push_error(&self, msg: impl Into<String>) {
        self.state_mut().errors.push(msg.into());
    }

    /// Returns all validation errors.
    pub fn pop_errors(&self) -> Vec<String> {
        core::mem::take(&mut self.state_mut().errors)
    }

    /// Adds a non-fatal validation warning. Logged when the kernel is built.
    pub fn push_warning(&self, msg: impl Into<String>) {
        self.state_mut().warnings.push(msg.into());
    }

    /// Returns all validation warnings.
    pub fn pop_warnings(&self) -> Vec<String> {
        core::mem::take(&mut self.state_mut().warnings)
    }

    /// Obtain the index-th buffer
    pub fn global(
        &self,
        buffer_pos: usize,
        ext_meta_pos: Option<usize>,
        value_ty: TypeHandle,
    ) -> Value {
        let entry_func = self.state().entry_func;
        let ctx = self.ctx_mut();

        let ty_arr = RuntimeArrayType::get(ctx, value_ty);
        let ty = PointerType::get(ctx, ty_arr.into(), AddressSpace::Global(buffer_pos));

        let id = entry_func.push_argument(ctx, ty.to_handle());
        entry_func.set_arg_attr(
            ctx,
            id,
            &ATTR_BUFFER_BINDING,
            Box::new(BufferBindingAttr::new(buffer_pos, ext_meta_pos)),
        );

        entry_func.get_entry_block(ctx).deref(ctx).get_argument(id)
    }

    /// Obtain the index-th buffer, in memory shared with peers (see [`ATTR_PEER_BINDING`]).
    pub fn peer_global(&self, buffer_pos: usize, value_ty: TypeHandle) -> Value {
        let value = self.global(buffer_pos, None, value_ty);
        let entry_func = self.state().entry_func;
        let ctx = self.ctx_mut();
        let id = entry_func
            .get_entry_block(ctx)
            .deref(ctx)
            .get_num_arguments()
            - 1;
        entry_func.set_arg_attr_unit(ctx, id, &ATTR_PEER_BINDING);
        value
    }

    /// Whether `ptr` points into a buffer in memory shared with peers.
    pub fn is_peer_pointer(&self, ptr: Value) -> bool {
        let ctx = self.ctx();
        let binding = match ptr.get_type(ctx).deref(ctx).downcast_ref::<PointerType>() {
            Some(PointerType {
                address_space: AddressSpace::Global(binding),
                ..
            }) => *binding,
            _ => return false,
        };
        peer_buffer_positions(ctx, self.state().entry_func).contains(&binding)
    }

    /// Obtain the index-th tensor map
    pub fn tensor_map(&self, buffer_pos: usize, ext_meta_pos: usize) -> Value {
        let entry_func = self.state().entry_func;
        let ctx = self.ctx();
        let ty = TensorMapType::get(ctx);

        let id = entry_func.push_argument(ctx, ty.to_handle());
        entry_func.set_arg_attr_unit(ctx, id, &ATTR_TENSOR_MAP_BINDING);
        entry_func.set_arg_attr(
            ctx,
            id,
            &ATTR_BUFFER_BINDING,
            Box::new(BufferBindingAttr::new(buffer_pos, Some(ext_meta_pos))),
        );
        entry_func.get_entry_block(ctx).deref(ctx).get_argument(id)
    }

    pub fn kernel_arg(&self, idx: usize) -> Value {
        let entry_block = self.state().entry_func.get_entry_block(self.ctx());
        entry_block.deref(self.ctx()).get_argument(idx)
    }

    pub fn extract_field(&self, aggregate: Value, field: usize) -> Value {
        let ctx = self.ctx_mut();
        let op = CompositeExtractOp::new(ctx, aggregate, field);
        self.register_with_result(&op)
    }

    pub fn const_usize(&self, value: usize) -> Value {
        let op = ConstantOp::new(self.ctx_mut(), Box::new(IndexAttr::new(value)));
        self.register_with_result(&op)
    }

    pub fn const_bool(&self, value: bool) -> Value {
        let op = ConstantOp::new(self.ctx_mut(), Box::new(BoolAttr::new(value)));
        self.register_with_result(&op)
    }

    pub fn into_context(self) -> Option<Context> {
        let entry = self.state().entry_func.get_entry_block(self.ctx());
        let term = entry.deref(self.ctx()).get_terminator(self.ctx());
        let is_yield = term.is_some_and(|term| term.is_op::<YieldOp>(self.ctx()));
        if let Some(term) = term
            && is_yield
        {
            self.inserter().set_insertion_point_to_block_end(entry);
            self.register(&ReturnOp::new(self.ctx_mut()));
            Operation::erase(term, self.ctx_mut());
        } else if term.is_none() {
            self.inserter().set_insertion_point_to_block_end(entry);
            self.register(&ReturnOp::new(self.ctx_mut()));
        }
        match self.ctx {
            CtxHandle::Rc(ctx) => Some(Rc::into_inner(ctx)?.into_inner()),
            CtxHandle::Ref(_) => None,
        }
    }
}

impl Display for Scope {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let ctx = self.ctx();
        let state = self.state();
        write!(f, "{}", state.module.disp(ctx))
    }
}

fn init_bool_flag(ctx: &mut Context, inserter: &mut dyn Inserter, name: &str) -> Value {
    let bool = TypeAttr::new(BoolType::get(ctx).to_handle());
    let r#true = BoolAttr::new(true).into();
    let flag = DeclareVariableOp::new(ctx, bool, AddressSpace::Local, 1, Some(r#true));
    inserter.append_op(ctx, &flag);
    set_operation_result_name(ctx, flag.get_operation(), 0, Some(ident(name)));
    flag.get_result(ctx)
}

/// A type's definition for a block, made by [`Scope::define_type`].
pub struct TypeDefinition<'a> {
    scope: &'a Scope,
    ty: TypeId,
    before: Option<ElemType>,
}

impl Drop for TypeDefinition<'_> {
    fn drop(&mut self) {
        if let Some(before) = self.before {
            self.scope.state_mut().typemap.insert(self.ty, before);
        }
    }
}

/// A size's definition for a block, made by [`Scope::define_size`].
pub struct SizeDefinition<'a> {
    scope: &'a Scope,
    size: TypeId,
    before: Option<usize>,
}

impl Drop for SizeDefinition<'_> {
    fn drop(&mut self) {
        if let Some(before) = self.before {
            self.scope.state_mut().sizemap.insert(self.size, before);
        }
    }
}