cubek_tile/algebra/
monoid.rs

1//! An identity and an associative operation.
2
3use cubecl::prelude::*;
4
5/// The ordering and arithmetic a [`Monoid`] needs of the values it combines.
6pub trait Carrier:
7    CubePartialOrd
8    + CubeAdd
9    + CubeMul
10    + core::ops::Add<Self, Output = Self>
11    + core::ops::Mul<Self, Output = Self>
12    + Sized
13{
14}
15
16impl<T> Carrier for T where
17    T: CubePartialOrd
18        + CubeAdd
19        + CubeMul
20        + core::ops::Add<Self, Output = Self>
21        + core::ops::Mul<Self, Output = Self>
22{
23}
24
25/// An identity and an associative operation, used by everything that merges values.
26#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
27pub enum Monoid {
28    /// `acc + val`, identity `0`.
29    Sum,
30    /// `acc * val`, identity `1`.
31    Prod,
32    /// `max(acc, val)`, identity the lowest value of the element.
33    Max,
34    /// `min(acc, val)`, identity the highest value of the element.
35    Min,
36}
37
38#[cube]
39impl Monoid {
40    /// This monoid's identity element; a masked read past an operand's extent must return it.
41    pub fn identity<E: Numeric>(#[comptime] monoid: Monoid) -> E {
42        match comptime!(monoid) {
43            Monoid::Sum => E::from_int(0),
44            Monoid::Prod => E::from_int(1),
45            Monoid::Max => E::min_value(),
46            Monoid::Min => E::max_value(),
47        }
48    }
49
50    /// `lhs ∗ rhs`, pointwise on vectors.
51    fn combine_of<T: Carrier>(lhs: T, rhs: T, #[comptime] monoid: Monoid) -> T {
52        match comptime!(monoid) {
53            Monoid::Sum => lhs + rhs,
54            Monoid::Prod => lhs * rhs,
55            Monoid::Max => max(lhs, rhs),
56            Monoid::Min => min(lhs, rhs),
57        }
58    }
59
60    /// A vector's first `width` components combined into one value.
61    pub fn reduce<E: Numeric, N: Size>(
62        v: Vector<E, N>,
63        #[comptime] width: usize,
64        #[comptime] monoid: Monoid,
65    ) -> E {
66        let mut acc = Monoid::identity::<E>(monoid);
67        #[unroll]
68        for j in 0..width {
69            acc = monoid.combine::<E>(acc, v.extract(j));
70        }
71        acc
72    }
73
74    /// `seed ∗ arr₀ ∗ … ∗ arr_{len-1}`: `len` elements of `arr` combined into `seed`.
75    pub fn reduce_array<E: Numeric>(
76        arr: &Array<E>,
77        #[comptime] len: usize,
78        seed: E,
79        #[comptime] monoid: Monoid,
80    ) -> E {
81        let mut acc = seed;
82        #[unroll]
83        for i in 0..len {
84            acc = monoid.combine::<E>(acc, arr[i]);
85        }
86        acc
87    }
88}
89
90/// `monoid.combine(a, b)`, hand-written since a comptime-only type has no `{Name}Expand`.
91impl Monoid {
92    pub fn combine<T: Carrier>(self, lhs: T, rhs: T) -> T {
93        Monoid::combine_of::<T>(lhs, rhs, self)
94    }
95
96    pub fn __expand_combine_method<T: Carrier>(
97        self,
98        scope: &Scope,
99        lhs: T::ExpandType,
100        rhs: T::ExpandType,
101    ) -> T::ExpandType {
102        Monoid::__expand_combine_of::<T>(scope, lhs, rhs, self)
103    }
104}
105
106/// Makes a comptime-only type expand as itself.
107macro_rules! comptime_only {
108    ($ty:ty) => {
109        impl CubeType for $ty {
110            type ExpandType = Self;
111        }
112
113        impl IntoExpand for $ty {
114            type Expand = Self;
115
116            fn into_expand(self, _scope: &Scope) -> Self {
117                self
118            }
119        }
120
121        impl IntoMut for $ty {
122            fn into_mut(self, _scope: &Scope) -> Self {
123                self
124            }
125        }
126
127        impl ExpandTypeClone for $ty {
128            fn clone_unchecked(&self) -> Self {
129                Clone::clone(self)
130            }
131        }
132
133        impl CubeDebug for $ty {}
134
135        impl AsRefExpand for $ty {
136            fn __expand_ref_method(&self, _scope: &Scope) -> &Self {
137                self
138            }
139        }
140
141        impl AsMutExpand for $ty {
142            fn __expand_ref_mut_method(&mut self, _scope: &Scope) -> &mut Self {
143                self
144            }
145        }
146    };
147}
148
149pub(crate) use comptime_only;
150
151comptime_only!(Monoid);
cubek_tile/algebra/monoid.rs

cubek_tile/algebra/
monoid.rs
