[PATCH 1/2] x86/fpu: Extend kernel_fpu_begin_mask() to initialize AMX state

[Chang S. Bae]

The In-Field Scan [1] test does not start its operation when the CPU is
unprepared. If the AMX state is uninitialized, the scan test will
immediately terminate. Therefore, a proper initialization of the AMX
state is necessary to run the test.

Although fpu_idle_fpregs() initializes the state, its usage should be
limited to specialized cases, primarily before entering the sleep state.
The restore_fpregs_from_fpstate() function offers a suitable mechanism
for initializing fpstate in general, which remains within the core code.

Extend kernel_fpu_begin_mask() to include AMX state initialization,
providing the in-field scan driver code access to the appropriate
initialization method while maintaining compliance with established FPU
API semantics.

[1] https://docs.kernel.org/arch/x86/ifs.html
Signed-off-by: Chang S. Bae <chang.seok.bae@xxxxxxxxx>
---
The necessity for AMX initialization is clarified in the Intel Software
Development Manual as of March 2024, particularly in Section 18.2
RECOMMENDATIONS FOR SYSTEM SOFTWARE of Vol. 1.

Side note: restore_fpregs_from_fpstate() also sets the x87 state to a
fixed value. However, this only applies to AMD CPUs with the FXSAVE_LEAK
quirk.
---
 arch/x86/include/asm/fpu/api.h | 1 +
 arch/x86/kernel/fpu/core.c     | 3 +++
 2 files changed, 4 insertions(+)

diff --git a/arch/x86/include/asm/fpu/api.h b/arch/x86/include/asm/fpu/api.h
index a2be3aefff9f..67887fc45c24 100644
--- a/arch/x86/include/asm/fpu/api.h
+++ b/arch/x86/include/asm/fpu/api.h
@@ -25,6 +25,7 @@
 /* Kernel FPU states to initialize in kernel_fpu_begin_mask() */
 #define KFPU_387	_BITUL(0)	/* 387 state will be initialized */
 #define KFPU_MXCSR	_BITUL(1)	/* MXCSR will be initialized */
+#define KFPU_AMX	_BITUL(2)	/* AMX will be initialized */
 
 extern void kernel_fpu_begin_mask(unsigned int kfpu_mask);
 extern void kernel_fpu_end(void);
diff --git a/arch/x86/kernel/fpu/core.c b/arch/x86/kernel/fpu/core.c
index 520deb411a70..0c0235b4a901 100644
--- a/arch/x86/kernel/fpu/core.c
+++ b/arch/x86/kernel/fpu/core.c
@@ -440,6 +440,9 @@ void kernel_fpu_begin_mask(unsigned int kfpu_mask)
 
 	if (unlikely(kfpu_mask & KFPU_387) && boot_cpu_has(X86_FEATURE_FPU))
 		asm volatile ("fninit");
+
+	if (unlikely(kfpu_mask & KFPU_AMX) && boot_cpu_has(X86_FEATURE_AMX_TILE))
+		restore_fpregs_from_fpstate(&init_fpstate, XFEATURE_MASK_XTILE);
 }
 EXPORT_SYMBOL_GPL(kernel_fpu_begin_mask);
 
-- 
2.34.1