diff --git a/Cargo.toml b/Cargo.toml index 20aee9a..0a60c01 100644 --- a/Cargo.toml +++ b/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.""" diff --git a/build.rs b/build.rs index d471f56..ac52cbf 100644 --- a/build.rs +++ b/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" { diff --git a/src/cgemm_kernel.rs b/src/cgemm_kernel.rs index daf3a47..1c0bf64 100644 --- a/src/cgemm_kernel.rs +++ b/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::*; diff --git a/src/dgemm_kernel.rs b/src/dgemm_kernel.rs index 20d56c8..1325eb0 100644 --- a/src/dgemm_kernel.rs +++ b/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::*; diff --git a/src/kernel_util.rs b/src/kernel_util.rs new file mode 100644 index 0000000..939703e --- /dev/null +++ b/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(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(ptr: *const T, i: usize) -> T { + *ptr.add(i) +} + diff --git a/src/lib.rs b/src/lib.rs index b49455d..23dd24e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -52,7 +52,7 @@ //! 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` @@ -60,8 +60,8 @@ //! - `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; diff --git a/src/sgemm_kernel.rs b/src/sgemm_kernel.rs index 0500716..26a0c4d 100644 --- a/src/sgemm_kernel.rs +++ b/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(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::*; diff --git a/src/zgemm_kernel.rs b/src/zgemm_kernel.rs index 5641f7c..71f4ce6 100644 --- a/src/zgemm_kernel.rs +++ b/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::*;