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