cubecl_tile/algebra/
monoid.rs

1//! An identity and an associative operation.
2
3use cubecl::prelude::*;
4
5use crate::comptime_only;
6
7/// The ordering and arithmetic a [`Monoid`] needs of the values it combines.
8pub trait Carrier:
9    CubePartialOrd
10    + CubeAdd
11    + CubeMul
12    + core::ops::Add<Self, Output = Self>
13    + core::ops::Mul<Self, Output = Self>
14    + Sized
15{
16}
17
18impl<T> Carrier for T where
19    T: CubePartialOrd
20        + CubeAdd
21        + CubeMul
22        + core::ops::Add<Self, Output = Self>
23        + core::ops::Mul<Self, Output = Self>
24{
25}
26
27/// An identity and an associative operation, used by everything that merges values.
28#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
29pub enum Monoid {
30    /// `acc + val`, identity `0`.
31    Sum,
32    /// `acc * val`, identity `1`.
33    Prod,
34    /// `max(acc, val)`, identity the lowest value of the element.
35    Max,
36    /// `min(acc, val)`, identity the highest value of the element.
37    Min,
38}
39
40#[cube]
41impl Monoid {
42    /// This monoid's identity element; a masked read past an operand's extent must return it.
43    pub fn identity<E: Numeric>(#[comptime] monoid: Monoid) -> E {
44        match comptime!(monoid) {
45            Monoid::Sum => E::from_int(0),
46            Monoid::Prod => E::from_int(1),
47            Monoid::Max => E::min_value(),
48            Monoid::Min => E::max_value(),
49        }
50    }
51
52    /// `lhs ∗ rhs`, pointwise on vectors.
53    fn combine_of<T: Carrier>(lhs: T, rhs: T, #[comptime] monoid: Monoid) -> T {
54        match comptime!(monoid) {
55            Monoid::Sum => lhs + rhs,
56            Monoid::Prod => lhs * rhs,
57            Monoid::Max => max(lhs, rhs),
58            Monoid::Min => min(lhs, rhs),
59        }
60    }
61
62    /// A vector's first `width` components combined into one value.
63    pub fn reduce<E: Numeric, N: Size>(
64        v: Vector<E, N>,
65        #[comptime] width: usize,
66        #[comptime] monoid: Monoid,
67    ) -> E {
68        let mut acc = Monoid::identity::<E>(monoid);
69        #[unroll]
70        for j in 0..width {
71            acc = monoid.combine::<E>(acc, v.extract(j));
72        }
73        acc
74    }
75}
76
77/// `monoid.combine(a, b)`, hand-written since a comptime-only type has no `{Name}Expand`.
78impl Monoid {
79    pub fn combine<T: Carrier>(self, lhs: T, rhs: T) -> T {
80        Monoid::combine_of::<T>(lhs, rhs, self)
81    }
82
83    pub fn __expand_combine_method<T: Carrier>(
84        self,
85        scope: &Scope,
86        lhs: T::ExpandType,
87        rhs: T::ExpandType,
88    ) -> T::ExpandType {
89        Monoid::__expand_combine_of::<T>(scope, lhs, rhs, self)
90    }
91}
92
93comptime_only!(Monoid);
cubecl_tile/algebra/monoid.rs

cubecl_tile/algebra/
monoid.rs
