cubecl_tile/tile/kind/procedural/
separable.rs

1//! A recipe stated as one factor per contracted axis.
2
3use cubecl::prelude::*;
4
5use super::{FactorReads, Reads, Recipe, RecipeCoords, RecipeExpand, Separable, SeparableExpand};
6
7/// The product of one factor per contracted axis, in contraction order: `K₀ ⊗ K₁ ⊗ … ⊗ Kₙ₋₁`.
8/// Each factor must read a distinct axis; this is not checked.
9#[derive(CubeType, Clone)]
10pub struct Factors<R: CubeType> {
11    pub factors: Sequence<R>,
12}
13
14impl<R: CubeType> FactorReads for FactorsExpand<R>
15where
16    R::ExpandType: Reads,
17{
18    fn factor_reads(&self, scope: &Scope, factor: usize, axis: crate::Axis) -> bool {
19        let factor = NativeExpand::from_lit(scope, factor);
20        self.factors
21            .__expand_index_method(scope, factor)
22            .reads(scope, axis)
23    }
24}
25
26#[cube]
27impl<R: CubeType> Factors<R> {
28    /// The product of `factors`, in contraction order.
29    pub fn new(factors: Sequence<R>) -> Factors<R> {
30        Factors::<R> { factors }
31    }
32
33    /// The sequence's length; panics when empty.
34    pub(crate) fn rank(&self) -> comptime_type!(usize) {
35        let rank = self.factors.len();
36        comptime!(assert!(
37            rank > 0,
38            "Factors: a separable recipe needs at least one factor"
39        ));
40        rank
41    }
42}
43
44#[cube]
45impl<T: Numeric, R: Recipe<T>> Recipe<T> for Factors<R> {
46    fn evaluate(&self, coordinates: &RecipeCoords) -> T {
47        let rank = self.rank();
48        let mut value = self.factors.index(0usize).evaluate(coordinates);
49
50        #[unroll]
51        for f in 1..rank {
52            value *= self.factors.index(f).evaluate(coordinates);
53        }
54
55        value
56    }
57}
58
59#[cube]
60impl<T: Numeric, R: Recipe<T>> Separable<T> for Factors<R> {
61    fn factors(&self) -> comptime_type!(usize) {
62        self.rank()
63    }
64
65    fn factor(&self, coordinates: &RecipeCoords, #[comptime] factor: usize) -> T {
66        self.factors.index(factor).evaluate(coordinates)
67    }
68}
cubecl_tile/tile/kind/procedural/separable.rs

cubecl_tile/tile/kind/procedural/
separable.rs
