Roll out complex contract validation

Cargo.toml

@@ -98,6 +98,9 @@ half = { version = "2.5", features = [
     "num-traits",
     "serde",
 ], default-features = false }
+num-complex = { version = "0.4", default-features = false, features = [
+    "bytemuck",
+] }
 num-integer = { version = "0.1", default-features = false }
 num-traits = { version = "0.2.19", default-features = false, features = [
     "libm",

crates/cubecl-core/Cargo.toml

@@ -56,6 +56,7 @@ enumset = { workspace = true }
 float-ord = { workspace = true }
 half = { workspace = true, features = ["bytemuck"] }
 hashbrown = { workspace = true }
+num-complex = { workspace = true }
 itertools = { workspace = true }
 num-integer = { workspace = true }
 num-traits = { workspace = true }

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

@@ -235,7 +235,9 @@ where
     }
 }
 
-impl<P: Scalar + Abs, N: Size> Abs for Vector<P, N> {}
+impl<P: Scalar + Abs, N: Size> Abs for Vector<P, N> {
+    type AbsElem = P::AbsElem;
+}
 impl<P: Scalar + Log, N: Size> Log for Vector<P, N> {}
 impl<P: Scalar + Log1p, N: Size> Log1p for Vector<P, N> {}
 impl<P: Scalar + Expm1, N: Size> Expm1 for Vector<P, N> {}

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

@@ -655,6 +655,8 @@ from_const!(e4m3);
 from_const!(e5m2);
 from_const!(ue8m0);
 from_const!(bool);
+from_const!(num_complex::Complex<f32>);
+from_const!(num_complex::Complex<f64>);
 
 macro_rules! tuple_cube_type {
     ($($P:ident),*) => {

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

@@ -0,0 +1,173 @@
+use core::ops::{Add, Div, Mul, Neg, Sub};
+
+use crate::{
+    Runtime,
+    ir::{ComplexKind, ElemType, Scope, StorageType, Type, Value},
+    prelude::{
+        CubeDebug, CubePrimitive, CubeType, IntoExpand, IntoRuntime, KernelBuilder, KernelLauncher,
+        LaunchArg, NativeAssign, NativeExpand, Scalar, ScalarArgSettings, impl_scalar_launch,
+    },
+    unexpanded,
+};
+use cubecl_ir::{Arithmetic, ConstantValue, Operator, features::ComplexUsage};
+use cubecl_runtime::client::ComputeClient;
+
+use crate::frontend::{
+    Abs,
+    operation::{unary_expand, unary_expand_fixed_output},
+};
+
+pub trait ComplexCore:
+    CubePrimitive
+    + Add<Output = Self>
+    + Sub<Output = Self>
+    + Mul<Output = Self>
+    + Div<Output = Self>
+    + Neg<Output = Self>
+    + Copy
+    + Clone
+    + PartialEq
+    + core::fmt::Debug
+    + Send
+    + Sync
+    + 'static
+{
+    type FloatElem: Scalar;
+
+    fn conj(self) -> Self {
+        unexpanded!()
+    }
+
+    fn real_val(self) -> Self::FloatElem {
+        unexpanded!()
+    }
+
+    fn imag_val(self) -> Self::FloatElem {
+        unexpanded!()
+    }
+
+    fn supported_complex_uses<R: Runtime>(
+        client: &ComputeClient<R>,
+    ) -> enumset::EnumSet<ComplexUsage> {
+        client
+            .properties()
+            .complex_usage(Self::as_type_native_unchecked().storage_type())
+    }
+}
+
+pub trait ComplexCompare: ComplexCore {}
+
+pub trait ComplexMath:
+    ComplexCore
+    + Abs<AbsElem = Self::FloatElem>
+    + crate::frontend::Exp
+    + crate::frontend::Log
+    + crate::frontend::Sin
+    + crate::frontend::Cos
+    + crate::frontend::Sqrt
+    + crate::frontend::Tanh
+    + crate::frontend::Powf
+{
+}
+
+pub trait ComplexCoreExpand {
+    fn __expand_conj_method(self, scope: &Scope) -> Self;
+    fn __expand_real_val_method(
+        self,
+        scope: &Scope,
+    ) -> NativeExpand<<Self as ComplexCoreExpand>::FloatElem>;
+    fn __expand_imag_val_method(
+        self,
+        scope: &Scope,
+    ) -> NativeExpand<<Self as ComplexCoreExpand>::FloatElem>;
+
+    type FloatElem: Scalar;
+}
+
+impl<T: ComplexCore> ComplexCoreExpand for NativeExpand<T> {
+    type FloatElem = T::FloatElem;
+
+    fn __expand_conj_method(self, scope: &Scope) -> Self {
+        unary_expand(scope, self.into(), Arithmetic::Conj).into()
+    }
+
+    fn __expand_real_val_method(self, scope: &Scope) -> NativeExpand<T::FloatElem> {
+        let expand_element: Value = self.into();
+        let item = <T::FloatElem as CubePrimitive>::__expand_as_type(scope);
+        unary_expand_fixed_output(scope, expand_element, item, Operator::Real).into()
+    }
+
+    fn __expand_imag_val_method(self, scope: &Scope) -> NativeExpand<T::FloatElem> {
+        let expand_element: Value = self.into();
+        let item = <T::FloatElem as CubePrimitive>::__expand_as_type(scope);
+        unary_expand_fixed_output(scope, expand_element, item, Operator::Imag).into()
+    }
+}
+
+macro_rules! impl_complex {
+    ($primitive:ty, $kind:ident, $float:ty) => {
+        impl CubeType for $primitive {
+            type ExpandType = NativeExpand<$primitive>;
+        }
+
+        impl CubeDebug for $primitive {}
+
+        impl CubePrimitive for $primitive {
+            type Scalar = Self;
+            type Size = crate::prelude::Const<1>;
+            type WithScalar<S: Scalar> = S;
+
+            fn as_type_native() -> Option<Type> {
+                Some(StorageType::Scalar(ElemType::Complex(ComplexKind::$kind)).into())
+            }
+
+            fn from_const_value(value: ConstantValue) -> Self {
+                let ConstantValue::Complex(re, im) = value else {
+                    unreachable!("expected Complex constant")
+                };
+                <$primitive>::new(re as $float, im as $float)
+            }
+        }
+
+        impl IntoRuntime for $primitive {
+            fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
+                self.into()
+            }
+        }
+
+        impl IntoExpand for $primitive {
+            type Expand = NativeExpand<$primitive>;
+
+            fn into_expand(self, _scope: &Scope) -> Self::Expand {
+                self.into()
+            }
+        }
+
+        impl NativeAssign for $primitive {}
+
+        impl crate::prelude::IntoMut for $primitive {
+            fn into_mut(self, _scope: &Scope) -> Self {
+                self
+            }
+        }
+
+        impl_scalar_launch!($primitive);
+
+        impl Scalar for $primitive {}
+
+        impl Abs for $primitive {
+            type AbsElem = $float;
+        }
+
+        impl ComplexCore for $primitive {
+            type FloatElem = $float;
+        }
+
+        impl ComplexCompare for $primitive {}
+
+        impl ComplexMath for $primitive {}
+    };
+}
+
+impl_complex!(num_complex::Complex<f32>, C32, f32);
+impl_complex!(num_complex::Complex<f64>, C64, f64);

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

@@ -2,6 +2,7 @@ mod atomic;
 mod base;
 mod bool;
 mod cast;
+mod complex;
 mod cube_elem;
 mod float;
 mod int;
@@ -13,6 +14,7 @@ pub use atomic::*;
 pub use base::*;
 pub use bool::*;
 pub use cast::*;
+pub use complex::*;
 pub use cube_elem::*;
 pub use float::*;
 pub use int::*;

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

@@ -20,7 +20,7 @@ use super::{LaunchArg, NativeAssign, NativeExpand};
 /// Used in kernels that should work for both.
 pub trait Numeric:
     Copy
-    + Abs
+    + Abs<AbsElem = Self>
     + VectorSum
     + ModFloor
     + Scalar

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

@@ -115,6 +115,7 @@ impl<Marker: 'static> Neg for DynamicScalar<Marker> {
             ConstantValue::Float(val) => (-val).into(),
             ConstantValue::UInt(val) => (-(val as i64)).into(),
             ConstantValue::Bool(val) => (!val).into(),
+            ConstantValue::Complex(_, _) => panic!("Complex values aren't supported"),
         })
     }
 }
@@ -315,7 +316,9 @@ impl<Marker: 'static> Recip for DynamicScalar<Marker> {}
 impl<Marker: 'static> Erf for DynamicScalar<Marker> {}
 impl<Marker: 'static> Exp for DynamicScalar<Marker> {}
 impl<Marker: 'static> ModFloor for DynamicScalar<Marker> {}
-impl<Marker: 'static> Abs for DynamicScalar<Marker> {}
+impl<Marker: 'static> Abs for DynamicScalar<Marker> {
+    type AbsElem = Self;
+}
 impl<Marker: 'static> Log for DynamicScalar<Marker> {}
 impl<Marker: 'static> Log1p for DynamicScalar<Marker> {}
 impl<Marker: 'static> Expm1 for DynamicScalar<Marker> {}
@@ -430,6 +433,7 @@ impl<Marker: 'static> Not for DynamicScalar<Marker> {
             ConstantValue::UInt(val) => (!val).into(),
             ConstantValue::Bool(val) => (!val).into(),
             ConstantValue::Float(val) => f64::from_bits(!val.to_bits()).into(),
+            ConstantValue::Complex(_, _) => panic!("Complex values aren't supported"),
         })
     }
 }
@@ -481,6 +485,7 @@ impl<Marker: 'static> Shl for DynamicScalar<Marker> {
             ConstantValue::Float(val) => f64::from_bits(val.to_bits() << rhs.val.as_u64()).into(),
             ConstantValue::UInt(val) => (val << rhs.val.as_u64()).into(),
             ConstantValue::Bool(_) => panic!("Invalid value"),
+            ConstantValue::Complex(_, _) => panic!("Complex values aren't supported"),
         })
     }
 }
@@ -493,6 +498,7 @@ impl<Marker: 'static> Shr for DynamicScalar<Marker> {
             ConstantValue::Float(val) => f64::from_bits(val.to_bits() >> rhs.val.as_u64()).into(),
             ConstantValue::UInt(val) => (val >> rhs.val.as_u64()).into(),
             ConstantValue::Bool(_) => panic!("Invalid value"),
+            ConstantValue::Complex(_, _) => panic!("Complex values aren't supported"),
         })
     }
 }
@@ -504,6 +510,7 @@ impl<Marker: 'static> ShrAssign<u32> for DynamicScalar<Marker> {
             ConstantValue::Float(val) => f64::from_bits(val.to_bits() >> rhs).into(),
             ConstantValue::UInt(val) => (val >> rhs).into(),
             ConstantValue::Bool(_) => panic!("Invalid value"),
+            ConstantValue::Complex(_, _) => panic!("Complex values aren't supported"),
         });
     }
 }
@@ -515,6 +522,7 @@ impl<Marker: 'static> ShlAssign<u32> for DynamicScalar<Marker> {
             ConstantValue::Float(val) => f64::from_bits(val.to_bits() << rhs).into(),
             ConstantValue::UInt(val) => (val << rhs).into(),
             ConstantValue::Bool(_) => panic!("Invalid value"),
+            ConstantValue::Complex(_, _) => panic!("Complex values aren't supported"),
         });
     }
 }
@@ -835,3 +843,11 @@ impl<Marker: 'static> Zero for DynamicScalar<Marker> {
         self.val.is_zero()
     }
 }
+
+impl<Marker: 'static> ComplexCore for DynamicScalar<Marker> {
+    type FloatElem = Self;
+}
+
+impl<Marker: 'static> ComplexCompare for DynamicScalar<Marker> {}
+
+impl<Marker: 'static> ComplexMath for DynamicScalar<Marker> {}