simd: encapsulate fpu amortization into nice functions

3 files changed, 66 insertions, 47 deletions

diff --git a/src/crypto/chacha20poly1305.c b/src/crypto/chacha20poly1305.c
index ccc6e1c..df1c5aa 100644
--- a/src/crypto/chacha20poly1305.c
+++ b/src/crypto/chacha20poly1305.c
@@ -6,6 +6,7 @@
 #include "chacha20poly1305.h"
 #include "chacha20.h"
 #include "poly1305.h"
+#include "simd.h"
 
 #include <linux/kernel.h>
 #include <crypto/scatterwalk.h>
@@ -65,9 +66,9 @@ void chacha20poly1305_encrypt(u8 *dst, const u8 *src, const size_t src_len,
 {
 	bool have_simd;
 
-	have_simd = chacha20poly1305_init_simd();
+	have_simd = simd_get();
 	__chacha20poly1305_encrypt(dst, src, src_len, ad, ad_len, nonce, key, have_simd);
-	chacha20poly1305_deinit_simd(have_simd);
+	simd_put(have_simd);
 }
 
 bool chacha20poly1305_encrypt_sg(struct scatterlist *dst, struct scatterlist *src, const size_t src_len,
@@ -176,9 +177,9 @@ bool chacha20poly1305_decrypt(u8 *dst, const u8 *src, const size_t src_len,
 {
 	bool have_simd, ret;
 
-	have_simd = chacha20poly1305_init_simd();
+	have_simd = simd_get();
 	ret = __chacha20poly1305_decrypt(dst, src, src_len, ad, ad_len, nonce, key, have_simd);
-	chacha20poly1305_deinit_simd(have_simd);
+	simd_put(have_simd);
 	return ret;
 }
 
@@ -253,13 +254,13 @@ void xchacha20poly1305_encrypt(u8 *dst, const u8 *src, const size_t src_len,
 			       const u8 nonce[XCHACHA20POLY1305_NONCELEN],
 			       const u8 key[CHACHA20POLY1305_KEYLEN])
 {
-	bool have_simd = chacha20poly1305_init_simd();
+	bool have_simd = simd_get();
 	u8 derived_key[CHACHA20POLY1305_KEYLEN] __aligned(16);
 
 	hchacha20(derived_key, nonce, key, have_simd);
 	__chacha20poly1305_encrypt(dst, src, src_len, ad, ad_len, le64_to_cpup((__le64 *)(nonce + 16)), derived_key, have_simd);
 	memzero_explicit(derived_key, CHACHA20POLY1305_KEYLEN);
-	chacha20poly1305_deinit_simd(have_simd);
+	simd_put(have_simd);
 }
 
 bool xchacha20poly1305_decrypt(u8 *dst, const u8 *src, const size_t src_len,
@@ -267,13 +268,13 @@ bool xchacha20poly1305_decrypt(u8 *dst, const u8 *src, const size_t src_len,
 			       const u8 nonce[XCHACHA20POLY1305_NONCELEN],
 			       const u8 key[CHACHA20POLY1305_KEYLEN])
 {
-	bool ret, have_simd = chacha20poly1305_init_simd();
+	bool ret, have_simd = simd_get();
 	u8 derived_key[CHACHA20POLY1305_KEYLEN] __aligned(16);
 
 	hchacha20(derived_key, nonce, key, have_simd);
 	ret = __chacha20poly1305_decrypt(dst, src, src_len, ad, ad_len, le64_to_cpup((__le64 *)(nonce + 16)), derived_key, have_simd);
 	memzero_explicit(derived_key, CHACHA20POLY1305_KEYLEN);
-	chacha20poly1305_deinit_simd(have_simd);
+	simd_put(have_simd);
 	return ret;
 }
 
diff --git a/src/crypto/chacha20poly1305.h b/src/crypto/chacha20poly1305.h
index 1b122ac..43b0a17 100644
--- a/src/crypto/chacha20poly1305.h
+++ b/src/crypto/chacha20poly1305.h
@@ -44,45 +44,6 @@ bool __must_check xchacha20poly1305_decrypt(u8 *dst, const u8 *src, const size_t
 					    const u8 nonce[XCHACHA20POLY1305_NONCELEN],
 					    const u8 key[CHACHA20POLY1305_KEYLEN]);
 
-#if defined(CONFIG_X86_64)
-#include <linux/version.h>
-#include <asm/fpu/api.h>
-#include <asm/simd.h>
-#elif IS_ENABLED(CONFIG_KERNEL_MODE_NEON)
-#include <asm/neon.h>
-#include <asm/simd.h>
-#endif
-
-static inline bool chacha20poly1305_init_simd(void)
-{
-	bool have_simd = false;
-#if defined(CONFIG_X86_64) && !defined(CONFIG_UML) && !defined(CONFIG_PREEMPT_RT_BASE)
-	have_simd = irq_fpu_usable();
-	if (have_simd)
-		kernel_fpu_begin();
-#elif IS_ENABLED(CONFIG_KERNEL_MODE_NEON) && !defined(CONFIG_PREEMPT_RT_BASE)
-#if defined(CONFIG_ARM64)
-	have_simd = true; /* ARM64 supports NEON in any context. */
-#elif defined(CONFIG_ARM)
-	have_simd = may_use_simd(); /* ARM doesn't support NEON in interrupt context. */
-#endif
-	if (have_simd)
-		kernel_neon_begin();
-#endif
-	return have_simd;
-}
-
-static inline void chacha20poly1305_deinit_simd(bool was_on)
-{
-#if defined(CONFIG_X86_64) && !defined(CONFIG_UML)
-	if (was_on)
-		kernel_fpu_end();
-#elif IS_ENABLED(CONFIG_KERNEL_MODE_NEON)
-	if (was_on)
-		kernel_neon_end();
-#endif
-}
-
 #ifdef DEBUG
 bool chacha20poly1305_selftest(void);
 #endif
diff --git a/src/crypto/simd.h b/src/crypto/simd.h
new file mode 100644
index 0000000..21e3c55
--- /dev/null
+++ b/src/crypto/simd.h
@@ -0,0 +1,57 @@
+/* SPDX-License-Identifier: GPL-2.0
+ *
+ * Copyright (C) 2015-2018 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved.
+ */
+
+#ifndef _WG_SIMD_H
+#define _WG_SIMD_H
+
+#if defined(CONFIG_X86_64)
+#include <linux/version.h>
+#include <asm/fpu/api.h>
+#include <asm/simd.h>
+#elif IS_ENABLED(CONFIG_KERNEL_MODE_NEON)
+#include <asm/neon.h>
+#include <asm/simd.h>
+#endif
+
+static inline bool simd_get(void)
+{
+	bool have_simd = false;
+#if defined(CONFIG_X86_64) && !defined(CONFIG_UML) && !defined(CONFIG_PREEMPT_RT_BASE)
+	have_simd = irq_fpu_usable();
+	if (have_simd)
+		kernel_fpu_begin();
+#elif IS_ENABLED(CONFIG_KERNEL_MODE_NEON) && !defined(CONFIG_PREEMPT_RT_BASE)
+#if defined(CONFIG_ARM64)
+	have_simd = true; /* ARM64 supports NEON in any context. */
+#elif defined(CONFIG_ARM)
+	have_simd = may_use_simd(); /* ARM doesn't support NEON in interrupt context. */
+#endif
+	if (have_simd)
+		kernel_neon_begin();
+#endif
+	return have_simd;
+}
+
+static inline void simd_put(bool was_on)
+{
+#if defined(CONFIG_X86_64) && !defined(CONFIG_UML) && !defined(CONFIG_PREEMPT_RT_BASE)
+	if (was_on)
+		kernel_fpu_end();
+#elif IS_ENABLED(CONFIG_KERNEL_MODE_NEON) && !defined(CONFIG_PREEMPT_RT_BASE)
+	if (was_on)
+		kernel_neon_end();
+#endif
+}
+
+static inline bool simd_relax(bool was_on)
+{
+	if (was_on && need_resched()) {
+		simd_put(true);
+		return simd_get();
+	}
+	return was_on;
+}
+
+#endif /* _WG_SIMD_H */