crypto: RSA Montgomery multiplication with mulx and two carry chains (parity R5)

lib/crypto/src/arch/mod.rs

@@ -54,6 +54,17 @@ pub(crate) fn value_barrier(value: u64) -> u64 {
     unsafe { core::ptr::read_volatile(&value) }
 }
 
+/// For a, b < m and n a multiple of four up to 64, by an instruction path:
+/// out and the returned top word are a b / 2^(64n) mod m plus at most one
+/// m, below 2m and unreduced; `None` where the CPU has no such path. RSA's
+/// data is public.
+pub(crate) fn rsa_mul(a: &[u64], b: &[u64], m: &[u64], n0: u64, n: usize, out: &mut [u64]) -> Option<u64> {
+    hosted!(x86_64::rsa_mul(a, b, m, n0, n, out), {
+        let _ = (a, b, m, n0, n, out);
+        None
+    })
+}
+
 /// GHASH key powers, (H^(i+1))' at index i in the instruction paths' form.
 pub(crate) type Powers = [u128; 16];
 

lib/crypto/src/arch/x86_64.rs

@@ -16,6 +16,8 @@ mod aesni;
 #[cfg(not(target_os = "uefi"))]
 mod chacha;
 #[cfg(not(target_os = "uefi"))]
+mod montgomery;
+#[cfg(not(target_os = "uefi"))]
 mod poly1305;
 #[cfg(not(target_os = "uefi"))]
 mod sha;
@@ -40,6 +42,7 @@ impl Features {
     const VPCLMUL: u32 = 1 << 8;
     const AVX512VL: u32 = 1 << 9;
     const AVX512BW: u32 = 1 << 10;
+    const ADX: u32 = 1 << 11;
 
     fn has(self, bits: u32) -> bool {
         self.0 & bits == bits
@@ -64,6 +67,11 @@ impl Features {
     pub(super) fn bmi2(self) -> bool {
         self.has(Self::BMI2)
     }
+    /// MULX and ADCX/ADOX.
+    #[cfg(not(target_os = "uefi"))]
+    pub(super) fn bmi2_adx(self) -> bool {
+        self.has(Self::BMI2 | Self::ADX)
+    }
     /// SHA-NI with the SSSE3 and SSE4.1 shuffles its message schedule takes.
     #[cfg(not(target_os = "uefi"))]
     pub(super) fn sha(self) -> bool {
@@ -86,7 +94,7 @@ impl Features {
 
     /// Intel SDM Vol. 2A, CPUID: leaf 1 ECX bits 1 (PCLMULQDQ), 9 (SSSE3),
     /// 19 (SSE4.1), 25 (AES), 27 (OSXSAVE) and 28 (AVX); leaf 7 sub-leaf 0
-    /// EBX bits 5 (AVX2), 8 (BMI2), 16 (AVX512F), 29 (SHA), 30 (AVX512BW)
+    /// EBX bits 5 (AVX2), 8 (BMI2), 16 (AVX512F), 19 (ADX), 29 (SHA), 30 (AVX512BW)
     /// and 31 (AVX512VL), ECX bits 9 (VAES) and 10 (VPCLMULQDQ). Vol. 1
     /// §13.3: AVX needs XCR0's SSE and AVX state (bits 1 and 2) enabled by
     /// the OS, AVX-512 its opmask, ZMM_Hi256 and Hi16_ZMM state (bits 5 to 7)
@@ -116,6 +124,7 @@ impl Features {
             | bit(leaf1, 9, Self::SSSE3)
             | bit(leaf1, 19, Self::SSE41)
             | bit(leaf7b, 8, Self::BMI2)
+            | bit(leaf7b, 19, Self::ADX)
             | bit(leaf7b, 29, Self::SHA);
         if ymm {
             bits |= bit(leaf7b, 5, Self::AVX2) | bit(leaf7c, 9, Self::VAES) | bit(leaf7c, 10, Self::VPCLMUL);
@@ -151,6 +160,18 @@ fn first_features() -> Features {
     read
 }
 
+/// RSA's Montgomery multiplication on BMI2 and ADX, when the CPU has them:
+/// out and the returned top word are the unreduced result, below 2m.
+#[cfg(not(target_os = "uefi"))]
+pub(super) fn rsa_mul(a: &[u64], b: &[u64], m: &[u64], n0: u64, n: usize, out: &mut [u64]) -> Option<u64> {
+    if !features().bmi2_adx() {
+        return None;
+    }
+    assert!(n.is_multiple_of(4) && n <= 64 && a.len() >= n && b.len() >= n && m.len() >= n && out.len() >= n);
+    // CPUID admitted BMI2 and ADX, and the lengths are checked above.
+    Some(unsafe { montgomery::mul(a, b, m, n0, n, out) })
+}
+
 /// A vector from 16 bytes by value. With no pointer there is no copy
 /// precondition, which debug assertions would check per access in hot loops.
 #[cfg(not(target_os = "uefi"))]
@@ -465,7 +486,7 @@ mod tests {
         for (bit, feature) in [(1, Features::CLMUL), (9, Features::SSSE3), (19, Features::SSE41), (25, Features::AES)] {
             assert_eq!(Features::decode(!(1 << bit), !0, !0, 0xe7).0 & feature, 0, "leaf 1 bit {bit}");
         }
-        for (bit, feature) in [(5, Features::AVX2), (8, Features::BMI2), (29, Features::SHA)] {
+        for (bit, feature) in [(5, Features::AVX2), (8, Features::BMI2), (19, Features::ADX), (29, Features::SHA)] {
             assert_eq!(Features::decode(!0, !(1 << bit), !0, 0xe7).0 & feature, 0, "leaf 7 EBX bit {bit}");
         }
         for (bit, feature) in [(9, Features::VAES), (10, Features::VPCLMUL)] {
@@ -481,6 +502,10 @@ mod tests {
         assert_eq!(f.clmul(), on && std::is_x86_feature_detected!("pclmulqdq"));
         assert_eq!(f.avx2(), on && std::is_x86_feature_detected!("avx2"));
         assert_eq!(f.bmi2(), on && std::is_x86_feature_detected!("bmi2"));
+        assert_eq!(
+            f.bmi2_adx(),
+            on && std::is_x86_feature_detected!("bmi2") && std::is_x86_feature_detected!("adx")
+        );
         assert_eq!(
             f.sha(),
             on && std::is_x86_feature_detected!("sha")

lib/crypto/src/arch/x86_64/montgomery.rs

@@ -0,0 +1,146 @@
+//! Montgomery multiplication for RSA moduli on BMI2 and ADX: `mulx` with two
+//! carry chains, CF through `adcx` and OF through `adox`, so each product
+//! row's low and high halves add without waiting on one flag (Intel SDM
+//! Vol. 2A, ADCX, ADOX and MULX). The data is public, as all RSA data here.
+
+use core::arch::asm;
+
+/// For a, b < m and n a multiple of four up to 64: out and the returned top
+/// word are a b / 2^(64n) mod m plus at most one m, below 2m and unreduced.
+/// For each word of b, one product row and one reduction row, each four
+/// words of a or m per loop step; `t` holds the n + 2 running words.
+///
+/// # Safety
+/// The CPU has BMI2 and ADX; n is a multiple of four up to 64, and `a`, `b`,
+/// `m` and `out` hold at least n words.
+#[target_feature(enable = "bmi2,adx")]
+pub(super) unsafe fn mul(a: &[u64], b: &[u64], m: &[u64], n0: u64, n: usize, out: &mut [u64]) -> u64 {
+    let mut t = [0u64; 66];
+    for &bi in &b[..n] {
+        // SAFETY: the rows read n words of `a` and `m` and write n + 2 of `t`;
+        // they write no other memory and leave the stack alone.
+        unsafe { row(a.as_ptr(), bi, m.as_ptr(), n0, n, t.as_mut_ptr()) };
+    }
+    out[..n].copy_from_slice(&t[..n]);
+    t[n]
+}
+
+/// t = (t + a bi + q m) / 2^64 for q = (t + a bi)_0 n0 mod 2^64.
+#[target_feature(enable = "bmi2,adx")]
+unsafe fn row(a: *const u64, bi: u64, m: *const u64, n0: u64, n: usize, t: *mut u64) {
+    unsafe {
+        asm!(
+            // Product row: t[j] += lo(a_j bi) on CF and hi(a_{j-1} bi) on OF.
+            "mov {a_ptr}, {a}",
+            "mov {t_ptr}, {t}",
+            "mov rcx, {quads}",
+            "xor {hi:e}, {hi:e}",
+            "2:",
+            "mulx {h0}, {lo}, [{a_ptr}]",
+            "adcx {lo}, [{t_ptr}]",
+            "adox {lo}, {hi}",
+            "mov [{t_ptr}], {lo}",
+            "mulx {hi}, {lo}, [{a_ptr} + 8]",
+            "adcx {lo}, [{t_ptr} + 8]",
+            "adox {lo}, {h0}",
+            "mov [{t_ptr} + 8], {lo}",
+            "mulx {h0}, {lo}, [{a_ptr} + 16]",
+            "adcx {lo}, [{t_ptr} + 16]",
+            "adox {lo}, {hi}",
+            "mov [{t_ptr} + 16], {lo}",
+            "mulx {hi}, {lo}, [{a_ptr} + 24]",
+            "adcx {lo}, [{t_ptr} + 24]",
+            "adox {lo}, {h0}",
+            "mov [{t_ptr} + 24], {lo}",
+            "lea {a_ptr}, [{a_ptr} + 32]",
+            "lea {t_ptr}, [{t_ptr} + 32]",
+            "lea rcx, [rcx - 1]",
+            "jrcxz 3f",
+            "jmp 2b",
+            "3:",
+            // t[n] += hi + CF + OF; t[n + 1] = the carries out of that.
+            "mov {lo}, 0",
+            "mov {h0}, [{t_ptr}]",
+            "adcx {h0}, {lo}",
+            "adox {h0}, {hi}",
+            "mov [{t_ptr}], {h0}",
+            "mov {h0}, 0",
+            "adcx {h0}, {lo}",
+            "adox {h0}, {lo}",
+            "mov [{t_ptr} + 8], {h0}",
+            // Reduction row: q = t_0 n0; t[j - 1] = t[j] + lo(q m_j) + hi(q m_{j-1}).
+            "mov rdx, [{t}]",
+            "imul rdx, {n0}",
+            "mov {a_ptr}, {m}",
+            "mov {t_ptr}, {t}",
+            "mov rcx, {quads}",
+            "xor {hi:e}, {hi:e}",
+            // Word 0 of the sum is zero by the choice of q; only its carries
+            // remain, so it is computed and not stored.
+            "mulx {hi}, {lo}, [{a_ptr}]",
+            "adcx {lo}, [{t_ptr}]",
+            "mulx {h0}, {lo}, [{a_ptr} + 8]",
+            "adcx {lo}, [{t_ptr} + 8]",
+            "adox {lo}, {hi}",
+            "mov [{t_ptr}], {lo}",
+            "mulx {hi}, {lo}, [{a_ptr} + 16]",
+            "adcx {lo}, [{t_ptr} + 16]",
+            "adox {lo}, {h0}",
+            "mov [{t_ptr} + 8], {lo}",
+            "mulx {h0}, {lo}, [{a_ptr} + 24]",
+            "adcx {lo}, [{t_ptr} + 24]",
+            "adox {lo}, {hi}",
+            "mov [{t_ptr} + 16], {lo}",
+            "lea {a_ptr}, [{a_ptr} + 32]",
+            "lea {t_ptr}, [{t_ptr} + 32]",
+            "lea rcx, [rcx - 1]",
+            "jrcxz 5f",
+            "4:",
+            "mulx {hi}, {lo}, [{a_ptr}]",
+            "adcx {lo}, [{t_ptr}]",
+            "adox {lo}, {h0}",
+            "mov [{t_ptr} - 8], {lo}",
+            "mulx {h0}, {lo}, [{a_ptr} + 8]",
+            "adcx {lo}, [{t_ptr} + 8]",
+            "adox {lo}, {hi}",
+            "mov [{t_ptr}], {lo}",
+            "mulx {hi}, {lo}, [{a_ptr} + 16]",
+            "adcx {lo}, [{t_ptr} + 16]",
+            "adox {lo}, {h0}",
+            "mov [{t_ptr} + 8], {lo}",
+            "mulx {h0}, {lo}, [{a_ptr} + 24]",
+            "adcx {lo}, [{t_ptr} + 24]",
+            "adox {lo}, {hi}",
+            "mov [{t_ptr} + 16], {lo}",
+            "lea {a_ptr}, [{a_ptr} + 32]",
+            "lea {t_ptr}, [{t_ptr} + 32]",
+            "lea rcx, [rcx - 1]",
+            "jrcxz 5f",
+            "jmp 4b",
+            "5:",
+            // t[n - 1] = t[n] + hi + CF + OF; t[n] = t[n + 1] + the carries.
+            "mov {lo}, [{t_ptr}]",
+            "adcx {lo}, {zero}",
+            "adox {lo}, {h0}",
+            "mov [{t_ptr} - 8], {lo}",
+            "mov {lo}, [{t_ptr} + 8]",
+            "adcx {lo}, {zero}",
+            "adox {lo}, {zero}",
+            "mov [{t_ptr}], {lo}",
+            a = in(reg) a,
+            m = in(reg) m,
+            t = in(reg) t,
+            n0 = in(reg) n0,
+            quads = in(reg) n / 4,
+            zero = in(reg) 0u64,
+            a_ptr = out(reg) _,
+            t_ptr = out(reg) _,
+            lo = out(reg) _,
+            hi = out(reg) _,
+            h0 = out(reg) _,
+            inout("rdx") bi => _,
+            out("rcx") _,
+            options(nostack),
+        );
+    }
+}

lib/crypto/src/rsa.rs

@@ -78,7 +78,9 @@ impl PublicKey {
         if bits < 2048 || e < 3 || e & 1 == 0 || modulus[modulus.len() - 1] & 1 == 0 {
             return Err(Error::InvalidKey);
         }
-        let words = modulus.len().div_ceil(8);
+        // Whole groups of four words, which the instruction path takes; the
+        // zero words above the modulus only enlarge R.
+        let words = modulus.len().div_ceil(32) * 4;
         let m = read(modulus);
         let mut key = Self {
             modulus: m,
@@ -119,11 +121,23 @@ impl PublicKey {
             subtract(x, &self.modulus, self.words);
         }
     }
-    /// out = a * b / R mod n, for a, b < n: word-serial Montgomery
-    /// multiplication with the reduction step interleaved (CIOS; Koc, Acar and
-    /// Kaliski, "Analyzing and comparing Montgomery multiplication
-    /// algorithms", IEEE Micro 1996).
+    /// out = a * b / R mod n, for a, b < n, by the instruction path where
+    /// the CPU has one.
     fn mul(&self, a: &[u64; WORDS], b: &[u64; WORDS], out: &mut [u64; WORDS]) {
+        let n = self.words;
+        match crate::arch::rsa_mul(a, b, &self.modulus, self.n0, n, out) {
+            Some(top) => {
+                if top != 0 || !less(out, &self.modulus, n) {
+                    subtract(out, &self.modulus, n);
+                }
+            }
+            None => self.mul_portable(a, b, out),
+        }
+    }
+    /// Word-serial Montgomery multiplication with the reduction step
+    /// interleaved (CIOS; Koc, Acar and Kaliski, "Analyzing and comparing
+    /// Montgomery multiplication algorithms", IEEE Micro 1996).
+    fn mul_portable(&self, a: &[u64; WORDS], b: &[u64; WORDS], out: &mut [u64; WORDS]) {
         let n = self.words;
         let (a, b, m) = (&a[..n], &b[..n], &self.modulus[..n]);
         let mut t = [0u64; WORDS + 2];
@@ -331,6 +345,53 @@ mod tests {
         }
     }
 
+    /// The instruction path, where the CPU has one, equals the portable
+    /// multiplication for the padded word counts 32 to 64 in steps of 8, and
+    /// at the largest residues.
+    #[test]
+    fn instruction_path_matches_the_portable_multiplication() {
+        // Which path this run compares, for a reader of `--nocapture` output.
+        let (zero, modulus) = ([0u64; WORDS], [u64::MAX; WORDS]);
+        let mut out = [0u64; WORDS];
+        let kernel = crate::arch::rsa_mul(&zero, &zero, &modulus, 1, 4, &mut out).is_some();
+        std::println!("instruction path: {}", if kernel { "the BMI2/ADX kernel" } else { "none; portable against itself" });
+        let mut seed = 0x2545f4914f6cdd1du64;
+        let mut next = move || {
+            seed ^= seed << 13;
+            seed ^= seed >> 7;
+            seed ^= seed << 17;
+            seed
+        };
+        for len in [256usize, 300, 384, 448, 512] {
+            let mut modulus: std::vec::Vec<u8> = (0..len).map(|_| next() as u8).collect();
+            modulus[0] |= 0x80;
+            modulus[len - 1] |= 1;
+            let key = PublicKey::from_components(&modulus, &[1, 0, 1]).unwrap();
+            let n = BigUint::from_bytes_be(&modulus);
+            for case in 0..64 {
+                let mut a = [0u64; WORDS];
+                let mut b = [0u64; WORDS];
+                let (x, y) = if case == 0 {
+                    (&n - 1u8, &n - 1u8)
+                } else {
+                    let x: std::vec::Vec<u8> = (0..len).map(|_| next() as u8).collect();
+                    let y: std::vec::Vec<u8> = (0..len).map(|_| next() as u8).collect();
+                    (BigUint::from_bytes_be(&x) % &n, BigUint::from_bytes_be(&y) % &n)
+                };
+                for (i, d) in x.to_u64_digits().iter().enumerate() {
+                    a[i] = *d;
+                }
+                for (i, d) in y.to_u64_digits().iter().enumerate() {
+                    b[i] = *d;
+                }
+                let (mut ours, mut portable) = ([0; WORDS], [0; WORDS]);
+                key.mul(&a, &b, &mut ours);
+                key.mul_portable(&a, &b, &mut portable);
+                assert_eq!(ours, portable, "{len} bytes, case {case}");
+            }
+        }
+    }
+
     fn integer(words: &[u64; WORDS]) -> BigUint {
         BigUint::from_bytes_le(&words.iter().flat_map(|w| w.to_le_bytes()).collect::<std::vec::Vec<_>>())
     }