cubecl-ir 0.12.0-pre.1

use core::fmt;

use alloc::{format, string::ToString, vec::Vec};

use cubecl_macros_internal::cube_op;
use derive_more::{Deref, From};
use derive_new::new;
use itertools::Itertools;
use pliron::{
    builtin::{
        attributes::{IdentifierAttr, TypeAttr},
        ops::FuncOp,
        types::{FunctionType, IntegerType, Signedness},
    },
    combine::{Parser, parser::char::char},
    derive::pliron_attr,
    identifier::Identifier,
    input_err,
    irfmt::parsers::{delimited_list_parser, spaced, ssa_opd_parse, ssa_opd_parser},
    location::Location,
    op::OpObj,
    parsable::{self, IntoParseResult, Parsable, ParseResult},
    printable::{self, Printable},
    symbol_table::SymbolTableCollection,
    r#type::TypedHandle,
    verify_err,
};

use crate::{
    AddressSpace, CanMaterialize, HasSideEffects, Pure,
    attributes::{BoolAttr, IndexAttr},
    dialect::{general::SymbolUserOpVerifyErr, synchronization::SyncScope},
    interfaces::{
        TypedExt,
        side_effects::{MemoryEffect, MemoryEffectsOp},
        synchronizes,
    },
    prelude::*,
    types::{
        ArrayType, LdMatrixForm, MatrixShape, PointerType, StMatrixForm, VectorType,
        matrix::{MatrixLayout, MatrixType},
    },
};

#[pliron_attr(name = "matrix.layout", format = "$0", verifier = "succ")]
#[derive(new, From, PartialEq, Eq, Clone, Debug, Hash, PartialOrd, Ord, Deref)]
pub struct MatrixLayoutAttr(pub MatrixLayout);

#[pliron_attr(name = "matrix.type", format = "$0", verifier = "succ")]
#[derive(new, From, Debug, Clone, PartialEq, Eq, Hash, Deref)]
pub struct MatrixTypeAttr(pub TypedHandle<MatrixType>);

#[pliron_attr(name = "matrix.type", format = "$0", verifier = "succ")]
#[derive(new, From, Debug, Clone, PartialEq, Eq, Hash, Deref)]
pub struct MatrixShapeAttr(pub MatrixShape);

#[pliron_attr(name = "matrix.ldmatrix_form", format = "$0", verifier = "succ")]
#[derive(new, From, Debug, Clone, PartialEq, Eq, Hash, Deref)]
pub struct LdMatrixFormAttr(pub LdMatrixForm);

#[pliron_attr(name = "matrix.stmatrix_form", format = "$0", verifier = "succ")]
#[derive(new, From, Debug, Clone, PartialEq, Eq, Hash, Deref)]
pub struct StMatrixFormAttr(pub StMatrixForm);

/// Fill a matrix with a scalar value.
/// Note: Unlike most matrix ops, this does not have implicit synchronization because there's no
/// coordination between threads.
#[cube_op(name = "matrix.fill")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
pub struct FillOp {
    #[operand(ptr_write)]
    pub matrix: Value,
    pub value: Value,
}

#[cube_op(name = "matrix.load")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
pub struct LoadOp {
    #[operand(ptr_write)]
    pub matrix: Value,
    #[operand(ptr_read)]
    pub source: Value,
    pub stride: Value,
    pub layout: MatrixLayoutAttr,
}
synchronizes!(LoadOp, SyncScope::Plane);

#[cube_op(name = "matrix.store")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
pub struct StoreOp {
    #[operand(ptr_read)]
    pub matrix: Value,
    #[operand(ptr_write)]
    pub destination: Value,
    pub stride: Value,
    pub layout: MatrixLayoutAttr,
}
synchronizes!(StoreOp, SyncScope::Plane);

#[cube_op(name = "matrix.multiply_accumulate")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
pub struct MultiplyAccumulateOp {
    #[operand(ptr_read)]
    pub mat_a: Value,
    #[operand(ptr_read)]
    pub mat_b: Value,
    #[operand(ptr_read)]
    pub mat_c: Value,
    #[operand(ptr_write)]
    pub mat_d: Value,
}
synchronizes!(MultiplyAccumulateOp, SyncScope::Plane);

/// Cast a matrix from one type to another.
/// Note: Unlike most matrix ops, this does not have implicit synchronization because there's no
/// coordination between threads.
#[cube_op(name = "matrix.cast")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
pub struct CastOp {
    #[operand(ptr_read)]
    pub input: Value,
    #[operand(ptr_write)]
    pub output: Value,
}

#[cube_op(name = "matrix.row_index")]
#[result_ty(fixed = IntegerType::get(ctx, 32, Signedness::Unsigned).into())]
#[op_traits(CanMaterialize, Pure)]
pub struct RowIndexOp {
    pub lane_id: Value,
    pub i: Value,
    pub matrix_ty: MatrixTypeAttr,
}

#[cube_op(name = "matrix.col_index")]
#[result_ty(fixed = IntegerType::get(ctx, 32, Signedness::Unsigned).into())]
#[op_traits(CanMaterialize, Pure)]
pub struct ColIndexOp {
    pub lane_id: Value,
    pub i: Value,
    pub matrix_ty: MatrixTypeAttr,
}

#[cube_op(name = "matrix.ldmatrix")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
pub struct LdMatrixOp {
    #[operand(ptr_read)]
    pub ptr: Value,
    #[operand(ptr_write)]
    pub out_arr: Value,
    pub factor: IndexAttr,
    pub transpose: BoolAttr,
    pub form: LdMatrixFormAttr,
}
synchronizes!(LdMatrixOp, SyncScope::Plane);

#[cube_op(name = "matrix.stmatrix")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
#[op_interfaces(OperandNOfType<0, ArrayType>, OperandNOfType<1, PointerType>)]
pub struct StMatrixOp {
    pub registers: Value,
    #[operand(ptr_write)]
    pub destination: Value,
    pub factor: IndexAttr,
    pub transpose: BoolAttr,
    pub form: StMatrixFormAttr,
}
synchronizes!(StMatrixOp, SyncScope::Plane);

#[cube_op(name = "matrix.mma_manual")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
#[op_interfaces(
    OperandNOfType<0, ArrayType>, OperandNOfType<1, ArrayType>, OperandNOfType<2, ArrayType>,
    OperandNOfType<3, PointerType>,
)]
pub struct MmaManualOp {
    pub registers_a: Value,
    pub registers_b: Value,
    pub registers_c: Value,
    #[operand(ptr_write)]
    pub registers_d: Value,
    pub shape: MatrixShapeAttr,
}
synchronizes!(MmaManualOp, SyncScope::Plane);

#[cube_op(name = "matrix.mma_manual_scaled")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
#[op_interfaces(
    OperandNOfType<0, ArrayType>, OperandNOfType<1, ArrayType>, OperandNOfType<2, ArrayType>,
    OperandNOfType<3, PointerType>, OperandNOfType<4, VectorType>, OperandNOfType<5, VectorType>,
)]
pub struct MmaManualScaledOp {
    pub registers_a: Value,
    pub registers_b: Value,
    pub registers_c: Value,
    #[operand(ptr_write)]
    pub registers_d: Value,
    pub scales_a: Value,
    pub scales_b: Value,
    pub scales_factor: IndexAttr,
    pub shape: MatrixShapeAttr,
}
synchronizes!(MmaManualScaledOp, SyncScope::Plane);

/// The swizzle a warpgroup MMA matrix descriptor declares for the tile it points at, named for
/// the width of the rows it permutes. The pattern repeats every 8 rows.
#[pliron_attr(name = "matrix.wgmma_swizzle", format, verifier = "succ")]
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug, PartialOrd, Ord)]
pub enum WgmmaSwizzle {
    /// No swizzle: the tile is made of 8x16-byte core matrices.
    None,
    B32,
    B64,
    B128,
}

impl WgmmaSwizzle {
    /// The bytes of a line the swizzle permutes the 16-byte chunks of. Without swizzle, a line
    /// is one chunk.
    pub fn width(&self) -> usize {
        match self {
            WgmmaSwizzle::None => 16,
            WgmmaSwizzle::B32 => 32,
            WgmmaSwizzle::B64 => 64,
            WgmmaSwizzle::B128 => 128,
        }
    }

    /// The alignment, in bytes, a tile with this swizzle must start at. A swizzled tile starts a
    /// repeat of its pattern, 8 rows of the swizzle's width, because the descriptor leaves the
    /// base offset that would shift the pattern at zero.
    pub fn alignment(&self) -> usize {
        match self {
            WgmmaSwizzle::None => 16,
            WgmmaSwizzle::B32 => 256,
            WgmmaSwizzle::B64 => 512,
            WgmmaSwizzle::B128 => 1024,
        }
    }
}

/// A warpgroup operation is `.sync.aligned` over four planes. No scope names a warpgroup, so it
/// declares the conservative end of each bound: at least a plane, so no pass takes it for a
/// cube barrier, and at most a cube, so no pass sinks it into control flow that only a plane
/// agrees on.
macro_rules! warpgroup_synchronizes {
    ($ty: ty) => {
        #[op_interface_impl]
        impl crate::interfaces::Synchronizes for $ty {
            fn minimum_scope(&self, _ctx: &Context) -> SyncScope {
                SyncScope::Plane
            }
            fn maximum_scope(&self, _ctx: &Context) -> SyncScope {
                SyncScope::Cube
            }
        }
    };
}

/// Units in a warpgroup: the four planes that issue one warpgroup MMA together.
pub const WARPGROUP_UNITS: usize = 128;
/// Rows every warpgroup MMA computes.
pub const WARPGROUP_M: usize = 64;
/// Bytes of K every warpgroup MMA reads.
pub const WARPGROUP_K_BYTES: usize = 32;

/// The elements of K of `elem_size` bytes every warpgroup MMA reads.
pub const fn warpgroup_k(elem_size: usize) -> usize {
    WARPGROUP_K_BYTES / elem_size
}

/// The elements each unit of a warpgroup holds of a `64 x cols` accumulator or `A` fragment.
pub const fn warpgroup_elems_per_unit(cols: usize) -> usize {
    WARPGROUP_M * cols / WARPGROUP_UNITS
}

/// Which dimension of a warpgroup MMA operand tile is contiguous in shared memory: K, as in a
/// row-major `A` or a column-major `B`, or M for `A` and N for `B`.
#[pliron_attr(name = "matrix.wgmma_major", format, verifier = "succ")]
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug, PartialOrd, Ord)]
pub enum WgmmaMajor {
    K,
    MN,
}

/// Builds the 64-bit matrix descriptor a warpgroup MMA reads a shared memory tile through.
/// `ptr` must point into shared memory, and the two offsets are in bytes.
#[cube_op(name = "matrix.wgmma_descriptor")]
#[result_ty(fixed = IntegerType::get(ctx, 64, Signedness::Unsigned).into())]
#[op_traits(CanMaterialize, Pure)]
pub struct WgmmaDescriptorOp {
    pub ptr: Value,
    pub leading_byte_offset: IndexAttr,
    pub stride_byte_offset: IndexAttr,
    pub swizzle: WgmmaSwizzle,
}

/// Orders register accesses before the warpgroup MMAs that follow it. Required before the first
/// MMA, and whenever registers an MMA reads were written by something else in between.
#[cube_op(name = "matrix.wgmma_fence")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
pub struct WgmmaFenceOp {}
warpgroup_synchronizes!(WgmmaFenceOp);

/// Commits every warpgroup MMA issued and not yet committed into one group.
#[cube_op(name = "matrix.wgmma_commit_group")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
pub struct WgmmaCommitGroupOp {}
warpgroup_synchronizes!(WgmmaCommitGroupOp);

/// Waits until at most `max_pending` committed groups of warpgroup MMAs are still running.
#[cube_op(name = "matrix.wgmma_wait_group")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
pub struct WgmmaWaitGroupOp {
    pub max_pending: IndexAttr,
}
warpgroup_synchronizes!(WgmmaWaitGroupOp);

/// Pins the registers of an array in place: no access to them moves across it. An MMA in flight
/// owns its accumulator until the wait that retires it, so the compiler must neither read the
/// registers early nor copy them elsewhere in between.
#[cube_op(name = "matrix.wgmma_fence_operand")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
#[op_interfaces(OperandNOfType<0, PointerType>)]
pub struct WgmmaFenceOperandOp {
    #[operand(ptr_read, ptr_write)]
    pub registers: Value,
}

/// Warpgroup MMA, `D = A * B + D`, issued asynchronously: the accumulator may only be read after
/// a [`WgmmaWaitGroupOp`] retires the group it was committed in.
///
/// `a` is either a `u64` matrix descriptor of a shared memory tile or an array of registers, and
/// `b` is always a descriptor, so the element types are carried as attributes. `a_major` is only
/// read for a descriptor: `A` in registers has no layout in memory.
#[cube_op(name = "matrix.wgmma")]
#[result_ty(none)]
#[op_traits(CanMaterialize, HasSideEffects)]
#[op_interfaces(OperandNOfType<2, PointerType>)]
pub struct WgmmaOp {
    pub a: Value,
    pub b: Value,
    pub accumulator: Value,
    pub a_ty: TypeAttr,
    pub b_ty: TypeAttr,
    pub shape: MatrixShapeAttr,
    pub a_major: WgmmaMajor,
    pub b_major: WgmmaMajor,
}
warpgroup_synchronizes!(WgmmaOp);

#[op_interface_impl]
impl MemoryEffectsOp for WgmmaOp {
    fn memory_effects(&self, ctx: &Context) -> Vec<MemoryEffect> {
        // The descriptors carry addresses the analyses cannot follow, so the operation reads
        // any of shared memory.
        let mut effects = vec![
            MemoryEffect::Read(self.accumulator(ctx)),
            MemoryEffect::Write(self.accumulator(ctx)),
            MemoryEffect::ReadAllInSpace(AddressSpace::Shared),
        ];
        // `A` in registers is read out of its array.
        let a = self.a(ctx);
        if a.is_ptr(ctx) {
            effects.push(MemoryEffect::Read(a));
        }
        effects
    }
}

/// Executes a closure for each element in the matrix.
/// Note: Unlike most matrix ops, this does not have implicit synchronization because there's no
/// coordination between threads.
#[pliron_op(
    name = "matrix.elementwise",
    attributes = (matrix_elementwise_closure: IdentifierAttr),
    verifier = "succ"
)]
#[op_traits(CanMaterialize, HasSideEffects)]
pub struct ElementwiseOp;

impl ElementwiseOp {
    pub fn new(
        ctx: &mut Context,
        matrix_in: Value,
        matrix_out: Value,
        closure: Identifier,
        captures: Vec<Value>,
    ) -> Self {
        let mut opds = vec![matrix_in, matrix_out];
        opds.extend(captures);
        let op = Self {
            op: Operation::new(ctx, Self::get_concrete_op_info(), vec![], opds, vec![], 0),
        };
        op.set_attr_matrix_elementwise_closure(ctx, IdentifierAttr::new(closure));
        op
    }

    pub fn matrix_in(&self, ctx: &Context) -> Value {
        self.get_operation().operand(ctx, 0)
    }

    pub fn matrix_out(&self, ctx: &Context) -> Value {
        self.get_operation().operand(ctx, 1)
    }

    pub fn closure(&self, ctx: &Context) -> Identifier {
        let attr = self.get_attr_matrix_elementwise_closure(ctx).unwrap();
        attr.clone().into()
    }

    pub fn closure_captures(&self, ctx: &Context) -> Vec<Value> {
        self.get_operation().deref(ctx).operands().skip(2).collect()
    }
}

impl ElementwiseOp {
    /// Callee type, including implicit args
    fn callee_type(&self, ctx: &Context) -> TypeHandle {
        let u32 = IntegerType::get(ctx, 32, Signedness::Unsigned).to_handle();
        let elem_ty = self
            .matrix_in(ctx)
            .unwrap_ptr(ctx)
            .element_ty(ctx)
            .scalar_ty(ctx);
        let mut args = vec![u32, u32, elem_ty];
        args.extend(self.closure_captures(ctx).iter().map(|it| it.get_type(ctx)));
        FunctionType::get(ctx, args, vec![elem_ty]).to_handle()
    }
}

impl Printable for ElementwiseOp {
    fn fmt(
        &self,
        ctx: &Context,
        _state: &printable::State,
        f: &mut fmt::Formatter<'_>,
    ) -> fmt::Result {
        let op_id = self.get_opid().disp(ctx).to_string();
        let mat_in = self.matrix_in(ctx).disp(ctx).to_string();
        let mat_out = self.matrix_out(ctx).disp(ctx).to_string();
        write!(f, "{op_id} ({mat_in}, {mat_out}) ")?;
        print_closure(ctx, &self.closure(ctx), &self.closure_captures(ctx), f)
    }
}
impl Parsable for ElementwiseOp {
    type Arg = Vec<(Identifier, Location)>;
    type Parsed = OpObj;
    fn parse<'a>(
        input: &mut parsable::StateStream<'a>,
        arg: Self::Arg,
    ) -> ParseResult<'a, Self::Parsed> {
        if !arg.is_empty() {
            return input_err!(input.loc(), "Expected no results").into_parse_result();
        }

        spaced(char('(')).parse_stream(input).into_result()?;
        let mat_in = ssa_opd_parse(input, ())?.0;
        spaced(char(',')).parse_stream(input).into_result()?;
        let mat_out = ssa_opd_parse(input, ())?.0;
        spaced(char(')')).parse_stream(input).into_result()?;

        let (closure, captures) = parse_closure(input)?.0;
        let ctx = &mut input.state.ctx;

        let op = ElementwiseOp::new(ctx, mat_in, mat_out, closure, captures);
        Ok(OpObj::new(op)).into_parse_result()
    }
}

#[op_interface_impl]
impl MemoryEffectsOp for ElementwiseOp {
    fn memory_effects(&self, ctx: &Context) -> Vec<MemoryEffect> {
        vec![
            MemoryEffect::Read(self.matrix_in(ctx)),
            MemoryEffect::Write(self.matrix_out(ctx)),
        ]
    }
}

#[op_interface_impl]
impl SymbolUserOpInterface for ElementwiseOp {
    fn verify_symbol_uses(
        &self,
        ctx: &Context,
        symbol_tables: &mut SymbolTableCollection,
    ) -> Result<()> {
        let callee_sym = self.closure(ctx);
        let Some(callee) =
            symbol_tables.lookup_symbol_in_nearest_table(ctx, self.get_operation(), &callee_sym)
        else {
            return verify_err!(
                self.loc(ctx),
                SymbolUserOpVerifyErr::SymbolNotFound(callee_sym.to_string())
            );
        };
        let Some(func_op) = (&*callee as &dyn Op).downcast_ref::<FuncOp>() else {
            return verify_err!(
                self.loc(ctx),
                SymbolUserOpVerifyErr::NotFunc(callee_sym.to_string())
            );
        };
        let func_op_ty = func_op.get_type(ctx);

        if func_op_ty != self.callee_type(ctx) {
            return verify_err!(
                self.loc(ctx),
                SymbolUserOpVerifyErr::FuncTypeErr(format!(
                    "expected {}, got {}",
                    func_op_ty.disp(ctx),
                    self.callee_type(ctx).disp(ctx)
                ))
            );
        }
        Ok(())
    }

    fn used_symbols(&self, ctx: &Context) -> Vec<Identifier> {
        vec![self.closure(ctx)]
    }
}

/// Reusable closure printer for maybe future ops, move this if it's used elsewhere
pub fn print_closure(
    ctx: &Context,
    closure: &Identifier,
    captures: &[Value],
    f: &mut fmt::Formatter<'_>,
) -> fmt::Result {
    let captures = captures
        .iter()
        .map(|it| it.disp(ctx).to_string())
        .join(", ");
    write!(f, "@{}({captures})", closure.disp(ctx))
}

/// Reusable closure parser for maybe future ops, move this if it's used elsewhere
pub fn parse_closure<'a>(
    input: &mut parsable::StateStream<'a>,
) -> ParseResult<'a, (Identifier, Vec<Value>)> {
    let mut parse_closure = char('@').with(Identifier::parser(()));
    let closure = parse_closure.parse_stream(input).into_result()?.0;
    let mut captures = delimited_list_parser('(', ')', ',', ssa_opd_parser());
    let captures = captures.parse_stream(input).into_result()?.0;
    Ok((closure, captures)).into_parse_result()
}