feat: add Trunc op (#956)

* feat: add `Trunc` op

* tests

Author: mooori

crates/cubecl-core/src/frontend/container/line/ops.rs

@@ -4,7 +4,7 @@ use num_traits::{NumCast, ToPrimitive};
 
 use crate::{
     self as cubecl,
-    prelude::{IsInf, IsNan, Powi, SaturatingAdd, SaturatingSub},
+    prelude::{IsInf, IsNan, Powi, SaturatingAdd, SaturatingSub, Trunc},
 };
 use crate::{
     frontend::{
@@ -255,6 +255,7 @@ impl<P: CubePrimitive + Remainder> Remainder for Line<P> {}
 impl<P: CubePrimitive + Round> Round for Line<P> {}
 impl<P: CubePrimitive + Floor> Floor for Line<P> {}
 impl<P: CubePrimitive + Ceil> Ceil for Line<P> {}
+impl<P: CubePrimitive + Trunc> Trunc for Line<P> {}
 impl<P: CubePrimitive + ReverseBits> ReverseBits for Line<P> {}
 impl<P: CubePrimitive + BitwiseNot> BitwiseNot for Line<P> {}
 impl<P: CubePrimitive + SaturatingAdd> SaturatingAdd for Line<P> {}

crates/cubecl-core/src/frontend/element/float.rs

@@ -32,6 +32,7 @@ pub trait Float:
     + Round
     + Floor
     + Ceil
+    + Trunc
     + Erf
     + Recip
     + Magnitude

crates/cubecl-core/src/frontend/element/float/typemap.rs

@@ -254,6 +254,7 @@ impl<const POS: u8> Sqrt for ElemExpand<POS> {}
 impl<const POS: u8> Round for ElemExpand<POS> {}
 impl<const POS: u8> Floor for ElemExpand<POS> {}
 impl<const POS: u8> Ceil for ElemExpand<POS> {}
+impl<const POS: u8> Trunc for ElemExpand<POS> {}
 impl<const POS: u8> IsNan for ElemExpand<POS> {}
 impl<const POS: u8> IsInf for ElemExpand<POS> {}
 

crates/cubecl-core/src/frontend/operation/unary.rs

@@ -235,6 +235,18 @@ impl_unary_func!(
     f32,
     f64
 );
+impl_unary_func!(
+    Trunc,
+    trunc,
+    __expand_trunc,
+    Arithmetic::Trunc,
+    f16,
+    bf16,
+    flex32,
+    tf32,
+    f32,
+    f64
+);
 impl_unary_func!(
     Erf,
     erf,

crates/cubecl-core/src/runtime_tests/unary.rs

@@ -305,6 +305,29 @@ test_unary_impl!(
     ]
 );
 
+test_unary_impl!(
+    test_trunc,
+    F,
+    F::trunc,
+    [{
+        input_vectorization: 1,
+        out_vectorization: 1,
+        input: as_type![F: -1.2, -1., -0., 0.],
+        expected: as_type![F: -1., -1., -0., 0.]
+    },
+    {
+        input_vectorization: 2,
+        out_vectorization: 2,
+        input: as_type![F: f32::NAN, 1., 1.2, 1.9],
+        expected: as_type![F: f32::NAN, 1., 1., 1.0]
+    },{
+        input_vectorization: 4,
+        out_vectorization: 4,
+        input: as_type![F: -0.9, 0.2, f32::NAN, 1.99],
+        expected: as_type![F: -0., 0., f32::NAN, 1.]
+    }]
+);
+
 test_unary_impl_fixed!(
     test_is_nan,
     F,
@@ -479,6 +502,7 @@ macro_rules! testgen_unary {
             add_test!(test_normalize);
             add_test!(test_magnitude);
             add_test!(test_abs);
+            add_test!(test_trunc);
             add_test!(test_is_nan);
             add_test!(test_is_inf);
         }

crates/cubecl-cpp/src/shared/base.rs

@@ -1012,6 +1012,9 @@ impl<D: Dialect> CppCompiler<D> {
             gpu::Arithmetic::Ceil(op) => {
                 instructions.push(Instruction::Ceil(self.compile_unary(op, out)))
             }
+            gpu::Arithmetic::Trunc(op) => {
+                instructions.push(Instruction::Trunc(self.compile_unary(op, out)))
+            }
             gpu::Arithmetic::Remainder(op) => {
                 instructions.push(Instruction::Remainder(self.compile_binary(op, out)))
             }

crates/cubecl-cpp/src/shared/instruction.rs

@@ -200,6 +200,7 @@ pub enum Instruction<D: Dialect> {
     },
     Round(UnaryInstruction<D>),
     Ceil(UnaryInstruction<D>),
+    Trunc(UnaryInstruction<D>),
     Floor(UnaryInstruction<D>),
     Warp(WarpInstruction<D>),
     Wmma(WmmaInstruction<D>),
@@ -538,6 +539,7 @@ for ({i_ty} {i} = {start}; {i} {cmp} {end}; {increment}) {{
             Instruction::ThreadFence => f.write_str("__threadfence();\n"),
             Instruction::Round(it) => Round::format(f, &it.input, &it.out),
             Instruction::Ceil(it) => Ceil::format(f, &it.input, &it.out),
+            Instruction::Trunc(it) => Trunc::format(f, &it.input, &it.out),
             Instruction::Floor(it) => Floor::format(f, &it.input, &it.out),
             Instruction::SliceLength { input, out } => {
                 let out = out.fmt_left();

crates/cubecl-cpp/src/shared/unary.rs

@@ -154,6 +154,7 @@ function!(Sin, "sin");
 function!(Sqrt, "sqrt");
 function!(Exp, "exp");
 function!(Ceil, "ceil");
+function!(Trunc, "trunc");
 function!(Floor, "floor");
 function!(Round, "rint");
 

crates/cubecl-cpu/src/compiler/visitor/operation/arithmetic.rs

@@ -40,6 +40,15 @@ impl<'a> Visitor<'a> {
                 ));
                 self.insert_variable(out, result);
             }
+            Arithmetic::Trunc(trunc) => {
+                let value = self.get_variable(trunc.input);
+                let result = self.append_operation_with_result(llvm_ods::intr_trunc(
+                    self.context,
+                    value,
+                    self.location,
+                ));
+                self.insert_variable(out, result);
+            }
             Arithmetic::Clamp(clamp) => {
                 let value = self.get_variable(clamp.input);
                 let mut min = self.get_variable(clamp.min_value);

crates/cubecl-ir/src/arithmetic.rs

@@ -32,6 +32,7 @@ pub enum Arithmetic {
     Round(UnaryOperator),
     Floor(UnaryOperator),
     Ceil(UnaryOperator),
+    Trunc(UnaryOperator),
     Erf(UnaryOperator),
     Recip(UnaryOperator),
     Clamp(ClampOperator),
@@ -73,6 +74,7 @@ impl Display for Arithmetic {
             Arithmetic::Round(op) => write!(f, "{}.round()", op.input),
             Arithmetic::Floor(op) => write!(f, "{}.floor()", op.input),
             Arithmetic::Ceil(op) => write!(f, "{}.ceil()", op.input),
+            Arithmetic::Trunc(op) => write!(f, "{}.trunc()", op.input),
             Arithmetic::Erf(op) => write!(f, "{}.erf()", op.input),
             Arithmetic::Recip(op) => write!(f, "{}.recip()", op.input),
             Arithmetic::Clamp(op) => {