Skip to main content

cubecl_tile/tile/kind/procedural/
phase.rs

1use cubecl::prelude::*;
2
3use crate::{Axis, Integer, IntegerExpand, floor_div_rem};
4
5use super::{Reads, Recipe, RecipeCoords, RecipeExpand};
6
7/// The fractional part a rational coordinate mapping leaves behind, scaled:
8/// `coefficient * frac((coord[axis] * numerator_scale + numerator_offset) / divisor)`.
9#[derive(CubeType, Clone)]
10pub struct Phase<T: Float> {
11    /// Multiplies the whole fraction.
12    pub coefficient: T,
13    pub numerator_scale: u32,
14    pub numerator_offset: i32,
15    pub divisor: u32,
16    #[cube(comptime)]
17    pub axis: Axis,
18}
19
20#[cube]
21impl<T: Float> Recipe<T> for Phase<T> {
22    fn evaluate(&self, coordinates: &RecipeCoords) -> T {
23        let divisor = self.divisor.constant();
24        // Checked here: a struct literal bypasses any constructor.
25        comptime!(assert!(
26            divisor != Some(0),
27            "Phase: divisor must be non-zero"
28        ));
29        let numerator = coordinates
30            .along(self.axis)
31            .times(self.numerator_scale)
32            .cast::<i32>()
33            .plus(self.numerator_offset);
34        let (_, residue) = floor_div_rem(numerator, self.divisor.cast::<i32>());
35        let residue = T::cast_from(residue);
36        // A constant divisor becomes a reciprocal literal to multiply by.
37        let fraction = if comptime!(divisor.is_some()) {
38            residue * T::new(comptime!(1.0 / divisor.unwrap() as f64))
39        } else {
40            residue / T::cast_from(self.divisor)
41        };
42        self.coefficient * fraction
43    }
44}
45
46impl<T: Float> Reads for PhaseExpand<T> {
47    fn reads(&self, _scope: &Scope, axis: Axis) -> bool {
48        self.axis == axis
49    }
50}