Smaller updates: edits on AVX-512 functionality, small refactor

Cargo.toml

@@ -13,7 +13,7 @@ repository = "https://github.com/bluss/matrixmultiply/"
 documentation = "https://docs.rs/matrixmultiply/"
 
 description = """
-General matrix multiplication for f32 and f64 matrices. Operates on matrices with general layout (they can use arbitrary row and column stride). Detects and uses AVX-512, AVX or SSE2 on x86 platforms transparently for higher performance. Uses a microkernel strategy, so that the implementation is easy to parallelize and optimize.
+General matrix multiplication for f32 and f64 matrices. Operates on matrices with general layout (they can use arbitrary row and column stride). Detects and uses SIMD features on x86/x86-64 and AArch64 transparently for higher performance. Uses a microkernel strategy, so that the implementation is easy to parallelize and optimize.
 
 Supports multithreading."""
 

build.rs

@@ -1,6 +1,8 @@
 fn main() {
+    // NOTE: from Rust 1.77: `cargo::` syntax. As long as before that is supported we use `cargo:`.
+
     println!("cargo:rerun-if-changed=build.rs");
-    let target_arch = std::env::var("CARGO_CFG_TARGET_ARCH").unwrap_or(String::new());
+    let target_arch = std::env::var("CARGO_CFG_TARGET_ARCH").unwrap_or_default();
 
     let ac = match autocfg::AutoCfg::new() {
         Ok(ac) => ac,
@@ -12,7 +14,8 @@ fn main() {
 
     // Avoid `unexpected_cfgs` lint from 1.80+ toolchains
     if ac.probe_rustc_version(1, 80) {
-        println!("cargo::rustc-check-cfg=cfg(has_avx512)");
+        println!("cargo:rustc-check-cfg=cfg(has_avx512)");
+        println!("cargo:rustc-check-cfg=cfg(has_aarch64_simd)");
     }
 
     if target_arch == "aarch64" {

src/cgemm_kernel.rs

@@ -11,6 +11,7 @@ use crate::kernel::GemmSelect;
 use crate::kernel::{U2, U4, c32, Element, c32_mul as mul};
 #[cfg(has_avx512)]
 use crate::kernel::U8;
+use crate::kernel_util::at;
 use crate::archparam;
 use crate::cgemm_common::pack_complex;
 use crate::packing::PackSlice;
@@ -300,11 +301,6 @@ kernel_fallback_impl_complex! {
     kernel_fallback_impl, T, TReal, KernelFallback::MR, KernelFallback::NR, 1
 }
 
-#[inline(always)]
-unsafe fn at(ptr: *const TReal, i: usize) -> TReal {
-    *ptr.add(i)
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

src/dgemm_kernel.rs

@@ -10,6 +10,9 @@ use crate::kernel::GemmKernel;
 use crate::kernel::GemmSelect;
 #[allow(unused)]
 use crate::kernel::{U4, U8};
+#[cfg(has_avx512)]
+use crate::kernel_util::preferential_transpose;
+use crate::kernel_util::at;
 use crate::archparam;
 
 #[cfg(target_arch="x86")]
@@ -951,10 +954,8 @@ unsafe fn kernel_target_avx512(k: usize, alpha: T, a: *const T, b: *const T,
 
     let mut ab = [_mm512_setzero_pd(); MR];
 
-    // Compute C in whichever orientation makes the output columns contiguous.
-    let prefer_row_major_c = rsc != 1;
-    let (mut a, mut b) = if prefer_row_major_c { (a, b) } else { (b, a) };
-    let (rsc, csc) = if prefer_row_major_c { (rsc, csc) } else { (csc, rsc) };
+    // Compute C in whichever orientation makes the output columns contiguous
+    let (mut a, mut b, rsc, csc) = preferential_transpose(MR, NR, a, b, rsc, csc);
 
     // Compute A B. The packed buffers are 64-byte aligned
     let mut bv = _mm512_load_pd(b);
@@ -1150,8 +1151,8 @@ unsafe fn kernel_fallback_impl(k: usize, alpha: T, a: *const T, b: *const T,
     unroll_by!(4 => k, {
         loop4!(i, loop4!(j, ab[i][j] += at(a, i) * at(b, j)));
 
-        a = a.offset(MR as isize);
-        b = b.offset(NR as isize);
+        a = a.add(MR);
+        b = b.add(NR);
     });
 
     macro_rules! c {
@@ -1162,11 +1163,6 @@ unsafe fn kernel_fallback_impl(k: usize, alpha: T, a: *const T, b: *const T,
     loop4!(j, loop4!(i, *c![i, j] = alpha * ab[i][j]));
 }
 
-#[inline(always)]
-unsafe fn at(ptr: *const T, i: usize) -> T {
-    *ptr.offset(i as isize)
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

src/kernel_util.rs

@@ -0,0 +1,36 @@
+/// If the microkernel tile is a square, mr == nr,
+/// then we can change the problem from C = A B to C^T = B^T A^T
+///
+/// Whenever rsc == 1, we switch the orientation around.
+/// *Any contiguous positive stride* C will exit this function with rsc arbitrary, csc == 1.
+///
+/// Input row stride (rs), column stride (cs) and outcomes
+/// rsc     csc     explanation
+/// M       1       no change => (M, 1)
+/// 1       Z       transpose => (Z, 1)
+/// M1      M2      no change => (M1, M2)
+///
+/// Where M signifies number M != 1 and Z any number (pos, neg, or zero)
+///
+/// + `mr`, `nr` kernel size; must be square
+/// + `a` pointer to packing buffer
+/// + `b` pointer to packing buffer
+/// + `rsc` row stride of c
+/// + `csc` colum stride of c
+#[inline(always)]
+#[allow(unused)]
+pub(crate) fn preferential_transpose<T>(mr: usize, nr: usize, a: *const T, b: *const T, rsc: isize, csc: isize)
+    -> (*const T, *const T, isize, isize)
+{
+    debug_assert_eq!(mr, nr, "Transpose requires that MR == NR");
+    // Compute C in whichever orientation makes the output columns contiguous
+    let prefer_row_major_c = rsc != 1;
+    if prefer_row_major_c { (a, b, rsc, csc) } else { (b, a, csc, rsc) }
+}
+
+/// Dereference pointer at offset `i`
+#[inline(always)]
+pub(crate) unsafe fn at<T: Copy>(ptr: *const T, i: usize) -> T {
+    *ptr.add(i)
+}
+

src/lib.rs

@@ -52,16 +52,16 @@
 //!   on all targets. These may depend on autovectorization to perform well.
 //!
 //! - *x86* and *x86-64* features can be detected at runtime by default or
-//!   compile time (if enabled), and the following kernel variants are
+//!   compile time, and the following kernel variants are
 //!   implemented:
 //!
 //!   - `fma`
 //!   - `avx`
 //!   - `sse2`
 //!   - `avx512f`
 //!
-//! - *aarch64* features can be detected at runtime by default or compile time
-//!   (if enabled), and the following kernel variants are implemented:
+//! - *aarch64* features can be detected at runtime by default or compile time,
+//!   and the following kernel variants are implemented:
 //!
 //!   - `neon`
 //!
@@ -97,14 +97,8 @@
 //!
 //! `avx512` is enabled by default.
 //!
-//! It compiles the AVX-512 kernels, which are then used at runtime on CPUs
-//! that support the `avx512f` (maybe more avx512 subsets support in the future) target feature.
-//! It requires Rust 1.89 or later and has no effect on older compilers.
-//! To disable it, use this in your `Cargo.toml`:
-//!
-//! ```toml
-//! matrixmultiply = { version = "0.3", default-features = false, features = ["std"] }
-//! ```
+//! It enables the AVX-512 kernels for x86/x86-64, which are then used at runtime on CPUs that
+//! support the `avx512f` target feature.
 //!
 //! ### `threading`
 //!
@@ -170,6 +164,7 @@ pub(crate) use archparam_defaults as archparam;
 
 mod gemm;
 mod kernel;
+mod kernel_util;
 mod packing;
 mod ptr;
 mod threading;

src/sgemm_kernel.rs

@@ -11,6 +11,9 @@ use crate::kernel::GemmSelect;
 use crate::kernel::{U4, U8};
 #[cfg(has_avx512)]
 use crate::kernel::U16;
+#[cfg(any(target_arch="x86", target_arch="x86_64", target_arch="aarch64", target_arch="wasm32"))]
+use crate::kernel_util::preferential_transpose;
+use crate::kernel_util::at;
 use crate::archparam;
 
 #[cfg(target_arch="x86")]
@@ -418,11 +421,8 @@ unsafe fn kernel_x86_avx<MA>(k: usize, alpha: T, a: *const T, b: *const T,
 
     let mut ab = [_mm256_setzero_ps(); MR];
 
-    // this kernel can operate in either transposition (C = A B or C^T = B^T A^T)
-    let prefer_row_major_c = rsc != 1;
-
-    let (mut a, mut b) = if prefer_row_major_c { (a, b) } else { (b, a) };
-    let (rsc, csc) = if prefer_row_major_c { (rsc, csc) } else { (csc, rsc) };
+    // Compute C in whichever orientation makes the output columns contiguous
+    let (mut a, mut b, rsc, csc) = preferential_transpose(MR, NR, a, b, rsc, csc);
 
     macro_rules! permute_mask {
         ($z:expr, $y:expr, $x:expr, $w:expr) => {
@@ -610,9 +610,7 @@ unsafe fn kernel_target_avx512(k: usize, alpha: T, a: *const T, b: *const T,
     let mut ab = [_mm512_setzero_ps(); MR];
 
     // Compute C in whichever orientation makes the output columns contiguous
-    let prefer_row_major_c = rsc != 1;
-    let (mut a, mut b) = if prefer_row_major_c { (a, b) } else { (b, a) };
-    let (rsc, csc) = if prefer_row_major_c { (rsc, csc) } else { (csc, rsc) };
+    let (mut a, mut b, rsc, csc) = preferential_transpose(MR, NR, a, b, rsc, csc);
 
     // Compute A B. The packed buffers are 64-byte aligned
     let mut bv = _mm512_load_ps(b);
@@ -672,7 +670,7 @@ unsafe fn kernel_target_neon(k: usize, alpha: T, a: *const T, b: *const T,
     const MR: usize = KernelNeon::MR;
     const NR: usize = KernelNeon::NR;
 
-    let (mut a, mut b, rsc, csc) = if rsc == 1 { (b, a, csc, rsc) } else { (a, b, rsc, csc) };
+    let (mut a, mut b, rsc, csc) = preferential_transpose(MR, NR, a, b, rsc, csc);
 
     // Kernel 8 x 8 (a x b)
     // Four quadrants of 4 x 4
@@ -819,7 +817,7 @@ unsafe fn kernel_target_wasm_simd(k: usize, alpha: T, a: *const T, b: *const T,
         f32x4_add(f32x4_mul(a, b), c)
     }
 
-    let (mut a, mut b, rsc, csc) = if rsc == 1 { (b, a, csc, rsc) } else { (a, b, rsc, csc) };
+    let (mut a, mut b, rsc, csc) = preferential_transpose(MR, NR, a, b, rsc, csc);
 
     // Kernel 8 x 8 (a x b)
     // Four quadrants of 4 x 4
@@ -961,8 +959,8 @@ unsafe fn kernel_fallback_impl(k: usize, alpha: T, a: *const T, b: *const T,
     unroll_by!(4 => k, {
         loop8!(i, loop4!(j, ab[i][j] += at(a, i) * at(b, j)));
 
-        a = a.offset(MR as isize);
-        b = b.offset(NR as isize);
+        a = a.add(MR);
+        b = b.add(NR);
     });
 
     macro_rules! c {
@@ -973,11 +971,6 @@ unsafe fn kernel_fallback_impl(k: usize, alpha: T, a: *const T, b: *const T,
     loop4!(j, loop8!(i, *c![i, j] = alpha * ab[i][j]));
 }
 
-#[inline(always)]
-unsafe fn at(ptr: *const T, i: usize) -> T {
-    *ptr.offset(i as isize)
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

src/zgemm_kernel.rs

@@ -9,6 +9,7 @@
 use crate::kernel::GemmKernel;
 use crate::kernel::GemmSelect;
 use crate::kernel::{U2, U4, c64, Element, c64_mul as mul};
+use crate::kernel_util::at;
 use crate::archparam;
 use crate::cgemm_common::pack_complex;
 use crate::packing::PackSlice;
@@ -278,11 +279,6 @@ kernel_fallback_impl_complex! {
     kernel_target_avx512, T, TReal, KernelAvx512::MR, KernelAvx512::NR, 4
 }
 
-#[inline(always)]
-unsafe fn at(ptr: *const TReal, i: usize) -> TReal {
-    *ptr.add(i)
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;