[librm] Enable the use of SSE instructions if supported by the CPU

Clear CR0.EM, clear CR0.TS, and set CR4.OSFXSR in order to allow the
execution of SSE instructions such as the AES-NI instructions for AES
hardware acceleration.

Note that we have to assert CR4.OSFXSR to allow these instructions to
execute, but our context switching logic (e.g. the protected-mode and
long-mode interrupt handlers) does not actually preserve the
FPU/MMX/SSE registers.

Our C code therefore cannot in general presume that the FPU/MMX/SSE
registers will be preserved across arbitrary context boundaries.
However, since C code executes with interrupts disabled, an individual
function may safely use temporary FPU/MMX/SSE registers provided that
it does not enable interrupts or otherwise relinquish the context.

For the same C code to also be usable under the UEFI IA-32 ABI, it
must preserve all registers other than %eax, %ecx, and %edx, including
preserving all MMX and XMM registers.

The practical upshot is therefore that C code may use SSE instructions
and may assume that SSE registers will not be changed arbitrarily
during execution (either because the ABI guarantees preservation, as
with UEFI or Linux, or because the runtime environment guarantees that
interrupts are disabled), but the C code must itself restore the
values of any modified FPU/MMX/SSE registers.

The "fxsave"/"fxrstor" performed by virt_call() would allow for a
slightly more relaxed constraint if support for the UEFI IA-32 ABI
were ever to be dropped in future.  A real-mode caller that is making
use of SSE must have already set OSFXSR, and so its non-64-bit
registers %xmm0-%xmm7 would already be saved and restored across
virt_call().  The tightest constraint would then become the UEFI X64
ABI, which defines %xmm0-%xmm5 as volatile (i.e. caller-saved): this
would allow C code in iPXE to use %xmm0-%xmm5 without needing to
explicitly save and restore their values.

Signed-off-by: Michael Brown <mcb30@ipxe.org>
diff --git a/src/arch/x86/transitions/librm.S b/src/arch/x86/transitions/librm.S
index c08d544..a550a27 100644
--- a/src/arch/x86/transitions/librm.S
+++ b/src/arch/x86/transitions/librm.S
@@ -16,12 +16,21 @@
 /* CR0: protection enabled */
 #define CR0_PE ( 1 << 0 )
 
+/* CR0: FPU emulation enabled */
+#define CR0_EM ( 1 << 2 )
+
+/* CR0: task switched */
+#define CR0_TS ( 1 << 3 )
+
 /* CR0: paging */
 #define CR0_PG ( 1 << 31 )
 
 /* CR4: physical address extensions */
 #define CR4_PAE ( 1 << 5 )
 
+/* CR4: OS support for FXSAVE/FXRSTOR across transitions */
+#define CR4_OSFXSR ( 1 << 9 )
+
 /* Extended feature enable MSR (EFER) */
 #define MSR_EFER 0xc0000080
 
@@ -426,6 +435,37 @@
 	/* Load protected-mode global descriptor table */
 	data32 lgdt gdtr
 
+	/* Enable SSE instruction execution, if supported.
+	 *
+	 * Note that we have to assert CR4.OSFXSR to allow these
+	 * instructions to execute, but our context switching logic
+	 * (e.g. the protected-mode and long-mode interrupt handlers)
+	 * does not actually preserve the FPU/MMX/SSE registers.
+	 *
+	 * Our C code therefore cannot in general presume that the
+	 * FPU/MMX/SSE registers will be preserved across arbitrary
+	 * context boundaries.  However, since C code executes with
+	 * interrupts disabled, an individual function may safely use
+	 * temporary FPU/MMX/SSE registers provided that it does not
+	 * enable interrupts or otherwise relinquish the context.
+	 *
+	 * For the same C code to also be usable under the UEFI IA-32
+	 * ABI, it must preserve all registers other than %eax, %ecx,
+	 * and %edx, including preserving all MMX and XMM registers.
+	 *
+	 * The practical upshot is therefore that C code may use SSE
+	 * instructions and may assume that SSE registers will not be
+	 * changed arbitrarily during execution (either because the
+	 * ABI guarantees preservation, as with UEFI or Linux, or
+	 * because the runtime environment guarantees that interrupts
+	 * are disabled), but the C code must itself restore the
+	 * values of any modified FPU/MMX/SSE registers.
+	 */
+.if32 ; testb	$0xff, fxsr_supported ; jz 1f ; .endif
+	movl	%cr4, %eax
+	orw	$CR4_OSFXSR, %ax
+	movl	%eax, %cr4
+1:
 	/* Zero segment registers.  This wastes around 12 cycles on
 	 * real hardware, but saves a substantial number of emulated
 	 * instructions under KVM.
@@ -440,7 +480,7 @@
 	/* Switch to protected mode with paging disabled */
 	cli
 	movl	%cr0, %eax
-	andl	$~CR0_PG, %eax
+	andl	$~( CR0_PG | CR0_TS | CR0_EM ), %eax
 	orb	$CR0_PE, %al
 	movl	%eax, %cr0
 	data32 ljmp	$VIRTUAL_CS, $VIRTUAL(r2p_pmode)