[PATCH 3/3] x86/fpu: Don't support kernel-mode FPU when irqs_disabled()

[Eric Biggers <ebiggers@kernel.org>]

From: Eric Biggers <ebiggers@google.com>

Make irq_fpu_usable() return false when irqs_disabled().  That makes the
irqs_disabled() checks in kernel_fpu_begin_mask() and kernel_fpu_end()
unnecessary, so also remove those.

Rationale:

- There's no known use case for kernel-mode FPU when irqs_disabled().
  arm64 and riscv already disallow kernel-mode FPU when irqs_disabled().
  __save_processor_state() previously did expect kernel_fpu_begin() and
  kernel_fpu_end() to work when irqs_disabled(), but this was a
  different use case and not actual kernel-mode FPU use.

- This is more efficient, since one call to irqs_disabled() replaces two
  irqs_disabled() and one in_hardirq().

- This fixes irq_fpu_usable() to correctly return false during CPU
  initialization.  Incorrectly returning true caused the SHA-256 library
  code, which is called when loading AMD microcode, to take a
  SIMD-optimized code path too early, causing a crash.  By correctly
  returning false from irq_fpu_usable(), the generic SHA-256 code
  correctly gets used instead.  (Note: SIMD-optimized SHA-256 doesn't
  get enabled until subsys_initcall, but CPU hotplug can happen later.)

Fixes: 11d7956d526f ("crypto: x86/sha256 - implement library instead of shash")
Reported-by: Ayush Jain <Ayush.Jain3@amd.com>
Closes: https://lore.kernel.org/r/20250516112217.GBaCcf6Yoc6LkIIryP@fat_crate.local
Signed-off-by: Eric Biggers <ebiggers@google.com>
---
 arch/x86/kernel/fpu/core.c | 49 ++++++++++++++------------------------
 1 file changed, 18 insertions(+), 31 deletions(-)

diff --git a/arch/x86/kernel/fpu/core.c b/arch/x86/kernel/fpu/core.c
index 476393b1d5e8f..9b3c5e17f86cd 100644
--- a/arch/x86/kernel/fpu/core.c
+++ b/arch/x86/kernel/fpu/core.c
@@ -66,42 +66,31 @@ struct fpu *x86_task_fpu(struct task_struct *task)
  * Can we use the FPU in kernel mode with the
  * whole "kernel_fpu_begin/end()" sequence?
  */
 bool irq_fpu_usable(void)
 {
-	if (WARN_ON_ONCE(in_nmi()))
-		return false;
-
 	/*
-	 * In kernel FPU usage already active?  This detects any explicitly
-	 * nested usage in task or softirq context, which is unsupported.  It
-	 * also detects attempted usage in a hardirq that has interrupted a
-	 * kernel-mode FPU section.
+	 * To ensure that (non-explicitly-nested) kernel-mode FPU is always
+	 * usable in task and softirq contexts, kernel_fpu_begin() disables
+	 * preemption and softirqs, and kernel_fpu_end() re-enables them.  That
+	 * is not compatible with hardirqs being disabled (including hardirq
+	 * context), or with NMI context.  Support for kernel-mode FPU is not
+	 * needed in those contexts anyway.  Return false in those contexts.
+	 *
+	 * Returning false when irqs_disabled() also eliminates the need to
+	 * explicitly check whether the FPU has been initialized yet during CPU
+	 * initialization.  Before then, hardirqs are still disabled.
 	 */
+	if (irqs_disabled() || WARN_ON_ONCE(in_nmi()))
+		return false;
+
+	/* Catch any explicitly nested usage, which should never happen. */
 	if (this_cpu_read(in_kernel_fpu)) {
-		WARN_ON_FPU(!in_hardirq());
+		WARN_ON_FPU(1);
 		return false;
 	}
-
-	/*
-	 * When not in NMI or hard interrupt context, FPU can be used in:
-	 *
-	 * - Task context except from within fpregs_lock()'ed critical
-	 *   regions.
-	 *
-	 * - Soft interrupt processing context which cannot happen
-	 *   while in a fpregs_lock()'ed critical region.
-	 */
-	if (!in_hardirq())
-		return true;
-
-	/*
-	 * In hard interrupt context it's safe when soft interrupts
-	 * are enabled, which means the interrupt did not hit in
-	 * a fpregs_lock()'ed critical region.
-	 */
-	return !softirq_count();
+	return true;
 }
 EXPORT_SYMBOL(irq_fpu_usable);
 
 /*
  * Track AVX512 state use because it is known to slow the max clock
@@ -443,12 +432,11 @@ static __always_inline void __fpu_save_state(void)
 	__cpu_invalidate_fpregs_state();
 }
 
 void kernel_fpu_begin_mask(unsigned int kfpu_mask)
 {
-	if (!irqs_disabled())
-		fpregs_lock();
+	fpregs_lock();
 
 	WARN_ON_FPU(!irq_fpu_usable());
 	WARN_ON_FPU(this_cpu_read(in_kernel_fpu));
 
 	this_cpu_write(in_kernel_fpu, true);
@@ -467,12 +455,11 @@ EXPORT_SYMBOL_GPL(kernel_fpu_begin_mask);
 void kernel_fpu_end(void)
 {
 	WARN_ON_FPU(!this_cpu_read(in_kernel_fpu));
 
 	this_cpu_write(in_kernel_fpu, false);
-	if (!irqs_disabled())
-		fpregs_unlock();
+	fpregs_unlock();
 }
 EXPORT_SYMBOL_GPL(kernel_fpu_end);
 
 #ifdef CONFIG_PM_SLEEP
 /*
-- 
2.49.0