Implemented magnitude

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

@@ -26,6 +26,7 @@ pub trait Float:
     + Ceil
     + Erf
     + Recip
+    + Magnitude
     + Normalize
     + Into<Self::ExpandType>
     + core::ops::Add<Output = Self>

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

@@ -170,6 +170,37 @@ where
     out
 }
 
+pub fn fixed_output_unary_expand<F>(
+    context: &mut CubeContext,
+    input: ExpandElement,
+    out_item: Item,
+    func: F,
+) -> ExpandElement
+where
+    F: Fn(UnaryOperator) -> Operator,
+{
+    let input_var: Variable = *input;
+
+    let input_item = input.item();
+
+    let out = if input.can_mut() && out_item == input_item {
+        input
+    } else {
+        context.create_local(out_item)
+    };
+
+    let out_var = *out;
+
+    let op = func(UnaryOperator {
+        input: input_var,
+        out: out_var,
+    });
+
+    context.register(op);
+
+    out
+}
+
 pub fn init_expand<F>(context: &mut CubeContext, input: ExpandElement, func: F) -> ExpandElement
 where
     F: Fn(UnaryOperator) -> Operator,

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

@@ -3,11 +3,11 @@ use half::{bf16, f16};
 use crate::{
     frontend::CubeContext,
     ir::Operator,
-    prelude::{CubePrimitive, ExpandElementTyped},
+    prelude::{CubePrimitive, ExpandElement, ExpandElementTyped},
     unexpanded,
 };
 
-use super::base::unary_expand;
+use super::base::{fixed_output_unary_expand, unary_expand};
 
 pub mod not {
     use super::*;
@@ -50,6 +50,26 @@ macro_rules! impl_unary_func {
     }
 }
 
+macro_rules! impl_fixed_out_vectorization_unary_func {
+    ($trait_name:ident, $method_name:ident, $method_name_expand:ident, $operator:expr, $out_vectorization: expr, $($type:ty),*) => {
+        pub trait $trait_name: CubePrimitive + Sized {
+            #[allow(unused_variables)]
+            fn $method_name(x: Self) -> Self {
+                unexpanded!()
+            }
+
+            fn $method_name_expand(context: &mut CubeContext, x: Self::ExpandType) -> ExpandElementTyped<Self> {
+                let expand_element: ExpandElement = x.into();
+                let mut item = expand_element.item();
+                item.vectorization = $out_vectorization;
+                fixed_output_unary_expand(context, expand_element, item, $operator).into()
+            }
+        }
+
+        $(impl $trait_name for $type {})*
+    }
+}
+
 impl_unary_func!(
     Abs,
     abs,
@@ -138,6 +158,17 @@ impl_unary_func!(
     f32,
     f64
 );
+impl_fixed_out_vectorization_unary_func!(
+    Magnitude,
+    magnitude,
+    __expand_magnitude,
+    Operator::Magnitude,
+    None,
+    f16,
+    bf16,
+    f32,
+    f64
+);
 impl_unary_func!(
     Normalize,
     normalize,

crates/cubecl-core/src/ir/operation.rs

@@ -82,6 +82,7 @@ pub enum Operator {
     AtomicOr(BinaryOperator),
     AtomicXor(BinaryOperator),
     AtomicCompareAndSwap(CompareAndSwapOperator),
+    Magnitude(UnaryOperator),
     Normalize(UnaryOperator),
 }
 

crates/cubecl-core/src/ir/processing.rs

@@ -223,6 +223,9 @@ impl ScopeProcessing {
                 Operator::AtomicXor(op) => {
                     sanitize_constant_scalar_ref_var(&mut op.rhs, &op.out);
                 }
+                Operator::Magnitude(op) => {
+                    sanitize_constant_scalar_ref_var(&mut op.input, &op.out);
+                }
                 Operator::Normalize(op) => {
                     sanitize_constant_scalar_ref_var(&mut op.input, &op.out);
                 }

crates/cubecl-core/src/ir/vectorization.rs

@@ -101,6 +101,7 @@ impl Operator {
             Operator::AtomicAnd(op) => Operator::AtomicAnd(op.vectorize(vectorization)),
             Operator::AtomicOr(op) => Operator::AtomicOr(op.vectorize(vectorization)),
             Operator::AtomicXor(op) => Operator::AtomicXor(op.vectorize(vectorization)),
+            Operator::Magnitude(op) => Operator::Magnitude(op.vectorize(vectorization)),
             Operator::Normalize(op) => Operator::Normalize(op.vectorize(vectorization)),
         }
     }

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

@@ -61,7 +61,7 @@ macro_rules! test_unary_impl {
                         CubeCount::Static(1, 1, 1),
                         CubeDim::new((input.len() / $input_vectorization as usize) as u32, 1, 1),
                         ArrayArg::from_raw_parts(&input_handle, input.len(), $input_vectorization),
-                        ArrayArg::from_raw_parts(&output_handle, input.len(), $out_vectorization),
+                        ArrayArg::from_raw_parts(&output_handle, $expected.len(), $out_vectorization),
                     )
                 };
 
@@ -72,6 +72,38 @@ macro_rules! test_unary_impl {
     };
 }
 
+test_unary_impl!(
+    test_magnitude,
+    F,
+    F::magnitude,
+    [
+        {
+            input_vectorization: 1,
+            out_vectorization: 1,
+            input: [-1., 23.1, -1.4, 5.1],
+            expected: [1., 23.1, 1.4, 5.1]
+        },
+        {
+            input_vectorization: 2,
+            out_vectorization: 1,
+            input: [-1., 0., 1., 5.],
+            expected: [1.0, 5.099]
+        },
+        {
+            input_vectorization: 4,
+            out_vectorization: 1,
+            input: [-1., 0., 1., 5.],
+            expected: [5.196]
+        },
+        {
+            input_vectorization: 4,
+            out_vectorization: 1,
+            input: [0., 0., 0., 0.],
+            expected: [0.]
+        }
+    ]
+);
+
 test_unary_impl!(
     test_normalize,
     F,
@@ -128,6 +160,7 @@ macro_rules! testgen_unary {
             }
 
             add_test!(test_normalize);
+            add_test!(test_magnitude);
         }
     };
 }

crates/cubecl-cuda/src/compiler/base.rs

@@ -554,6 +554,9 @@ impl CudaCompiler {
             gpu::Operator::Normalize(op) => {
                 instructions.push(Instruction::Normalize(self.compile_unary(op)))
             }
+            gpu::Operator::Magnitude(op) => {
+                instructions.push(Instruction::Magnitude(self.compile_unary(op)))
+            }
         };
     }
 

crates/cubecl-cuda/src/compiler/instruction.rs

@@ -140,6 +140,7 @@ pub enum Instruction {
         out: Variable,
     },
     Negate(UnaryInstruction),
+    Magnitude(UnaryInstruction),
     Normalize(UnaryInstruction),
 }
 
@@ -390,6 +391,7 @@ for ({i_ty} {i} = {start}; {i} {cmp} {end}; {increment}) {{
                 f.write_fmt(format_args!("{out} = !{input};\n"))
             }
             Instruction::Normalize(inst) => Normalize::format(f, &inst.input, &inst.out),
+            Instruction::Magnitude(inst) => Magnitude::format(f, &inst.input, &inst.out),
         }
     }
 }
@@ -478,6 +480,28 @@ impl Remainder {
     }
 }
 
+struct Magnitude;
+
+impl Magnitude {
+    fn format(
+        f: &mut core::fmt::Formatter<'_>,
+        input: &Variable,
+        out: &Variable,
+    ) -> core::fmt::Result {
+        let num = input.item().vectorization;
+        let elem = input.elem();
+
+        f.write_fmt(format_args!("{out} = 0.0;\n"))?;
+
+        for i in 0..num {
+            let input_i = input.index(i);
+            f.write_fmt(format_args!("{out} += {input_i} * {input_i};\n"))?;
+        }
+
+        Sqrt::format_unary(f, out, out, elem)
+    }
+}
+
 struct Normalize;
 
 impl Normalize {

crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs

@@ -754,6 +754,10 @@ impl WgslCompiler {
                 input: self.compile_variable(op.input),
                 out: self.compile_variable(op.out),
             },
+            cube::Operator::Magnitude(op) => wgsl::Instruction::Magnitude {
+                input: self.compile_variable(op.input),
+                out: self.compile_variable(op.out),
+            },
             cube::Operator::Normalize(op) => wgsl::Instruction::Normalize {
                 input: self.compile_variable(op.input),
                 out: self.compile_variable(op.out),

crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs

@@ -313,6 +313,10 @@ pub enum Instruction {
         input: Variable,
         out: Variable,
     },
+    Magnitude {
+        input: Variable,
+        out: Variable,
+    },
     Normalize {
         input: Variable,
         out: Variable,
@@ -689,6 +693,9 @@ for (var {i}: {i_ty} = {start}; {i} {cmp} {end}; {increment}) {{
                 "{out} = atomicCompareExchangeWeak({lhs}, {cmp}, {value}).old_value;\n"
             )),
             Instruction::Negate { input, out } => f.write_fmt(format_args!("{out} = -{input};\n")),
+            Instruction::Magnitude { input, out } => {
+                f.write_fmt(format_args!("{out} = length({input});\n"))
+            }
             Instruction::Normalize { input, out } => {
                 if input.item().vectorization_factor() == 1 {
                     // We need a check for vectorization factor 1 here, for compatibility with cuda.

examples/normalization/src/lib.rs

@@ -1,33 +1,44 @@
 use cubecl::prelude::*;
 
 #[cube(launch_unchecked)]
-fn norm_test<F: Float>(input: &Array<F>, output: &mut Array<F>) {
+fn norm_test<F: Float>(input: &Array<F>, output_a: &mut Array<F>, output_b: &mut Array<F>) {
     if ABSOLUTE_POS < input.len() {
-        output[ABSOLUTE_POS] = F::normalize(input[ABSOLUTE_POS]);
+        output_a[ABSOLUTE_POS] = F::normalize(input[ABSOLUTE_POS]);
+        output_b[ABSOLUTE_POS] = input[ABSOLUTE_POS] / F::magnitude(input[ABSOLUTE_POS]);
     }
 }
 
 pub fn launch<R: Runtime>(device: &R::Device) {
     let client = R::client(device);
     let input = &[-1., 0., 1., 5.];
-    let output_handle = client.empty(input.len() * core::mem::size_of::<f32>());
     let input_handle = client.create(f32::as_bytes(input));
+    let output_a_handle = client.empty(input.len() * core::mem::size_of::<f32>());
+    let output_b_handle = client.empty(input.len() * core::mem::size_of::<f32>());
 
     unsafe {
         norm_test::launch_unchecked::<f32, R>(
             &client,
             CubeCount::Static(1, 1, 1),
             CubeDim::new(input.len() as u32, 1, 1),
             ArrayArg::from_raw_parts(&input_handle, input.len(), 4),
-            ArrayArg::from_raw_parts(&output_handle, input.len(), 4),
+            ArrayArg::from_raw_parts(&output_a_handle, input.len(), 4),
+            ArrayArg::from_raw_parts(&output_b_handle, input.len(), 4),
         )
     };
 
-    let bytes = client.read(output_handle.binding());
+    let bytes = client.read(output_a_handle.binding());
     let output = f32::from_bytes(&bytes);
 
     println!(
         "Executed normalize with runtime {:?} => {output:?}",
         R::name()
     );
+
+    let bytes = client.read(output_b_handle.binding());
+    let output = f32::from_bytes(&bytes);
+
+    println!(
+        "Executed normalize using magnitude with runtime {:?} => {output:?}",
+        R::name()
+    );
 }