Merge branch 'add-avx512-vpclmulqdq'

onethumb · onethumb · commit cc6683301ddc · 2024-12-27T11:32:39.000-08:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -16,6 +16,7 @@ rust-version = "1.70.0"
 
 [dependencies]
 crc = "3"
+lazy_static = { version = "1.4.0", optional = true }
 
 [dev-dependencies]
 crc = "3"
@@ -25,6 +26,7 @@ rand = "0.8"
 
 [features]
 pmull = [] # deprecated, no longer have any effect.
+vpclmulqdq = ["lazy_static"]
 fake-simd = []
 
 [[bench]]
diff --git a/README.md b/README.md
@@ -8,7 +8,7 @@ crc64fast-nvme
 SIMD-accelerated carryless-multiplication [CRC-64/NVME](https://reveng.sourceforge.io/crc-catalogue/all.htm#crc.cat.crc-64-nvme) checksum computation
 (similar to [crc32fast](https://crates.io/crates/crc32fast) and forked from [crc64fast](https://github.com/tikv/crc64fast) which calculates [CRC-64/XZ](https://reveng.sourceforge.io/crc-catalogue/all.htm#crc.cat.crc-64-xz) [a.k.a `CRC-64/GO-ECMA`]).
 
-`CRC-64/NVME` comes from the [NVM Express® NVM Command Set Specification](https://nvmexpress.org/wp-content/uploads/NVM-Express-NVM-Command-Set-Specification-1.0d-2023.12.28-Ratified.pdf) (Revision 1.0d, December 2023) and has also been implemented in the [Linux kernel](https://github.com/torvalds/linux/blob/786c8248dbd33a5a7a07f7c6e55a7bfc68d2ca48/lib/crc64.c#L66-L73) (where it's called `CRC-64/Rocksoft`) and [AWS S3's recommended checksum option](https://docs.aws.amazon.com/AmazonS3/latest/userguide/checking-object-integrity.html) as `CRC64-NVME`. (Note that the Check value in the spec uses incorrect endianness (Section 5.2.1.3.4, Figure 120, page 83).
+`CRC-64/NVME` comes from the [NVM Express® NVM Command Set Specification](https://nvmexpress.org/wp-content/uploads/NVM-Express-NVM-Command-Set-Specification-1.0d-2023.12.28-Ratified.pdf) (Revision 1.0d, December 2023) and has also been implemented in the [Linux kernel](https://github.com/torvalds/linux/blob/786c8248dbd33a5a7a07f7c6e55a7bfc68d2ca48/lib/crc64.c#L66-L73) (where it's called `CRC-64/Rocksoft`) and is [AWS S3's recommended checksum option](https://docs.aws.amazon.com/AmazonS3/latest/userguide/checking-object-integrity.html) as `CRC64-NVME`. (Note that the Check value in the spec uses incorrect endianness [Section 5.2.1.3.4, Figure 120, page 83]).
 
 SIMD-accelerated carryless-multiplication is based on the Intel [Fast CRC Computation for Generic Polynomials Using PCLMULQDQ Instruction](https://web.archive.org/web/20131224125630/https://www.intel.com/content/dam/www/public/us/en/documents/white-papers/fast-crc-computation-generic-polynomials-pclmulqdq-paper.pdf) paper.
 
@@ -48,14 +48,28 @@ be chosen based on CPU feature at runtime.
     * using PCLMULQDQ + SSE 4.1 on x86/x86_64
     * using PMULL + NEON on AArch64 (64-bit ARM)
 
-| Algorithm              | Throughput (x86_64) | Throughput (aarch64) |
-|:-----------------------|--------------------:|---------------------:|
-| [crc 3.0.1]            |  0.5 GiB/s          |  0.3 GiB/s           |
-| crc64fast-nvme (table) |  2.3 GiB/s          |  1.8 GiB/s           |
-| crc64fast-nvme (simd)  | 28.2 GiB/s          | 20.0 GiB/s           |
+| Algorithm                   | Throughput (x86_64) | Throughput (aarch64) |
+|:----------------------------|--------------------:|---------------------:|
+| [crc 3.0.1]                 |           0.5 GiB/s |            0.3 GiB/s |
+| crc64fast-nvme (table)      |           2.3 GiB/s |            1.8 GiB/s |
+| crc64fast-nvme (SIMD)       |          28.2 GiB/s |           20.0 GiB/s |
+| crc64fast-nvme (VPCLMULQDQ) |            52 GiB/s |                 n/a  |
 
 [crc 3.0.1]: https://docs.rs/crc/3.0.1/crc/index.html
 
+## Experimental "Vector Carry-Less Multiplication of Quadwords" (VPCLMULQDQ) support
+
+Using Rust's support for [AVX512 intrinsics](https://github.com/rust-lang/rust/issues/111137), specifically [VPCLMULQDQ](https://doc.rust-lang.org/src/core/stdarch/crates/core_arch/src/x86/vpclmulqdq.rs.html), we can massively improve throughput for x86_64 processors which support them (Intel Ice Lake+ and AMD Zen4+).
+
+Specifically, on an `m7i.8xlarge` EC2 instance (4th gen Xeon, aka Sapphire Rapids), throughput approximately _doubles_ from ~26GiB/s to ~52GiB/s.
+
+Since these are currently marked as unstable features in Rust, you'll need to build with `nightly` and enable the `vpclmulqdq` feature:
+
+``` 
+rustup toolchain install nightly
+cargo +nightly build --features="vpclmulqdq" -r
+```
+
 ## References
 
 * [crc32-fast](https://crates.io/crates/crc32fast) - Original `crc32` implementation in Rust.
@@ -70,6 +84,7 @@ be chosen based on CPU feature at runtime.
 * [StackOverflow PCLMULQDQ CRC32 question](https://stackoverflow.com/questions/21171733/calculating-constants-for-crc32-using-pclmulqdq) - Insightful question & answer to CRC32 implementation details.
 * [AWS S3 announcement about CRC64-NVME support](https://aws.amazon.com/blogs/aws/introducing-default-data-integrity-protections-for-new-objects-in-amazon-s3/)
 * [AWS S3 docs on checking object integrity using CRC64-NVME](https://docs.aws.amazon.com/AmazonS3/latest/userguide/checking-object-integrity.html)
+* [Vector Carry-Less Multiplication of Quadwords (VPCLMULQDQ) details](https://en.wikichip.org/wiki/x86/vpclmulqdq)
 
 ## License
 
diff --git a/src/lib.rs b/src/lib.rs
@@ -17,11 +17,22 @@
 //! let checksum = c.sum64();
 //! assert_eq!(checksum, 0xd9160d1fa8e418e3);
 //! ```
+//!
+//! Tracking links for unstable features used here (which require nightly builds):
+//!
+//! - simd_ffi: https://github.com/rust-lang/rust/issues/27731
+//! - link_llvm_intrinsics: https://github.com/rust-lang/rust/issues/29602
+//! - avx512_target_feature: https://github.com/rust-lang/rust/issues/111137
+
+#![cfg_attr(
+    feature = "vpclmulqdq",
+    feature(avx512_target_feature, stdarch_x86_avx512)
+)]
 
 mod pclmulqdq;
 mod table;
 
-type UpdateFn = fn(u64, &[u8]) -> u64;
+type UpdateFn = unsafe fn(u64, &[u8]) -> u64;
 
 /// Represents an in-progress CRC-64 computation.
 #[derive(Clone)]
@@ -52,7 +63,9 @@ impl Digest {
 
     /// Writes some data into the digest.
     pub fn write(&mut self, bytes: &[u8]) {
-        self.state = (self.computer)(self.state, bytes);
+        unsafe {
+            self.state = (self.computer)(self.state, bytes);
+        }
     }
 
     /// Computes the current CRC-64/NVME value.
diff --git a/src/pclmulqdq/mod.rs b/src/pclmulqdq/mod.rs
@@ -7,22 +7,24 @@
 //!
 //! [white paper]: https://web.archive.org/web/20131224125630/https://www.intel.com/content/dam/www/public/us/en/documents/white-papers/fast-crc-computation-generic-polynomials-pclmulqdq-paper.pdf
 
+use std::{
+    fmt::Debug,
+    ops::{BitXor, BitXorAssign},
+};
+
+use super::table;
+
+use self::arch::Simd;
+
 #[cfg(not(feature = "fake-simd"))]
-#[cfg_attr(target_arch = "x86_64", path = "x86_64.rs")]
+#[cfg_attr(target_arch = "x86_64", path = "x86_64/mod.rs")]
 #[cfg_attr(target_arch = "aarch64", path = "aarch64.rs")]
 #[cfg_attr(target_arch = "x86", path = "x86.rs")]
 mod arch;
 
 #[cfg(feature = "fake-simd")]
 mod arch;
 
-use self::arch::Simd;
-use super::table;
-use std::{
-    fmt::Debug,
-    ops::{BitXor, BitXorAssign},
-};
-
 /// This trait must be implemented on `self::arch::Simd` to provide the
 /// platform-specific SIMD implementations.
 trait SimdExt: Copy + Debug + BitXor {
@@ -71,24 +73,47 @@ impl BitXorAssign for Simd {
 }
 
 pub fn get_update() -> super::UpdateFn {
+    #[cfg(feature = "vpclmulqdq")]
+    {
+        use arch::vpclmulqdq::*;
+        if Simd256::is_supported() {
+            return update_256_batch;
+        }
+    }
+
     if Simd::is_supported() {
-        update
+        update_128_batch
     } else {
         table::update
     }
 }
 
-fn update(mut state: u64, bytes: &[u8]) -> u64 {
-    let (left, middle, right) = unsafe { bytes.align_to::<[Simd; 8]>() };
+// This function is unsafe because it uses platform dependent functions.
+unsafe fn update_128_batch(mut state: u64, bytes: &[u8]) -> u64 {
+    let (left, middle, right) = bytes.align_to::<[Simd; 8]>();
     if let Some((first, rest)) = middle.split_first() {
         state = table::update(state, left);
-        state = unsafe { update_simd(state, first, rest) };
+        state = update_simd(state, first, rest);
         table::update(state, right)
     } else {
         table::update(state, bytes)
     }
 }
 
+#[cfg(feature = "vpclmulqdq")]
+#[target_feature(enable = "avx2", enable = "vpclmulqdq")]
+unsafe fn update_256_batch(mut state: u64, bytes: &[u8]) -> u64 {
+    use arch::vpclmulqdq::*;
+    let (left, middle, right) = bytes.align_to::<[[Simd256; 4]; 2]>();
+    if let Some((first, rest)) = middle.split_first() {
+        state = update_128_batch(state, left);
+        state = update_vpclmulqdq(state, first, rest);
+        update_128_batch(state, right)
+    } else {
+        update_128_batch(state, bytes)
+    }
+}
+
 #[cfg_attr(
     any(target_arch = "x86", target_arch = "x86_64"),
     target_feature(enable = "pclmulqdq", enable = "sse2", enable = "sse4.1")
@@ -112,6 +137,11 @@ unsafe fn update_simd(state: u64, first: &[Simd; 8], rest: &[[Simd; 8]]) -> u64
         }
     }
 
+    fold_tail(x)
+}
+
+#[inline(always)]
+unsafe fn fold_tail(x: [Simd; 8]) -> u64 {
     let coeffs = [
         Simd::new(table::K_895, table::K_959), // fold by distance of 112 bytes
         Simd::new(table::K_767, table::K_831), // fold by distance of 96 bytes
diff --git a/src/pclmulqdq/x86_64/mod.rs b/src/pclmulqdq/x86_64/mod.rs
@@ -1,11 +1,14 @@
 // Copyright 2020 TiKV Project Authors. Licensed under MIT or Apache-2.0.
 
-//! x86_64 implementation of the PCLMULQDQ-based CRC calculation.
+//! x86 (32-bit) implementation of the PCLMULQDQ-based CRC calculation.
 
 #[cfg(target_arch = "x86_64")]
 use std::arch::x86_64::*;
 use std::ops::BitXor;
 
+#[cfg(feature = "vpclmulqdq")]
+pub mod vpclmulqdq;
+
 #[repr(transparent)]
 #[derive(Copy, Clone, Debug)]
 pub struct Simd(__m128i);
@@ -14,41 +17,48 @@ impl super::SimdExt for Simd {
     fn is_supported() -> bool {
         is_x86_feature_detected!("pclmulqdq") // _mm_clmulepi64_si128
             && is_x86_feature_detected!("sse2") // (all other _mm_*)
-            && is_x86_feature_detected!("sse4.1") // _mm_extract_epi64
+            && is_x86_feature_detected!("sse4.1")
     }
 
     #[inline]
     #[target_feature(enable = "sse2")]
     unsafe fn new(high: u64, low: u64) -> Self {
-        Self(_mm_set_epi64x(high as i64, low as i64))
+        // On 32-bit systems, we need to split u64 into low and high 32-bit parts
+        let high_low = (high & 0xFFFFFFFF) as i32;
+        let high_high = ((high >> 32) & 0xFFFFFFFF) as i32;
+        let low_low = (low & 0xFFFFFFFF) as i32;
+        let low_high = ((low >> 32) & 0xFFFFFFFF) as i32;
+
+        // Create the 128-bit vector using 32-bit parts
+        Self(_mm_set_epi32(high_high, high_low, low_high, low_low))
     }
 
     #[inline]
     #[target_feature(enable = "sse2", enable = "pclmulqdq")]
     unsafe fn fold_16(self, coeff: Self) -> Self {
-        let h = Self(_mm_clmulepi64_si128(self.0, coeff.0, 0x11));
-        let l = Self(_mm_clmulepi64_si128(self.0, coeff.0, 0x00));
+        let h = Self(_mm_clmulepi64_si128::<0x11>(self.0, coeff.0));
+        let l = Self(_mm_clmulepi64_si128::<0x00>(self.0, coeff.0));
         h ^ l
     }
 
     #[inline]
     #[target_feature(enable = "sse2", enable = "pclmulqdq")]
     unsafe fn fold_8(self, coeff: u64) -> Self {
         let coeff = Self::new(0, coeff);
-        let h = Self(_mm_clmulepi64_si128(self.0, coeff.0, 0x00));
-        let l = Self(_mm_srli_si128(self.0, 8));
+        let h = Self(_mm_clmulepi64_si128::<0x00>(self.0, coeff.0));
+        let l = Self(_mm_srli_si128::<8>(self.0));
         h ^ l
     }
 
     #[inline]
     #[target_feature(enable = "sse2", enable = "sse4.1", enable = "pclmulqdq")]
     unsafe fn barrett(self, poly: u64, mu: u64) -> u64 {
         let polymu = Self::new(poly, mu);
-        let t1 = _mm_clmulepi64_si128(self.0, polymu.0, 0x00);
-        let h = Self(_mm_slli_si128(t1, 8));
-        let l = Self(_mm_clmulepi64_si128(t1, polymu.0, 0x10));
+        let t1 = _mm_clmulepi64_si128::<0x00>(self.0, polymu.0);
+        let h = Self(_mm_slli_si128::<8>(t1));
+        let l = Self(_mm_clmulepi64_si128::<0x10>(t1, polymu.0));
         let reduced = h ^ l ^ self;
-        _mm_extract_epi64(reduced.0, 1) as u64
+        _mm_extract_epi64::<1>(reduced.0) as u64
     }
 }
 
diff --git a/src/pclmulqdq/x86_64/vpclmulqdq.rs b/src/pclmulqdq/x86_64/vpclmulqdq.rs
