cubecl_tile/tile/kind/procedural/
normalization.rs

1//! The divisor of a normalized separable factor run: which taps it sums, and its zero guard.
2
3use cubecl::prelude::*;
4
5use crate::Space;
6
7/// What a normalized factor run sums and divides by, over the complete runs of `over`.
8#[derive(Clone, PartialEq, Debug)]
9pub(crate) struct Normalization {
10    pub taps: TapSupport,
11    pub guard: DivGuard,
12    pub over: Space,
13}
14
15impl Normalization {
16    pub(crate) fn new(taps: TapSupport, guard: DivGuard, over: Space) -> Self {
17        Normalization { taps, guard, over }
18    }
19}
20
21/// How division handles a denominator too small to divide by. The default maps it to zero.
22#[derive(Clone, Copy, PartialEq, Debug, Default)]
23pub struct DivGuard {
24    pub epsilon: f32,
25    pub fallback: f32,
26}
27
28impl DivGuard {
29    /// Falls back to `fallback` where the denominator's magnitude is at most `epsilon`.
30    pub fn new(epsilon: f32, fallback: f32) -> Self {
31        assert!(
32            epsilon.is_finite() && epsilon >= 0.0,
33            "DivGuard: epsilon must be finite and non-negative"
34        );
35        assert!(fallback.is_finite(), "DivGuard: fallback must be finite");
36        DivGuard { epsilon, fallback }
37    }
38}
39
40/// Which taps contribute to the sum of a normalized separable filter factor.
41#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
42pub enum TapSupport {
43    /// Sum only taps whose projected input sample is in bounds.
44    /// The contraction must read the source window in place; a shared-memory stage is rejected.
45    #[default]
46    InBounds,
47    /// Sum the filter's whole support, preserving the fade of zero padding at an edge.
48    Whole,
49}
50
51/// A guarded reciprocal that preserves the sign of a valid denominator; NaN takes the fallback.
52#[cube]
53pub(crate) fn guarded_recip<E: Numeric>(d: E, #[comptime] guard: DivGuard) -> E {
54    let epsilon = E::cast_from(comptime!(guard.epsilon));
55    let fallback = E::cast_from(comptime!(guard.fallback));
56    let valid = d.abs() > epsilon;
57    let safe = select(valid, d, E::from_int(1));
58    select(valid, E::from_int(1) / safe, fallback)
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64
65    #[test]
66    fn division_guard_accepts_finite_non_negative_thresholds() {
67        DivGuard::new(0.0, 0.0);
68        DivGuard::new(1.0e-7, -1.0);
69    }
70
71    #[test]
72    fn division_guard_rejects_invalid_thresholds() {
73        for epsilon in [-1.0, f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
74            assert!(
75                std::panic::catch_unwind(|| DivGuard::new(epsilon, 0.0)).is_err(),
76                "epsilon {epsilon:?} should be rejected"
77            );
78        }
79        for fallback in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
80            assert!(
81                std::panic::catch_unwind(|| DivGuard::new(1.0e-7, fallback)).is_err(),
82                "fallback {fallback:?} should be rejected"
83            );
84        }
85    }
86}
cubecl_tile/tile/kind/procedural/normalization.rs

cubecl_tile/tile/kind/procedural/
normalization.rs
