1use core::ops::{Add, Div, Mul, Neg, Sub};
2
3use cubecl_ir::{
4 ComplexKind, ConstantValue, ElemType, ExpandValue, FloatKind, Scope,
5 dialect::{
6 cmp::{FEqualOp, FNotEqualOp},
7 math::*,
8 },
9 features::ComplexUsage,
10 interfaces::TypedExt,
11 pliron::{
12 builtin::op_interfaces::OneResultInterface, context::Context, op::Op, r#type::TypeHandle,
13 value::Value,
14 },
15 types::{
16 VectorType,
17 scalar::{Complex32Type, Complex64Type},
18 },
19};
20use cubecl_runtime::client::Client;
21
22use crate::{
23 frontend::{
24 ConstantTyped, Cos, CosNativeExpand, Exp, ExpNativeExpand, Log, LogNativeExpand, Powf,
25 PowfNativeExpand, ScalarArgSettings, Sin, SinNativeExpand, Sqrt, SqrtNativeExpand, Tanh,
26 TanhNativeExpand,
27 operation::{
28 AddNativeExpand, DivNativeExpand, MulNativeExpand, NegNativeExpand,
29 PartialEqNativeExpand, SubNativeExpand, binary_expand, unary_expand,
30 },
31 require_complex_usage,
32 },
33 prelude::{
34 CubeDebug, CubePrimitive, CubeType, IntoExpand, IntoRuntime, KernelBuilder, KernelLauncher,
35 LaunchArg, NativeAssign, NativeExpand, Scalar, impl_scalar_launch,
36 },
37 unexpanded,
38};
39
40pub trait ComplexCore:
41 Scalar
42 + IntoRuntime
43 + CubePrimitive<
44 Scalar: ComplexNativeExpand<FloatElem = Self::FloatElem>
45 + AddNativeExpand
46 + SubNativeExpand
47 + MulNativeExpand
48 + DivNativeExpand
49 + NegNativeExpand
50 + PartialEqNativeExpand,
51 > + Add<Output = Self>
52 + Sub<Output = Self>
53 + Mul<Output = Self>
54 + Div<Output = Self>
55 + Neg<Output = Self>
56 + Copy
57 + Clone
58 + PartialEq
59 + core::fmt::Debug
60 + Send
61 + Sync
62 + 'static
63{
64 type FloatElem: Scalar;
65
66 fn conj(self) -> Self {
67 unexpanded!()
68 }
69
70 fn real_val(self) -> Self::FloatElem {
71 unexpanded!()
72 }
73
74 fn imag_val(self) -> Self::FloatElem {
75 unexpanded!()
76 }
77
78 fn supported_complex_uses(client: &Client) -> cubecl_ir::EnumSet<ComplexUsage> {
79 client.properties().complex_usage(Self::elem_type_native())
80 }
81}
82
83pub trait ComplexCompare: ComplexCore {}
84
85pub trait ComplexAbs:
86 ComplexCore + CubePrimitive<Scalar: ComplexAbsNativeExpand<FloatElem = Self::FloatElem>>
87{
88 fn abs(self) -> Self::FloatElem {
89 unexpanded!()
90 }
91
92 fn norm(self) -> Self::FloatElem {
93 unexpanded!()
94 }
95
96 fn __expand_abs(scope: &Scope, x: NativeExpand<Self>) -> NativeExpand<Self::FloatElem> {
97 x.__expand_abs_method(scope)
98 }
99
100 fn __expand_norm(scope: &Scope, x: NativeExpand<Self>) -> NativeExpand<Self::FloatElem> {
101 x.__expand_norm_method(scope)
102 }
103}
104
105pub trait ComplexMath: ComplexAbs + Exp + Log + Sin + Cos + Sqrt + Tanh + Powf {}
106
107pub trait ComplexAbsExpand {
108 type FloatElem: Scalar;
109 fn __expand_abs_method(self, scope: &Scope) -> NativeExpand<Self::FloatElem>;
110 fn __expand_norm_method(self, scope: &Scope) -> NativeExpand<Self::FloatElem>;
111}
112
113pub trait ComplexAbsNativeExpand {
114 type FloatElem: Scalar;
115 fn __expand_native_abs(scope: &Scope, input: ExpandValue) -> ExpandValue;
116}
117
118impl<T: ComplexAbs> ComplexAbsExpand for NativeExpand<T> {
119 type FloatElem = T::FloatElem;
120
121 fn __expand_abs_method(self, scope: &Scope) -> NativeExpand<T::FloatElem> {
122 T::Scalar::__expand_native_abs(scope, self.into()).into()
123 }
124
125 fn __expand_norm_method(self, scope: &Scope) -> NativeExpand<T::FloatElem> {
126 T::Scalar::__expand_native_abs(scope, self.into()).into()
127 }
128}
129
130pub trait ComplexNativeExpand {
131 type FloatElem: Scalar;
132 fn __expand_native_conj(scope: &Scope, input: ExpandValue) -> ExpandValue;
133 fn __expand_native_real(scope: &Scope, input: ExpandValue) -> ExpandValue;
134 fn __expand_native_imag(scope: &Scope, input: ExpandValue) -> ExpandValue;
135}
136
137pub trait ComplexCoreExpand {
138 type FloatElem: Scalar;
139 fn __expand_conj_method(self, scope: &Scope) -> Self;
140 fn __expand_real_val_method(self, scope: &Scope) -> NativeExpand<Self::FloatElem>;
141 fn __expand_imag_val_method(self, scope: &Scope) -> NativeExpand<Self::FloatElem>;
142}
143
144impl<T: ComplexCore> ComplexCoreExpand for NativeExpand<T> {
145 type FloatElem = T::FloatElem;
146
147 fn __expand_conj_method(self, scope: &Scope) -> Self {
148 T::Scalar::__expand_native_conj(scope, self.into()).into()
149 }
150
151 fn __expand_real_val_method(self, scope: &Scope) -> NativeExpand<T::FloatElem> {
152 T::Scalar::__expand_native_real(scope, self.into()).into()
153 }
154
155 fn __expand_imag_val_method(self, scope: &Scope) -> NativeExpand<T::FloatElem> {
156 T::Scalar::__expand_native_imag(scope, self.into()).into()
157 }
158}
159
160fn complex_component<O>(
161 scope: &Scope,
162 input: ExpandValue,
163 out_scalar: TypeHandle,
164 op: impl FnOnce(&mut Context, Value, TypeHandle) -> O,
165) -> ExpandValue
166where
167 O: Op + OneResultInterface,
168{
169 let input = input.read_value(scope);
170 let vector_size = input.vector_size(scope.ctx());
171 let out_ty = if vector_size == 1 {
172 out_scalar
173 } else {
174 VectorType::get(scope.ctx(), out_scalar, vector_size).into()
175 };
176 let operation = op(scope.ctx_mut(), input, out_ty);
177 scope.register_with_result(&operation).into()
178}
179
180macro_rules! impl_complex_unary {
181 ($primitive:ty, $trait:ident, $native:ident, $method:ident, $op:ty, $name:literal) => {
182 impl $trait for $primitive {}
183 impl $native for $primitive {
184 fn $method(scope: &Scope, input: ExpandValue) -> ExpandValue {
185 require_complex_usage(scope, Self::elem_type_native(), ComplexUsage::Math, $name);
186 unary_expand(scope, input, <$op>::new)
187 }
188 }
189 };
190}
191
192macro_rules! impl_complex {
193 ($primitive:ty, $kind:ident, $float:ty, $ir_ty:ty, $float_kind:ident) => {
194 impl CubeType for $primitive {
195 type ExpandType = NativeExpand<Self>;
196 }
197
198 impl CubeDebug for $primitive {}
199
200 impl Scalar for $primitive {
201 fn elem_type_native() -> ElemType {
202 ElemType::Complex(ComplexKind::$kind)
203 }
204 }
205
206 impl CubePrimitive for $primitive {
207 type Scalar = Self;
208 type Size = crate::prelude::Const<1>;
209 type WithScalar<S: Scalar> = S;
210
211 fn from_const_value(value: ConstantValue) -> Self {
212 let ConstantValue::Complex(re, im) = value else {
213 unreachable!("expected complex constant")
214 };
215 <$primitive>::new(re as $float, im as $float)
216 }
217 }
218
219 impl ConstantTyped for $primitive {
220 fn __expand_as_type(scope: &Scope) -> TypeHandle {
221 <$ir_ty>::get(scope.ctx()).into()
222 }
223 }
224
225 impl IntoRuntime for $primitive {
226 fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
227 self.into()
228 }
229 }
230
231 impl IntoExpand for $primitive {
232 type Expand = NativeExpand<Self>;
233 fn into_expand(self, _scope: &Scope) -> Self::Expand {
234 self.into()
235 }
236 }
237
238 impl NativeAssign for $primitive {}
239 impl_scalar_launch!($primitive);
240
241 impl ComplexNativeExpand for $primitive {
242 type FloatElem = $float;
243
244 fn __expand_native_conj(scope: &Scope, input: ExpandValue) -> ExpandValue {
245 require_complex_usage(scope, Self::elem_type_native(), ComplexUsage::Core, "conj");
246 unary_expand(scope, input, CConjOp::new)
247 }
248
249 fn __expand_native_real(scope: &Scope, input: ExpandValue) -> ExpandValue {
250 require_complex_usage(
251 scope,
252 Self::elem_type_native(),
253 ComplexUsage::Core,
254 "real_val",
255 );
256 complex_component(
257 scope,
258 input,
259 FloatKind::$float_kind.to_type(scope.ctx()),
260 |ctx, input, ty| CRealOp::new(ctx, ty, input),
261 )
262 }
263
264 fn __expand_native_imag(scope: &Scope, input: ExpandValue) -> ExpandValue {
265 require_complex_usage(
266 scope,
267 Self::elem_type_native(),
268 ComplexUsage::Core,
269 "imag_val",
270 );
271 complex_component(
272 scope,
273 input,
274 FloatKind::$float_kind.to_type(scope.ctx()),
275 |ctx, input, ty| CImagOp::new(ctx, ty, input),
276 )
277 }
278 }
279
280 impl ComplexCore for $primitive {
281 type FloatElem = $float;
282 }
283 impl ComplexCompare for $primitive {}
284 impl ComplexMath for $primitive {}
285
286 impl AddNativeExpand for $primitive {
287 fn __expand_native_add(
288 scope: &Scope,
289 lhs: ExpandValue,
290 rhs: ExpandValue,
291 ) -> ExpandValue {
292 require_complex_usage(scope, Self::elem_type_native(), ComplexUsage::Core, "+");
293 binary_expand(scope, lhs, rhs, FAddOp::new)
294 }
295 }
296 impl SubNativeExpand for $primitive {
297 fn __expand_native_sub(
298 scope: &Scope,
299 lhs: ExpandValue,
300 rhs: ExpandValue,
301 ) -> ExpandValue {
302 require_complex_usage(scope, Self::elem_type_native(), ComplexUsage::Core, "-");
303 binary_expand(scope, lhs, rhs, FSubOp::new)
304 }
305 }
306 impl MulNativeExpand for $primitive {
307 fn __expand_native_mul(
308 scope: &Scope,
309 lhs: ExpandValue,
310 rhs: ExpandValue,
311 ) -> ExpandValue {
312 require_complex_usage(scope, Self::elem_type_native(), ComplexUsage::Core, "*");
313 binary_expand(scope, lhs, rhs, FMulOp::new)
314 }
315 }
316 impl DivNativeExpand for $primitive {
317 fn __expand_native_div(
318 scope: &Scope,
319 lhs: ExpandValue,
320 rhs: ExpandValue,
321 ) -> ExpandValue {
322 require_complex_usage(scope, Self::elem_type_native(), ComplexUsage::Core, "/");
323 binary_expand(scope, lhs, rhs, FDivOp::new)
324 }
325 }
326 impl NegNativeExpand for $primitive {
327 fn __expand_native_neg(scope: &Scope, input: ExpandValue) -> ExpandValue {
328 require_complex_usage(scope, Self::elem_type_native(), ComplexUsage::Core, "neg");
329 unary_expand(scope, input, FNegOp::new)
330 }
331 }
332 impl PartialEqNativeExpand for $primitive {
333 fn __expand_native_eq(
334 scope: &Scope,
335 lhs: ExpandValue,
336 rhs: ExpandValue,
337 ) -> ExpandValue {
338 require_complex_usage(scope, Self::elem_type_native(), ComplexUsage::Compare, "==");
339 binary_expand(scope, lhs, rhs, FEqualOp::new)
340 }
341 fn __expand_native_ne(
342 scope: &Scope,
343 lhs: ExpandValue,
344 rhs: ExpandValue,
345 ) -> ExpandValue {
346 require_complex_usage(scope, Self::elem_type_native(), ComplexUsage::Compare, "!=");
347 binary_expand(scope, lhs, rhs, FNotEqualOp::new)
348 }
349 }
350
351 impl ComplexAbs for $primitive {}
352 impl ComplexAbsNativeExpand for $primitive {
353 type FloatElem = $float;
354 fn __expand_native_abs(scope: &Scope, input: ExpandValue) -> ExpandValue {
355 require_complex_usage(scope, Self::elem_type_native(), ComplexUsage::Math, "abs");
356 complex_component(
357 scope,
358 input,
359 FloatKind::$float_kind.to_type(scope.ctx()),
360 |ctx, input, ty| CAbsOp::new(ctx, ty, input),
361 )
362 }
363 }
364
365 impl_complex_unary!(
366 $primitive,
367 Exp,
368 ExpNativeExpand,
369 __expand_native_exp,
370 ExpOp,
371 "exp"
372 );
373 impl_complex_unary!(
374 $primitive,
375 Log,
376 LogNativeExpand,
377 __expand_native_ln,
378 LogOp,
379 "log"
380 );
381 impl_complex_unary!(
382 $primitive,
383 Sin,
384 SinNativeExpand,
385 __expand_native_sin,
386 SinOp,
387 "sin"
388 );
389 impl_complex_unary!(
390 $primitive,
391 Cos,
392 CosNativeExpand,
393 __expand_native_cos,
394 CosOp,
395 "cos"
396 );
397 impl_complex_unary!(
398 $primitive,
399 Sqrt,
400 SqrtNativeExpand,
401 __expand_native_sqrt,
402 SqrtOp,
403 "sqrt"
404 );
405 impl_complex_unary!(
406 $primitive,
407 Tanh,
408 TanhNativeExpand,
409 __expand_native_tanh,
410 TanhOp,
411 "tanh"
412 );
413
414 impl Powf for $primitive {}
415 impl PowfNativeExpand for $primitive {
416 fn __expand_native_powf(
417 scope: &Scope,
418 lhs: ExpandValue,
419 rhs: ExpandValue,
420 ) -> ExpandValue {
421 require_complex_usage(scope, Self::elem_type_native(), ComplexUsage::Math, "powf");
422 binary_expand(scope, lhs, rhs, PowfOp::new)
423 }
424 }
425 };
426}
427
428impl_complex!(num_complex::Complex<f32>, C32, f32, Complex32Type, F32);
429impl_complex!(num_complex::Complex<f64>, C64, f64, Complex64Type, F64);