1use super::blas::{Op, gemm};
4use crate::math::complex::{Complex, Float};
5use crate::math::numbers::Number;
6use core::fmt;
7use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
8
9pub(crate) mod sealed {
11 use super::*;
12
13 pub trait Element: Clone + PartialEq + fmt::Debug {
15 fn zero() -> Self;
17 fn one() -> Self;
19 fn is_zero(&self) -> bool;
21 fn conj(&self) -> Self;
23 fn add(a: &Self, b: &Self) -> Self;
25 fn sub(a: &Self, b: &Self) -> Self;
27 fn mul(a: &Self, b: &Self) -> Self;
29 fn neg(a: &Self) -> Self;
31 fn add_assign(a: &mut Self, b: &Self);
33 fn sub_assign(a: &mut Self, b: &Self);
35 fn mul_assign(a: &mut Self, b: &Self);
37 fn mul_add(c: &mut Self, a: &Self, b: &Self);
39
40 fn product(c: &mut [Self], a: &[Self], b: &[Self], m: usize, k: usize, n: usize) {
43 for j in 0..n {
44 let c = &mut c[j * m..(j + 1) * m];
45 for p in 0..k {
46 let x = &b[p + j * k];
47 if x.is_zero() {
48 continue;
49 }
50 for (c, a) in c.iter_mut().zip(&a[p * m..(p + 1) * m]) {
51 Self::mul_add(c, a, x);
52 }
53 }
54 }
55 }
56 }
57
58 pub trait Scalar:
60 Element
61 + Copy
62 + Default
63 + Add<Output = Self>
64 + Sub<Output = Self>
65 + Mul<Output = Self>
66 + Div<Output = Self>
67 + Neg<Output = Self>
68 + AddAssign
69 + SubAssign
70 + MulAssign
71 + DivAssign
72 + Send
73 + Sync
74 + 'static
75 {
76 type R: Float;
78 type C: super::Scalar<Real = Self::R>;
80 type Re: super::Scalar<Real = Self::R>;
83 const IS_COMPLEX: bool;
85 fn re(self) -> Self::R;
87 fn im(self) -> Self::R;
89 fn from_real(re: Self::R) -> Self;
91 fn from_parts(re: Self::R, im: Self::R) -> Self;
93 fn abs(self) -> Self::R;
95 fn abs2(self) -> Self::R;
97 fn abs1(self) -> Self::R;
99 fn scale(self, r: Self::R) -> Self;
101 fn mul_sub(self, a: Self, b: Self) -> Self;
103 fn sqrt(self) -> Self;
105 fn to_complex(self) -> Complex<Self::R>;
107 fn to_c(self) -> Self::C;
109 fn from_c(c: Self::C) -> Self;
111 fn is_finite(self) -> bool;
113 #[allow(clippy::too_many_arguments)]
116 fn gemm_real(
117 m: usize,
118 n: usize,
119 k: usize,
120 alpha: Self::R,
121 a: &[Self::R],
122 lda: usize,
123 b: &[Self::R],
124 ldb: usize,
125 beta: Self::R,
126 c: &mut [Self::R],
127 ldc: usize,
128 );
129 #[allow(clippy::too_many_arguments)]
131 fn gemm(
132 m: usize,
133 n: usize,
134 k: usize,
135 alpha: Self,
136 a: &[Self],
137 lda: usize,
138 opa: Op,
139 b: &[Self],
140 ldb: usize,
141 opb: Op,
142 beta: Self,
143 c: &mut [Self],
144 ldc: usize,
145 );
146 }
147}
148
149pub trait Element: sealed::Element {}
152
153pub trait Scalar: Element + sealed::Scalar<R = <Self as Scalar>::Real> {
157 type Real: Float;
160}
161
162impl Element for Number {}
163
164impl sealed::Element for Number {
165 #[inline]
166 fn zero() -> Self {
167 Number::default()
168 }
169 #[inline]
170 fn one() -> Self {
171 Number::from(1u8)
172 }
173 #[inline]
174 fn is_zero(&self) -> bool {
175 Number::is_zero(self)
176 }
177 #[inline]
178 fn conj(&self) -> Self {
179 self.clone()
180 }
181 #[inline]
182 fn add(a: &Self, b: &Self) -> Self {
183 a + b
184 }
185 #[inline]
186 fn sub(a: &Self, b: &Self) -> Self {
187 a - b
188 }
189 #[inline]
190 fn mul(a: &Self, b: &Self) -> Self {
191 a * b
192 }
193 #[inline]
194 fn neg(a: &Self) -> Self {
195 -a
196 }
197 #[inline]
198 fn add_assign(a: &mut Self, b: &Self) {
199 *a += b;
200 }
201 #[inline]
202 fn sub_assign(a: &mut Self, b: &Self) {
203 *a -= b;
204 }
205 #[inline]
206 fn mul_assign(a: &mut Self, b: &Self) {
207 *a *= b;
208 }
209 #[inline]
210 fn mul_add(c: &mut Self, a: &Self, b: &Self) {
211 if !a.is_zero() {
212 *c += a * b;
213 }
214 }
215}
216
217macro_rules! gemm {
219 ($f:ident) => {
220 #[inline]
221 #[track_caller]
222 fn gemm(
223 m: usize,
224 n: usize,
225 k: usize,
226 alpha: Self,
227 a: &[Self],
228 lda: usize,
229 opa: Op,
230 b: &[Self],
231 ldb: usize,
232 opb: Op,
233 beta: Self,
234 c: &mut [Self],
235 ldc: usize,
236 ) {
237 super::gemm::$f(m, n, k, alpha, a, lda, opa, b, ldb, opb, beta, c, ldc);
238 }
239 };
240}
241
242macro_rules! copy_element {
244 ($zero:expr, $one:expr, $conj:expr) => {
245 #[inline]
246 fn zero() -> Self {
247 $zero
248 }
249 #[inline]
250 fn one() -> Self {
251 $one
252 }
253 #[inline]
254 fn is_zero(&self) -> bool {
255 *self == $zero
256 }
257 #[inline]
258 fn conj(&self) -> Self {
259 $conj(*self)
260 }
261 #[inline]
262 fn add(a: &Self, b: &Self) -> Self {
263 *a + *b
264 }
265 #[inline]
266 fn sub(a: &Self, b: &Self) -> Self {
267 *a - *b
268 }
269 #[inline]
270 fn mul(a: &Self, b: &Self) -> Self {
271 *a * *b
272 }
273 #[inline]
274 fn neg(a: &Self) -> Self {
275 -*a
276 }
277 #[inline]
278 fn add_assign(a: &mut Self, b: &Self) {
279 *a += *b;
280 }
281 #[inline]
282 fn sub_assign(a: &mut Self, b: &Self) {
283 *a -= *b;
284 }
285 #[inline]
286 fn mul_assign(a: &mut Self, b: &Self) {
287 *a *= *b;
288 }
289 #[inline]
290 fn mul_add(c: &mut Self, a: &Self, b: &Self) {
291 *c += *a * *b;
292 }
293 fn product(c: &mut [Self], a: &[Self], b: &[Self], m: usize, k: usize, n: usize) {
294 gemm(m, n, k, $one, a, m, Op::N, b, k, Op::N, $zero, c, m);
295 }
296 };
297}
298
299macro_rules! real {
300 ($t:ident) => {
301 impl Element for $t {}
302 impl Scalar for $t {
303 type Real = $t;
304 }
305
306 impl sealed::Element for $t {
307 copy_element!(0.0, 1.0, core::convert::identity);
308 }
309
310 impl sealed::Scalar for $t {
311 type R = $t;
312 type C = Complex<$t>;
313 type Re = $t;
314 const IS_COMPLEX: bool = false;
315 #[inline]
316 fn re(self) -> $t {
317 self
318 }
319 #[inline]
320 fn im(self) -> $t {
321 0.0
322 }
323 #[inline]
324 fn from_real(re: $t) -> Self {
325 re
326 }
327 #[inline]
328 fn from_parts(re: $t, _: $t) -> Self {
329 re
330 }
331 #[inline]
332 fn abs(self) -> $t {
333 $t::abs(self)
334 }
335 #[inline]
336 fn abs2(self) -> $t {
337 self * self
338 }
339 #[inline]
340 fn abs1(self) -> $t {
341 $t::abs(self)
342 }
343 #[inline]
344 fn scale(self, r: $t) -> Self {
345 self * r
346 }
347 #[inline(always)]
348 fn mul_sub(self, a: Self, b: Self) -> Self {
349 $t::mul_add(-a, b, self)
350 }
351 #[inline]
352 fn sqrt(self) -> Self {
353 $t::sqrt(self)
354 }
355 #[inline]
356 fn to_complex(self) -> Complex<$t> {
357 Complex::new(self, 0.0)
358 }
359 #[inline]
360 fn to_c(self) -> Complex<$t> {
361 Complex::new(self, 0.0)
362 }
363 #[inline]
364 fn from_c(c: Complex<$t>) -> Self {
365 c.re
366 }
367 #[inline]
368 fn is_finite(self) -> bool {
369 $t::is_finite(self)
370 }
371 #[inline]
372 fn gemm_real(
373 m: usize,
374 n: usize,
375 k: usize,
376 alpha: $t,
377 a: &[$t],
378 lda: usize,
379 b: &[$t],
380 ldb: usize,
381 beta: $t,
382 c: &mut [$t],
383 ldc: usize,
384 ) {
385 gemm(m, n, k, alpha, a, lda, Op::N, b, ldb, Op::N, beta, c, ldc);
386 }
387 gemm!(real);
388 }
389
390 impl Element for Complex<$t> {}
391 impl Scalar for Complex<$t> {
392 type Real = $t;
393 }
394
395 impl sealed::Element for Complex<$t> {
396 copy_element!(Complex::<$t>::ZERO, Complex::<$t>::ONE, Complex::conj);
397 }
398
399 impl sealed::Scalar for Complex<$t> {
400 type R = $t;
401 type C = Self;
402 type Re = $t;
403 const IS_COMPLEX: bool = true;
404 #[inline]
405 fn re(self) -> $t {
406 self.re
407 }
408 #[inline]
409 fn im(self) -> $t {
410 self.im
411 }
412 #[inline]
413 fn from_real(re: $t) -> Self {
414 Complex::new(re, 0.0)
415 }
416 #[inline]
417 fn from_parts(re: $t, im: $t) -> Self {
418 Complex::new(re, im)
419 }
420 #[inline]
421 fn abs(self) -> $t {
422 Complex::abs(self)
423 }
424 #[inline]
425 fn abs2(self) -> $t {
426 self.norm_sqr()
427 }
428 #[inline]
429 fn abs1(self) -> $t {
430 self.re.abs() + self.im.abs()
431 }
432 #[inline]
433 fn scale(self, r: $t) -> Self {
434 self * r
435 }
436 #[inline(always)]
437 fn mul_sub(self, a: Self, b: Self) -> Self {
438 let re = $t::mul_add(a.im, b.im, $t::mul_add(-a.re, b.re, self.re));
439 let im = $t::mul_add(-a.im, b.re, $t::mul_add(-a.re, b.im, self.im));
440 Complex::new(re, im)
441 }
442 #[inline]
443 fn sqrt(self) -> Self {
444 Complex::sqrt(self)
445 }
446 #[inline]
447 fn to_complex(self) -> Complex<$t> {
448 self
449 }
450 #[inline]
451 fn to_c(self) -> Self {
452 self
453 }
454 #[inline]
455 fn from_c(c: Self) -> Self {
456 c
457 }
458 #[inline]
459 fn is_finite(self) -> bool {
460 Complex::is_finite(self)
461 }
462 #[inline]
463 fn gemm_real(
464 m: usize,
465 n: usize,
466 k: usize,
467 alpha: $t,
468 a: &[$t],
469 lda: usize,
470 b: &[$t],
471 ldb: usize,
472 beta: $t,
473 c: &mut [$t],
474 ldc: usize,
475 ) {
476 gemm(m, n, k, alpha, a, lda, Op::N, b, ldb, Op::N, beta, c, ldc);
477 }
478 gemm!(complex);
479 }
480 };
481}
482
483real!(f32);
484real!(f64);