cubecl_tile/tile/kind/procedural/
phase.rs1use cubecl::prelude::*;
2
3use crate::{Axis, Integer, IntegerExpand, floor_div_rem};
4
5use super::{Reads, Recipe, RecipeCoords, RecipeExpand};
6
7#[derive(CubeType, Clone)]
10pub struct Phase<T: Float> {
11 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 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 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}