[PATCH] lib/crypto: x86/chacha: Use bzhi to simplify mask generation

[Eric Biggers <ebiggers@kernel.org>]

When XOR'ing the bytes from a partial register at the end, use the bzhi
instruction (from BMI2) to generate the needed mask a bit more
efficiently.  All real processors with AVX-512 also have BMI2, so we
might as well make use of it.  This follows the example of the AVX-512
optimized AES-GCM code which already does this.

Signed-off-by: Eric Biggers <ebiggers@kernel.org>
---
 lib/crypto/x86/chacha-avx512-x86_64.S | 30 +++++++++++----------------
 lib/crypto/x86/chacha.h               |  3 ++-
 2 files changed, 14 insertions(+), 19 deletions(-)

diff --git a/lib/crypto/x86/chacha-avx512-x86_64.S b/lib/crypto/x86/chacha-avx512-x86_64.S
index 92babb1d3844..422479631d5b 100644
--- a/lib/crypto/x86/chacha-avx512-x86_64.S
+++ b/lib/crypto/x86/chacha-avx512-x86_64.S
@@ -189,15 +189,13 @@ SYM_FUNC_START(chacha_2block_xor_avx512vl)
 
 .Lxorpart2:
 	# xor remaining bytes from partial register into output
-	mov		%rcx,%rax
+	mov		%rcx,%r9
+	and		$~0xf,%r9
 	and		$0xf,%rcx
 	jz		.Ldone2
-	mov		%rax,%r9
-	and		$~0xf,%r9
 
-	mov		$1,%rax
-	shld		%cl,%rax,%rax
-	sub		$1,%rax
+	mov		$-1,%rax
+	bzhi		%rcx,%rax,%rax
 	kmovq		%rax,%k1
 
 	vmovdqu8	(%rdx,%r9),%xmm1{%k1}{z}
@@ -455,15 +453,13 @@ SYM_FUNC_START(chacha_4block_xor_avx512vl)
 
 .Lxorpart4:
 	# xor remaining bytes from partial register into output
-	mov		%rcx,%rax
+	mov		%rcx,%r9
+	and		$~0xf,%r9
 	and		$0xf,%rcx
 	jz		.Ldone4
-	mov		%rax,%r9
-	and		$~0xf,%r9
 
-	mov		$1,%rax
-	shld		%cl,%rax,%rax
-	sub		$1,%rax
+	mov		$-1,%rax
+	bzhi		%rcx,%rax,%rax
 	kmovq		%rax,%k1
 
 	vmovdqu8	(%rdx,%r9),%xmm1{%k1}{z}
@@ -977,15 +973,13 @@ SYM_FUNC_END(chacha_4block_xor_avx512vl)
 
 .Lxorpart\@:
 	# xor remaining bytes from partial register into output
-	mov		%rcx,%rax
+	mov		%rcx,%r9
+	and		$~(VL - 1),%r9
 	and		$(VL - 1),%rcx
 	jz		.Ldone\@
-	mov		%rax,%r9
-	and		$~(VL - 1),%r9
 
-	mov		$1,%rax
-	shld		%cl,%rax,%rax
-	sub		$1,%rax
+	mov		$-1,%rax
+	bzhi		%rcx,%rax,%rax
 	kmovq		%rax,%k1
 
 	vmovdqu8	(%rdx,%r9),V1{%k1}{z}
diff --git a/lib/crypto/x86/chacha.h b/lib/crypto/x86/chacha.h
index fdf01b07742e..ed39b3fdda84 100644
--- a/lib/crypto/x86/chacha.h
+++ b/lib/crypto/x86/chacha.h
@@ -191,7 +191,8 @@ static void chacha_mod_init_arch(void)
 		static_branch_enable(&chacha_use_avx2);
 
 		if (boot_cpu_has(X86_FEATURE_AVX512VL) &&
-		    boot_cpu_has(X86_FEATURE_AVX512BW)) { /* kmovq */
+		    boot_cpu_has(X86_FEATURE_AVX512BW) &&
+		    boot_cpu_has(X86_FEATURE_BMI2)) {
 			static_branch_enable(&chacha_use_avx512vl);
 
 			if (!boot_cpu_has(X86_FEATURE_PREFER_YMM))

base-commit: 9fbbbaefd5f95551da18426c200cf41e1b4255a1
-- 
2.56.0