diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b34f3b2..26a8773 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -211,3 +211,31 @@ jobs: - name: Miri run: ci/miri.sh --features cgemm + # run the test suite under Intel SDE to ensure AVX-512 kernels are available + avx512_test: + runs-on: ubuntu-latest + strategy: + matrix: + include: + - rust: stable + target: x86_64-unknown-linux-gnu + name: avx512_test/${{ matrix.target }}/${{ matrix.rust }} + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + with: + toolchain: ${{ matrix.rust }} + targets: ${{ matrix.target }} + - name: Set up Intel SDE + # Sets SDE_PATH to the directory containing the SDE binaries + uses: petarpetrovt/setup-sde@v5.0 + - name: Tests under SDE (force + ensure avx512f) + run: | + export PATH="$SDE_PATH:$PATH" + cargo test -v --tests --lib --release --no-fail-fast --features cgemm + env: + # Emulate a CPU with avx512f + CARGO_TARGET_X86_64_UNKNOWN_LINUX_GNU_RUNNER: "sde64 -spr --" + MMTEST_FEATURE: avx512f + MMTEST_ENSUREFEATURE: 1 + MMTEST_FAST_TEST: 1 diff --git a/Cargo.toml b/Cargo.toml index ede9561..20aee9a 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 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 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. Supports multithreading.""" @@ -43,13 +43,14 @@ bencher = "0.1.2" itertools = "0.8" [features] -default = ["std"] +default = ["std", "avx512"] # support for complex f32, complex f64 cgemm = [] threading = ["thread-tree", "std", "once_cell", "num_cpus"] std = [] +avx512 = [] # support for compile-time configuration constconf = [] diff --git a/README.rst b/README.rst index 0cdf47e..ca6544c 100644 --- a/README.rst +++ b/README.rst @@ -51,6 +51,11 @@ __ https://bluss.github.io/rust/2016/03/28/a-gemmed-rabbit-hole/ Recent Changes -------------- +- Unreleased + + - Add AVX-512 microkernels for sgemm (16×16) and dgemm (8×8) + - Add AVX-512 microkernels for cgemm (8×4) and zgemm (4×4) + - 0.3.10 - sgemm: Reduce unnecessary AVX register permutations by `@SongXiaoXi `_ `#88 `_ diff --git a/build.rs b/build.rs index 0fa2751..d471f56 100644 --- a/build.rs +++ b/build.rs @@ -1,12 +1,32 @@ fn main() { println!("cargo:rerun-if-changed=build.rs"); - if std::env::var("CARGO_CFG_TARGET_ARCH").unwrap_or(String::new()) == "aarch64" { - match autocfg::AutoCfg::new() { - // From 1.61 aarch64 intrinsics and #[target_feature] - Ok(ac) => if ac.probe_rustc_version(1, 61) { - println!("cargo:rustc-cfg=has_aarch64_simd"); - } - Err(err) => println!("cargo:warning={}", err), + let target_arch = std::env::var("CARGO_CFG_TARGET_ARCH").unwrap_or(String::new()); + + let ac = match autocfg::AutoCfg::new() { + Ok(ac) => ac, + Err(err) => { + println!("cargo:warning={}", err); + return; + } + }; + + // Avoid `unexpected_cfgs` lint from 1.80+ toolchains + if ac.probe_rustc_version(1, 80) { + println!("cargo::rustc-check-cfg=cfg(has_avx512)"); + } + + if target_arch == "aarch64" { + // From 1.61 aarch64 intrinsics and #[target_feature] + if ac.probe_rustc_version(1, 61) { + println!("cargo:rustc-cfg=has_aarch64_simd"); + } + } + if target_arch == "x86" || target_arch == "x86_64" { + // From 1.89 AVX-512 intrinsics ("avx512f") + if ac.probe_rustc_version(1, 89) + && std::env::var_os("CARGO_FEATURE_AVX512").is_some() + { + println!("cargo:rustc-cfg=has_avx512"); } } } diff --git a/src/cgemm_kernel.rs b/src/cgemm_kernel.rs index 795704e..daf3a47 100644 --- a/src/cgemm_kernel.rs +++ b/src/cgemm_kernel.rs @@ -9,10 +9,14 @@ use crate::kernel::GemmKernel; use crate::kernel::GemmSelect; use crate::kernel::{U2, U4, c32, Element, c32_mul as mul}; +#[cfg(has_avx512)] +use crate::kernel::U8; use crate::archparam; use crate::cgemm_common::pack_complex; use crate::packing::PackSlice; +#[cfg(has_avx512)] +struct KernelAvx512; #[cfg(any(target_arch="x86", target_arch="x86_64"))] struct KernelAvx2; #[cfg(any(target_arch="x86", target_arch="x86_64"))] @@ -37,6 +41,12 @@ pub(crate) fn detect(selector: G) where G: GemmSelect { // dispatch to specific compiled versions #[cfg(any(target_arch="x86", target_arch="x86_64"))] { + #[cfg(has_avx512)] + { + if is_x86_feature_detected_!("avx512f") { + return selector.select(KernelAvx512); + } + } if is_x86_feature_detected_!("fma") { if is_x86_feature_detected_!("avx2") { return selector.select(KernelAvx2); @@ -54,6 +64,40 @@ pub(crate) fn detect(selector: G) where G: GemmSelect { return selector.select(KernelFallback); } +#[cfg(has_avx512)] +impl GemmKernel for KernelAvx512 { + type Elem = T; + + type MRTy = U8; + type NRTy = U4; + + #[inline(always)] + fn align_to() -> usize { 32 } + + #[inline(always)] + fn always_masked() -> bool { KernelFallback::always_masked() } + + #[inline(always)] + fn nc() -> usize { archparam::C_NC } + #[inline(always)] + fn kc() -> usize { archparam::C_KC } + #[inline(always)] + fn mc() -> usize { archparam::C_MC } + + pack_methods!{} + + #[inline(always)] + unsafe fn kernel( + k: usize, + alpha: T, + a: *const T, + b: *const T, + beta: T, + c: *mut T, rsc: isize, csc: isize) { + kernel_target_avx512(k, alpha, a, b, beta, c, rsc, csc) + } +} + #[cfg(any(target_arch="x86", target_arch="x86_64"))] impl GemmKernel for KernelAvx2 { type Elem = T; @@ -190,6 +234,19 @@ impl GemmKernel for KernelFallback { } } +// Kernel AVX-512 +#[cfg(has_avx512)] +macro_rules! loop_m { ($i:ident, $e:expr) => { loop8!($i, $e) }; } +#[cfg(has_avx512)] +macro_rules! loop_n { ($j:ident, $e:expr) => { loop4!($j, $e) }; } + +#[cfg(has_avx512)] +kernel_fallback_impl_complex! { + // instantiate separately + [inline target_feature(enable="avx512f")] [fma_yes] + kernel_target_avx512, T, TReal, KernelAvx512::MR, KernelAvx512::NR, 4 +} + // Kernel AVX2 #[cfg(any(target_arch="x86", target_arch="x86_64"))] macro_rules! loop_m { ($i:ident, $e:expr) => { loop4!($i, $e) }; } @@ -312,5 +369,10 @@ mod tests { "fma", fma, KernelFma, "avx2", avx2, KernelAvx2 } + + #[cfg(has_avx512)] + test_arch_kernels_x86! { + "avx512f", avx512f, KernelAvx512 + } } } diff --git a/src/dgemm_kernel.rs b/src/dgemm_kernel.rs index aa46150..20d56c8 100644 --- a/src/dgemm_kernel.rs +++ b/src/dgemm_kernel.rs @@ -30,6 +30,8 @@ struct KernelFmaAvx2; struct KernelFma; #[cfg(any(target_arch="x86", target_arch="x86_64"))] struct KernelSse2; +#[cfg(has_avx512)] +struct KernelAvx512; #[cfg(target_arch="aarch64")] #[cfg(has_aarch64_simd)] @@ -49,6 +51,12 @@ pub(crate) fn detect(selector: G) where G: GemmSelect { // dispatch to specific compiled versions #[cfg(any(target_arch="x86", target_arch="x86_64"))] { + #[cfg(has_avx512)] + { + if is_x86_feature_detected_!("avx512f") { + return selector.select(KernelAvx512); + } + } if is_x86_feature_detected_!("fma") { if is_x86_feature_detected_!("avx2") { return selector.select(KernelFmaAvx2); @@ -235,6 +243,54 @@ impl GemmKernel for KernelSse2 { } } +#[cfg(has_avx512)] +impl GemmKernel for KernelAvx512 { + type Elem = T; + + type MRTy = U8; + type NRTy = U8; + + #[inline(always)] + fn align_to() -> usize { 64 } + + #[inline(always)] + fn always_masked() -> bool { false } + + #[inline(always)] + fn nc() -> usize { archparam::D_NC } + #[inline(always)] + fn kc() -> usize { archparam::D_KC } + #[inline(always)] + fn mc() -> usize { archparam::D_MC } + + #[inline] + unsafe fn pack_mr(kc: usize, mc: usize, pack: PackSlice, + a: *const Self::Elem, rsa: isize, csa: isize) + { + // safety: avx512f is enabled + crate::packing::pack_avx512::(kc, mc, pack, a, rsa, csa) + } + + #[inline] + unsafe fn pack_nr(kc: usize, mc: usize, pack: PackSlice, + a: *const Self::Elem, rsa: isize, csa: isize) + { + // safety: avx512f is enabled + crate::packing::pack_avx512::(kc, mc, pack, a, rsa, csa) + } + + #[inline(always)] + unsafe fn kernel( + k: usize, + alpha: T, + a: *const T, + b: *const T, + beta: T, + c: *mut T, rsc: isize, csc: isize) { + kernel_target_avx512(k, alpha, a, b, beta, c, rsc, csc) + } +} + #[cfg(target_arch="aarch64")] #[cfg(has_aarch64_simd)] impl GemmKernel for KernelNeon { @@ -883,6 +939,71 @@ unsafe fn kernel_x86_avx(k: usize, alpha: T, a: *const T, b: *const T, } } +// no inline for unmasked kernels +#[cfg(has_avx512)] +#[target_feature(enable="avx512f")] +unsafe fn kernel_target_avx512(k: usize, alpha: T, a: *const T, b: *const T, + beta: T, c: *mut T, rsc: isize, csc: isize) +{ + const MR: usize = KernelAvx512::MR; + const NR: usize = KernelAvx512::NR; + debug_assert_ne!(k, 0); + + 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 A B. The packed buffers are 64-byte aligned + let mut bv = _mm512_load_pd(b); + unroll_by_with_last!(4 => k, is_last, { + loop8!(i, ab[i] = _mm512_fmadd_pd(_mm512_set1_pd(*a.add(i)), bv, ab[i])); + if !is_last { + a = a.add(MR); + b = b.add(NR); + bv = _mm512_load_pd(b); + } + }); + + macro_rules! c { + ($i:expr, $j:expr) => (c.offset(rsc * $i as isize + csc * $j as isize)); + } + + // C <- alpha (A B) + beta C, in a single epilogue pass + // Fold alpha into the final FMA + // When beta == 0 the kernel must not read C + let alphav = _mm512_set1_pd(alpha); + if beta != 0. { + let betav = _mm512_set1_pd(beta); + if csc == 1 { + loop8!(i, { + let cv = _mm512_mul_pd(_mm512_loadu_pd(c![i, 0]), betav); + _mm512_storeu_pd(c![i, 0], _mm512_fmadd_pd(alphav, ab[i], cv)); + }); + } else { + loop8!(i, { + let mut tmp = [0.; NR]; + for j in 0..NR { tmp[j] = *c![i, j]; } + let cv = _mm512_mul_pd(_mm512_loadu_pd(tmp.as_ptr()), betav); + _mm512_storeu_pd(tmp.as_mut_ptr(), _mm512_fmadd_pd(alphav, ab[i], cv)); + for j in 0..NR { *c![i, j] = tmp[j]; } + }); + } + } else { + if csc == 1 { + loop8!(i, _mm512_storeu_pd(c![i, 0], _mm512_mul_pd(alphav, ab[i]))); + } else { + loop8!(i, { + let mut tmp = [0.; NR]; + _mm512_storeu_pd(tmp.as_mut_ptr(), _mm512_mul_pd(alphav, ab[i])); + for j in 0..NR { *c![i, j] = tmp[j]; } + }); + } + } +} + #[cfg(target_arch="aarch64")] #[cfg(has_aarch64_simd)] #[target_feature(enable="neon")] @@ -1124,5 +1245,10 @@ mod tests { "avx", avx, KernelAvx, "sse2", sse2, KernelSse2 } + + #[cfg(has_avx512)] + test_arch_kernels_x86! { + "avx512f", avx512f, KernelAvx512 + } } } diff --git a/src/gemm.rs b/src/gemm.rs index fa2284a..cebb8e8 100644 --- a/src/gemm.rs +++ b/src/gemm.rs @@ -220,6 +220,15 @@ impl GemmSelect for GemmParameters { } } +#[cfg(not(has_avx512))] +const KERNEL_MAX_MR: usize = 8; +#[cfg(not(has_avx512))] +const KERNEL_MAX_NR: usize = 8; +// Widen MR and NR when AVX-512 is enabled +#[cfg(has_avx512)] +const KERNEL_MAX_MR: usize = 16; +#[cfg(has_avx512)] +const KERNEL_MAX_NR: usize = 16; /// Ensure that GemmKernel parameters are supported /// (alignment, microkernel size). @@ -233,10 +242,10 @@ fn ensure_kernel_params() let nr = K::NR; // These are current limitations, // can change if corresponding code in gemm_loop is updated. - assert!(mr > 0 && mr <= 8); - assert!(nr > 0 && nr <= 8); - assert!(mr * nr * size_of::() <= 8 * 4 * 8); - assert!(K::align_to() <= 32); + assert!(mr > 0 && mr <= KERNEL_MAX_MR); + assert!(nr > 0 && nr <= KERNEL_MAX_NR); + assert!(mr * nr * size_of::() <= KERNEL_MAX_SIZE); + assert!(K::align_to() <= KERNEL_MAX_ALIGN); // one row/col of the kernel is limiting the max align we can provide let max_align = size_of::() * min(mr, nr); assert!(K::align_to() <= max_align); @@ -337,8 +346,15 @@ unsafe fn gemm_loop( } // set up buffer for masked (redirected output of) kernel +#[cfg(not(has_avx512))] const KERNEL_MAX_SIZE: usize = 8 * 8 * 4; +#[cfg(has_avx512)] +const KERNEL_MAX_SIZE: usize = 16 * 16 * 4; +#[cfg(not(has_avx512))] const KERNEL_MAX_ALIGN: usize = 32; +// The AVX-512 kernels load a full 64-byte ZMM register +#[cfg(has_avx512)] +const KERNEL_MAX_ALIGN: usize = 64; const MASK_BUF_SIZE: usize = KERNEL_MAX_SIZE + KERNEL_MAX_ALIGN - 1; // Pointers into buffer will be manually aligned anyway, due to @@ -489,7 +505,7 @@ unsafe fn align_ptr(mut align_to: usize, mut ptr: *mut T) -> *mut T { } /// Call the GEMM kernel with a "masked" output C. -/// +/// /// Simply redirect the MR by NR kernel output to the passed /// in `mask_buf`, and copy the non masked region to the real /// C. diff --git a/src/kernel.rs b/src/kernel.rs index 5ea4e38..190689a 100644 --- a/src/kernel.rs +++ b/src/kernel.rs @@ -199,11 +199,15 @@ pub(crate) trait ConstNum { pub(crate) struct U2; pub(crate) struct U4; pub(crate) struct U8; +#[cfg(has_avx512)] +pub(crate) struct U16; #[cfg(feature = "cgemm")] impl ConstNum for U2 { const VALUE: usize = 2; } impl ConstNum for U4 { const VALUE: usize = 4; } impl ConstNum for U8 { const VALUE: usize = 8; } +#[cfg(has_avx512)] +impl ConstNum for U16 { const VALUE: usize = 16; } #[cfg(test)] diff --git a/src/lib.rs b/src/lib.rs index 9f1d132..b49455d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -58,6 +58,7 @@ //! - `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: @@ -92,6 +93,19 @@ //! [`target-feature`](https://doc.rust-lang.org/rustc/codegen-options/index.html#target-feature) //! option to `rustc`.) //! +//! ### `avx512` +//! +//! `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"] } +//! ``` +//! //! ### `threading` //! //! `threading` is an optional crate feature @@ -128,7 +142,8 @@ //! considered upgrade policy, where updating the minimum Rust version is not a breaking //! change. //! -//! Some features are enabled with later versions: from Rust 1.61 AArch64 NEON support. +//! Some features are enabled with later versions: from Rust 1.61 AArch64 NEON +//! support, and from Rust 1.89 x86/x86-64 AVX-512 support. #![doc(html_root_url = "https://docs.rs/matrixmultiply/0.3/")] #![cfg_attr(not(feature = "std"), no_std)] diff --git a/src/loopmacros.rs b/src/loopmacros.rs index ac210e8..fca4b6e 100644 --- a/src/loopmacros.rs +++ b/src/loopmacros.rs @@ -66,6 +66,37 @@ macro_rules! loop8 { }} } +#[allow(unused)] +#[cfg(debug_assertions)] +macro_rules! loop16 { + ($i:ident, $e:expr) => { + for $i in 0..16 { $e } + } +} + +#[allow(unused)] +#[cfg(not(debug_assertions))] +macro_rules! loop16 { + ($i:ident, $e:expr) => {{ + let $i = 0; $e; + let $i = 1; $e; + let $i = 2; $e; + let $i = 3; $e; + let $i = 4; $e; + let $i = 5; $e; + let $i = 6; $e; + let $i = 7; $e; + let $i = 8; $e; + let $i = 9; $e; + let $i = 10; $e; + let $i = 11; $e; + let $i = 12; $e; + let $i = 13; $e; + let $i = 14; $e; + let $i = 15; $e; + }} +} + #[cfg(debug_assertions)] macro_rules! unroll_by { ($by:tt => $ntimes:expr, $e:expr) => { diff --git a/src/packing.rs b/src/packing.rs index 167222d..4fd69d6 100644 --- a/src/packing.rs +++ b/src/packing.rs @@ -78,6 +78,18 @@ pub(crate) unsafe fn pack_avx2(kc: usize, mc: usize, pack: PackSlice, pack_impl::(kc, mc, pack, a, rsa, csa) } +/// Specialized for AVX-512 +/// Safety: Requires AVX-512F +#[cfg(has_avx512)] +#[target_feature(enable="avx512f")] +pub(crate) unsafe fn pack_avx512(kc: usize, mc: usize, pack: PackSlice, + a: *const T, rsa: isize, csa: isize) + where T: Element, + MR: ConstNum, +{ + pack_impl::(kc, mc, pack, a, rsa, csa) +} + /// Pack implementation, see pack above for docs. /// /// Uses inline(always) so that it can be instantiated for different target features. diff --git a/src/sgemm_kernel.rs b/src/sgemm_kernel.rs index dbbcf98..0500716 100644 --- a/src/sgemm_kernel.rs +++ b/src/sgemm_kernel.rs @@ -9,6 +9,8 @@ use crate::kernel::GemmKernel; use crate::kernel::GemmSelect; use crate::kernel::{U4, U8}; +#[cfg(has_avx512)] +use crate::kernel::U16; use crate::archparam; #[cfg(target_arch="x86")] @@ -29,6 +31,8 @@ struct KernelFmaAvx2; struct KernelFma; #[cfg(any(target_arch="x86", target_arch="x86_64"))] struct KernelSse2; +#[cfg(has_avx512)] +struct KernelAvx512; #[cfg(target_arch="aarch64")] #[cfg(has_aarch64_simd)] @@ -49,6 +53,12 @@ pub(crate) fn detect(selector: G) where G: GemmSelect { // dispatch to specific compiled versions #[cfg(any(target_arch="x86", target_arch="x86_64"))] { + #[cfg(has_avx512)] + { + if is_x86_feature_detected_!("avx512f") { + return selector.select(KernelAvx512); + } + } if is_x86_feature_detected_!("fma") { if is_x86_feature_detected_!("avx2") { return selector.select(KernelFmaAvx2); @@ -224,6 +234,53 @@ impl GemmKernel for KernelSse2 { } } +#[cfg(has_avx512)] +impl GemmKernel for KernelAvx512 { + type Elem = T; + + type MRTy = U16; + type NRTy = U16; + + #[inline(always)] + fn align_to() -> usize { 64 } + + #[inline(always)] + fn always_masked() -> bool { false } + + #[inline(always)] + fn nc() -> usize { archparam::S_NC } + #[inline(always)] + fn kc() -> usize { archparam::S_KC } + #[inline(always)] + fn mc() -> usize { archparam::S_MC } + + #[inline] + unsafe fn pack_mr(kc: usize, mc: usize, pack: PackSlice, + a: *const Self::Elem, rsa: isize, csa: isize) + { + // safety: avx512f is enabled + crate::packing::pack_avx512::(kc, mc, pack, a, rsa, csa) + } + + #[inline] + unsafe fn pack_nr(kc: usize, mc: usize, pack: PackSlice, + a: *const Self::Elem, rsa: isize, csa: isize) + { + // safety: avx512f is enabled + crate::packing::pack_avx512::(kc, mc, pack, a, rsa, csa) + } + + #[inline(always)] + unsafe fn kernel( + k: usize, + alpha: T, + a: *const T, + b: *const T, + beta: T, + c: *mut T, rsc: isize, csc: isize) { + kernel_target_avx512(k, alpha, a, b, beta, c, rsc, csc) + } +} #[cfg(target_arch="aarch64")] #[cfg(has_aarch64_simd)] @@ -540,6 +597,71 @@ unsafe fn kernel_x86_avx(k: usize, alpha: T, a: *const T, b: *const T, } } +// no inline for unmasked kernels +#[cfg(has_avx512)] +#[target_feature(enable="avx512f")] +unsafe fn kernel_target_avx512(k: usize, alpha: T, a: *const T, b: *const T, + beta: T, c: *mut T, rsc: isize, csc: isize) +{ + const MR: usize = KernelAvx512::MR; + const NR: usize = KernelAvx512::NR; + debug_assert_ne!(k, 0); + + 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) }; + + // Compute A B. The packed buffers are 64-byte aligned + let mut bv = _mm512_load_ps(b); + unroll_by_with_last!(4 => k, is_last, { + loop16!(i, ab[i] = _mm512_fmadd_ps(_mm512_set1_ps(*a.add(i)), bv, ab[i])); + if !is_last { + a = a.add(MR); + b = b.add(NR); + bv = _mm512_load_ps(b); + } + }); + + macro_rules! c { + ($i:expr, $j:expr) => (c.offset(rsc * $i as isize + csc * $j as isize)); + } + + // C <- alpha (A B) + beta C, in a single epilogue pass + // Fold alpha into the final FMA + // When beta == 0 the kernel must not read C + let alphav = _mm512_set1_ps(alpha); + if beta != 0. { + let betav = _mm512_set1_ps(beta); + if csc == 1 { + loop16!(i, { + let cv = _mm512_mul_ps(_mm512_loadu_ps(c![i, 0]), betav); + _mm512_storeu_ps(c![i, 0], _mm512_fmadd_ps(alphav, ab[i], cv)); + }); + } else { + loop16!(i, { + let mut tmp = [0.; NR]; + for j in 0..NR { tmp[j] = *c![i, j]; } + let cv = _mm512_mul_ps(_mm512_loadu_ps(tmp.as_ptr()), betav); + _mm512_storeu_ps(tmp.as_mut_ptr(), _mm512_fmadd_ps(alphav, ab[i], cv)); + for j in 0..NR { *c![i, j] = tmp[j]; } + }); + } + } else { + if csc == 1 { + loop16!(i, _mm512_storeu_ps(c![i, 0], _mm512_mul_ps(alphav, ab[i]))); + } else { + loop16!(i, { + let mut tmp = [0.; NR]; + _mm512_storeu_ps(tmp.as_mut_ptr(), _mm512_mul_ps(alphav, ab[i])); + for j in 0..NR { *c![i, j] = tmp[j]; } + }); + } + } +} + #[cfg(target_arch="aarch64")] #[cfg(has_aarch64_simd)] #[target_feature(enable="neon")] @@ -947,6 +1069,11 @@ mod tests { "sse2", sse2, KernelSse2 } + #[cfg(has_avx512)] + test_arch_kernels_x86! { + "avx512f", avx512f, KernelAvx512 + } + #[test] fn ensure_target_features_tested() { // If enabled, this test ensures that the requested feature actually @@ -963,6 +1090,7 @@ mod tests { "avx" => is_x86_feature_detected_!("avx"), "fma" => is_x86_feature_detected_!("fma"), "sse2" => is_x86_feature_detected_!("sse2"), + "avx512f" => is_x86_feature_detected_!("avx512f"), _ => false, }; assert!(detected, "Feature {:?} was not detected, so it could not be tested", diff --git a/src/zgemm_kernel.rs b/src/zgemm_kernel.rs index a689a63..5641f7c 100644 --- a/src/zgemm_kernel.rs +++ b/src/zgemm_kernel.rs @@ -13,6 +13,8 @@ use crate::archparam; use crate::cgemm_common::pack_complex; use crate::packing::PackSlice; +#[cfg(has_avx512)] +struct KernelAvx512; #[cfg(any(target_arch="x86", target_arch="x86_64"))] struct KernelAvx2; #[cfg(any(target_arch="x86", target_arch="x86_64"))] @@ -37,6 +39,12 @@ pub(crate) fn detect(selector: G) where G: GemmSelect { // dispatch to specific compiled versions #[cfg(any(target_arch="x86", target_arch="x86_64"))] { + #[cfg(has_avx512)] + { + if is_x86_feature_detected_!("avx512f") { + return selector.select(KernelAvx512); + } + } if is_x86_feature_detected_!("fma") { if is_x86_feature_detected_!("avx2") { return selector.select(KernelAvx2); @@ -57,6 +65,40 @@ pub(crate) fn detect(selector: G) where G: GemmSelect { macro_rules! loop_m { ($i:ident, $e:expr) => { loop4!($i, $e) }; } macro_rules! loop_n { ($j:ident, $e:expr) => { loop2!($j, $e) }; } +#[cfg(has_avx512)] +impl GemmKernel for KernelAvx512 { + type Elem = T; + + type MRTy = U4; + type NRTy = U4; + + #[inline(always)] + fn align_to() -> usize { 32 } + + #[inline(always)] + fn always_masked() -> bool { KernelFallback::always_masked() } + + #[inline(always)] + fn nc() -> usize { archparam::Z_NC } + #[inline(always)] + fn kc() -> usize { archparam::Z_KC } + #[inline(always)] + fn mc() -> usize { archparam::Z_MC } + + pack_methods!{} + + #[inline(always)] + unsafe fn kernel( + k: usize, + alpha: T, + a: *const T, + b: *const T, + beta: T, + c: *mut T, rsc: isize, csc: isize) { + kernel_target_avx512(k, alpha, a, b, beta, c, rsc, csc) + } +} + #[cfg(any(target_arch="x86", target_arch="x86_64"))] impl GemmKernel for KernelAvx2 { type Elem = T; @@ -223,6 +265,19 @@ kernel_fallback_impl_complex! { kernel_fallback_impl, T, TReal, KernelFallback::MR, KernelFallback::NR, 1 } +// 4x4 loop macros for AVX-512 only +#[cfg(has_avx512)] +macro_rules! loop_m { ($i:ident, $e:expr) => { loop4!($i, $e) }; } +#[cfg(has_avx512)] +macro_rules! loop_n { ($j:ident, $e:expr) => { loop4!($j, $e) }; } + +#[cfg(has_avx512)] +kernel_fallback_impl_complex! { + // instantiate separately + [inline target_feature(enable="avx512f")] [fma_yes] + kernel_target_avx512, T, TReal, KernelAvx512::MR, KernelAvx512::NR, 4 +} + #[inline(always)] unsafe fn at(ptr: *const TReal, i: usize) -> TReal { *ptr.add(i) @@ -292,5 +347,10 @@ mod tests { "fma", fma, KernelFma, "avx2", avx2, KernelAvx2 } + + #[cfg(has_avx512)] + test_arch_kernels_x86! { + "avx512f", avx512f, KernelAvx512 + } } }