1#![allow(clippy::items_after_statements)]
47
48#[cfg(feature = "bench-internals")]
53pub mod scalar;
54#[cfg(not(feature = "bench-internals"))]
55pub(crate) mod scalar;
56
57#[cfg(test)]
58mod proptest;
59
60#[allow(
66 clippy::cast_possible_wrap,
67 clippy::cast_ptr_alignment,
68 clippy::incompatible_msrv,
69 clippy::inline_always,
70 clippy::ptr_as_ptr,
71 clippy::unreadable_literal,
72 clippy::verbose_bit_mask,
73 clippy::wildcard_imports
74)]
75#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
76pub(crate) mod x86;
77
78#[cfg(target_arch = "aarch64")]
79pub(crate) mod arm;
80
81pub(crate) mod wasm;
82
83pub(crate) type AxpyFn = unsafe fn(c: u8, x: &[u8], y: &mut [u8]);
89pub(crate) type ScaleFn = unsafe fn(c: u8, x: &[u8], y: &mut [u8]);
91pub(crate) type ScaleInplaceFn = unsafe fn(c: u8, y: &mut [u8]);
93
94pub(crate) struct KernelSet {
96 pub(crate) axpy: AxpyFn,
98 pub(crate) scale: ScaleFn,
100 pub(crate) scale_inplace: ScaleInplaceFn,
102 pub(crate) name: &'static str,
104}
105
106#[cfg(feature = "std")]
111mod runtime {
112 #[cfg(target_arch = "aarch64")]
113 use super::arm;
114 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
115 use super::x86;
116 #[cfg(not(target_arch = "aarch64"))]
117 use super::{scalar_axpy_wrapper, scalar_scale_inplace_wrapper, scalar_scale_wrapper};
118 use super::{AxpyFn, KernelSet, ScaleFn, ScaleInplaceFn};
119 use std::sync::OnceLock;
120
121 static KERNEL: OnceLock<KernelSet> = OnceLock::new();
122
123 pub(crate) fn get() -> &'static KernelSet {
125 KERNEL.get_or_init(detect)
126 }
127
128 fn detect() -> KernelSet {
130 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
132 {
133 if is_x86_feature_detected!("gfni")
135 && is_x86_feature_detected!("avx512f")
136 && is_x86_feature_detected!("avx512bw")
137 {
138 return KernelSet {
139 axpy: x86::gfni_avx512::axpy_gfni_avx512 as AxpyFn,
140 scale: x86::gfni_avx512::scale_gfni_avx512 as ScaleFn,
141 scale_inplace: x86::gfni_avx512::scale_inplace_gfni_avx512 as ScaleInplaceFn,
142 name: "gfni+avx512 (tier1)",
143 };
144 }
145
146 if is_x86_feature_detected!("gfni") && is_x86_feature_detected!("avx2") {
148 return KernelSet {
149 axpy: x86::gfni_avx2::axpy_gfni_avx2 as AxpyFn,
150 scale: x86::gfni_avx2::scale_gfni_avx2 as ScaleFn,
151 scale_inplace: x86::gfni_avx2::scale_inplace_gfni_avx2 as ScaleInplaceFn,
152 name: "gfni+avx2 (tier2)",
153 };
154 }
155
156 if is_x86_feature_detected!("gfni") && is_x86_feature_detected!("sse4.2") {
158 return KernelSet {
159 axpy: x86::gfni_sse::axpy_gfni_sse as AxpyFn,
160 scale: x86::gfni_sse::scale_gfni_sse as ScaleFn,
161 scale_inplace: x86::gfni_sse::scale_inplace_gfni_sse as ScaleInplaceFn,
162 name: "gfni+sse4.2 (tier3)",
163 };
164 }
165
166 if is_x86_feature_detected!("avx512f")
168 && is_x86_feature_detected!("avx512bw")
169 && is_x86_feature_detected!("ssse3")
170 {
171 return KernelSet {
172 axpy: x86::avx512_ssse3::axpy_avx512_ssse3 as AxpyFn,
173 scale: x86::avx512_ssse3::scale_avx512_ssse3 as ScaleFn,
174 scale_inplace: x86::avx512_ssse3::scale_inplace_avx512_ssse3 as ScaleInplaceFn,
175 name: "avx512+ssse3 (tier4)",
176 };
177 }
178
179 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("ssse3") {
181 return KernelSet {
182 axpy: x86::avx2_ssse3::axpy_avx2_ssse3 as AxpyFn,
183 scale: x86::avx2_ssse3::scale_avx2_ssse3 as ScaleFn,
184 scale_inplace: x86::avx2_ssse3::scale_inplace_avx2_ssse3 as ScaleInplaceFn,
185 name: "avx2+ssse3 (tier5)",
186 };
187 }
188
189 if is_x86_feature_detected!("ssse3") {
191 return KernelSet {
192 axpy: x86::ssse3::axpy_ssse3 as AxpyFn,
193 scale: x86::ssse3::scale_ssse3 as ScaleFn,
194 scale_inplace: x86::ssse3::scale_inplace_ssse3 as ScaleInplaceFn,
195 name: "ssse3 (tier6)",
196 };
197 }
198 }
199
200 #[cfg(target_arch = "aarch64")]
202 {
203 return KernelSet {
205 axpy: arm::neon::axpy_neon as AxpyFn,
206 scale: arm::neon::scale_neon as ScaleFn,
207 scale_inplace: arm::neon::scale_inplace_neon as ScaleInplaceFn,
208 name: "neon (tier7)",
209 };
210 }
211
212 #[cfg(not(target_arch = "aarch64"))]
219 KernelSet {
220 axpy: scalar_axpy_wrapper as AxpyFn,
221 scale: scalar_scale_wrapper as ScaleFn,
222 scale_inplace: scalar_scale_inplace_wrapper as ScaleInplaceFn,
223 name: "scalar (tier9)",
224 }
225 }
226}
227
228#[cfg(all(feature = "std", not(target_arch = "aarch64")))]
233unsafe fn scalar_axpy_wrapper(c: u8, x: &[u8], y: &mut [u8]) {
234 scalar::axpy(c, x, y);
235}
236
237#[cfg(all(feature = "std", not(target_arch = "aarch64")))]
238unsafe fn scalar_scale_wrapper(c: u8, x: &[u8], y: &mut [u8]) {
239 scalar::scale(c, x, y);
240}
241
242#[cfg(all(feature = "std", not(target_arch = "aarch64")))]
243unsafe fn scalar_scale_inplace_wrapper(c: u8, y: &mut [u8]) {
244 scalar::scale_inplace(c, y);
245}
246
247#[inline]
267pub fn axpy(c: u8, x: &[u8], y: &mut [u8]) {
268 assert_eq!(
269 x.len(),
270 y.len(),
271 "rlnc_simdx::kernel::axpy: length mismatch"
272 );
273 assert!(
274 !ranges_overlap(x.as_ptr(), x.len(), y.as_ptr(), y.len()),
275 "rlnc_simdx::kernel::axpy: overlapping buffers are not allowed"
276 );
277 #[cfg(feature = "std")]
278 unsafe {
281 (runtime::get().axpy)(c, x, y);
282 }
283
284 #[cfg(not(feature = "std"))]
285 axpy_static(c, x, y)
286}
287
288#[inline]
299pub fn scale(c: u8, x: &[u8], y: &mut [u8]) {
300 assert_eq!(
301 x.len(),
302 y.len(),
303 "rlnc_simdx::kernel::scale: length mismatch"
304 );
305 assert!(
306 !ranges_overlap(x.as_ptr(), x.len(), y.as_ptr(), y.len()),
307 "rlnc_simdx::kernel::scale: overlapping buffers are not allowed; use scale_inplace for in-place"
308 );
309 #[cfg(feature = "std")]
310 unsafe {
311 (runtime::get().scale)(c, x, y);
312 }
313
314 #[cfg(not(feature = "std"))]
315 scale_static(c, x, y)
316}
317
318#[inline]
324pub fn scale_inplace(c: u8, y: &mut [u8]) {
325 #[cfg(feature = "std")]
326 unsafe {
328 (runtime::get().scale_inplace)(c, y);
329 }
330
331 #[cfg(not(feature = "std"))]
332 scale_inplace_static(c, y)
333}
334
335#[inline]
350pub fn axpy_multi(coeffs: &[u8], sources: &[&[u8]], y: &mut [u8]) {
351 assert_eq!(
352 coeffs.len(),
353 sources.len(),
354 "axpy_multi: coeffs/sources len"
355 );
356
357 for (i, src) in sources.iter().enumerate() {
360 assert_eq!(
361 src.len(),
362 y.len(),
363 "axpy_multi: source[{i}] length mismatch"
364 );
365 }
366 for (i, src) in sources.iter().enumerate() {
367 assert!(
368 !ranges_overlap(src.as_ptr(), src.len(), y.as_ptr(), y.len()),
369 "axpy_multi: source[{i}] overlaps destination"
370 );
371 }
372
373 #[cfg(feature = "std")]
374 let axpy_kernel = runtime::get().axpy;
375
376 const BLOCK: usize = 4096;
379 let n = y.len();
380 let mut off = 0usize;
381 while off < n {
382 let end = (off + BLOCK).min(n);
383 for (i, &c) in coeffs.iter().enumerate() {
384 if c != 0 {
385 let source_block = &sources[i][off..end];
386 let destination_block = &mut y[off..end];
387
388 #[cfg(feature = "std")]
389 unsafe {
394 axpy_kernel(c, source_block, destination_block);
395 }
396
397 #[cfg(not(feature = "std"))]
398 axpy_static(c, source_block, destination_block);
399 }
400 }
401 off = end;
402 }
403}
404
405#[inline]
408fn ranges_overlap(a: *const u8, a_len: usize, b: *const u8, b_len: usize) -> bool {
409 if a_len == 0 || b_len == 0 {
410 return false;
411 }
412 let a0 = a as usize;
413 let b0 = b as usize;
414 let a1 = a0 + a_len;
415 let b1 = b0 + b_len;
416 a0 < b1 && b0 < a1
417}
418
419#[inline]
424#[must_use]
425pub fn dot(a: &[u8], b: &[u8]) -> u8 {
426 assert_eq!(a.len(), b.len(), "rlnc_simdx::kernel::dot: length mismatch");
427 scalar::dot(a, b)
428}
429
430#[must_use]
432pub fn active_kernel_name() -> &'static str {
433 #[cfg(feature = "std")]
434 {
435 runtime::get().name
436 }
437
438 #[cfg(not(feature = "std"))]
439 {
440 active_kernel_name_static()
441 }
442}
443
444#[cfg(not(feature = "std"))]
450#[inline]
451fn axpy_static(c: u8, x: &[u8], y: &mut [u8]) {
452 #[cfg(all(target_family = "wasm", target_feature = "simd128"))]
454 return unsafe { wasm::simd128::axpy_wasm(c, x, y) };
456
457 #[cfg(target_arch = "aarch64")]
459 return unsafe { arm::neon::axpy_neon(c, x, y) };
460
461 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
463 {
464 #[cfg(all(
465 target_feature = "gfni",
466 target_feature = "avx512f",
467 target_feature = "avx512bw"
468 ))]
469 return unsafe { x86::gfni_avx512::axpy_gfni_avx512(c, x, y) };
470 #[cfg(all(
471 target_feature = "gfni",
472 target_feature = "avx2",
473 not(all(target_feature = "avx512f", target_feature = "avx512bw"))
474 ))]
475 return unsafe { x86::gfni_avx2::axpy_gfni_avx2(c, x, y) };
476 #[cfg(all(
477 target_feature = "avx2",
478 target_feature = "ssse3",
479 not(target_feature = "gfni")
480 ))]
481 return unsafe { x86::avx2_ssse3::axpy_avx2_ssse3(c, x, y) };
482 #[cfg(all(target_feature = "ssse3", not(target_feature = "avx2")))]
483 return unsafe { x86::ssse3::axpy_ssse3(c, x, y) };
484 }
485
486 scalar::axpy(c, x, y)
488}
489
490#[cfg(not(feature = "std"))]
491#[inline]
492fn scale_static(c: u8, x: &[u8], y: &mut [u8]) {
493 #[cfg(all(target_family = "wasm", target_feature = "simd128"))]
494 return unsafe { wasm::simd128::scale_wasm(c, x, y) };
496
497 #[cfg(target_arch = "aarch64")]
498 return unsafe { arm::neon::scale_neon(c, x, y) };
499
500 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
501 {
502 #[cfg(all(
503 target_feature = "gfni",
504 target_feature = "avx512f",
505 target_feature = "avx512bw"
506 ))]
507 return unsafe { x86::gfni_avx512::scale_gfni_avx512(c, x, y) };
508 #[cfg(all(
509 target_feature = "gfni",
510 target_feature = "avx2",
511 not(all(target_feature = "avx512f", target_feature = "avx512bw"))
512 ))]
513 return unsafe { x86::gfni_avx2::scale_gfni_avx2(c, x, y) };
514 #[cfg(all(
515 target_feature = "avx2",
516 target_feature = "ssse3",
517 not(target_feature = "gfni")
518 ))]
519 return unsafe { x86::avx2_ssse3::scale_avx2_ssse3(c, x, y) };
520 #[cfg(all(target_feature = "ssse3", not(target_feature = "avx2")))]
521 return unsafe { x86::ssse3::scale_ssse3(c, x, y) };
522 }
523
524 scalar::scale(c, x, y)
525}
526
527#[cfg(not(feature = "std"))]
528#[inline]
529fn scale_inplace_static(c: u8, y: &mut [u8]) {
530 #[cfg(all(target_family = "wasm", target_feature = "simd128"))]
531 return unsafe { wasm::simd128::scale_inplace_wasm(c, y) };
533
534 #[cfg(target_arch = "aarch64")]
535 return unsafe { arm::neon::scale_inplace_neon(c, y) };
536
537 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
538 {
539 #[cfg(all(
540 target_feature = "gfni",
541 target_feature = "avx512f",
542 target_feature = "avx512bw"
543 ))]
544 return unsafe { x86::gfni_avx512::scale_inplace_gfni_avx512(c, y) };
545 #[cfg(all(
546 target_feature = "gfni",
547 target_feature = "avx2",
548 not(all(target_feature = "avx512f", target_feature = "avx512bw"))
549 ))]
550 return unsafe { x86::gfni_avx2::scale_inplace_gfni_avx2(c, y) };
551 #[cfg(all(
552 target_feature = "avx2",
553 target_feature = "ssse3",
554 not(target_feature = "gfni")
555 ))]
556 return unsafe { x86::avx2_ssse3::scale_inplace_avx2_ssse3(c, y) };
557 #[cfg(all(target_feature = "ssse3", not(target_feature = "avx2")))]
558 return unsafe { x86::ssse3::scale_inplace_ssse3(c, y) };
559 }
560
561 scalar::scale_inplace(c, y)
562}
563
564#[cfg(not(feature = "std"))]
565fn active_kernel_name_static() -> &'static str {
566 #[cfg(all(target_family = "wasm", target_feature = "simd128"))]
567 return "wasm-simd128 (tier8)";
568 #[cfg(target_arch = "aarch64")]
569 return "neon (tier7)";
570 #[cfg(all(target_feature = "gfni", target_feature = "avx512f"))]
571 return "gfni+avx512 (tier1)";
572 #[cfg(all(
573 target_feature = "gfni",
574 target_feature = "avx2",
575 not(target_feature = "avx512f")
576 ))]
577 return "gfni+avx2 (tier2)";
578 #[cfg(all(
579 target_feature = "avx2",
580 target_feature = "ssse3",
581 not(target_feature = "gfni")
582 ))]
583 return "avx2+ssse3 (tier5)";
584 #[cfg(all(target_feature = "ssse3", not(target_feature = "avx2")))]
585 return "ssse3 (tier6)";
586 "scalar (tier9)"
587}
588
589#[cfg(test)]
594mod tests {
595 use super::*;
596
597 #[test]
598 fn axpy_round_trip() {
599 let c = 0x53u8;
600 let x: Vec<u8> = (0u8..128).collect();
601 let mut y = vec![0u8; 128];
602 axpy(c, &x, &mut y);
603 axpy(c, &x, &mut y);
604 assert_eq!(y, vec![0u8; 128]);
605 }
606
607 #[test]
608 fn scale_axpy_consistency() {
609 let c = 0xC7u8;
610 let x: Vec<u8> = (1u8..=64).collect();
611 let mut y_scale = vec![0u8; 64];
612 let mut y_axpy = vec![0u8; 64];
613 scale(c, &x, &mut y_scale);
614 axpy(c, &x, &mut y_axpy);
615 assert_eq!(y_scale, y_axpy);
616 }
617
618 #[test]
619 fn kernel_name_not_empty() {
620 let name = active_kernel_name();
621 assert!(!name.is_empty());
622 println!("Active kernel: {name}");
623 }
624
625 #[test]
626 #[cfg(feature = "std")]
627 fn runtime_dispatch_selects_best() {
628 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
630 {
631 let name = active_kernel_name();
632 if is_x86_feature_detected!("gfni") && is_x86_feature_detected!("avx512f") {
633 assert!(name.contains("tier1"), "Expected tier1, got: {name}");
634 } else if is_x86_feature_detected!("gfni") && is_x86_feature_detected!("avx2") {
635 assert!(name.contains("tier2"), "Expected tier2, got: {name}");
636 } else if is_x86_feature_detected!("ssse3") {
637 assert!(
638 name.contains("tier3")
639 || name.contains("tier4")
640 || name.contains("tier5")
641 || name.contains("tier6"),
642 "Expected SSSE3-tier, got: {name}"
643 );
644 }
645 }
646
647 #[cfg(target_arch = "aarch64")]
648 assert_eq!(active_kernel_name(), "neon (tier7)");
649
650 #[cfg(not(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64")))]
651 assert!(!active_kernel_name().is_empty());
652 }
653}