Skip to main content

cubecl_tile/tile/kind/procedural/
affine.rs

1use cubecl::prelude::*;
2
3use crate::Axis;
4
5use super::{Reads, Recipe, RecipeCoords, RecipeExpand};
6
7/// A one-dimensional affine coordinate expression, `offset + coefficient * coordinate[axis]`.
8#[derive(CubeType, Clone)]
9pub struct AffineCoordinate<T: Numeric> {
10    pub offset: T,
11    pub coefficient: T,
12    #[cube(comptime)]
13    pub axis: Axis,
14}
15
16/// Construct an [`AffineCoordinate`] along the specified axis.
17#[cube]
18pub fn affine_along<T: Numeric>(
19    #[comptime] axis: Axis,
20    offset: T,
21    coefficient: T,
22) -> AffineCoordinate<T> {
23    AffineCoordinate::<T> {
24        offset,
25        coefficient,
26        axis,
27    }
28}
29
30#[cube]
31impl<T: Numeric> Recipe<T> for AffineCoordinate<T> {
32    fn evaluate(&self, coordinates: &RecipeCoords) -> T {
33        self.offset + self.coefficient * T::cast_from(coordinates.along(self.axis))
34    }
35}
36
37impl<T: Numeric> Reads for AffineCoordinateExpand<T> {
38    fn reads(&self, _scope: &Scope, axis: Axis) -> bool {
39        self.axis == axis
40    }
41}