1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

use crate::{
    ir::{ClampOperator, Operator},
    prelude::{CubeContext, CubePrimitive, ExpandElement, UInt, BF16, F16, F32, F64, I32, I64},
    unexpanded,
};

use super::unary_expand;

pub trait Clamp: CubePrimitive + Sized {
    /// Clamp the input value between the max and min values provided.
    #[allow(unused_variables)]
    fn clamp(input: Self, min_value: Self, max_value: Self) -> Self {
        unexpanded!()
    }
    fn __expand_clamp(
        context: &mut CubeContext,
        input: Self::ExpandType,
        min_value: Self::ExpandType,
        max_value: Self::ExpandType,
    ) -> Self::ExpandType {
        let input: ExpandElement = input.into();
        let min_value: ExpandElement = min_value.into();
        let max_value: ExpandElement = max_value.into();

        unary_expand(context, input, |op| {
            Operator::Clamp(ClampOperator {
                input: op.input,
                min_value: *min_value,
                max_value: *max_value,
                out: op.out,
            })
        })
        .into()
    }
}

impl Clamp for F16 {}
impl Clamp for BF16 {}
impl Clamp for F32 {}
impl Clamp for F64 {}
impl Clamp for I32 {}
impl Clamp for I64 {}
impl Clamp for UInt {}