From a86746d34ecc0dbbc8ac3970578328b3b004174f Mon Sep 17 00:00:00 2001 From: GiggleLiu Date: Sat, 29 Aug 2026 16:29:34 +0800 Subject: [PATCH 1/4] fix(cpp): lower CUDA complex casts with cuComplex helpers --- crates/cubecl-cpp/src/shared/unary.rs | 23 +++++++-- crates/cubecl-cpp/src/shared/unary/tests.rs | 52 +++++++++++++++++++++ 2 files changed, 72 insertions(+), 3 deletions(-) create mode 100644 crates/cubecl-cpp/src/shared/unary/tests.rs diff --git a/crates/cubecl-cpp/src/shared/unary.rs b/crates/cubecl-cpp/src/shared/unary.rs index ed8372fc3f..d583ab4cf9 100644 --- a/crates/cubecl-cpp/src/shared/unary.rs +++ b/crates/cubecl-cpp/src/shared/unary.rs @@ -413,9 +413,23 @@ impl Unary 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, 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) => 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 +497,6 @@ impl Unary for IsInf { true } } + +#[cfg(test)] +mod tests; diff --git a/crates/cubecl-cpp/src/shared/unary/tests.rs b/crates/cubecl-cpp/src/shared/unary/tests.rs new file mode 100644 index 0000000000..513d4fa483 --- /dev/null +++ b/crates/cubecl-cpp/src/shared/unary/tests.rs @@ -0,0 +1,52 @@ +use crate::{ + cuda::{CudaDialect, mma::PtxWmmaCompiler}, + shared::{Elem, Instruction, Item, UnaryInstruction, Variable}, +}; +use cubecl_core::ir::ConstantValue; + +type Dialect = CudaDialect; + +fn cast_variable(input: Variable, output: Elem) -> String { + Instruction::Assign(UnaryInstruction { + input, + out: Variable::Named { + name: "output", + item: Item::scalar(output, false), + }, + }) + .to_string() +} + +fn cast(input: Elem, output: Elem) -> 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"); +} From f274a76d8589d39fbbbcb4916b4f354f1c2447db Mon Sep 17 00:00:00 2001 From: GiggleLiu Date: Sat, 29 Aug 2026 16:45:18 +0800 Subject: [PATCH 2/4] fix(core): satisfy drain-collect lint --- crates/cubecl-core/src/post_processing/unroll.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/crates/cubecl-core/src/post_processing/unroll.rs b/crates/cubecl-core/src/post_processing/unroll.rs index 92490bea66..ce19ac21ae 100644 --- a/crates/cubecl-core/src/post_processing/unroll.rs +++ b/crates/cubecl-core/src/post_processing/unroll.rs @@ -534,19 +534,19 @@ 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, ); } @@ -554,27 +554,27 @@ impl UnrollProcessor { 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, ); } From 5939d8e3b1c479ed6db31de5f01cb6d432178f2d Mon Sep 17 00:00:00 2001 From: GiggleLiu Date: Sat, 29 Aug 2026 16:55:28 +0800 Subject: [PATCH 3/4] fix(std): remove redundant prelude import --- crates/cubecl-std/src/tensor/view/launch.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/cubecl-std/src/tensor/view/launch.rs b/crates/cubecl-std/src/tensor/view/launch.rs index 3bf02c1368..3acf8ed250 100644 --- a/crates/cubecl-std/src/tensor/view/launch.rs +++ b/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}, }; From 5d664af894505c9fab9b4f6cd4959ae5796e4d79 Mon Sep 17 00:00:00 2001 From: GiggleLiu Date: Sat, 29 Aug 2026 23:27:27 +0800 Subject: [PATCH 4/4] fix(cpp): preserve complex truthiness in bool casts --- crates/cubecl-cpp/src/shared/unary.rs | 6 +++++- crates/cubecl-cpp/src/shared/unary/tests.rs | 14 +++++++++++++- crates/cubecl-ir/src/variable.rs | 5 ++++- 3 files changed, 22 insertions(+), 3 deletions(-) diff --git a/crates/cubecl-cpp/src/shared/unary.rs b/crates/cubecl-cpp/src/shared/unary.rs index d583ab4cf9..bf07b28095 100644 --- a/crates/cubecl-cpp/src/shared/unary.rs +++ b/crates/cubecl-cpp/src/shared/unary.rs @@ -414,7 +414,8 @@ impl Unary for Assign { // Cast only when necessary. if elem != input.elem() { // cuComplex types are C structs without converting constructors. Complex-to-real - // casts use the real component, matching `ConstantValue::cast_to`. + // 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)") @@ -427,6 +428,9 @@ impl Unary for Assign { (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})"), diff --git a/crates/cubecl-cpp/src/shared/unary/tests.rs b/crates/cubecl-cpp/src/shared/unary/tests.rs index 513d4fa483..037c229782 100644 --- a/crates/cubecl-cpp/src/shared/unary/tests.rs +++ b/crates/cubecl-cpp/src/shared/unary/tests.rs @@ -2,7 +2,7 @@ use crate::{ cuda::{CudaDialect, mma::PtxWmmaCompiler}, shared::{Elem, Instruction, Item, UnaryInstruction, Variable}, }; -use cubecl_core::ir::ConstantValue; +use cubecl_core::ir::{ConstantValue, ElemType}; type Dialect = CudaDialect; @@ -49,4 +49,16 @@ fn cuda_complex_casts_use_cucomplex_components_and_constructors() { "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) + ); } diff --git a/crates/cubecl-ir/src/variable.rs b/crates/cubecl-ir/src/variable.rs index 646da93efa..a0fd376a48 100644 --- a/crates/cubecl-ir/src/variable.rs +++ b/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)