cubecl-ir 0.11.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,
        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::{
    CanMaterialize, HasSideEffects, Pure,
    attributes::{BoolAttr, IndexAttr},
    dialect::{general::SymbolUserOpVerifyErr, synchronization::SyncScope},
    interfaces::{
        TypedExt,
        side_effects::{MemoryEffect, MemoryEffectsOp},
        synchronizes,
    },
    prelude::*,
    types::{
        ArrayType, MatrixShape, PointerType, 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);

/// 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,
}
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,
}
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);

/// 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()
}