cubecl-core 0.11.0-pre.3

use cubecl_ir::ExpandValue;
use half::{bf16, f16};

use super::{CubePrimitive, Vector};
use crate::prelude::*;
use crate::{
    ir::{Scope, attributes::IndexAttr, dialect::plane::*},
    unexpanded,
};

pub trait PlaneNumeric {
    fn __expand_native_sum(scope: &Scope, value: ExpandValue) -> ExpandValue;
    fn __expand_native_inclusive_sum(scope: &Scope, value: ExpandValue) -> ExpandValue;
    fn __expand_native_exclusive_sum(scope: &Scope, value: ExpandValue) -> ExpandValue;

    fn __expand_native_prod(scope: &Scope, value: ExpandValue) -> ExpandValue;
    fn __expand_native_inclusive_prod(scope: &Scope, value: ExpandValue) -> ExpandValue;
    fn __expand_native_exclusive_prod(scope: &Scope, value: ExpandValue) -> ExpandValue;

    fn __expand_native_plane_min(scope: &Scope, value: ExpandValue) -> ExpandValue;
    fn __expand_native_plane_max(scope: &Scope, value: ExpandValue) -> ExpandValue;
}

macro_rules! plane_numeric {
    ($($ty: ty),*; $sum: ty, $inc_sum: ty, $exc_sum: ty, $prod: ty, $inc_prod: ty, $exc_prod: ty, $min: ty, $max: ty) => {
        $(impl PlaneNumeric for $ty {
            fn __expand_native_sum(scope: &Scope, value: ExpandValue) -> ExpandValue {
                unary_expand(scope, value, <$sum>::new)
            }
            fn __expand_native_inclusive_sum(scope: &Scope, value: ExpandValue) -> ExpandValue {
                unary_expand(scope, value, <$inc_sum>::new)
            }
            fn __expand_native_exclusive_sum(scope: &Scope, value: ExpandValue) -> ExpandValue {
                unary_expand(scope, value, <$exc_sum>::new)
            }

            fn __expand_native_prod(scope: &Scope, value: ExpandValue) -> ExpandValue {
                unary_expand(scope, value, <$prod>::new)
            }
            fn __expand_native_inclusive_prod(scope: &Scope, value: ExpandValue) -> ExpandValue {
                unary_expand(scope, value, <$inc_prod>::new)
            }
            fn __expand_native_exclusive_prod(scope: &Scope, value: ExpandValue) -> ExpandValue {
                unary_expand(scope, value, <$exc_prod>::new)
            }

            fn __expand_native_plane_min(scope: &Scope, value: ExpandValue) -> ExpandValue {
                unary_expand(scope, value, <$min>::new)
            }
            fn __expand_native_plane_max(scope: &Scope, value: ExpandValue) -> ExpandValue {
                unary_expand(scope, value, <$max>::new)
            }
        })*
    };
}

plane_numeric!(i8, i16, i32, i64, isize; ISumOp, InclusiveISumOp, ExclusiveISumOp, IProdOp, InclusiveIProdOp, ExclusiveIProdOp, SMinOp, SMaxOp);
plane_numeric!(u8, u16, u32, u64, usize; ISumOp, InclusiveISumOp, ExclusiveISumOp, IProdOp, InclusiveIProdOp, ExclusiveIProdOp, UMinOp, UMaxOp);
plane_numeric!(f16, bf16, f32, flex32, tf32, f64; FSumOp, InclusiveFSumOp, ExclusiveFSumOp, FProdOp, InclusiveFProdOp, ExclusiveFProdOp, FMinOp, FMaxOp);

/// Returns true if the cube unit has the lowest `plane_unit_id` among active unit in the plane
pub fn plane_elect() -> bool {
    unexpanded!()
}

/// Module containing the expand function for [`plane_elect()`].
pub mod plane_elect {
    use super::*;

    /// Expand method of [`plane_elect()`].
    pub fn expand(scope: &Scope) -> NativeExpand<bool> {
        let op = ElectOp::new(scope.ctx_mut());
        scope.register_with_result(&op).into()
    }
}

/// Broadcasts the value from the specified plane unit at the given index
/// to all active units within that plane. Requires a constant index. For non-constant indices,
/// use [`plane_shuffle()`].
#[allow(unused_variables)]
pub fn plane_broadcast<E: CubePrimitive>(value: E, index: u32) -> E {
    unexpanded!()
}

/// Module containing the expand function for [`plane_broadcast()`].
pub mod plane_broadcast {
    use super::*;

    /// Expand method of [`plane_broadcast()`].
    pub fn expand<E: CubePrimitive>(
        scope: &Scope,
        value: NativeExpand<E>,
        id: u32,
    ) -> NativeExpand<E> {
        let value = value.read_value(scope);
        let op = BroadcastOp::new(scope.ctx_mut(), value, IndexAttr::new(id as usize));
        scope.register_with_result(&op).into()
    }
}

/// Perform an arbitrary lane shuffle operation across the plane.
/// Each unit reads the value from the specified source lane.
///
/// # Example
/// `plane_shuffle(value, 0)` - all lanes read from lane 0 (same as broadcast)
/// `plane_shuffle(value, lane_id ^ 1)` - butterfly pattern (same as `shuffle_xor`)
#[allow(unused_variables)]
pub fn plane_shuffle<E: CubePrimitive>(value: E, src_lane: u32) -> E {
    unexpanded!()
}

/// Module containing the expand function for [`plane_shuffle()`].
pub mod plane_shuffle {
    use super::*;

    /// Expand method of [`plane_shuffle()`].
    pub fn expand<E: CubePrimitive>(
        scope: &Scope,
        value: NativeExpand<E>,
        src_lane: NativeExpand<u32>,
    ) -> NativeExpand<E> {
        let value = value.read_value(scope);
        let src_lane = src_lane.read_value(scope);
        let op = ShuffleOp::new(scope.ctx_mut(), value, src_lane);
        scope.register_with_result(&op).into()
    }
}

/// Perform a shuffle XOR operation across the plane.
/// Each unit exchanges its value with another unit at an index determined by XOR with the mask.
/// This is useful for butterfly reduction patterns.
///
/// # Example
/// For a 32-lane warp with mask=1:
/// - Lane 0 gets value from lane 1, lane 1 gets value from lane 0
/// - Lane 2 gets value from lane 3, lane 3 gets value from lane 2
/// - etc.
#[allow(unused_variables)]
pub fn plane_shuffle_xor<E: CubePrimitive>(value: E, mask: u32) -> E {
    unexpanded!()
}

/// Module containing the expand function for [`plane_shuffle_xor()`].
pub mod plane_shuffle_xor {
    use super::*;

    /// Expand method of [`plane_shuffle_xor()`].
    pub fn expand<E: CubePrimitive>(
        scope: &Scope,
        value: NativeExpand<E>,
        mask: NativeExpand<u32>,
    ) -> NativeExpand<E> {
        let value = value.read_value(scope);
        let mask = mask.read_value(scope);
        let op = ShuffleXorOp::new(scope.ctx_mut(), value, mask);
        scope.register_with_result(&op).into()
    }
}

/// Perform a shuffle up operation across the plane.
/// Each unit reads the value from a unit with a lower lane ID (`current_id` - delta).
/// Units with `lane_id` < delta will read from themselves (no change).
///
/// # Example
/// For delta=1: `[a, b, c, d] -> [a, a, b, c]`
#[allow(unused_variables)]
pub fn plane_shuffle_up<E: CubePrimitive>(value: E, delta: u32) -> E {
    unexpanded!()
}

/// Module containing the expand function for [`plane_shuffle_up()`].
pub mod plane_shuffle_up {
    use super::*;

    /// Expand method of [`plane_shuffle_up()`].
    pub fn expand<E: CubePrimitive>(
        scope: &Scope,
        value: NativeExpand<E>,
        delta: NativeExpand<u32>,
    ) -> NativeExpand<E> {
        let value = value.read_value(scope);
        let delta = delta.read_value(scope);
        let op = ShuffleUpOp::new(scope.ctx_mut(), value, delta);
        scope.register_with_result(&op).into()
    }
}

/// Perform a shuffle down operation across the plane.
/// Each unit reads the value from a unit with a higher lane ID (`current_id` + delta).
/// Units at the end will read from themselves if (`lane_id` + delta >= `plane_dim`).
///
/// # Example
/// For delta=1: `[a, b, c, d] -> [b, c, d, d]`
#[allow(unused_variables)]
pub fn plane_shuffle_down<E: CubePrimitive>(value: E, delta: u32) -> E {
    unexpanded!()
}

/// Module containing the expand function for [`plane_shuffle_down()`].
pub mod plane_shuffle_down {
    use super::*;

    /// Expand method of [`plane_shuffle_down()`].
    pub fn expand<E: CubePrimitive>(
        scope: &Scope,
        value: NativeExpand<E>,
        delta: NativeExpand<u32>,
    ) -> NativeExpand<E> {
        let value = value.read_value(scope);
        let delta = delta.read_value(scope);
        let op = ShuffleDownOp::new(scope.ctx_mut(), value, delta);
        scope.register_with_result(&op).into()
    }
}

/// Perform a reduce sum operation across all units in a plane.
#[allow(unused_variables)]
pub fn plane_sum<E: CubePrimitive<Scalar: PlaneNumeric>>(value: E) -> E {
    unexpanded!()
}

/// Module containing the expand function for [`plane_sum()`].
pub mod plane_sum {
    use super::*;

    /// Expand method of [`plane_sum()`].
    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
        scope: &Scope,
        elem: NativeExpand<E>,
    ) -> NativeExpand<E> {
        E::Scalar::__expand_native_sum(scope, elem.into()).into()
    }
}

/// Perform an inclusive sum operation across all units in a plane.
/// This sums all values to the "left" of the unit, including this unit's value.
/// Also known as "prefix sum" or "inclusive scan".
///
/// # Example
/// `inclusive_sum([1, 2, 3, 4, 5]) == [1, 3, 6, 10, 15]`
#[allow(unused_variables)]
pub fn plane_inclusive_sum<E: CubePrimitive<Scalar: PlaneNumeric>>(value: E) -> E {
    unexpanded!()
}

/// Module containing the expand function for [`plane_inclusive_sum()`].
pub mod plane_inclusive_sum {
    use super::*;

    /// Expand method of [`plane_inclusive_sum()`].
    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
        scope: &Scope,
        elem: NativeExpand<E>,
    ) -> NativeExpand<E> {
        E::Scalar::__expand_native_inclusive_sum(scope, elem.into()).into()
    }
}

/// Perform an exclusive sum operation across all units in a plane.
/// This sums all values to the "left" of the unit, excluding this unit's value. The 0th unit will
/// be set to `E::zero()`.
/// Also known as "exclusive prefix sum" or "exclusive scan".
///
/// # Example
/// `exclusive_sum([1, 2, 3, 4, 5]) == [0, 1, 3, 6, 10]`
#[allow(unused_variables)]
pub fn plane_exclusive_sum<E: CubePrimitive<Scalar: PlaneNumeric>>(value: E) -> E {
    unexpanded!()
}

/// Module containing the expand function for [`plane_exclusive_sum()`].
pub mod plane_exclusive_sum {
    use super::*;

    /// Expand method of [`plane_exclusive_sum()`].
    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
        scope: &Scope,
        elem: NativeExpand<E>,
    ) -> NativeExpand<E> {
        E::Scalar::__expand_native_exclusive_sum(scope, elem.into()).into()
    }
}

/// Perform a reduce prod operation across all units in a plane.
pub fn plane_prod<E: CubePrimitive<Scalar: PlaneNumeric>>(_elem: E) -> E {
    unexpanded!()
}

/// Module containing the expand function for [`plane_prod()`].
pub mod plane_prod {
    use super::*;

    /// Expand method of [`plane_prod()`].
    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
        scope: &Scope,
        elem: NativeExpand<E>,
    ) -> NativeExpand<E> {
        E::Scalar::__expand_native_prod(scope, elem.into()).into()
    }
}

/// Perform an inclusive product operation across all units in a plane.
/// This multiplies all values to the "left" of the unit, including this unit's value.
/// Also known as "prefix product" or "inclusive scan".
///
/// # Example
/// `exclusive_prod([1, 2, 3, 4, 5]) == [1, 2, 6, 24, 120]`
#[allow(unused_variables)]
pub fn plane_inclusive_prod<E: CubePrimitive<Scalar: PlaneNumeric>>(value: E) -> E {
    unexpanded!()
}

/// Module containing the expand function for [`plane_inclusive_prod()`].
pub mod plane_inclusive_prod {
    use super::*;

    /// Expand method of [`plane_inclusive_prod()`].
    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
        scope: &Scope,
        elem: NativeExpand<E>,
    ) -> NativeExpand<E> {
        E::Scalar::__expand_native_inclusive_prod(scope, elem.into()).into()
    }
}

/// Perform an exclusive product operation across all units in a plane.
/// This multiplies all values to the "left" of the unit, excluding this unit's value. The 0th unit
/// will be set to `E::one()`.
/// Also known as "exclusive prefix product" or "exclusive scan".
///
/// # Example
/// `exclusive_prod([1, 2, 3, 4, 5]) == [1, 1, 2, 6, 24]`
#[allow(unused_variables)]
pub fn plane_exclusive_prod<E: CubePrimitive<Scalar: PlaneNumeric>>(value: E) -> E {
    unexpanded!()
}

/// Module containing the expand function for [`plane_exclusive_prod()`].
pub mod plane_exclusive_prod {
    use super::*;

    /// Expand method of [`plane_exclusive_prod()`].
    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
        scope: &Scope,
        elem: NativeExpand<E>,
    ) -> NativeExpand<E> {
        E::Scalar::__expand_native_exclusive_prod(scope, elem.into()).into()
    }
}

/// Perform a reduce max operation across all units in a plane.
pub fn plane_max<E: CubePrimitive<Scalar: PlaneNumeric>>(_elem: E) -> E {
    unexpanded!()
}

/// Module containing the expand function for [`plane_max()`].
pub mod plane_max {
    use super::*;

    /// Expand method of [`plane_max()`].
    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
        scope: &Scope,
        elem: NativeExpand<E>,
    ) -> NativeExpand<E> {
        E::Scalar::__expand_native_plane_max(scope, elem.into()).into()
    }
}

/// Perform a reduce min operation across all units in a plane.
pub fn plane_min<E: CubePrimitive<Scalar: PlaneNumeric>>(_elem: E) -> E {
    unexpanded!()
}

/// Module containing the expand function for [`plane_min()`].
pub mod plane_min {
    use super::*;

    /// Expand method of [`plane_min()`].
    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
        scope: &Scope,
        elem: NativeExpand<E>,
    ) -> NativeExpand<E> {
        E::Scalar::__expand_native_plane_min(scope, elem.into()).into()
    }
}

/// Perform a reduce all operation across all units in a plane.
pub fn plane_all(_elem: bool) -> bool {
    unexpanded!()
}

/// Module containing the expand function for [`plane_all()`].
pub mod plane_all {
    use super::*;

    /// Expand method of [`plane_all()`].
    pub fn expand(scope: &Scope, elem: NativeExpand<bool>) -> NativeExpand<bool> {
        let value = elem.read_value(scope);
        let op = AllOp::new(scope.ctx_mut(), value);
        scope.register_with_result(&op).into()
    }
}

/// Perform a reduce any operation across all units in a plane.
pub fn plane_any(_elem: bool) -> bool {
    unexpanded!()
}

/// Module containing the expand function for [`plane_any()`].
pub mod plane_any {
    use super::*;

    /// Expand method of [`plane_any()`].
    pub fn expand(scope: &Scope, elem: NativeExpand<bool>) -> NativeExpand<bool> {
        let value = elem.read_value(scope);
        let op = AnyOp::new(scope.ctx_mut(), value);
        scope.register_with_result(&op).into()
    }
}

/// Perform a ballot operation across all units in a plane.
/// Returns a set of 32-bit bitfields as a [`Vector`], with each element containing the value from 32
/// invocations.
/// Note that vector size will always be set to 4 even for `PLANE_DIM <= 64`, because we can't
/// retrieve the actual plane size at expand time. Use the runtime`PLANE_DIM` to index appropriately.
pub fn plane_ballot(_elem: bool) -> Vector<u32, Const<4>> {
    unexpanded!()
}

/// Module containing the expand function for [`plane_ballot()`].
pub mod plane_ballot {
    use super::*;

    /// Expand method of [`plane_ballot()`].
    pub fn expand(scope: &Scope, elem: NativeExpand<bool>) -> NativeExpand<Vector<u32, Const<4>>> {
        let value = elem.read_value(scope);
        let op = BallotOp::new(scope.ctx_mut(), value);
        scope.register_with_result(&op).into()
    }
}