[PATCH 03/20] lib/crypto: x86/aes: Clean up aes-aesni.S in preparation for AES modes

From: Eric Biggers

Date: Mon Sep 21 2026 - 01:17:58 EST

Various miscellaneous updates to the assembly code in preparation for
adding implementations of AES modes to the same file:

- Define macros for the function argument registers, function prologues,
and function epilogues to centralize some of the handling of 32-bit vs
64-bit.

- Refactor the actual AES encryption and AES decryption logic into
macros _do_aes and _do_aes_ecb so that some of the modes can reuse it.

- Rename mask to expandkey_mask to differentiate it from the bswap_mask
that will be added.

- Update the prototypes of aes_encrypt_aesni() and aes_decrypt_aesni()
to be dst, src, key so that they will match the mode functions.

Note that this means passing a pointer to the key struct instead of a
(nrounds, rndkeys) pair, similar to what arch/x86/crypto/aes*.S do.
Although this makes the assembly code depend on the format of the key
struct, having one fewer argument makes it easier to accommodate
32-bit mode, and the C glue code becomes slightly simpler.

Signed-off-by: Eric Biggers <ebiggers@xxxxxxxxxx>
---
lib/crypto/x86/aes-aesni.S | 249 +++++++++++++++++++++++++------------
lib/crypto/x86/aes.h | 28 +++--
2 files changed, 185 insertions(+), 92 deletions(-)

diff --git a/lib/crypto/x86/aes-aesni.S b/lib/crypto/x86/aes-aesni.S
index b8c3e104a3be..90a3765d35b8 100644
--- a/lib/crypto/x86/aes-aesni.S
+++ b/lib/crypto/x86/aes-aesni.S
@@ -8,25 +8,126 @@
// AVX. It does use up to SSE4.1, which all CPUs with AES-NI have.
#include <linux/linkage.h>

-.section .rodata
#ifdef __x86_64__
#define RODATA(label) label(%rip)
-#else
+
+#define ARG0 %rdi
+#define ARG1 %rsi
+#define ARG2 %rdx
+#define ARG2_32 %edx
+#define ARG3 %rcx
+#define ARG4 %r8
+#define TMP %rax
+#define TMP_32 %eax
+#define TMP_16 %ax
+#define TMP_8 %al
+
+#else // __x86_64__
+
#define RODATA(label) label
-#endif

+// Caller-save GPRs and the first 3 function arguments, assuming -mregparm=3
+#define ARG0 %eax
+#define ARG1 %edx
+#define ARG2 %ecx
+#define ARG2_32 %ecx
+
+// *Callee*-save GPRs.
+#define ARG3 %edi
+#define ARG3_32 %edi
+#define ARG4 %esi
+#define TMP %ebx
+#define TMP_32 %ebx
+#define TMP_16 %bx
+#define TMP_8 %bl
+#endif // !__x86_64__
+
+// Offsets in struct aes_key
+#define OFFSETOF_NROUNDS 4
+#define OFFSETOF_ROUNDKEYS 16
+#define OFFSETOF_INVROUNDKEYS 256
+
+.section .rodata
+.p2align 4
+.Lexpandkey_mask:
// A mask for pshufb that extracts the last dword, rotates it right by 8
// bits, and copies the result to all four dwords.
-.p2align 4
-.Lmask:
.byte 13, 14, 15, 12, 13, 14, 15, 12, 13, 14, 15, 12, 13, 14, 15, 12

- // The AES round constants, used during key expansion
.Lrcon:
+ // The AES round constants, used during key expansion
.long 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36

.text

+// In 32-bit mode, push certain callee-saved GPRs and optionally load function
+// arguments from the stack into them. Do nothing in 64-bit mode.
+//
+// TMP is always made available as a temporary register.
+//
+// \uses_arg3 is 0 to not touch ARG3, 1 to make it available as a temporary
+// register only, or 2 to actually load it as an argument from the stack.
+// Likewise for \uses_arg4 and ARG4.
+.macro _prologue uses_arg3=0, uses_arg4=0
+#ifdef __i386__
+ .set ARG3_OFFSET, 4
+ .set ARG3_PUSHED, \uses_arg3
+ .set ARG4_PUSHED, \uses_arg4
+.if \uses_arg3
+ push ARG3
+ .set ARG3_OFFSET, ARG3_OFFSET + 4
+ .if \uses_arg3 == 2
+ mov ARG3_OFFSET(%esp), ARG3
+ .endif
+.endif
+.if \uses_arg4
+ push ARG4
+ .set ARG3_OFFSET, ARG3_OFFSET + 4
+ .if \uses_arg4 == 2
+ mov ARG3_OFFSET+4(%esp), ARG4
+ .endif
+.endif
+ push TMP
+ .set ARG3_OFFSET, ARG3_OFFSET + 4
+#endif // __i386__
+.endm
+
+.macro _reload_arg3
+#ifdef __i386__
+ mov ARG3_OFFSET(%esp), ARG3
+#endif
+.endm
+
+// Undo any pushes that _prologue did, then return.
+.macro _epilogue
+#ifdef __i386__
+ pop TMP
+.if ARG4_PUSHED
+ pop ARG4
+.endif
+.if ARG3_PUSHED
+ pop ARG3
+.endif
+#endif
+ RET
+.endm
+
+.macro _aesenc enc, rndkey, data
+.if \enc
+ aesenc \rndkey, \data
+.else
+ aesdec \rndkey, \data
+.endif
+.endm
+
+.macro _aesenclast enc, rndkeylast, data
+.if \enc
+ aesenclast \rndkeylast, \data
+.else
+ aesdeclast \rndkeylast, \data
+.endif
+.endm
+
// Transform four dwords [a0, a1, a2, a3] in \a into
// [a0, a0^a1, a0^a1^a2, a0^a1^a2^a3]. \tmp is a temporary xmm register.
//
@@ -71,32 +172,18 @@
.endm

.macro _aes_expandkey_aesni is_aes128
-#ifdef __x86_64__
// Arguments
- .set RNDKEYS, %rdi
- .set INV_RNDKEYS, %rsi
- .set IN_KEY, %rdx
+ .set RNDKEYS, ARG0
+ .set INV_RNDKEYS, ARG1
+ .set IN_KEY, ARG2

// Other local variables
- .set RCON_PTR, %rcx
- .set COUNTER, %eax
-#else
- // Arguments, assuming -mregparm=3
- .set RNDKEYS, %eax
- .set INV_RNDKEYS, %edx
- .set IN_KEY, %ecx
-
- // Other local variables
- .set RCON_PTR, %ebx
- .set COUNTER, %esi
-#endif
+ .set RCON_PTR, ARG3
+ .set COUNTER, TMP_32
.set RCON, %xmm6
.set MASK, %xmm7

-#ifdef __i386__
- push %ebx
- push %esi
-#endif
+ _prologue uses_arg3=1

.if \is_aes128
// AES-128: the first round key is simply a copy of the raw key.
@@ -112,7 +199,7 @@
.endif

// Generate the remaining round keys.
- movdqa RODATA(.Lmask), MASK
+ movdqa RODATA(.Lexpandkey_mask), MASK
.if \is_aes128
lea RODATA(.Lrcon), RCON_PTR
mov $10, COUNTER
@@ -176,11 +263,7 @@
movdqu %xmm0, 16(INV_RNDKEYS) // => Last inverse round key

.Ldone\@:
-#ifdef __i386__
- pop %esi
- pop %ebx
-#endif
- RET
+ _epilogue
.endm

// void aes128_expandkey_aesni(u32 rndkeys[], u32 *inv_rndkeys,
@@ -195,67 +278,73 @@ SYM_FUNC_START(aes256_expandkey_aesni)
_aes_expandkey_aesni 0
SYM_FUNC_END(aes256_expandkey_aesni)

-.macro _aes_crypt_aesni enc
-#ifdef __x86_64__
- .set RNDKEYS, %rdi
- .set NROUNDS, %esi
- .set OUT, %rdx
- .set IN, %rcx
-#else
- // Assuming -mregparm=3
- .set RNDKEYS, %eax
- .set NROUNDS, %edx
- .set OUT, %ecx
- .set IN, %ebx // Passed on stack
-#endif
-
-#ifdef __i386__
- push %ebx
- mov 8(%esp), %ebx
-#endif
-
- // Zero-th round
- movdqu (IN), %xmm0
- movdqu (RNDKEYS), %xmm1
- pxor %xmm1, %xmm0
-
- // Normal rounds
- add $16, RNDKEYS
+// AES-encrypt (\enc=1) or decrypt (\enc=0) the AESDATA registers specified in
+// \vecs using the aes_enckey or aes_key pointed to by KEY. RNDKEY must be set
+// to a temporary XMM register, NROUNDS to a temporary 32-bit GPR, and
+// RNDKEY_PTR to a temporary full-size GPR.
+.macro _do_aes enc, vecs:vararg
+ movl OFFSETOF_NROUNDS(KEY), NROUNDS
dec NROUNDS
-.Lnext_round\@:
- movdqu (RNDKEYS), %xmm1
.if \enc
- aesenc %xmm1, %xmm0
+ .set rndkey0_offs, OFFSETOF_ROUNDKEYS
.else
- aesdec %xmm1, %xmm0
+ .set rndkey0_offs, OFFSETOF_INVROUNDKEYS
.endif
- add $16, RNDKEYS
+
+ // Do the zero-th AES round.
+ movdqu rndkey0_offs(KEY), RNDKEY
+.irp i, \vecs
+ pxor RNDKEY, AESDATA\i
+.endr
+ // Do the regular AES rounds.
+ lea rndkey0_offs+16(KEY), RNDKEY_PTR
+.Lnext_round\@:
+ movdqu (RNDKEY_PTR), RNDKEY
+ add $16, RNDKEY_PTR
+.irp i, \vecs
+ _aesenc \enc, RNDKEY, AESDATA\i
+.endr
dec NROUNDS
- jne .Lnext_round\@
+ jnz .Lnext_round\@
+ // Do the last AES round.
+ movdqu (RNDKEY_PTR), RNDKEY
+.irp i, \vecs
+ _aesenclast \enc, RNDKEY, AESDATA\i
+.endr
+.endm

- // Last round
- movdqu (RNDKEYS), %xmm1
-.if \enc
- aesenclast %xmm1, %xmm0
-.else
- aesdeclast %xmm1, %xmm0
-.endif
- movdqu %xmm0, (OUT)
+.macro _do_aes_ecb enc, vecs:vararg
+.irp i, \vecs
+ movdqu \i*16(SRC), AESDATA\i
+.endr
+ _do_aes \enc, \vecs
+.irp i, \vecs
+ movdqu AESDATA\i, \i*16(DST)
+.endr
+.endm

-#ifdef __i386__
- pop %ebx
-#endif
- RET
+.macro _aes_crypt_aesni enc
+ .set DST, ARG0
+ .set SRC, ARG1
+ .set KEY, ARG2
+ .set RNDKEY_PTR, ARG3 // Temporary register for _do_aes
+ .set NROUNDS, TMP_32 // Temporary register for _do_aes
+ .set AESDATA0, %xmm0
+ .set RNDKEY, %xmm1 // Temporary register for _do_aes
+
+ _prologue uses_arg3=1
+ _do_aes_ecb \enc, 0
+ _epilogue
.endm

-// void aes_encrypt_aesni(const u32 rndkeys[], int nrounds,
-// u8 out[AES_BLOCK_SIZE], const u8 in[AES_BLOCK_SIZE]);
+// void aes_encrypt_aesni(u8 dst[AES_BLOCK_SIZE], const u8 src[AES_BLOCK_SIZE],
+// const struct aes_enckey *key);
SYM_FUNC_START(aes_encrypt_aesni)
_aes_crypt_aesni 1
SYM_FUNC_END(aes_encrypt_aesni)

-// void aes_decrypt_aesni(const u32 inv_rndkeys[], int nrounds,
-// u8 out[AES_BLOCK_SIZE], const u8 in[AES_BLOCK_SIZE]);
+// void aes_decrypt_aesni(u8 dst[AES_BLOCK_SIZE], const u8 src[AES_BLOCK_SIZE],
+// const struct aes_key *key);
SYM_FUNC_START(aes_decrypt_aesni)
_aes_crypt_aesni 0
SYM_FUNC_END(aes_decrypt_aesni)
diff --git a/lib/crypto/x86/aes.h b/lib/crypto/x86/aes.h
index b047dee94f57..06146fef06be 100644
--- a/lib/crypto/x86/aes.h
+++ b/lib/crypto/x86/aes.h
@@ -7,16 +7,21 @@

#include <asm/fpu/api.h>

-static __ro_after_init DEFINE_STATIC_KEY_FALSE(have_aes);
+static __ro_after_init DEFINE_STATIC_KEY_FALSE(have_aesni);
+
+/* The assembly code assumes the following offsets. */
+static_assert(offsetof(struct aes_enckey, nrounds) == 4);
+static_assert(offsetof(struct aes_enckey, k.rndkeys) == 16);
+static_assert(offsetof(struct aes_key, inv_k.inv_rndkeys) == 256);

void aes128_expandkey_aesni(u32 rndkeys[], u32 *inv_rndkeys,
const u8 in_key[AES_KEYSIZE_128]);
void aes256_expandkey_aesni(u32 rndkeys[], u32 *inv_rndkeys,
const u8 in_key[AES_KEYSIZE_256]);
-void aes_encrypt_aesni(const u32 rndkeys[], int nrounds,
- u8 out[AES_BLOCK_SIZE], const u8 in[AES_BLOCK_SIZE]);
-void aes_decrypt_aesni(const u32 inv_rndkeys[], int nrounds,
- u8 out[AES_BLOCK_SIZE], const u8 in[AES_BLOCK_SIZE]);
+void aes_encrypt_aesni(u8 dst[AES_BLOCK_SIZE], const u8 src[AES_BLOCK_SIZE],
+ const struct aes_enckey *key);
+void aes_decrypt_aesni(u8 dst[AES_BLOCK_SIZE], const u8 src[AES_BLOCK_SIZE],
+ const struct aes_key *key);

/*
* Expand an AES key using AES-NI if supported and usable or generic code
@@ -36,7 +41,7 @@ static void aes_preparekey_arch(union aes_enckey_arch *k,
u32 *rndkeys = k->rndkeys;
u32 *inv_rndkeys = inv_k ? inv_k->inv_rndkeys : NULL;

- if (static_branch_likely(&have_aes) && key_len != AES_KEYSIZE_192 &&
+ if (static_branch_likely(&have_aesni) && key_len != AES_KEYSIZE_192 &&
irq_fpu_usable()) {
kernel_fpu_begin();
if (key_len == AES_KEYSIZE_128)
@@ -53,9 +58,9 @@ static void aes_encrypt_arch(const struct aes_enckey *key,
u8 out[AES_BLOCK_SIZE],
const u8 in[AES_BLOCK_SIZE])
{
- if (static_branch_likely(&have_aes) && irq_fpu_usable()) {
+ if (static_branch_likely(&have_aesni) && irq_fpu_usable()) {
kernel_fpu_begin();
- aes_encrypt_aesni(key->k.rndkeys, key->nrounds, out, in);
+ aes_encrypt_aesni(out, in, key);
kernel_fpu_end();
} else {
aes_encrypt_generic(key->k.rndkeys, key->nrounds, out, in);
@@ -66,10 +71,9 @@ static void aes_decrypt_arch(const struct aes_key *key,
u8 out[AES_BLOCK_SIZE],
const u8 in[AES_BLOCK_SIZE])
{
- if (static_branch_likely(&have_aes) && irq_fpu_usable()) {
+ if (static_branch_likely(&have_aesni) && irq_fpu_usable()) {
kernel_fpu_begin();
- aes_decrypt_aesni(key->inv_k.inv_rndkeys, key->nrounds,
- out, in);
+ aes_decrypt_aesni(out, in, key);
kernel_fpu_end();
} else {
aes_decrypt_generic(key->inv_k.inv_rndkeys, key->nrounds,
@@ -81,5 +85,5 @@ static void aes_decrypt_arch(const struct aes_key *key,
static void aes_mod_init_arch(void)
{
if (boot_cpu_has(X86_FEATURE_AES))
- static_branch_enable(&have_aes);
+ static_branch_enable(&have_aesni);
}
--
2.55.0