cubecl_tile/tile/kind/procedural/
base.rs

1//! Recipes: memory-free sources evaluated at logical coordinates, and their factorization.
2
3use cubecl::ir::Scope;
4use cubecl::prelude::*;
5
6use crate::{Axis, Coords, Space};
7
8/// Coordinate dependence of a recipe, queried while a kernel is expanded.
9pub trait Reads {
10    fn reads(&self, scope: &Scope, axis: Axis) -> bool;
11}
12
13/// Per-factor coordinate dependencies of a separable recipe, queried during expansion.
14pub trait FactorReads {
15    fn factor_reads(&self, scope: &Scope, factor: usize, axis: Axis) -> bool;
16}
17
18/// The absolute logical coordinates a [`Recipe`] is evaluated at: `origin` plus read offset.
19#[derive(CubeType, Clone)]
20#[expand(derive(Clone))]
21pub struct RecipeCoords {
22    origin: Coords<u32>,
23    offset: Coords<u32>,
24    #[cube(comptime)]
25    pub space: Space,
26}
27
28#[cube]
29impl RecipeCoords {
30    pub(crate) fn new(
31        origin: &Coords<u32>,
32        offset: &Coords<u32>,
33        #[comptime] space: Space,
34    ) -> Self {
35        RecipeCoords {
36            origin: origin.clone(),
37            offset: offset.clone(),
38            space,
39        }
40    }
41
42    /// The absolute coordinate along the axis at comptime position `p`.
43    pub fn at(&self, #[comptime] p: usize) -> u32 {
44        self.origin.at(p) + self.offset.at(p)
45    }
46
47    /// The absolute coordinate along the specified `axis`.
48    pub fn along(&self, #[comptime] axis: Axis) -> u32 {
49        self.at(comptime!(self.space.position(axis)))
50    }
51}
52
53/// An N-dimensional scalar field evaluated at absolute logical coordinates.
54#[cube(expand_base_traits = "ExpandTypeClone")]
55pub trait Recipe<T: Numeric> {
56    fn evaluate(&self, coordinates: &RecipeCoords) -> T;
57}
58
59/// A [`Recipe`] that factorizes into one factor per contracted axis, `R = ∏ᵢ Rᵢ`.
60/// Factor `i` must vary only along the `i`-th contracted axis.
61#[cube]
62pub trait Separable<T: Numeric>: Recipe<T> {
63    fn factors(&self) -> comptime_type!(usize);
64    /// The factor at comptime position `factor`, in the contraction's walk order of axes.
65    fn factor(&self, coordinates: &RecipeCoords, #[comptime] factor: usize) -> T;
66}
cubecl_tile/tile/kind/procedural/base.rs

cubecl_tile/tile/kind/procedural/
base.rs
