1use crate::{
9 Float, agreement::{Mapping, Walked, expanded_items}, buffer::{Element, Source, read_all}, convert, error::Error, float::{FUSED, sealed::Core},
10 parallel, slice
11};
12use num_complex::Complex;
13use num_traits::{One, Zero};
14use std::{borrow::Cow, mem::MaybeUninit, sync::LazyLock};
15
16#[cfg(accelerate)]
17mod accelerate;
18mod factor;
19mod inner;
20mod scalar;
21#[cfg(accelerate)]
23pub(crate) use accelerate::Accelerated as Vendor;
24#[cfg(not(accelerate))]
25pub(crate) trait Vendor {}
26#[cfg(not(accelerate))]
27impl<T> Vendor for T {}
28
29fn switched() -> bool {
31 static NATIVE: LazyLock<bool> = LazyLock::new(|| std::env::var_os("RUSTYMATH_NATIVE").is_some_and(|v| v == "1"));
32 cfg!(all(accelerate, native)) && *NATIVE
33}
34
35pub(crate) fn accelerated() -> bool { cfg!(accelerate) && !switched() }
37
38macro_rules! dispatch {
40 ($vendor:expr, $native:expr) => {{
41 #[cfg(all(accelerate, native))]
42 let r = if accelerated() { $vendor } else { $native };
43 #[cfg(all(accelerate, not(native)))]
44 let r = $vendor;
45 #[cfg(not(accelerate))]
46 let r = $native;
47 r
48 }};
49}
50macro_rules! from_order {
52 ($n:expr, $from:expr, $vendor:expr, $native:expr) => {{
53 #[cfg(accelerate)]
54 if $n >= $from && accelerated() {
55 return $vendor;
56 }
57 $native
58 }};
59}
60pub mod exact;
61pub use factor::*;
62pub use inner::inner_with;
63pub use scalar::Scalar;
64
65pub(crate) fn identity<T: num_traits::Zero + num_traits::One>(n: usize) -> Vec<T> {
67 let mut x = Vec::new();
68 identity_into(&mut x, n);
69 x
70}
71
72pub(crate) fn identity_into<T: num_traits::Zero + num_traits::One>(x: &mut Vec<T>, n: usize) {
74 x.clear();
75 x.extend(std::iter::repeat_with(T::zero).take(n * n));
76 for i in 0..n { x[i * (n + 1)] = T::one() }
77}
78
79#[inline]
81pub fn transpose_into<S: Copy, D>(src: &[S], ls: usize, dst: &mut [D], ld: usize, [rows, cols]: [usize; 2], f: impl Fn(S) -> D) {
82 if rows == 0 || cols == 0 { return; }
83 assert!(src.len() >= (rows - 1) * ls + cols && dst.len() >= (cols - 1) * ld + rows, "both matrices lie in their slices");
84 let (whole_rows, whole_cols) = (rows / 8 * 8, cols / 8 * 8);
85 match 16 / size_of::<S>().max(1) {
87 8.. => tiles::<8, S, D>(src, ls, dst, ld, [whole_rows, whole_cols], &f), 4.. => tiles::<4, S, D>(src, ls, dst, ld, [whole_rows, whole_cols], &f),
88 2.. => tiles::<2, S, D>(src, ls, dst, ld, [whole_rows, whole_cols], &f), _ => tiles::<1, S, D>(src, ls, dst, ld, [whole_rows, whole_cols], &f)
89 }
90 if whole_cols < cols { for r in 0..whole_rows { for c in whole_cols..cols { dst[c * ld + r] = f(src[r * ls + c]) } } }
91 for r in whole_rows..rows { for c in 0..cols { dst[c * ld + r] = f(src[r * ls + c]) } }
92}
93
94#[inline(always)]
96fn tiles<const P: usize, S: Copy, D>(src: &[S], ls: usize, dst: &mut [D], ld: usize, [rows, cols]: [usize; 2], f: &impl Fn(S) -> D) {
97 let (from, to) = (src.as_ptr(), dst.as_mut_ptr());
98 for r0 in (0..rows).step_by(8) {
99 for c0 in (0..cols).step_by(8) {
100 for i in (0..8).step_by(P) {
101 let read = |k: usize| unsafe { from.add((r0 + i + k) * ls + c0).cast::<[S; 8]>().read() };
103 let mut runs = [read(0); P];
104 for k in 1..P { runs[k] = read(k) }
105 for j in 0..8 { let out = unsafe { &mut *to.add((c0 + j) * ld + r0 + i).cast::<[D; P]>() }; for k in 0..P { out[k] = f(runs[k][j]) } }
106 }
107 }
108 }
109}
110
111pub trait Inner<T> {
113 type Out: Copy + Default;
115 fn inner_into(&self, a: &[T], b: &[T], size: [usize; 3], out: &mut [Self::Out]);
117 fn inner(&self, a: &[T], b: &[T], size @ [m, _, n]: [usize; 3]) -> Vec<Self::Out> {
119 let mut c = filled(m * n, Self::Out::default());
120 self.inner_into(a, b, size, &mut c);
121 c
122 }
123 #[inline]
125 fn checked_into(&self, a: &[T], b: &[T], size: [usize; 3], out: &mut [Self::Out]) -> bool { self.inner_into(a, b, size, out); true }
126 fn inner_ops(&self, a: &[T], b: &[T], [oa, ob]: [Orientation; 2], [m, k, n]: [usize; 3]) -> Vec<Self::Out> where T: Copy {
128 self.inner(&oriented(a, oa, [m, k]), &oriented(b, ob, [k, n]), [m, k, n])
129 }
130 fn inner_mapped(&self, a: Mapped<'_, T>, b: Mapped<'_, T>, orientations: [Orientation; 2], size @ [m, _, n]: [usize; 3]) -> Vec<Self::Out> where T: Element
132 {
133 let mut out = filled(m * n, Self::Out::default());
134 self.inner_mapped_into(a, b, orientations, size, &mut out);
135 out
136 }
137 fn inner_mapped_into(&self, a: Mapped<'_, T>, b: Mapped<'_, T>, [oa, ob]: [Orientation; 2], size @ [m, k, n]: [usize; 3], out: &mut [Self::Out]) where
139 T: Element {
140 let (mut ab, mut bb) = (Vec::new(), Vec::new());
141 self.inner_into(&oriented(a.read(0, m * k, &mut ab), oa, [m, k]), &oriented(b.read(0, k * n, &mut bb), ob, [k, n]), size, out)
142 }
143}
144
145fn oriented<T: Copy>(x: &[T], o: Orientation, [rows, cols]: [usize; 2]) -> std::borrow::Cow<'_, [T]> {
146 match o {
147 Orientation::Normal => std::borrow::Cow::Borrowed(x), Orientation::Transpose => (0..rows * cols).map(|ij| x[ij % cols * rows + ij / cols]).collect()
148 }
149}
150
151#[derive(Default)]
153pub struct PlusTimes;
154#[derive(Default)]
156pub struct MaxPlus;
157#[derive(Default)]
159pub struct MinPlus;
160#[derive(Default)]
162pub struct OrAnd;
163#[derive(Default)]
165pub struct AndOr;
166#[derive(Default)]
168pub struct PlusEqual;
169#[derive(Default)]
171pub struct PlusNotEqual;
172
173#[derive(Clone, Copy, Debug, PartialEq, Eq)]
175pub enum Product { PlusTimes, MaxPlus, MinPlus, OrAnd, AndOr, PlusEqual, PlusNotEqual }
176
177impl Product { pub fn compares_exactly(self) -> bool { matches!(self, Self::PlusEqual | Self::PlusNotEqual) } }
178
179#[derive(Clone, Debug)]
181pub struct ProductPlan { pub layout: crate::layout::Layout, pub iteration: crate::layout::Layout, pub mappings: [Mapping; 2], pub contracted: usize }
182
183#[derive(Clone, Copy, Debug)]
184pub struct MatrixProduct { pub frame_rank: usize, pub row_rank: usize, pub size: [usize; 3] }
185impl MatrixProduct {
186 fn of(shape: &[usize], owners: &[u8]) -> Option<Self> {
187 let frame_rank = owners.iter().take_while(|&&owner| owner == 3).count();
188 let row_rank = owners[frame_rank..].iter().take_while(|&&owner| owner == 1).count();
189 if !owners[frame_rank + row_rank..].iter().all(|&owner| owner == 2) { return None; }
190 let sizes = [&shape[frame_rank..frame_rank + row_rank], &shape[owners.len()..], &shape[frame_rank + row_rank..owners.len()]]
191 .map(|shape| crate::layout::element_count(shape).ok());
192 match sizes { [Some(m), Some(k), Some(n)] => Some(Self { frame_rank, row_rank, size: [m, k, n] }), _ => None }
193 }
194}
195
196impl ProductPlan {
197 pub fn contracted_len(&self) -> usize { let shape = self.iteration.shape(); crate::layout::count(&shape[shape.len() - self.contracted..]) }
199 pub fn folded(iteration: crate::layout::Layout, mappings: [Mapping; 2], contracted: usize) -> Result<Self, Error> {
200 let rank = iteration.shape().len();
201 if contracted > rank { return Err(Error::new(crate::error::Kind::Rank, "too many contracted axes")); }
202 let layout = iteration.axes(0..rank - contracted);
203 Ok(Self { layout, iteration, mappings, contracted })
204 }
205
206 pub fn matrix_for(&self, inputs: [crate::agreement::Read<'_>; 2]) -> Option<MatrixProduct> {
208 let shape = self.iteration.shape();
209 let mut owners = vec![0; shape.len() - self.contracted];
210 let mut nominal = owners.clone();
211 for (side, input) in inputs.iter().enumerate() {
212 let plan = input.plan.remapped(self.mappings[side].clone(), self.iteration.count(), crate::agreement::Missing::Invalid);
213 let map =
214 crate::agreement::Read { sources: input.sources, plan: &plan }.positional(shape).map_or_else(|| self.mappings[side].clone(), |(_, map)| map);
215 let (_, steps) = map.affine(shape)?;
216 for (axis, owner) in owners.iter_mut().enumerate() { if steps[axis] != 0 { *owner |= 1 << side; } }
217 if let Some((_, free)) = self.mappings[side].affine(shape) {
218 for (named, &step) in nominal.iter_mut().zip(&free) { if step != 0 { *named |= 1 << side; } }
219 }
220 }
221 let mut last = 3;
223 for ((owner, named), &len) in owners.iter_mut().zip(nominal).zip(shape) {
224 if *owner == 0 { *owner = if len == 1 { last } else if named != 0 { named } else { 3 }; }
225 last = *owner;
226 }
227 MatrixProduct::of(shape, &owners)
228 }
229 pub fn dot([left, right]: [&crate::layout::Layout; 2]) -> Result<Self, Error> {
230 let rank = left.shape().len();
231 if rank == 0 || right.shape().is_empty() { return Err(Error::new(crate::error::Kind::Rank, "an inner product needs vectors or matrices")); }
232 Self::paired([left, right], &[(rank - 1, 0)], None)
233 }
234 pub fn matrices(inputs: [&crate::layout::Layout; 2], frame: usize) -> Result<Self, Error> {
235 let (contracted, destinations) = Self::matrix_axes(inputs, frame)?;
236 Self::paired(inputs, &[contracted], Some(destinations))
237 }
238 pub fn matrix_axes([left, right]: [&crate::layout::Layout; 2], frame: usize) -> Result<((usize, usize), [Vec<usize>; 2]), Error> {
240 use crate::agreement::Agreement;
241 let [lr, rr] = [left, right].map(|l| l.shape().len());
242 if lr <= frame || rr <= frame { return Err(Error::new(crate::error::Kind::Rank, "a matrix product needs vector or matrix cells")); }
243 let [lc, rc] = [(lr - frame).min(2), (rr - frame).min(2)];
244 let [lf, rf] = [left.axes(0..lr - lc), right.axes(0..rr - rc)];
245 let [mut xa, mut ya] = Agreement::axis_pairs(&lf, &rf);
246 let first = xa.iter().chain(&ya).max().map_or(0, |a| a + 1);
247 if lc == 2 { xa.push(first); }
248 if rc == 2 { ya.push(first + usize::from(lc == 2)); }
249 Ok(((lr - 1, rr - rc), [xa, ya]))
250 }
251 pub fn paired(inputs @ [left, right]: [&crate::layout::Layout; 2], pairs: &[(usize, usize)], destinations: Option<[Vec<usize>; 2]>) -> Result<Self, Error> {
253 use crate::{agreement::Agreement, error::Kind, layout::{element_count, generated_len}};
254 let contracted: [Vec<usize>; 2] = [pairs.iter().map(|p| p.0).collect(), pairs.iter().map(|p| p.1).collect()];
255 for (input, axes) in inputs.iter().zip(&contracted) {
256 if axes.iter().enumerate().any(|(i, a)| *a >= input.shape().len() || axes[..i].contains(a)) {
257 return Err(Error::new(Kind::Rank, "contracted axes must be distinct axes of their inputs"));
258 }
259 }
260 for &(a, b) in pairs {
261 let (x, y) = (left.shape()[a], right.shape()[b]);
262 if x != y {
263 return Err(Error::new(Kind::Length, format!("an inner product needs matching contracted lengths, not {x} and {y}")));
264 }
265 Mapping::contract(left, a, right, b).map_err(|_| Error::new(Kind::Length, "an inner product needs the same keys on its contracted axes"))?;
266 }
267 let free: [Vec<usize>; 2] = std::array::from_fn(|i| (0..inputs[i].shape().len()).filter(|a| !contracted[i].contains(a)).collect());
268 let free_layouts = std::array::from_fn::<_, 2, _>(|i| inputs[i].axes(free[i].iter().copied()));
269 let total = free[0].len() + free[1].len();
270 let destinations = destinations.unwrap_or_else(|| [(0..free[0].len()).collect(), (free[0].len()..total).collect()]);
271 for (axes, positions) in free.iter().zip(&destinations) {
272 if axes.len() != positions.len() || positions.iter().enumerate().any(|(i, p)| *p >= total || positions[..i].contains(p)) {
273 return Err(Error::new(Kind::Rank, "free axes need distinct output destinations"));
274 }
275 }
276 let agreement = Agreement::mapped(&free_layouts[0], &free_layouts[1], &destinations[0], &destinations[1], false)?;
277 if agreement.left.fills() || agreement.right.fills() { return Err(Error::new(Kind::Length, "keyed frames need a cell that an argument lacks")); }
278 let layout = agreement.layout;
279 generated_len(layout.shape())?;
280 let first = layout.shape().len();
281 let iteration = layout.concat(&left.axes(contracted[0].iter().copied()));
282 let items = element_count(iteration.shape())?;
283 let mappings = std::array::from_fn(|i| {
284 let mut axes = vec![0; inputs[i].shape().len()];
285 for (&a, &p) in free[i].iter().zip(&destinations[i]) { axes[a] = p; }
286 for (p, &a) in contracted[i].iter().enumerate() { axes[a] = first + p; }
287 Mapping::new(inputs[i], &axes, &iteration, items)
288 });
289 Ok(Self { layout, iteration, mappings, contracted: pairs.len() })
290 }
291}
292
293trait Pair<T> {
295 type Out: Copy + Send;
296 const UNIT: Self::Out;
298 fn g(a: T, b: T) -> Self::Out;
299 fn f(x: Self::Out, y: Self::Out) -> Self::Out;
300 #[inline(always)]
302 fn step(acc: Self::Out, a: T, b: T) -> Self::Out { Self::f(acc, Self::g(a, b)) }
303}
304
305macro_rules! pairs {
307 ($($obj:ty: $t:ty => $out:ty = $unit:expr, |$a:ident, $b:ident| $g:expr, |$x:ident, $y:ident| $f:expr $(, |$c:ident, $p:ident, $q:ident| $step:expr)?;)+) => {$(
308 impl Pair<$t> for $obj {
309 type Out = $out;
310 const UNIT: $out = $unit;
311 #[inline(always)]
312 fn g($a: $t, $b: $t) -> $out { $g }
313 #[inline(always)]
314 fn f($x: $out, $y: $out) -> $out { $f }
315 $(#[inline(always)]
316 fn step($c: $out, $p: $t, $q: $t) -> $out { $step })?
317 }
318 )+};
319}
320macro_rules! floats {
321 ($($t:ty),+) => {$(pairs! {
322 PlusTimes: $t => $t = 0.0, |a, b| a * b, |x, y| x + y, |acc, a, b| a.mla::<FUSED>(b, acc);
323 MaxPlus: $t => $t = <$t>::NEG_INFINITY, |a, b| a + b, |acc, term| Core::maximum(term, acc);
324 MinPlus: $t => $t = <$t>::INFINITY, |a, b| a + b, |acc, term| Core::minimum(term, acc);
325 PlusTimes: Complex<$t> => Complex<$t> = Complex::new(0.0, 0.0), |a, b| a * b, |x, y| x + y, |acc, a, b| a.mul_acc(b, acc);
326 })+};
327}
328floats!(f32, f64);
329
330pairs! {
331 OrAnd: bool => bool = false, |a, b| a & b, |x, y| x | y;
332 AndOr: bool => bool = true, |a, b| a | b, |x, y| x & y;
333}
334
335impl Inner<bool> for OrAnd {
336 type Out = bool;
337 #[inline]
338 fn inner_into(&self, a: &[bool], b: &[bool], size: [usize; 3], out: &mut [bool]) {
339 if size[0] * size[1] * size[2] <= DIRECT_WORK { direct::<bool, bool, OrAnd>(a, b, size, out) } else { packed::<false>(a, b, size, out) }
340 }
341}
342
343impl Inner<bool> for AndOr {
344 type Out = bool;
345 #[inline]
346 fn inner_into(&self, a: &[bool], b: &[bool], size: [usize; 3], out: &mut [bool]) {
347 if size[0] * size[1] * size[2] <= DIRECT_WORK { direct::<bool, bool, AndOr>(a, b, size, out) } else { packed::<true>(a, b, size, out) }
348 }
349}
350
351fn packed<const NOT: bool>(a: &[bool], b: &[bool], [m, k, n]: [usize; 3], c: &mut [bool]) {
353 const WORDS: usize = 4;
354 assert!(a.len() == m * k && b.len() == k * n && c.len() == m * n, "`a` must hold `m`×`k` items, `b` `k`×`n` and `c` `m`×`n`");
355 if k == 0 || n == 0 { return c.fill(NOT); }
356 let blocks = n.div_ceil(64 * WORDS);
357 let mut packed = vec![[0u64; WORDS]; k * blocks];
358 for (row, bits) in b.chunks_exact(n).zip(packed.chunks_exact_mut(blocks)) {
359 for (j, chunk) in row.chunks(64).enumerate() {
360 bits[j / WORDS][j % WORDS] = chunk.iter().enumerate().fold(0, |w, (t, &v)| w | u64::from(v ^ NOT) << t);
361 }
362 }
363 for (row, out) in a.chunks_exact(k).zip(c.chunks_exact_mut(n)) {
364 for (block, out) in out.chunks_mut(64 * WORDS).enumerate() {
365 let mut acc = [0u64; WORDS];
366 for (&x, bits) in row.iter().zip(packed[block..].iter().step_by(blocks)) {
367 let mask = u64::from(x ^ NOT).wrapping_neg();
368 for t in 0..WORDS { acc[t] |= bits[t] & mask }
369 }
370 for (j, o) in out.iter_mut().enumerate() { *o = (acc[j / 64] >> (j % 64) & 1 != 0) ^ NOT }
371 }
372 }
373}
374pairs! {
376 PlusTimes: i64 => i64 = 0, |a, b| a.wrapping_mul(b), |x, y| x.wrapping_add(y);
377 PlusTimes: i32 => i32 = 0, |a, b| a.wrapping_mul(b), |x, y| x.wrapping_add(y);
378 MaxPlus: i64 => i64 = i64::MIN, |a, b| a.wrapping_add(b), |x, y| x.max(y);
379 MinPlus: i64 => i64 = i64::MAX, |a, b| a.wrapping_add(b), |x, y| x.min(y);
380 PlusTimes: i128 => Option<i128> = Some(0), |a, b| a.checked_mul(b), |x, y| x?.checked_add(y?);
381 MaxPlus: i128 => Option<i128> = Some(i128::MIN), |a, b| a.checked_add(b), |x, y| Some(x?.max(y?));
382 MinPlus: i128 => Option<i128> = Some(i128::MAX), |a, b| a.checked_add(b), |x, y| Some(x?.min(y?));
383}
384macro_rules! counts {
385 ($($t:ty => $out:ty),+) => {$(pairs! { PlusEqual: $t => $out = 0, |a, b| <$out>::from(a == b), |x, y| x + y; })+};
386}
387counts!(i64 => i64, i32 => i32, i16 => i16, f64 => i64, f32 => i32);
388
389#[doc(hidden)]
391pub const KC: usize = 256;
392#[doc(hidden)]
394pub const MC: usize = 128;
395#[doc(hidden)]
397pub const NC: usize = 2048;
398
399#[cfg(target_arch = "aarch64")]
401mod tiles {
402 pub(super) const BYTES: usize = 16;
403 #[cfg(native)]
404 pub(super) const TIMES64: [usize; 2] = [8, 4];
405 pub(super) const TIMES32: [usize; 2] = [8, 8];
406 pub(super) const COMPARES64: [usize; 2] = [6, 8];
407 pub(super) const COMPARES32: [usize; 2] = [6, 16];
408 pub(super) const COMPARES16: [usize; 2] = [6, 32];
409}
410#[cfg(all(target_arch = "x86_64", target_feature = "avx2", target_feature = "fma"))]
411mod tiles {
412 pub(super) const BYTES: usize = 32;
413 #[cfg(native)]
414 pub(super) const TIMES64: [usize; 2] = [6, 8];
415 pub(super) const TIMES32: [usize; 2] = [6, 16];
416 pub(super) const COMPARES64: [usize; 2] = [4, 8];
417 pub(super) const COMPARES32: [usize; 2] = [4, 16];
418 pub(super) const COMPARES16: [usize; 2] = [4, 32];
419}
420#[cfg(not(any(target_arch = "aarch64", all(target_arch = "x86_64", target_feature = "avx2", target_feature = "fma"))))]
421mod tiles {
422 pub(super) const BYTES: usize = 16;
423 pub(super) const TIMES64: [usize; 2] = [4, 4];
424 pub(super) const TIMES32: [usize; 2] = [4, 8];
425 pub(super) const COMPARES64: [usize; 2] = [4, 4];
426 pub(super) const COMPARES32: [usize; 2] = [4, 8];
427 pub(super) const COMPARES16: [usize; 2] = [4, 16];
428}
429
430const PLAIN: [usize; 2] = [4, 8];
432
433macro_rules! inner {
434 ($($object:ident at $t:ty => $out:ty: $tile:expr, $kernel:ident;)+) => {$(
435 impl Inner<$t> for $object {
436 type Out = $out;
437 fn inner_into(&self, a: &[$t], b: &[$t], size: [usize; 3], out: &mut [Self::Out]) {
438 const MR: usize = $tile[0];
439 const NR: usize = $tile[1];
440 tiled_into::<$t, $t, $object, MR, NR>(a, b, size, out, $kernel::<MR, NR, { NR * size_of::<$t>() / tiles::BYTES }>)
441 }
442 }
443 )+};
444}
445fn filled<T: Clone>(n: usize, x: T) -> Vec<T> {
447 let mut v = Vec::with_capacity(n);
448 v.resize(n, x);
449 v
450}
451#[inline(always)]
453fn tiled_into<S: Source<T> + Copy + Sync, T: Copy + Default + Send + Sync, P: Pair<T>, const MR: usize, const NR: usize>(
454 a: &[S], b: &[S], [m, k, n]: [usize; 3], c: &mut [P::Out], kernel: impl Fn(&[T], &[T], &mut [[P::Out; NR]; MR], bool) + Sync) {
455 assert!(a.len() == m * k && b.len() == k * n && c.len() == m * n, "`a` must hold `m`×`k` items, `b` `k`×`n` and `c` `m`×`n`");
456 if m * k * n <= DIRECT_WORK { return direct::<S, T, P>(a, b, [m, k, n], c); }
457 let f = <S as Source<T>>::read;
458 product::<T, P, MR, NR>(
459 c, n, |rows, p0, out| strip(a, k, false, rows, p0, out, f), |columns, p0, out| strip(b, n, true, columns, p0, out, f), [m, k, n], true, Part::Whole,
460 kernel);
461}
462
463const DIRECT_WORK: usize = 64;
465
466#[inline]
469fn direct<S: Source<T> + Copy, T: Copy, P: Pair<T>>(a: &[S], b: &[S], [m, k, n]: [usize; 3], c: &mut [P::Out]) {
470 let f = <S as Source<T>>::read;
471 const SIDE: usize = 4;
472 assert!(a.len() >= m * k && b.len() >= k * n && c.len() >= m * n, "`a` must hold `m`×`k` items, `b` `k`×`n` and `c` `m`×`n`");
473 for i in 0..m {
475 let mut j = 0;
476 while j + SIDE <= n {
477 let mut acc = [P::UNIT; SIDE];
478 for p in 0..k {
479 let x = f(unsafe { *a.get_unchecked(i * k + p) });
480 for s in 0..SIDE { acc[s] = P::step(acc[s], x, f(unsafe { *b.get_unchecked(p * n + j + s) })) }
481 }
482 for s in 0..SIDE { unsafe { *c.get_unchecked_mut(i * n + j + s) = acc[s] } }
483 j += SIDE;
484 }
485 for j in j..n {
486 let mut acc = P::UNIT;
487 for p in 0..k { acc = P::step(acc, f(unsafe { *a.get_unchecked(i * k + p) }), f(unsafe { *b.get_unchecked(p * n + j) })) }
488 unsafe { *c.get_unchecked_mut(i * n + j) = acc }
489 }
490 }
491}
492inner! {
493 MaxPlus at f64 => f64: tiles::COMPARES64, max_plus64;
494 MaxPlus at f32 => f32: tiles::COMPARES32, max_plus32;
495 MinPlus at f64 => f64: tiles::COMPARES64, min_plus64;
496 MinPlus at f32 => f32: tiles::COMPARES32, min_plus32;
497 PlusEqual at i32 => i32: tiles::COMPARES32, equal_i32;
498 PlusEqual at i16 => i16: tiles::COMPARES16, equal_i16;
499 PlusEqual at f64 => i64: tiles::COMPARES64, equal64;
500 PlusEqual at f32 => i32: tiles::COMPARES32, equal32;
501}
502
503impl Inner<i64> for PlusEqual {
505 type Out = i64;
506 fn inner_into(&self, a: &[i64], b: &[i64], size: [usize; 3], out: &mut [i64]) {
507 if let Some(c) = windowed(a, b, size) { return out.iter_mut().zip(&c).for_each(|(o, &x)| *o = i64::from(x)); }
508 let narrow = |x: &[i64]| x.iter().map(|&v| i32::try_from(v).ok()).collect::<Option<Vec<_>>>();
509 if let (Some(a), Some(b)) = (narrow(a), narrow(b)) { return out.iter_mut().zip(&self.inner(&a, &b, size)).for_each(|(o, &x)| *o = i64::from(x)); }
510 const MR: usize = tiles::COMPARES64[0];
511 const NR: usize = tiles::COMPARES64[1];
512 tiled_into::<i64, i64, PlusEqual, MR, NR>(a, b, size, out, equal_i64::<MR, NR, { NR * size_of::<i64>() / tiles::BYTES }>)
513 }
514}
515
516fn windowed(a: &[i64], b: &[i64], size: [usize; 3]) -> Option<Vec<i16>> {
518 if size[1] > i16::MAX as usize { return None; }
519 let ((la, ha), (lb, hb)) = (crate::reduce::extremes(a), crate::reduce::extremes(b));
520 let (lo, hi) = (la.min(lb), ha.max(hb));
521 if i128::from(hi) - i128::from(lo) > i128::from(u16::MAX) { return None; }
522 let shifted = |x: &[i64]| x.iter().map(|&v| v.wrapping_sub(lo) as u16 as i16).collect::<Vec<_>>();
523 Some(PlusEqual.inner(&shifted(a), &shifted(b), size))
524}
525
526macro_rules! differences {
528 ($($t:ty => $out:ty),+) => {$(
529 impl Inner<$t> for PlusNotEqual {
530 type Out = $out;
531 fn inner_into(&self, a: &[$t], b: &[$t], size: [usize; 3], out: &mut [$out]) {
532 PlusEqual.inner_into(a, b, size, out);
533 let k = size[1] as $out;
534 for x in out { *x = k - *x }
535 }
536 }
537 )+};
538}
539differences!(i64 => i64, i32 => i32, i16 => i16, f64 => i64, f32 => i32);
540
541fn magnitude<T: Copy>(x: &[T]) -> u128 where i64: From<T> { u128::from(x.iter().fold(0, |m, &v| m.max(i64::from(v).unsigned_abs()))) }
543
544macro_rules! integers {
546 ($($object:ident: $bound:expr;)+) => {$(
547 impl Inner<i64> for $object {
548 type Out = i64;
549 fn inner_into(&self, a: &[i64], b: &[i64], size: [usize; 3], out: &mut [i64]) {
550 assert!(self.checked_into(a, b, size, out), "an integer inner product fits i64");
551 }
552 fn checked_into(&self, a: &[i64], b: &[i64], size: [usize; 3], out: &mut [i64]) -> bool {
553 let bound = $bound(magnitude(a), magnitude(b), size[1] as u128);
554 if bound <= i64::MAX as u128 { tiered(self, a, b, size, bound, out); return true; }
555 wide::<Self>(a, b, size).map(|c| out.copy_from_slice(&c)).is_some()
556 }
557 }
558 )+};
559}
560integers! {
561 PlusTimes: |a: u128, b: u128, k: u128| a.saturating_mul(b).saturating_mul(k);
562 MaxPlus: |a: u128, b: u128, _| a + b;
563 MinPlus: |a: u128, b: u128, _| a + b;
564}
565
566fn wide<P: Pair<i128, Out = Option<i128>>>(a: &[i64], b: &[i64], [m, k, n]: [usize; 3]) -> Option<Vec<i64>> {
568 (0..m * n).map(|c| {
569 let (i, j) = (c / n, c % n);
570 let total = (0..k).fold(P::UNIT, |t, p| P::step(t, a[i * k + p].into(), b[p * n + j].into()))?;
571 i64::try_from(total).ok()
572 }).collect()
573}
574
575fn tiered<P>(pair: &P, a: &[i64], b: &[i64], [m, k, n]: [usize; 3], bound: u128, out: &mut [i64]) where
577 P: Inner<f32, Out = f32> + Inner<f64, Out = f64> + Pair<i64, Out = i64> {
578 assert!(a.len() == m * k && b.len() == k * n, "`a` must hold `m`×`k` items and `b` `k`×`n`");
579 if k == 0 { return out.fill(<P as Pair<i64>>::UNIT); }
580 if exactly(pair, a, b, [m, k, n], bound, true, out, |x| x) { return; }
581 tiled_into::<i64, i64, P, { PLAIN[0] }, { PLAIN[1] }>(a, b, [m, k, n], out, micro::<i64, P, { PLAIN[0] }, { PLAIN[1] }>)
582}
583
584fn exactly<P, S, O>(pair: &P, a: &[S], b: &[S], size: [usize; 3], bound: u128, wide: bool, out: &mut [O], whole: impl Fn(i64) -> O) -> bool where
586 P: Inner<f32, Out = f32> + Inner<f64, Out = f64>, S: Source<f32> + Source<f64> {
587 fn at<F: Float + Element, P: Inner<F, Out = F>, S: Source<F>, O>(
588 pair: &P, a: &[S], b: &[S], size: [usize; 3], out: &mut [O], whole: impl Fn(i64) -> O) -> bool {
589 for (o, x) in out.iter_mut().zip(pair.inner(&read_all(a), &read_all(b), size)) { *o = whole(convert::shifted(x.wide())) }
590 true
591 }
592 if bound <= 1 << 24 { at::<f32, P, S, O>(pair, a, b, size, out, whole) }
593 else { wide && bound <= convert::WIDE_LIMIT as u128 && at::<f64, P, S, O>(pair, a, b, size, out, whole) }
594}
595
596fn times_i32_into<S: Source<i32>>(a: &[S], b: &[S], size: [usize; 3], out: &mut [i32]) {
598 const MR: usize = tiles::TIMES32[0];
599 const NR: usize = tiles::TIMES32[1];
600 tiled_into::<S, i32, PlusTimes, MR, NR>(a, b, size, out, times_i32::<MR, NR, { NR * size_of::<i32>() / tiles::BYTES }>)
601}
602
603macro_rules! wrapped_narrow {
605 ($($t:ty: $largest:expr;)+) => {$(
606 impl Inner<$t> for crate::Wrapped<PlusTimes> {
607 type Out = i32;
608 fn inner_into(&self, a: &[$t], b: &[$t], size @ [m, k, n]: [usize; 3], out: &mut [i32]) {
609 assert!(a.len() == m * k && b.len() == k * n && out.len() == m * n, "`a` must hold `m`×`k` items, `b` `k`×`n` and `out` `m`×`n`");
610 if m * k * n > DIRECT_WORK && exactly(&PlusTimes, a, b, size, $largest(a, b) * k as u128, accelerated(), out, |x| x as i32) { return }
611 times_i32_into(a, b, size, out)
612 }
613 }
614 )+};
615}
616wrapped_narrow! {
617 u8: |_, _| 255 * 255;
618 i16: |_, _| 1 << 30;
619 i32: |a, b| magnitude(a) * magnitude(b);
620}
621
622impl Inner<half::f16> for PlusTimes {
624 type Out = f32;
625 fn inner_into(&self, a: &[half::f16], b: &[half::f16], size @ [m, k, n]: [usize; 3], c: &mut [f32]) {
626 assert!(a.len() == m * k && b.len() == k * n && c.len() == m * n, "`a` must hold `m`×`k` items, `b` `k`×`n` and `c` `m`×`n`");
627 if m * k * n <= DIRECT_WORK { return direct::<half::f16, f32, PlusTimes>(a, b, size, c); }
628 dispatch!(half_products_accelerated(a, b, size, c), half_products_native(a, b, size, c))
629 }
630}
631
632#[cfg(accelerate)]
634fn half_products_accelerated(a: &[half::f16], b: &[half::f16], size @ [_, k, n]: [usize; 3], c: &mut [f32]) {
635 f32::gemm(c, n, Matrix::rows(&read_all(a), k), Matrix::rows(&read_all(b), n), size, 1.0, 0.0)
636}
637
638#[cfg(native)]
640fn half_products_native(a: &[half::f16], b: &[half::f16], size: [usize; 3], c: &mut [f32]) {
641 const MR: usize = tiles::TIMES32[0];
642 const NR: usize = tiles::TIMES32[1];
643 tiled_into::<half::f16, f32, PlusTimes, MR, NR>(a, b, size, c, times32::<MR, NR, { NR * size_of::<f32>() / tiles::BYTES }>)
644}
645
646#[derive(Clone, Copy)]
648pub(crate) struct Matrix<'a, T> { pub data: &'a [T], pub ld: usize, pub op: Op }
649#[derive(Clone, Copy)]
650pub(crate) enum Op { Normal, Transpose, Adjoint, Conjugate }
651#[derive(Clone, Copy, Debug, PartialEq, Eq)]
653pub enum Orientation { Normal, Transpose }
654#[derive(Clone, Copy)]
656pub struct Mapped<'a, T> { pub data: &'a [T], pub mapping: &'a Walked<'a> }
657
658impl<'a, T: Element> Mapped<'a, T> {
659 fn matrix(&self, o: Orientation, [rows, cols]: [usize; 2]) -> Option<Matrix<'a, T>> {
661 let [r, c] = if o == Orientation::Normal { [rows, cols] } else { [cols, rows] };
662 let (first, strides) = self.mapping.affine(&[r, c])?;
663 let sc = if c == 1 { 1 } else { strides[1] };
665 let sr = if r > 1 { strides[0] } else if sc == 1 { c as isize } else { 1 };
666 let data = self.data.get(first..)?;
667 let stored = match (sr, sc) {
668 (s, 1) if s >= c as isize => Matrix::rows(data, s as usize), (1, s) if s >= r as isize => Matrix { data, ld: s as usize, op: Op::Transpose },
669 _ => return None
670 };
671 Some(if o == Orientation::Normal { stored } else { stored.transpose() })
672 }
673 pub fn in_place(&self, o: Orientation, size: [usize; 2]) -> bool { self.matrix(o, size).is_some() }
675 pub fn read<'b>(&'b self, start: usize, n: usize, buf: &'b mut Vec<T>) -> &'b [T] { expanded_items(self.mapping, self.data, start, n, buf) }
677}
678#[derive(Clone, Copy)]
679pub(crate) enum Side { Left, Right }
680#[derive(Clone, Copy)]
682pub(crate) enum Triangle { Lower, UnitUpper }
683impl Op {
684 #[inline(always)]
685 fn transposed(self) -> bool { matches!(self, Self::Transpose | Self::Adjoint) }
686}
687impl<'a, T: Copy> Matrix<'a, T> {
688 pub(crate) fn rows(data: &'a [T], ld: usize) -> Self { Self { data, ld, op: Op::Normal } }
689 pub(crate) fn transpose(self) -> Self {
690 Self { op: match self.op {
691 Op::Normal => Op::Transpose, Op::Transpose => Op::Normal, Op::Adjoint => Op::Conjugate, Op::Conjugate => Op::Adjoint
692 }, ..self }
693 }
694 pub(crate) fn offset(self, row: usize, col: usize) -> Self {
696 let at = if self.op.transposed() { col * self.ld + row } else { row * self.ld + col };
697 Self { data: &self.data[at..], ..self }
698 }
699 pub(crate) fn oriented(data: &'a [T], o: Orientation, [rows, cols]: [usize; 2]) -> Self {
701 match o { Orientation::Normal => Self::rows(data, cols), Orientation::Transpose => Self { data, ld: rows, op: Op::Transpose } }
702 }
703 fn check(self, mut rows: usize, mut cols: usize) {
705 if self.op.transposed() { std::mem::swap(&mut rows, &mut cols) }
706 if rows == 0 || cols == 0 { return; }
707 let end = (rows - 1).checked_mul(self.ld).and_then(|n| n.checked_add(cols));
708 assert!(self.ld >= cols && end.is_some_and(|n| n <= self.data.len()), "matrix view must hold its rows and columns");
709 }
710}
711impl<T: Scalar> Matrix<'_, T> {
712 pub(crate) fn conjugate(self) -> Self {
713 if !T::COMPLEX { return self; }
715 Self { op: match self.op {
716 Op::Normal => Op::Conjugate, Op::Transpose => Op::Adjoint, Op::Adjoint => Op::Transpose, Op::Conjugate => Op::Normal
717 }, ..self }
718 }
719 pub(crate) fn adjoint(self) -> Self { self.transpose().conjugate() }
720 #[cfg(native)]
721 #[inline(always)]
722 pub(crate) fn at(self, i: usize, j: usize) -> T {
723 let x = self.data[if self.op.transposed() { j * self.ld + i } else { i * self.ld + j }];
724 if matches!(self.op, Op::Adjoint | Op::Conjugate) { x.conj() } else { x }
725 }
726 #[cfg(native)]
728 #[inline(always)]
729 fn strip<const W: usize>(self, rows: [usize; 2], p0: usize, scale: T, out: &mut [[MaybeUninit<T>; W]]) where T: Default {
730 let (transposed, one) = (self.op.transposed(), scale == T::one());
731 if !T::COMPLEX || !matches!(self.op, Op::Adjoint | Op::Conjugate) {
732 if one { strip(self.data, self.ld, transposed, rows, p0, out, |x| x) } else { strip(self.data, self.ld, transposed, rows, p0, out, |x| scale * x) }
733 } else if one { strip(self.data, self.ld, transposed, rows, p0, out, |x: T| x.conj()) }
734 else { strip(self.data, self.ld, transposed, rows, p0, out, |x: T| scale * x.conj()) }
735 }
736}
737#[inline(always)]
739fn strip<S: Copy, T: Default, const W: usize>(
740 data: &[S], ld: usize, transposed: bool, [r0, lines]: [usize; 2], p0: usize, out: &mut [[MaybeUninit<T>; W]], f: impl Fn(S) -> T) {
741 assert!(lines <= W);
742 let kc = out.len();
743 let rows: [&[S]; W] = std::array::from_fn(|l| if transposed || l >= lines { &[] } else { &data[(r0 + l) * ld + p0..][..kc] });
744 match (transposed, lines == W) {
745 (true, true) => {
746 for (p, step) in out.iter_mut().enumerate() {
747 let run: &[S; W] = data[(p0 + p) * ld + r0..].first_chunk().expect("a whole step lies in `data`");
748 for l in 0..W { step[l] = MaybeUninit::new(f(run[l])) }
749 }
750 }
751 (true, false) => {
752 for (p, step) in out.iter_mut().enumerate() {
753 let run = &data[(p0 + p) * ld + r0..][..lines];
754 for l in 0..W { step[l] = MaybeUninit::new(if l < lines { f(run[l]) } else { T::default() }) }
755 }
756 }
757 (false, true) => { for (p, step) in out.iter_mut().enumerate() { for l in 0..W { step[l] = MaybeUninit::new(f(rows[l][p])) } } }
758 (false, false) => {
759 for (p, step) in out.iter_mut().enumerate() { for l in 0..W { step[l] = MaybeUninit::new(if l < lines { f(rows[l][p]) } else { T::default() }) } }
760 }
761 }
762}
763pub(crate) trait Tiled: Scalar + Default + Vendor {
765 #[inline]
767 fn multiply(a: Matrix<'_, Self>, b: Matrix<'_, Self>, size: [usize; 3]) -> Vec<Self> {
768 let mut c = filled(size[0] * size[2], Self::zero());
769 Self::gemm(&mut c, size[2], a, b, size, Self::one(), Self::zero());
770 c
771 }
772 #[inline]
774 fn gemm(c: &mut [Self], ldc: usize, a: Matrix<'_, Self>, b: Matrix<'_, Self>, size: [usize; 3], alpha: Self, beta: Self) {
775 gemm(c, ldc, a, b, size, alpha, beta, Part::Whole)
776 }
777 #[inline]
779 fn gram(c: &mut [Self], ldc: usize, x: Matrix<'_, Self>, [n, k]: [usize; 2], alpha: Self::Real, beta: Self::Real) {
780 gemm(c, ldc, x, x.adjoint(), [n, k, n], Self::from_real(alpha), Self::from_real(beta), Part::Lower)
781 }
782 fn triangular(x: &mut [Self], ldx: usize, a: Matrix<'_, Self>, [m, n]: [usize; 2], side: Side, triangle: Triangle) {
784 let order = match side { Side::Left => m, Side::Right => n };
785 a.check(order, order);
786 Matrix::rows(x, ldx).check(m, n);
787 if m == 0 || n == 0 { return; }
788 dispatch!(accelerate::triangular(x, ldx, a, [m, n], side, triangle), factor::triangular_native(x, ldx, a, [m, n], side, triangle))
789 }
790 #[cfg(native)]
792 fn native_product(c: &mut [Self], ldc: usize, a: Matrix<'_, Self>, alpha: Self, b: Matrix<'_, Self>, size: [usize; 3], start: bool, part: Part);
793 fn svd(a: &[Self], size: [usize; 2], vectors: bool) -> Option<(Vec<Self>, Vec<Self::Real>, Vec<Self>)>;
795 fn eigh(a: &[Self], n: usize, vectors: bool) -> Option<Eigh<Self>>;
797 fn eig(a: &[Self], n: usize, vectors: bool) -> Option<Eig<Self::Real>>;
799}
800
801fn gemm<T: Tiled>(c: &mut [T], ldc: usize, a: Matrix<'_, T>, b: Matrix<'_, T>, [m, k, n]: [usize; 3], alpha: T, beta: T, part: Part) {
802 a.check(m, k);
803 b.check(k, n);
804 Matrix::rows(c, ldc).check(m, n);
805 if m == 0 || n == 0 { return; }
806 if k == 0 || alpha == T::zero() { return scale(c, ldc, [m, n], beta, part); }
807 if by_vector(c, ldc, a, alpha, b, [m, k, n], beta, part) { return; }
808 dispatch!(accelerate::product(c, ldc, a, b, [m, k, n], alpha, beta, part), native_gemm(c, ldc, a, b, [m, k, n], alpha, beta, part))
809}
810
811const BLOCK_ITEMS: usize = 1 << 18;
813const PIECE_MATRIX: usize = 1 << 17;
815const WHOLE_ITEMS: usize = 1 << 24;
817
818fn block_rows(cols: usize, least: usize) -> usize { least.max(BLOCK_ITEMS).div_ceil(cols) }
820
821fn blocks<T: Element>(op: Mapped<'_, T>, rows: std::ops::Range<usize>, cols: usize, least: usize, mut f: impl FnMut(usize, usize, &[T])) {
823 let (step, mut block) = (block_rows(cols, least), Vec::new());
824 for first in rows.clone().step_by(step) {
825 let len = step.min(rows.end - first);
826 f(first, len, op.read(first * cols, len * cols, &mut block));
827 }
828}
829
830fn contracted<T: Tiled + Element>(op: Mapped<'_, T>, [rows, cols]: [usize; 2], c: &mut [T], f: &(dyn Fn(usize, usize, &[T], &mut [T], T) + Sync)) {
832 let beta = |first: usize| if first == 0 { T::zero() } else { T::one() };
833 if c.len() > BLOCK_ITEMS { return blocks(op, 0..rows, cols, c.len(), |first, len, block| f(first, len, block, c, beta(first))); }
834 parallel::summed(rows, block_rows(cols, 0), c, &mut Vec::new(), &|range, sums| {
835 blocks(op, range, cols, 0, |first, len, block| f(first, len, block, sums, beta(first)))
836 });
837}
838
839const COLUMN_ROWS: usize = 4096;
841const COLUMN_STRIP: usize = 32;
843const TAIL_STRIP: usize = 8;
845const COLUMN_RUN: usize = 64;
847
848fn by_columns<T: Tiled + Element>(a: Mapped<'_, T>, b: Mapped<'_, T>, oa: Orientation, [m, k]: [usize; 2], c: &mut [T]) -> bool {
850 use crate::gather::Steps;
851 let Mapping::Steps { base, steps } = &**a.mapping else { return false };
852 if m * k == 0 { return false; }
853 let lead = if oa == Orientation::Normal { m } else { k };
855 let (mut split, mut size) = (0, 1);
856 while size < lead && split < steps.len() {
857 size *= steps[split].len();
858 split += 1;
859 }
860 if size != lead { return false; }
861 let (outer, inner) = steps.split_at(split);
862 let (rows, cols) = if oa == Orientation::Normal { (outer, inner) } else { (inner, outer) };
863 let Some(&Steps::Stride { step: 1, len: width, .. }) = rows.iter().rev().find(|s| s.len() > 1) else { return false };
864 if width < COLUMN_RUN { return false; }
865 let mut whole = Vec::new();
866 let weights = b.read(0, k, &mut whole);
867 let columns: Vec<(T, usize)> = (0..k).filter_map(|p| Some((weights[p], Steps::reach(0, cols, p)?))).collect();
868 parallel::split(c, COLUMN_ROWS, &|first, part| {
869 let mut r = first;
870 while r < first + part.len() {
871 let len = (width - r % width).min(COLUMN_ROWS).min(first + part.len() - r);
872 let out = &mut part[r - first..r - first + len];
874 if let Some(at) = Steps::reach(*base, rows, r) { scaled_runs(out, a.data, at, &columns, true) }
875 else { out.fill(T::zero()) }
876 r += len;
877 }
878 true
879 });
880 true
881}
882
883fn scaled_runs<T: Tiled>(out: &mut [T], data: &[T], at: usize, columns: &[(T, usize)], start: bool) {
885 let wide = scaled_strips::<T, COLUMN_STRIP>(out, data, at, columns, start);
886 let done = wide + scaled_strips::<T, TAIL_STRIP>(&mut out[wide..], data, at + wide, columns, start);
887 let at = at + done;
888 for (t, y) in out[done..].iter_mut().enumerate() {
889 *y = columns.iter().fold(if start { T::zero() } else { *y }, |acc, &(w, o)| w.mul_acc(data[at + o + t], acc))
890 }
891}
892
893fn scaled_strips<T: Tiled, const N: usize>(out: &mut [T], data: &[T], at: usize, columns: &[(T, usize)], start: bool) -> usize {
895 let (strips, _) = out.as_chunks_mut::<N>();
896 for (s, strip) in strips.iter_mut().enumerate() {
897 let (at, mut acc) = (at + s * N, if start { [T::zero(); N] } else { *strip });
898 for &(w, o) in columns { let x = &data[at + o..][..N]; for t in 0..N { acc[t] = w.mul_acc(x[t], acc[t]) } }
899 *strip = acc;
900 }
901 strips.len() * N
902}
903
904fn dots<T: Tiled>(out: &mut [T], data: &[T], ld: usize, first: usize, x: &[T]) {
906 let k = x.len();
907 let (quads, tail) = out.as_chunks_mut::<4>();
908 let (xc, xr) = x.as_chunks::<8>();
909 let sum = |s: &[T; 8]| ((s[0] + s[4]) + (s[2] + s[6])) + ((s[1] + s[5]) + (s[3] + s[7]));
910 for (q, quad) in quads.iter_mut().enumerate() {
911 let rows: [&[T]; 4] = std::array::from_fn(|r| &data[(first + 4 * q + r) * ld..][..k]);
912 let mut s = [[T::zero(); 8]; 4];
913 for (c, xs) in xc.iter().enumerate() { for r in 0..4 { let a = &rows[r][c * 8..][..8]; for l in 0..8 { s[r][l] = a[l].mul_acc(xs[l], s[r][l]) } } }
914 let at = xc.len() * 8;
915 for r in 0..4 { quad[r] = xr.iter().enumerate().fold(sum(&s[r]), |t, (p, &b)| rows[r][at + p].mul_acc(b, t)) }
916 }
917 let first = first + 4 * quads.len();
918 for (i, y) in tail.iter_mut().enumerate() { *y = factor::dot(&data[(first + i) * ld..][..k], x) }
919}
920
921const NARROW: usize = 8;
923
924fn narrow_dots<T: Tiled, const K: usize>(out: &mut [T], data: &[T], ld: usize, first: usize, x: &[T]) {
926 let x: &[T; K] = x.try_into().expect("`x` must hold `K` items");
927 let sums = |i: usize| {
928 let a = &data[i * ld..][..K];
929 let mut s = [T::zero(); 2];
930 for l in 0..K { s[l % 2] = a[l].mul_acc(x[l], s[l % 2]) }
931 s
932 };
933 let last = first + out.len().saturating_sub(1);
934 let (pairs, odd) = out.as_chunks_mut::<2>();
935 for (p, pair) in pairs.iter_mut().enumerate() {
936 let (s, t) = (sums(first + 2 * p), sums(first + 2 * p + 1));
937 *pair = [s[0] + s[1], t[0] + t[1]];
938 }
939 if let [y] = odd { let s = sums(last); *y = s[0] + s[1] }
940}
941
942fn narrow_sums<T: Tiled>(out: &mut [T], data: &[T], ld: usize, w: &[T]) {
944 let n = out.len();
945 let whole = if data.len() < NARROW { 0 } else { ((data.len() - NARROW) / ld.max(1) + 1).min(w.len()) };
946 let mut acc = [[T::zero(); NARROW]; 4];
947 let (quads, rest) = w[..whole].as_chunks::<4>();
948 for (q, ws) in quads.iter().enumerate() {
949 for s in 0..4 { let a = &data[(4 * q + s) * ld..][..NARROW]; for t in 0..NARROW { acc[s][t] = a[t].mul_acc(ws[s], acc[s][t]) } }
950 }
951 for (p, &wp) in rest.iter().enumerate() {
952 let a = &data[(4 * quads.len() + p) * ld..][..NARROW];
953 for t in 0..NARROW { acc[0][t] = a[t].mul_acc(wp, acc[0][t]) }
954 }
955 for (p, &wp) in w.iter().enumerate().skip(whole) { for t in 0..n { acc[0][t] = data[p * ld + t].mul_acc(wp, acc[0][t]) } }
956 for (t, y) in out.iter_mut().enumerate() { *y = (acc[0][t] + acc[2][t]) + (acc[1][t] + acc[3][t]) }
957}
958
959fn blocked<T: Tiled + Element>(a: Mapped<'_, T>, b: Mapped<'_, T>, [oa, ob]: [Orientation; 2], [m, k, n]: [usize; 3], c: &mut [T]) {
961 assert_eq!(c.len(), m * n);
962 if m == 0 || k == 0 || n == 0 { c.fill(T::zero()); return; }
963 let (ma, mb) = (a.matrix(oa, [m, k]), b.matrix(ob, [k, n]));
964 if let (Some(a), Some(b)) = (ma, mb) { return T::gemm(c, n, a, b, [m, k, n], T::one(), T::zero()); }
965 if ma.is_none() && n == 1 && by_columns(a, b, oa, [m, k], c) { crate::events::taken(crate::events::Route::MappedProduct, m * k); return; }
966 let read = if ma.is_none() { m * k } else { k * n };
967 if read > m * n { crate::events::taken(crate::events::Route::MappedProduct, read); }
968 let (one, zero) = (T::one(), T::zero());
969 let mut whole = Vec::new();
970 match ma {
971 None => {
972 let b = match mb { Some(b) => b, None => Matrix::oriented(b.read(0, k * n, &mut whole), ob, [k, n]) };
973 match oa {
974 Orientation::Normal => {
975 let (rows, least) =
976 if parallel::threads() == 1 || k * n <= PIECE_MATRIX { (m.div_ceil(parallel::threads()), k * n) } else { (m, WHOLE_ITEMS) };
977 parallel::pieces_of(c, rows * n, usize::MAX, &|start, part| {
978 let from = start / n;
979 blocks(a, from..from + part.len() / n, k, least, |first, len, block| {
980 T::gemm(&mut part[(first - from) * n..], n, Matrix::rows(block, k), b, [len, k, n], one, zero)
981 });
982 true
983 });
984 }
985 Orientation::Transpose => contracted(a, [k, m], c, &|first, len, part, c, beta| {
986 T::gemm(c, n, Matrix::oriented(part, oa, [m, len]), b.offset(first, 0), [m, len, n], one, beta)
987 }),
988 }
989 }
990 Some(a) => match ob {
991 Orientation::Normal => {
992 contracted(b, [k, n], c, &|first, len, part, c, beta| T::gemm(c, n, a.offset(0, first), Matrix::rows(part, n), [m, len, n], one, beta))
993 }
994 Orientation::Transpose => {
995 blocks(b, 0..n, k, m * n, |first, len, block| T::gemm(&mut c[first..], n, a, Matrix::oriented(block, ob, [k, len]), [m, k, len], one, zero))
996 }
997 },
998 }
999}
1000
1001fn scale<T: Tiled>(c: &mut [T], ldc: usize, [m, n]: [usize; 2], beta: T, part: Part) {
1003 if beta == T::one() { return; }
1004 for i in 0..m {
1005 for j in part.columns(i, n) {
1006 let x = &mut c[i * ldc + j];
1007 *x = if beta == T::zero() { T::zero() } else { beta * *x };
1008 }
1009 }
1010}
1011
1012const VECTOR_WORK: usize = 1 << 15;
1014const VECTOR_STRIP: usize = 2048;
1016const VECTOR_ROWS: usize = 8;
1018const VECTOR_PIECE: usize = 2 * THREAD_WORK;
1020
1021fn by_vector<T: Tiled>(c: &mut [T], ldc: usize, a: Matrix<'_, T>, alpha: T, b: Matrix<'_, T>, [m, k, n]: [usize; 3], beta: T, part: Part) -> bool {
1023 let (matrix, vector, rows, step) = match (m, n, part) {
1024 (_, 1, Part::Whole) => (a, b.transpose(), m, ldc), (1, _, Part::Whole) => (b.transpose(), a, n, 1), _ => return false
1025 };
1026 let narrow = if matrix.op.transposed() { rows } else { k } <= NARROW;
1027 if accelerated() && !narrow { return false; }
1028 let conj = matches!(matrix.op, Op::Adjoint | Op::Conjugate);
1030 let flip = T::COMPLEX && conj != matches!(vector.op, Op::Adjoint | Op::Conjugate);
1031 let stride = if vector.op.transposed() { vector.ld } else { 1 };
1032 let x: Cow<'_, [T]> = if stride == 1 && !flip { Cow::Borrowed(&vector.data[..k]) }
1033 else { (0..k).map(|p| if flip { vector.data[p * stride].conj() }
1034 else { vector.data[p * stride] }).collect() };
1035 let (data, ld, unit) = (matrix.data, matrix.ld, (VECTOR_WORK / k).max(1).next_multiple_of(4));
1036 let (one, start) = (alpha == T::one(), beta == T::zero());
1037 let direct = step == 1 && one && start && !conj;
1038 let mut kept = if direct { Vec::new() } else { filled(rows, T::zero()) };
1039 let y = if direct { &mut c[..rows] } else { &mut kept[..] };
1040 if matrix.op.transposed() {
1041 let piece = (VECTOR_PIECE / rows).max(1);
1043 if narrow {
1044 parallel::summed(k, piece, y, &mut Vec::new(), &|range, part| narrow_sums(part, &data[range.start * ld..], ld, &x[range]));
1045 } else {
1046 let columns: Vec<(T, usize)> = x.iter().enumerate().map(|(p, &w)| (w, p * ld)).collect();
1047 if rows <= VECTOR_STRIP {
1048 parallel::summed(columns.len(), piece, y, &mut Vec::new(), &|range, part| {
1049 for (i, block) in columns[range].chunks(VECTOR_ROWS).enumerate() { scaled_runs(part, data, 0, block, i == 0) }
1050 });
1051 } else {
1052 parallel::split(y, unit, &|first, part| {
1053 for (s, strip) in part.chunks_mut(VECTOR_STRIP).enumerate() {
1054 for (i, block) in columns.chunks(VECTOR_ROWS).enumerate() { scaled_runs(strip, data, first + s * VECTOR_STRIP, block, i == 0) }
1055 }
1056 true
1057 });
1058 }
1059 }
1060 } else {
1061 let dots = match k {
1062 1 => narrow_dots::<T, 1>, 2 => narrow_dots::<T, 2>, 3 => narrow_dots::<T, 3>, 4 => narrow_dots::<T, 4>, 5 => narrow_dots::<T, 5>,
1063 6 => narrow_dots::<T, 6>, 7 => narrow_dots::<T, 7>, 8 => narrow_dots::<T, 8>, _ => dots::<T>
1064 };
1065 parallel::split(y, unit, &|first, part| { dots(part, data, ld, first, &x); true });
1066 }
1067 if direct { return true; }
1068 for (i, &y) in kept.iter().enumerate() {
1070 let y = if conj { y.conj() } else { y };
1071 let (y, z) = (if one { y } else { alpha * y }, &mut c[i * step]);
1072 *z = if start { y } else { beta.mul_acc(*z, y) };
1073 }
1074 true
1075}
1076
1077#[cfg(native)]
1079fn native_gemm<T: Tiled>(c: &mut [T], ldc: usize, a: Matrix<'_, T>, b: Matrix<'_, T>, [m, k, n]: [usize; 3], alpha: T, beta: T, part: Part) {
1080 let start = beta == T::zero();
1081 if !start { scale(c, ldc, [m, n], beta, part) }
1082 T::native_product(c, ldc, a, alpha, b, [m, k, n], start, part)
1083}
1084#[derive(Clone, Copy)]
1086pub(crate) enum Part { Whole, Lower }
1087impl Part {
1088 #[inline(always)]
1090 fn columns(self, i: usize, n: usize) -> std::ops::Range<usize> { match self { Part::Whole => 0..n, Part::Lower => 0..(i + 1).min(n) } }
1091}
1092macro_rules! tiled {
1093 (@real $t:ty, $tile:ident, $kernel:ident; $c:ident, $ldc:ident, $a:ident, $alpha:ident, $b:ident, $size:ident, $start:ident, $part:ident) => {{
1094 const MR: usize = tiles::$tile[0];
1095 const NR: usize = tiles::$tile[1];
1096 let one = <$t>::one();
1097 product::<$t, PlusTimes, MR, NR>($c, $ldc, |rows, p0, out| $a.strip(rows, p0, $alpha, out), |columns, p0, out| $b.transpose().strip(columns, p0, one, out), $size, $start, $part, $kernel::<MR, NR, { NR * size_of::<$t>() / tiles::BYTES }>)
1098 }};
1099 (@complex $r:ty; $c:ident, $ldc:ident, $a:ident, $alpha:ident, $b:ident, $size:ident, $start:ident, $part:ident) => {
1100 complex_product::<$r>($c, $ldc, $a, $alpha, $b, $size, $start, $part)
1101 };
1102 (@eig real; $a:ident, $n:ident, $vectors:ident) => { from_order!($n, accelerate::GEEV_FROM, accelerate::eig($a, $n, $vectors), factor::general::native($a, $n, $vectors)) };
1103 (@eig complex; $a:ident, $n:ident, $vectors:ident) => { factor::general::native($a, $n, $vectors) };
1104 ($($t:ty: $kind:ident($($args:tt)*);)+) => {$(
1105 impl Tiled for $t {
1106 #[cfg(native)]
1107 #[inline]
1108 fn native_product(c: &mut [$t], ldc: usize, a: Matrix<'_, $t>, alpha: $t, b: Matrix<'_, $t>, size: [usize; 3], start: bool, part: Part) {
1109 tiled!(@$kind $($args)*; c, ldc, a, alpha, b, size, start, part)
1110 }
1111 fn svd(a: &[$t], size: [usize; 2], vectors: bool) -> Option<(Vec<$t>, Vec<<$t as Scalar>::Real>, Vec<$t>)> {
1112 from_order!(size[0].max(size[1]), accelerate::GESDD_FROM, accelerate::svd(a, size, vectors), factor::svd::native(a, size, vectors))
1113 }
1114 fn eigh(a: &[$t], n: usize, vectors: bool) -> Option<Eigh<$t>> {
1115 from_order!(n, accelerate::HEEVD_FROM, accelerate::eigh(a, n, vectors), factor::eigen::native(a, n, vectors))
1116 }
1117 fn eig(a: &[$t], n: usize, vectors: bool) -> Option<Eig<<$t as Scalar>::Real>> {
1118 tiled!(@eig $kind; a, n, vectors)
1119 }
1120 }
1121 impl Inner<$t> for PlusTimes {
1122 type Out = $t;
1123 fn inner_ops(&self, a: &[$t], b: &[$t], [oa, ob]: [Orientation; 2], [m, k, n]: [usize; 3]) -> Vec<$t> {
1124 assert!(a.len() == m * k && b.len() == k * n, "`a` must hold `m`×`k` items and `b` `k`×`n`");
1125 <$t>::multiply(Matrix::oriented(a, oa, [m, k]), Matrix::oriented(b, ob, [k, n]), [m, k, n])
1126 }
1127 #[inline]
1128 fn inner_into(&self, a: &[$t], b: &[$t], [m, k, n]: [usize; 3], c: &mut [$t]) {
1129 assert!(a.len() == m * k && b.len() == k * n && c.len() == m * n, "`a` must hold `m`×`k` items, `b` `k`×`n` and `c` `m`×`n`");
1130 if m * k * n <= DIRECT_WORK { return direct::<$t, $t, PlusTimes>(a, b, [m, k, n], c) }
1131 <$t>::gemm(c, n, Matrix::rows(a, k), Matrix::rows(b, n), [m, k, n], <$t>::one(), <$t>::zero())
1132 }
1133 fn inner_mapped_into(&self, a: Mapped<'_, $t>, b: Mapped<'_, $t>, orientations: [Orientation; 2], size: [usize; 3], out: &mut [$t]) { blocked(a, b, orientations, size, out) }
1134 }
1135 )+};
1136}
1137tiled! {
1138 f32: real(f32, TIMES32, times32);
1139 f64: real(f64, TIMES64, times64);
1140 Complex<f32>: complex(f32);
1141 Complex<f64>: complex(f64);
1142}
1143#[cfg(native)]
1145fn complex_product<R: crate::Native + Tiled>(
1146 c: &mut [Complex<R>], ldc: usize, a: Matrix<'_, Complex<R>>, alpha: Complex<R>, b: Matrix<'_, Complex<R>>, size: [usize; 3], start: bool, part: Part) {
1147 let [m, k, n] = size;
1148 if m == 0 || k == 0 || n == 0 { return; }
1149 let ([ar, ai], [br, bi]) = (parts(a, [m, k], alpha), parts(b, [k, n], Complex::new(R::one(), R::zero())));
1150 let real = |c: &mut [R], x: &[R], y: &[R], scale: R, start: bool| R::native_product(c, n, Matrix::rows(x, k), scale, Matrix::rows(y, n), size, start, part);
1151 let (mut re, mut im) = (filled(m * n, R::zero()), filled(m * n, R::zero()));
1152 real(&mut re, &ar, &br, R::one(), true);
1153 real(&mut re, &ai, &bi, -R::one(), false);
1154 real(&mut im, &ar, &bi, R::one(), true);
1155 real(&mut im, &ai, &br, R::one(), false);
1156 for i in 0..m {
1157 for j in part.columns(i, n) {
1158 let (t, x) = (i * n + j, &mut c[i * ldc + j]);
1159 *x = if start { Complex::new(re[t], im[t]) } else { Complex::new(x.re + re[t], x.im + im[t]) };
1160 }
1161 }
1162}
1163
1164#[cfg(native)]
1166fn parts<R: crate::Native + Tiled>(x: Matrix<'_, Complex<R>>, [rows, cols]: [usize; 2], scale: Complex<R>) -> [Vec<R>; 2] {
1167 let (one, conj) = (scale == Complex::new(R::one(), R::zero()), matches!(x.op, Op::Adjoint | Op::Conjugate));
1168 let value = |v: Complex<R>| { let v = if conj { v.conj() } else { v }; if one { v } else { scale * v } };
1169 let mut z = filled(rows * cols, Complex::default());
1170 if x.op.transposed() { transpose_into(x.data, x.ld, &mut z, cols, [cols, rows], value) } else {
1171 for (row, out) in z.chunks_exact_mut(cols).enumerate() { for (o, &v) in out.iter_mut().zip(&x.data[row * x.ld..][..cols]) { *o = value(v) } }
1172 }
1173 [z.iter().map(|v| v.re).collect(), z.iter().map(|v| v.im).collect()]
1174}
1175
1176#[cfg(not(any(
1178 target_arch = "aarch64", all(target_arch = "wasm32", target_feature = "simd128"),
1179 all(target_arch = "x86_64", target_feature = "avx2", target_feature = "fma")
1180)))]
1181mod plain {
1182 use super::*;
1183 macro_rules! plain {
1184 ($($name:ident: $t:ty, $object:ident;)+) => {$(
1185 #[inline(always)]
1186 pub(super) fn $name<const MR: usize, const NR: usize, const V: usize>(
1187 a: &[$t], b: &[$t], t: &mut [[<$object as Pair<$t>>::Out; NR]; MR], first: bool,
1188 ) {
1189 micro::<$t, $object, MR, NR>(a, b, t, first)
1190 }
1191 )+};
1192 }
1193 plain! {
1194 times64: f64, PlusTimes;
1195 times32: f32, PlusTimes;
1196 times_i32: i32, PlusTimes;
1197 max_plus64: f64, MaxPlus;
1198 max_plus32: f32, MaxPlus;
1199 min_plus64: f64, MinPlus;
1200 min_plus32: f32, MinPlus;
1201 equal64: f64, PlusEqual;
1202 equal32: f32, PlusEqual;
1203 equal_i64: i64, PlusEqual;
1204 equal_i32: i32, PlusEqual;
1205 equal_i16: i16, PlusEqual;
1206 }
1207}
1208#[cfg(not(any(
1209 target_arch = "aarch64", all(target_arch = "wasm32", target_feature = "simd128"),
1210 all(target_arch = "x86_64", target_feature = "avx2", target_feature = "fma")
1211)))]
1212use plain::*;
1213
1214#[cfg(any(
1216 target_arch = "aarch64", all(target_arch = "wasm32", target_feature = "simd128"),
1217 all(target_arch = "x86_64", target_feature = "avx2", target_feature = "fma")
1218))]
1219mod lanes {
1220 #[cfg(target_arch = "aarch64")]
1221 use core::arch::aarch64::*;
1222 #[cfg(target_arch = "wasm32")]
1223 use core::arch::wasm32::*;
1224 #[cfg(target_arch = "x86_64")]
1225 use core::arch::x86_64::*;
1226 macro_rules! kernels {
1227 ($($name:ident: $t:ty => $out:ty, $lanes:literal, $splat:ident, $load:ident, $start:expr, $load_out:ident, $store_out:ident, |$a:ident, $b:ident, $c:ident| $step:expr;)+) => {$(
1228 #[inline(always)]
1229 pub(super) fn $name<const MR: usize, const NR: usize, const V: usize>(a: &[$t], b: &[$t], t: &mut [[$out; NR]; MR], first: bool) {
1230 assert!(NR == $lanes * V && a.len() / MR == b.len() / NR);
1231 let steps = a.len() / MR;
1232 unsafe {
1234 let (ap, bp) = (a.as_ptr(), b.as_ptr());
1235 let mut acc = [[$start; V]; MR];
1236 if !first { for r in 0..MR { for v in 0..V { acc[r][v] = $load_out(t[r].as_ptr().add($lanes * v).cast()) } } }
1237 for p in 0..steps {
1238 let y: [_; V] = core::array::from_fn(|v| $load(bp.add(p * NR + $lanes * v).cast()));
1239 for r in 0..MR {
1240 let $a = $splat(*ap.add(p * MR + r) as _);
1241 for v in 0..V { acc[r][v] = { let ($c, $b) = (acc[r][v], y[v]); $step } }
1242 }
1243 }
1244 for r in 0..MR { for v in 0..V { $store_out(t[r].as_mut_ptr().add($lanes * v).cast(), acc[r][v]) } }
1245 }
1246 }
1247 )+};
1248 }
1249 macro_rules! lanes {
1251 ($($name:ident: $t:ty, $pair:ident, $lanes:literal, $splat:ident, $load:ident, $store:ident, |$a:ident, $b:ident, $c:ident| $step:expr;)+) => {
1252 kernels! { $($name: $t => $t, $lanes, $splat, $load, $splat(<super::$pair as super::Pair<$t>>::UNIT), $load, $store, |$a, $b, $c| $step;)+ }
1253 };
1254 }
1255 #[cfg(all(target_arch = "aarch64", native))]
1257 lanes! {
1258 times64: f64, PlusTimes, 2, vdupq_n_f64, vld1q_f64, vst1q_f64, |a, b, c| vfmaq_f64(c, a, b);
1259 times32: f32, PlusTimes, 4, vdupq_n_f32, vld1q_f32, vst1q_f32, |a, b, c| vfmaq_f32(c, a, b);
1260 }
1261 #[cfg(target_arch = "aarch64")]
1262 lanes! {
1263 times_i32: i32, PlusTimes, 4, vdupq_n_s32, vld1q_s32, vst1q_s32, |a, b, c| vmlaq_s32(c, a, b);
1265 max_plus64: f64, MaxPlus, 2, vdupq_n_f64, vld1q_f64, vst1q_f64, |a, b, c| vmaxnmq_f64(c, vaddq_f64(a, b));
1266 max_plus32: f32, MaxPlus, 4, vdupq_n_f32, vld1q_f32, vst1q_f32, |a, b, c| vmaxnmq_f32(c, vaddq_f32(a, b));
1267 min_plus64: f64, MinPlus, 2, vdupq_n_f64, vld1q_f64, vst1q_f64, |a, b, c| vminnmq_f64(c, vaddq_f64(a, b));
1268 min_plus32: f32, MinPlus, 4, vdupq_n_f32, vld1q_f32, vst1q_f32, |a, b, c| vminnmq_f32(c, vaddq_f32(a, b));
1269 }
1270 #[cfg(target_arch = "wasm32")]
1272 lanes! {
1273 times64: f64, PlusTimes, 2, f64x2_splat, v128_load, v128_store, |a, b, c| f64x2_add(c, f64x2_mul(a, b));
1274 times32: f32, PlusTimes, 4, f32x4_splat, v128_load, v128_store, |a, b, c| f32x4_add(c, f32x4_mul(a, b));
1275 times_i32: i32, PlusTimes, 4, i32x4_splat, v128_load, v128_store, |a, b, c| i32x4_add(c, i32x4_mul(a, b));
1276 max_plus64: f64, MaxPlus, 2, f64x2_splat, v128_load, v128_store, |a, b, c| f64x2_pmax(c, f64x2_add(a, b));
1277 max_plus32: f32, MaxPlus, 4, f32x4_splat, v128_load, v128_store, |a, b, c| f32x4_pmax(c, f32x4_add(a, b));
1278 min_plus64: f64, MinPlus, 2, f64x2_splat, v128_load, v128_store, |a, b, c| f64x2_pmin(c, f64x2_add(a, b));
1279 min_plus32: f32, MinPlus, 4, f32x4_splat, v128_load, v128_store, |a, b, c| f32x4_pmin(c, f32x4_add(a, b));
1280 }
1281 #[cfg(all(target_arch = "x86_64", native))]
1283 lanes! {
1284 times64: f64, PlusTimes, 4, _mm256_set1_pd, _mm256_loadu_pd, _mm256_storeu_pd, |a, b, c| _mm256_fmadd_pd(a, b, c);
1285 times32: f32, PlusTimes, 8, _mm256_set1_ps, _mm256_loadu_ps, _mm256_storeu_ps, |a, b, c| _mm256_fmadd_ps(a, b, c);
1286 }
1287 #[cfg(target_arch = "x86_64")]
1288 lanes! {
1289 times_i32: i32, PlusTimes, 8, _mm256_set1_epi32, _mm256_loadu_si256, _mm256_storeu_si256, |a, b, c| _mm256_add_epi32(c, _mm256_mullo_epi32(a, b));
1291 max_plus64: f64, MaxPlus, 4, _mm256_set1_pd, _mm256_loadu_pd, _mm256_storeu_pd, |a, b, c| _mm256_max_pd(_mm256_add_pd(a, b), c);
1293 max_plus32: f32, MaxPlus, 8, _mm256_set1_ps, _mm256_loadu_ps, _mm256_storeu_ps, |a, b, c| _mm256_max_ps(_mm256_add_ps(a, b), c);
1294 min_plus64: f64, MinPlus, 4, _mm256_set1_pd, _mm256_loadu_pd, _mm256_storeu_pd, |a, b, c| _mm256_min_pd(_mm256_add_pd(a, b), c);
1295 min_plus32: f32, MinPlus, 8, _mm256_set1_ps, _mm256_loadu_ps, _mm256_storeu_ps, |a, b, c| _mm256_min_ps(_mm256_add_ps(a, b), c);
1296 }
1297 #[cfg(target_arch = "aarch64")]
1299 kernels! {
1300 equal64: f64 => i64, 2, vdupq_n_f64, vld1q_f64, vdupq_n_s64(0), vld1q_s64, vst1q_s64, |a, b, c| vsubq_s64(c, vreinterpretq_s64_u64(vceqq_f64(a, b)));
1301 equal32: f32 => i32, 4, vdupq_n_f32, vld1q_f32, vdupq_n_s32(0), vld1q_s32, vst1q_s32, |a, b, c| vsubq_s32(c, vreinterpretq_s32_u32(vceqq_f32(a, b)));
1302 equal_i64: i64 => i64, 2, vdupq_n_s64, vld1q_s64, vdupq_n_s64(0), vld1q_s64, vst1q_s64, |a, b, c| vsubq_s64(c, vreinterpretq_s64_u64(vceqq_s64(a, b)));
1303 equal_i32: i32 => i32, 4, vdupq_n_s32, vld1q_s32, vdupq_n_s32(0), vld1q_s32, vst1q_s32, |a, b, c| vsubq_s32(c, vreinterpretq_s32_u32(vceqq_s32(a, b)));
1304 equal_i16: i16 => i16, 8, vdupq_n_s16, vld1q_s16, vdupq_n_s16(0), vld1q_s16, vst1q_s16, |a, b, c| vsubq_s16(c, vreinterpretq_s16_u16(vceqq_s16(a, b)));
1305 }
1306 #[cfg(target_arch = "wasm32")]
1307 kernels! {
1308 equal64: f64 => i64, 2, f64x2_splat, v128_load, i64x2_splat(0), v128_load, v128_store, |a, b, c| i64x2_sub(c, f64x2_eq(a, b));
1309 equal32: f32 => i32, 4, f32x4_splat, v128_load, i32x4_splat(0), v128_load, v128_store, |a, b, c| i32x4_sub(c, f32x4_eq(a, b));
1310 equal_i64: i64 => i64, 2, i64x2_splat, v128_load, i64x2_splat(0), v128_load, v128_store, |a, b, c| i64x2_sub(c, i64x2_eq(a, b));
1311 equal_i32: i32 => i32, 4, i32x4_splat, v128_load, i32x4_splat(0), v128_load, v128_store, |a, b, c| i32x4_sub(c, i32x4_eq(a, b));
1312 equal_i16: i16 => i16, 8, i16x8_splat, v128_load, i16x8_splat(0), v128_load, v128_store, |a, b, c| i16x8_sub(c, i16x8_eq(a, b));
1313 }
1314 #[cfg(target_arch = "x86_64")]
1315 kernels! {
1316 equal64: f64 => i64, 4, _mm256_set1_pd, _mm256_loadu_pd, _mm256_setzero_si256(), _mm256_loadu_si256, _mm256_storeu_si256, |a, b, c| {
1317 _mm256_sub_epi64(c, _mm256_castpd_si256(_mm256_cmp_pd::<_CMP_EQ_OQ>(a, b)))
1318 };
1319 equal32: f32 => i32, 8, _mm256_set1_ps, _mm256_loadu_ps, _mm256_setzero_si256(), _mm256_loadu_si256, _mm256_storeu_si256, |a, b, c| {
1320 _mm256_sub_epi32(c, _mm256_castps_si256(_mm256_cmp_ps::<_CMP_EQ_OQ>(a, b)))
1321 };
1322 equal_i64: i64 => i64, 4, _mm256_set1_epi64x, _mm256_loadu_si256, _mm256_setzero_si256(), _mm256_loadu_si256, _mm256_storeu_si256, |a, b, c| {
1323 _mm256_sub_epi64(c, _mm256_cmpeq_epi64(a, b))
1324 };
1325 equal_i32: i32 => i32, 8, _mm256_set1_epi32, _mm256_loadu_si256, _mm256_setzero_si256(), _mm256_loadu_si256, _mm256_storeu_si256, |a, b, c| {
1326 _mm256_sub_epi32(c, _mm256_cmpeq_epi32(a, b))
1327 };
1328 equal_i16: i16 => i16, 16, _mm256_set1_epi16, _mm256_loadu_si256, _mm256_setzero_si256(), _mm256_loadu_si256, _mm256_storeu_si256, |a, b, c| {
1329 _mm256_sub_epi16(c, _mm256_cmpeq_epi16(a, b))
1330 };
1331 }
1332}
1333#[cfg(any(
1334 target_arch = "aarch64", all(target_arch = "wasm32", target_feature = "simd128"),
1335 all(target_arch = "x86_64", target_feature = "avx2", target_feature = "fma")
1336))]
1337use lanes::*;
1338
1339const BAND_WORK: usize = 1 << 19;
1341
1342const THREAD_WORK: usize = 1 << 17;
1344
1345const PANEL: usize = 1 << 22;
1347const SERIAL_PANEL: usize = 2 << 20;
1349const ROWS: usize = 16;
1351const SWEEP: usize = 128;
1353#[inline(always)]
1355fn product<T: Copy + Default + Send + Sync, P: Pair<T>, const MR: usize, const NR: usize>(
1356 c: &mut [P::Out], ldc: usize, a: impl Fn([usize; 2], usize, &mut [[MaybeUninit<T>; MR]]) + Sync,
1357 b: impl Fn([usize; 2], usize, &mut [[MaybeUninit<T>; NR]]) + Sync, [m, k, n]: [usize; 3], start: bool, part: Part,
1358 kernel: impl Fn(&[T], &[T], &mut [[P::Out; NR]; MR], bool) + Sync) {
1359 if m == 0 || k == 0 || n == 0 { return; }
1360 assert!(ldc >= n && c.len() >= (m - 1) * ldc + n, "`c` must hold `m` rows of `n` items, `ldc` apart");
1361 let c = &mut c[..(m - 1) * ldc + n];
1362 let width = NC.min(n).next_multiple_of(NR);
1363 let panel = if parallel::threads() == 1 { SERIAL_PANEL / size_of::<T>() } else { PANEL };
1364 let depth = (panel / (width * KC)).clamp(1, k.div_ceil(KC)) * KC;
1365 let mut bp = Vec::with_capacity(width * depth.min(k));
1366 for jc in (0..n).step_by(NC) {
1367 let nc = NC.min(n - jc);
1368 for pp in (0..k).step_by(depth) {
1369 let kp = depth.min(k - pp);
1370 let unit = NR * kp;
1371 let work = m * kp * nc;
1372 let threads = if work < 2 * BAND_WORK { 1 } else { parallel::threads().min((work / THREAD_WORK).isqrt()) };
1373 let band_rows = if threads == 1 { m } else { m.div_ceil((work / BAND_WORK).min(4 * threads)).next_multiple_of(MR).min(MC) };
1374 let pack_b = |at: usize, units: &mut [MaybeUninit<T>]| {
1375 for p0 in (0..kp).step_by(ROWS) {
1376 let rows = ROWS.min(kp - p0);
1377 for (u, strip) in units.chunks_exact_mut(unit).enumerate() {
1378 let s = at / unit + u;
1379 let (j, lines) = (jc + s * NR, NR.min(nc - s * NR));
1380 pack::<_, NR>(&mut strip[p0 * NR..][..rows * NR], rows, lines, |first, ln, steps| b([j + first, ln], pp + p0, steps));
1381 }
1382 }
1383 true
1384 };
1385 bp.clear();
1386 unsafe {
1387 slice::appended(&mut bp, nc.div_ceil(NR) * unit, |slots| if threads == 1 { pack_b(0, slots) } else { parallel::split(slots, unit, &pack_b) })
1388 };
1389 let (a, bp, kernel) = (&a, &bp[..], &kernel);
1390 let lower = matches!(part, Part::Lower);
1391 parallel::pieces_of(c, band_rows * ldc, threads, &move |first, band| {
1392 let (i0, i1) = (first / ldc, (first / ldc + band_rows).min(m));
1393 let mut slots = Box::<[T]>::new_uninit_slice(MC.min(i1 - i0).next_multiple_of(MR) * KC.min(kp));
1394 for step in 0..kp.div_ceil(KC) {
1395 let pc = pp + step * KC;
1396 let (kc, load) = (KC.min(k - pc), pc > 0 || !start);
1397 for ic in (i0..i1).step_by(MC) {
1398 let mc = MC.min(i1 - ic);
1399 if lower && ic + mc <= jc { continue; }
1400 let ap = pack::<_, MR>(&mut slots, kc, mc, |first, lines, steps| a([ic + first, lines], pc, steps));
1401 for jb in (0..nc).step_by(SWEEP) {
1402 for ir in (0..mc).step_by(MR) {
1403 for jr in (jb..nc.min(jb + SWEEP)).step_by(NR) {
1404 let bj = &bp[jr / NR * unit + step * NR * KC..][..NR * kc];
1405 if lower && ic + ir + MR.min(mc - ir) <= jc + jr { continue; }
1406 let (rows, cols) = (MR.min(mc - ir), NR.min(nc - jr));
1407 let at = |r: usize| (ic - i0 + ir + r) * ldc + jc + jr;
1408 let mut t = [[P::UNIT; NR]; MR];
1409 if load {
1411 for r in 0..rows {
1412 if cols == NR { t[r].copy_from_slice(&band[at(r)..][..NR]) } else { for s in 0..cols { t[r][s] = band[at(r) + s] } }
1413 }
1414 }
1415 kernel(&ap[ir * kc..(ir + MR) * kc], bj, &mut t, !load);
1416 for r in 0..rows {
1417 if cols == NR { band[at(r)..][..NR].copy_from_slice(&t[r]) } else { for s in 0..cols { band[at(r) + s] = t[r][s] } }
1418 }
1419 }
1420 }
1421 }
1422 }
1423 }
1424 true
1425 });
1426 }
1427 }
1428}
1429
1430fn pack<T: Copy + Default, const W: usize>(
1432 buffer: &mut [MaybeUninit<T>], kc: usize, len: usize, strip: impl Fn(usize, usize, &mut [[MaybeUninit<T>; W]])) -> &[T] {
1433 let strips = len.div_ceil(W);
1434 for (s, chunk) in buffer.chunks_exact_mut(W * kc).take(strips).enumerate() { strip(s * W, W.min(len - s * W), chunk.as_chunks_mut::<W>().0); }
1435 unsafe { slice::written(&mut buffer[..strips * W * kc]) }
1436}
1437
1438#[inline(always)]
1440fn micro<T: Copy, P: Pair<T>, const MR: usize, const NR: usize>(a: &[T], b: &[T], t: &mut [[P::Out; NR]; MR], first: bool) {
1441 let mut acc = if first { [[P::UNIT; NR]; MR] } else { *t };
1442 for (a, b) in a.as_chunks::<MR>().0.iter().zip(b.as_chunks::<NR>().0) { for r in 0..MR { for s in 0..NR { acc[r][s] = P::step(acc[r][s], a[r], b[s]) } } }
1443 *t = acc;
1444}