Add AVX-512 kernel

.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

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 = []

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 <https://github.com/SongXiaoXi>`_ `#88 <https://github.com/bluss/matrixmultiply/pull/88>`_

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");
         }
     }
 }

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<G>(selector: G) where G: GemmSelect<T> {
     // 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<G>(selector: G) where G: GemmSelect<T> {
     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
+        }
     }
 }

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<G>(selector: G) where G: GemmSelect<T> {
     // 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<Self::Elem>,
+                      a: *const Self::Elem, rsa: isize, csa: isize)
+    {
+        // safety: avx512f is enabled
+        crate::packing::pack_avx512::<Self::MRTy, T>(kc, mc, pack, a, rsa, csa)
+    }
+
+    #[inline]
+    unsafe fn pack_nr(kc: usize, mc: usize, pack: PackSlice<Self::Elem>,
+                      a: *const Self::Elem, rsa: isize, csa: isize)
+    {
+        // safety: avx512f is enabled
+        crate::packing::pack_avx512::<Self::NRTy, T>(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<MA>(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
+        }
     }
 }