fix(cpp): lower CUDA complex casts with cuComplex helpers

crates/cubecl-core/src/post_processing/unroll.rs

@@ -534,47 +534,47 @@ impl UnrollProcessor {
                     Branch::If(op) => {
                         op.scope.instructions = self.transform_instructions(
                             allocator,
-                            op.scope.instructions.drain(..).collect(),
+                            core::mem::take(&mut op.scope.instructions),
                             mappings,
                         );
                     }
                     Branch::IfElse(op) => {
                         op.scope_if.instructions = self.transform_instructions(
                             allocator,
-                            op.scope_if.instructions.drain(..).collect(),
+                            core::mem::take(&mut op.scope_if.instructions),
                             mappings,
                         );
                         op.scope_else.instructions = self.transform_instructions(
                             allocator,
-                            op.scope_else.instructions.drain(..).collect(),
+                            core::mem::take(&mut op.scope_else.instructions),
                             mappings,
                         );
                     }
                     Branch::Switch(op) => {
                         for (_, case) in &mut op.cases {
                             case.instructions = self.transform_instructions(
                                 allocator,
-                                case.instructions.drain(..).collect(),
+                                core::mem::take(&mut case.instructions),
                                 mappings,
                             );
                         }
                         op.scope_default.instructions = self.transform_instructions(
                             allocator,
-                            op.scope_default.instructions.drain(..).collect(),
+                            core::mem::take(&mut op.scope_default.instructions),
                             mappings,
                         );
                     }
                     Branch::RangeLoop(op) => {
                         op.scope.instructions = self.transform_instructions(
                             allocator,
-                            op.scope.instructions.drain(..).collect(),
+                            core::mem::take(&mut op.scope.instructions),
                             mappings,
                         );
                     }
                     Branch::Loop(op) => {
                         op.scope.instructions = self.transform_instructions(
                             allocator,
-                            op.scope.instructions.drain(..).collect(),
+                            core::mem::take(&mut op.scope.instructions),
                             mappings,
                         );
                     }

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

@@ -413,9 +413,27 @@ impl<D: Dialect> Unary<D> for Assign {
     {
         // Cast only when necessary.
         if elem != input.elem() {
-            match elem {
-                Elem::TF32 => write!(f, "nvcuda::wmma::__float_to_tf32({input})"),
-                elem => write!(f, "{elem}({input})"),
+            // cuComplex types are C structs without converting constructors. Complex-to-real
+            // casts use the real component, while complex-to-bool checks both components,
+            // matching `ConstantValue::cast_to`.
+            match (input.elem(), elem) {
+                (Elem::CF32 | Elem::CF64, Elem::CF32) => {
+                    write!(f, "make_cuFloatComplex({input}.x, {input}.y)")
+                }
+                (Elem::CF32 | Elem::CF64, Elem::CF64) => {
+                    write!(f, "make_cuDoubleComplex({input}.x, {input}.y)")
+                }
+                (_, Elem::CF32) => write!(f, "make_cuFloatComplex({input}, 0.0f)"),
+                (_, Elem::CF64) => write!(f, "make_cuDoubleComplex({input}, 0.0)"),
+                (Elem::CF32 | Elem::CF64, Elem::TF32) => {
+                    write!(f, "nvcuda::wmma::__float_to_tf32({input}.x)")
+                }
+                (Elem::CF32 | Elem::CF64, Elem::Bool) => {
+                    write!(f, "({input}.x != 0 || {input}.y != 0)")
+                }
+                (Elem::CF32 | Elem::CF64, elem) => write!(f, "{elem}({input}.x)"),
+                (_, Elem::TF32) => write!(f, "nvcuda::wmma::__float_to_tf32({input})"),
+                (_, elem) => write!(f, "{elem}({input})"),
             }
         } else {
             write!(f, "{input}")
@@ -483,3 +501,6 @@ impl<D: Dialect> Unary<D> for IsInf {
         true
     }
 }
+
+#[cfg(test)]
+mod tests;

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

@@ -0,0 +1,64 @@
+use crate::{
+    cuda::{CudaDialect, mma::PtxWmmaCompiler},
+    shared::{Elem, Instruction, Item, UnaryInstruction, Variable},
+};
+use cubecl_core::ir::{ConstantValue, ElemType};
+
+type Dialect = CudaDialect<PtxWmmaCompiler>;
+
+fn cast_variable(input: Variable<Dialect>, output: Elem<Dialect>) -> String {
+    Instruction::Assign(UnaryInstruction {
+        input,
+        out: Variable::Named {
+            name: "output",
+            item: Item::scalar(output, false),
+        },
+    })
+    .to_string()
+}
+
+fn cast(input: Elem<Dialect>, output: Elem<Dialect>) -> String {
+    cast_variable(
+        Variable::Named {
+            name: "input",
+            item: Item::scalar(input, false),
+        },
+        output,
+    )
+}
+
+#[test]
+fn cuda_complex_casts_use_cucomplex_components_and_constructors() {
+    assert_eq!(
+        cast_variable(
+            Variable::Constant(ConstantValue::UInt(0), Item::scalar(Elem::U32, false)),
+            Elem::CF64,
+        ),
+        "output = make_cuDoubleComplex(uint32(0), 0.0);\n"
+    );
+    assert_eq!(
+        cast(Elem::F64, Elem::CF32),
+        "output = make_cuFloatComplex(input, 0.0f);\n"
+    );
+    assert_eq!(
+        cast(Elem::CF64, Elem::CF32),
+        "output = make_cuFloatComplex(input.x, input.y);\n"
+    );
+    assert_eq!(
+        cast(Elem::CF32, Elem::CF64),
+        "output = make_cuDoubleComplex(input.x, input.y);\n"
+    );
+    assert_eq!(cast(Elem::CF32, Elem::F64), "output = double(input.x);\n");
+    assert_eq!(
+        cast(Elem::CF32, Elem::Bool),
+        "output = (input.x != 0 || input.y != 0);\n"
+    );
+    assert_eq!(
+        ConstantValue::Complex(0.0, 1.0).cast_to(ElemType::Bool),
+        ConstantValue::Bool(true)
+    );
+    assert_eq!(
+        ConstantValue::Complex(0.0, 0.0).cast_to(ElemType::Bool),
+        ConstantValue::Bool(false)
+    );
+}

crates/cubecl-ir/src/variable.rs

@@ -505,7 +505,10 @@ impl ConstantValue {
                     }
                 }
                 .into(),
-                ElemType::Bool => self.as_bool().into(),
+                ElemType::Bool => match self {
+                    ConstantValue::Complex(re, im) => (*re != 0.0 || *im != 0.0).into(),
+                    _ => self.as_bool().into(),
+                },
                 ElemType::Complex(kind) => match (self, kind) {
                     (ConstantValue::Complex(re, im), ComplexKind::C32) => {
                         ConstantValue::Complex(*re as f32 as f64, *im as f32 as f64)

crates/cubecl-std/src/tensor/view/launch.rs

@@ -15,7 +15,6 @@ mod layout {
         self as cubecl,
         format::DebugRaw,
         hash::{StableHash, StableHasher},
-        prelude::*,
         zspace::{Shape, Strides, metadata::Metadata},
     };