cubecl_tile/algebra/
semiring.rs

1//! The pair of monoids a contraction runs under.
2
3use cubecl::prelude::*;
4
5use super::monoid::{Carrier, Monoid};
6use crate::comptime_only;
7
8/// A product monoid and the accumulation monoid its products sum into.
9#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
10pub struct Semiring {
11    add: Monoid,
12    mul: Monoid,
13}
14
15impl Semiring {
16    /// `(+, *)`: the ordinary matmul.
17    pub const SUM_PROD: Self = Self {
18        add: Monoid::Sum,
19        mul: Monoid::Prod,
20    };
21    /// `(min, +)`: shortest paths, DTW.
22    pub const MIN_SUM: Self = Self {
23        add: Monoid::Min,
24        mul: Monoid::Sum,
25    };
26    /// `(max, +)`: Viterbi in log space.
27    pub const MAX_SUM: Self = Self {
28        add: Monoid::Max,
29        mul: Monoid::Sum,
30    };
31
32    /// The monoid products accumulate into, and every drain sums partials under.
33    pub const fn add(self) -> Monoid {
34        self.add
35    }
36
37    /// The product a pair of operands forms before it is summed in.
38    pub const fn mul(self) -> Monoid {
39        self.mul
40    }
41}
42
43#[cube]
44impl Semiring {
45    /// One accumulation step, `acc + (lhs * rhs)`, in `fma`'s argument order.
46    /// [`SUM_PROD`](Semiring::SUM_PROD) over floats must stay a single `fma`. An `fma` is a float
47    /// instruction, so integers take the product and the sum apart, which is exact.
48    fn step_of<T: Carrier + CubePrimitive>(
49        lhs: T,
50        rhs: T,
51        acc: T,
52        #[comptime] semiring: Semiring,
53    ) -> T {
54        let element = elem_type_of::<T>();
55        if comptime!(semiring == Semiring::SUM_PROD && element.is_float()) {
56            fma(lhs, rhs, acc)
57        } else {
58            let product = semiring.mul().combine::<T>(lhs, rhs);
59            semiring.add().combine::<T>(product, acc)
60        }
61    }
62}
63
64/// `semiring.step(a, b, acc)`, the pair [`Monoid::combine`] documents.
65impl Semiring {
66    pub fn step<T: Carrier + CubePrimitive>(self, lhs: T, rhs: T, acc: T) -> T {
67        Semiring::step_of::<T>(lhs, rhs, acc, self)
68    }
69
70    pub fn __expand_step_method<T: Carrier + CubePrimitive>(
71        self,
72        scope: &Scope,
73        lhs: T::ExpandType,
74        rhs: T::ExpandType,
75        acc: T::ExpandType,
76    ) -> T::ExpandType {
77        Semiring::__expand_step_of::<T>(scope, lhs, rhs, acc, self)
78    }
79}
80
81comptime_only!(Semiring);
cubecl_tile/algebra/semiring.rs

cubecl_tile/algebra/
semiring.rs
