arm: support aes_armv8 on Rust 1.59+ using asm!

tarcieri · tarcieri · commit 4655fe470dd7 · 2023-06-15T13:30:53.000-06:00
Adds "polyfills" for the unstable ARMv8 AES intrinsics using the `asm!`
macro which was stabilized in Rust 1.59.

Based on benchmarks this has no effect on performance, although it was
necessary to place AESE/AESMC and AESD/AESIMC into a single `asm!` block
in order to ensure that instructions fuse properly, as they did when
using the proper intrinsics.
diff --git a/.github/workflows/aes.yml b/.github/workflows/aes.yml
@@ -219,15 +219,15 @@ jobs:
           cross test --package aes --target ${{ matrix.target }}
           cross test --package aes --target ${{ matrix.target }} --features hazmat
 
-  # ARMv8 cross-compiled tests for AES intrinsics (nightly-only)
+  # ARMv8 cross-compiled tests for AES intrinsics
   armv8:
     env:
       RUSTFLAGS: "-Dwarnings --cfg aes_armv8"
     strategy:
       matrix:
         include:
           - target: aarch64-unknown-linux-gnu
-            rust: nightly
+            rust: 1.59.0 # MSRV for inline `asm!`
     runs-on: ubuntu-latest
     # Cross mounts only current package, i.e. by default it ignores workspace's Cargo.toml
     defaults:
diff --git a/aes/src/armv8.rs b/aes/src/armv8.rs
@@ -14,6 +14,7 @@ pub(crate) mod hazmat;
 
 mod encdec;
 mod expand;
+mod intrinsics;
 #[cfg(test)]
 mod test_expand;
 
diff --git a/aes/src/armv8/encdec.rs b/aes/src/armv8/encdec.rs
@@ -4,6 +4,12 @@ use crate::{Block, Block8};
 use cipher::inout::InOut;
 use core::arch::aarch64::*;
 
+// Stable "polyfills" for unstable core::arch::aarch64 intrinsics
+// TODO(tarcieri): remove when these intrinsics have been stabilized
+use super::intrinsics::{
+    vaesdq_u8, vaesdq_u8_and_vaesimcq_u8, vaeseq_u8, vaeseq_u8_and_vaesmcq_u8,
+};
+
 /// Perform AES encryption using the given expanded keys.
 #[target_feature(enable = "aes")]
 #[target_feature(enable = "neon")]
@@ -19,11 +25,8 @@ pub(super) unsafe fn encrypt1<const N: usize>(
     let mut state = vld1q_u8(in_ptr as *const u8);
 
     for k in expanded_keys.iter().take(rounds - 1) {
-        // AES single round encryption
-        state = vaeseq_u8(state, *k);
-
-        // AES mix columns
-        state = vaesmcq_u8(state);
+        // AES single round encryption and mix columns
+        state = vaeseq_u8_and_vaesmcq_u8(state, *k);
     }
 
     // AES single round encryption
@@ -62,11 +65,8 @@ pub(super) unsafe fn encrypt8<const N: usize>(
 
     for k in expanded_keys.iter().take(rounds - 1) {
         for i in 0..8 {
-            // AES single round encryption
-            state[i] = vaeseq_u8(state[i], *k);
-
-            // AES mix columns
-            state[i] = vaesmcq_u8(state[i]);
+            // AES single round encryption and mix columns
+            state[i] = vaeseq_u8_and_vaesmcq_u8(state[i], *k);
         }
     }
 
@@ -95,11 +95,8 @@ pub(super) unsafe fn decrypt1<const N: usize>(
     let mut state = vld1q_u8(in_ptr as *const u8);
 
     for k in expanded_keys.iter().take(rounds - 1) {
-        // AES single round decryption
-        state = vaesdq_u8(state, *k);
-
-        // AES inverse mix columns
-        state = vaesimcq_u8(state);
+        // AES single round decryption and inverse mix columns
+        state = vaesdq_u8_and_vaesimcq_u8(state, *k);
     }
 
     // AES single round decryption
@@ -138,11 +135,8 @@ pub(super) unsafe fn decrypt8<const N: usize>(
 
     for k in expanded_keys.iter().take(rounds - 1) {
         for i in 0..8 {
-            // AES single round decryption
-            state[i] = vaesdq_u8(state[i], *k);
-
-            // AES inverse mix columns
-            state[i] = vaesimcq_u8(state[i]);
+            // AES single round decryption and inverse mix columns
+            state[i] = vaesdq_u8_and_vaesimcq_u8(state[i], *k);
         }
     }
 
diff --git a/aes/src/armv8/expand.rs b/aes/src/armv8/expand.rs
@@ -2,6 +2,10 @@
 
 use core::{arch::aarch64::*, mem, slice};
 
+// Stable "polyfills" for unstable core::arch::aarch64 intrinsics
+// TODO(tarcieri): remove when these intrinsics have been stabilized
+use super::intrinsics::{vaeseq_u8, vaesimcq_u8};
+
 /// There are 4 AES words in a block.
 const BLOCK_WORDS: usize = 4;
 
diff --git a/aes/src/armv8/hazmat.rs b/aes/src/armv8/hazmat.rs
@@ -7,6 +7,9 @@
 use crate::{Block, Block8};
 use core::arch::aarch64::*;
 
+// Stable "polyfills" for unstable core::arch::aarch64 intrinsics
+use super::intrinsics::{vaesdq_u8, vaeseq_u8, vaesimcq_u8, vaesmcq_u8};
+
 /// AES cipher (encrypt) round function.
 #[allow(clippy::cast_ptr_alignment)]
 #[target_feature(enable = "aes")]
diff --git a/aes/src/armv8/intrinsics.rs b/aes/src/armv8/intrinsics.rs
@@ -0,0 +1,85 @@
+//! Stable "polyfills" for unstable `core::arch::aarch64` intrinsics which use
+//! `asm!` internally to allow use on stable Rust.
+// TODO(tarcieri): remove when these intrinsics have been stabilized
+
+use core::arch::{aarch64::uint8x16_t, asm};
+
+/// AES single round encryption.
+#[inline(always)]
+pub(super) unsafe fn vaeseq_u8(mut data: uint8x16_t, key: uint8x16_t) -> uint8x16_t {
+    asm!(
+        "AESE {:v}.16B, {:v}.16B",
+        inout(vreg) data, in(vreg) key,
+        options(pure, nomem, nostack, preserves_flags)
+    );
+    data
+}
+
+/// AES single round decryption.
+#[inline(always)]
+pub(super) unsafe fn vaesdq_u8(mut data: uint8x16_t, key: uint8x16_t) -> uint8x16_t {
+    asm!(
+        "AESD {:v}.16B, {:v}.16B",
+        inout(vreg) data, in(vreg) key,
+        options(pure, nomem, nostack, preserves_flags)
+    );
+    data
+}
+
+/// AES mix columns.
+#[cfg(feature = "hazmat")]
+#[inline(always)]
+pub(super) unsafe fn vaesmcq_u8(mut data: uint8x16_t) -> uint8x16_t {
+    asm!(
+        "AESMC {:v}.16B, {:v}.16B",
+        out(vreg) data, in(vreg) data,
+        options(pure, nomem, nostack, preserves_flags)
+    );
+    data
+}
+
+/// AES inverse mix columns.
+#[inline(always)]
+pub(super) unsafe fn vaesimcq_u8(mut data: uint8x16_t) -> uint8x16_t {
+    asm!(
+        "AESIMC {:v}.16B, {:v}.16B",
+        out(vreg) data, in(vreg) data,
+        options(pure, nomem, nostack, preserves_flags)
+    );
+    data
+}
+
+/// AES single round encryption combined with mix columns.
+///
+/// These two instructions are combined into a single assembly block to ensure
+/// that instructions fuse properly.
+#[inline(always)]
+pub(super) unsafe fn vaeseq_u8_and_vaesmcq_u8(mut data: uint8x16_t, key: uint8x16_t) -> uint8x16_t {
+    asm!(
+        "AESE {d:v}.16B, {k:v}.16B",
+        "AESMC {d:v}.16B, {d:v}.16B",
+        d = inout(vreg) data,
+        k = in(vreg) key,
+        options(pure, nomem, nostack, preserves_flags)
+    );
+    data
+}
+
+/// AES single round decryption combined with mix columns.
+///
+/// These two instructions are combined into a single assembly block to ensure
+/// that instructions fuse properly.
+#[inline(always)]
+pub(super) unsafe fn vaesdq_u8_and_vaesimcq_u8(
+    mut data: uint8x16_t,
+    key: uint8x16_t,
+) -> uint8x16_t {
+    asm!(
+        "AESD {d:v}.16B, {k:v}.16B",
+        "AESIMC {d:v}.16B, {d:v}.16B",
+        d = inout(vreg) data,
+        k = in(vreg) key,
+        options(pure, nomem, nostack, preserves_flags)
+    );
+    data
+}
diff --git a/aes/src/lib.rs b/aes/src/lib.rs
@@ -26,11 +26,11 @@
 //! backend at the cost of decreased performance (using a modified form of
 //! the fixslicing technique called "semi-fixslicing").
 //!
-//! ## ARMv8 intrinsics (nightly-only)
+//! ## ARMv8 intrinsics (Rust 1.59+)
 //! On `aarch64` targets including `aarch64-apple-darwin` (Apple M1) and Linux
 //! targets such as `aarch64-unknown-linux-gnu` and `aarch64-unknown-linux-musl`,
 //! support for using AES intrinsics provided by the ARMv8 Cryptography Extensions
-//! is available when using the nightly compiler, and can be enabled using the
+//! is available when using Rust 1.59 or above, and can be enabled using the
 //! `aes_armv8` configuration flag.
 //!
 //! On Linux and macOS, when the `aes_armv8` flag is enabled support for AES
@@ -99,7 +99,7 @@
 //!
 //! You can modify crate using the following configuration flags:
 //!
-//! - `aes_armv8`: enable ARMv8 AES intrinsics (nightly-only).
+//! - `aes_armv8`: enable ARMv8 AES intrinsics (Rust 1.59+).
 //! - `aes_force_soft`: force software implementation.
 //! - `aes_compact`: reduce code size at the cost of slower performance
 //! (affects only software backend).
@@ -119,7 +119,6 @@
 )]
 #![cfg_attr(docsrs, feature(doc_cfg))]
 #![warn(missing_docs, rust_2018_idioms)]
-#![cfg_attr(all(aes_armv8, target_arch = "aarch64"), feature(stdsimd))]
 
 #[cfg(feature = "hazmat")]
 #[cfg_attr(docsrs, doc(cfg(feature = "hazmat")))]
