[LibOS] Save/restore FP/SSE/AVX/... control words on syscalls

Previously, Graphene saved and restored only the GP registers of the
CPU context on entering and leaving syscall emulation in syscalldb().
In reality, the CPU context must also contain FP control word (fpcw)
and the SSE/AVX/... control word (mxcsr). This commit preserves these
control words across app-to-Graphene context switches.

During syscall emulation, Graphene may clobber the FP/SSE/AVX/...
state (except the control words). We rely on the fact that apps do
*not* assume that this state is preserved across syscalls (except
the control words). Thus, it is enough to save/restore only the
control words on each syscall. This commit also removes previous
hack of performing expensive xsave/xrstor instructions on clone().

LibOS/shim/include/shim_tcb.h

@@ -11,6 +11,8 @@
 
 struct shim_context {
     struct shim_regs* regs;
+    uint16_t          fpcw;    /* FPU Control Word (for x87) */
+    uint32_t          mxcsr;   /* MXCSR control/status register (for SSE/AVX/...) */
     uint64_t          fs_base;
     struct atomic_int preempt;
 };

LibOS/shim/src/generated-offsets.c

@@ -7,6 +7,8 @@
 __attribute__((__used__)) static void dummy(void) {
     OFFSET_T(SHIM_TCB_OFFSET, PAL_TCB, libos_tcb);
     OFFSET_T(TCB_REGS, shim_tcb_t, context.regs);
+    OFFSET_T(TCB_FPCW, shim_tcb_t, context.fpcw);
+    OFFSET_T(TCB_MXCSR, shim_tcb_t, context.mxcsr);
     OFFSET(SHIM_REGS_RSP, shim_regs, rsp);
     OFFSET(SHIM_REGS_R15, shim_regs, r15);
     OFFSET(SHIM_REGS_RIP, shim_regs, rip);

LibOS/shim/src/shim_context-x86_64.c

@@ -149,7 +149,9 @@ noreturn void restore_child_context_after_clone(struct shim_context* context) {
     shim_tcb_t* tcb = shim_get_tcb();
     __enable_preempt(tcb);
 
-    __asm__ volatile("movq %0, %%rsp\r\n"
+    __asm__ volatile("fldcw (%0)\r\n" /* restore FP (fpcw) and SSE/AVX/... (mxcsr) control words */
+                     "ldmxcsr (%1)\r\n"
+                     "movq %2, %%rsp\r\n"
                      "addq $2 * 8, %%rsp\r\n"    /* skip orig_rax and rsp */
                      "popq %%r15\r\n"
                      "popq %%r14\r\n"
@@ -169,7 +171,7 @@ noreturn void restore_child_context_after_clone(struct shim_context* context) {
                      "movq "XSTRINGIFY(SHIM_REGS_RSP)" - "XSTRINGIFY(SHIM_REGS_RIP)"(%%rsp), %%rsp\r\n"
                      "movq $0, %%rax\r\n"
                      "jmp *-"XSTRINGIFY(RED_ZONE_SIZE)"-8(%%rsp)\r\n"
-                     :: "g"(&regs) : "memory");
+                     :: "g"(&context->fpcw), "g"(&context->mxcsr), "g"(&regs) : "memory");
 
     __builtin_unreachable();
 }

LibOS/shim/src/sys/shim_clone.c

@@ -30,7 +30,8 @@ struct shim_clone_args {
     void* stack;
     unsigned long fs_base;
     struct shim_regs regs;
-    void* xstate_extended;
+    uint16_t fpcw;
+    uint32_t mxcsr;
 };
 
 /*
@@ -76,12 +77,6 @@ static int clone_implementation_wrapper(struct shim_clone_args* arg) {
 
     debug("set fs_base to 0x%lx\n", tcb->context.fs_base);
 
-    /* FIXME: The below XSAVE area restore is not really correct but rather a dummy and will be
-     * fixed later. Now it restores the extended state from within LibOS rather than the app. In
-     * reality, XSAVE area should be part of shim_regs, and XRSTOR should happen during
-     * restore_context(). */
-    shim_xstate_restore(arg->xstate_extended);
-
     if (my_thread->set_child_tid) {
         *(my_thread->set_child_tid) = my_thread->tid;
         my_thread->set_child_tid = NULL;
@@ -101,6 +96,10 @@ static int clone_implementation_wrapper(struct shim_clone_args* arg) {
 
     add_thread(my_thread);
 
+    /* New thread inherits FP (fpcw) and SSE/AVX/... (mxcsr) control words of parent thread. */
+    tcb->context.fpcw  = arg->fpcw;
+    tcb->context.mxcsr = arg->mxcsr;
+
     /* Copy regs before we let the parent release them. */
     struct shim_regs regs = arg->regs;
 
@@ -164,6 +163,10 @@ static long do_clone_new_vm(unsigned long flags, struct shim_thread* thread, uns
      * shared with the parent. */
     shim_tcb.context.regs = self->shim_tcb->context.regs;
 
+    /* new process inherits FP (fpcw) and SSE/AVX/... (mxcsr) control words */
+    shim_tcb.context.fpcw  = self->shim_tcb->context.fpcw;
+    shim_tcb.context.mxcsr = self->shim_tcb->context.mxcsr;
+
     if (flags & CLONE_SETTLS) {
         shim_tcb.context.fs_base = fs_base;
     } else {
@@ -416,16 +419,8 @@ long shim_do_clone(unsigned long flags, unsigned long user_stack_addr, int* pare
     new_args.stack   = (void*)(user_stack_addr ?: shim_context_get_sp(&self->shim_tcb->context));
     new_args.fs_base = fs_base;
     new_args.regs    = *self->shim_tcb->context.regs;
-
-    /* FIXME: The below XSAVE area save is not really correct but rather a dummy and will be fixed
-     * later. Now it saves the extended state from within LibOS rather than the app. In reality,
-     * XSAVE area should be part of shim_regs, and XSAVE should happen during syscalldb().
-     * Also note that we require up to 4KB of stack space for XSAVE -- this is wrong for e.g. Go
-     * because its goroutines start with 2KB stack size; but we'll remove XSAVE here anyway. */
-    size_t xstate_extended_size = g_shim_xsave_size + SHIM_FP_XSTATE_MAGIC2_SIZE;
-    new_args.xstate_extended    = ALIGN_DOWN_PTR(new_args.stack - xstate_extended_size,
-                                                 SHIM_XSTATE_ALIGN);
-    shim_xstate_save(new_args.xstate_extended);
+    new_args.fpcw    = self->shim_tcb->context.fpcw;
+    new_args.mxcsr   = self->shim_tcb->context.mxcsr;
 
     // Invoke DkThreadCreate to spawn off a child process using the actual
     // "clone" system call. DkThreadCreate allocates a stack for the child

LibOS/shim/src/syscallas-x86_64.S

@@ -2,9 +2,20 @@
 /* Copyright (C) 2014 Stony Brook University */
 
 /*
- * syscallas.S
+ * This file contains the entry point of system call table in library OS (the function syscalldb()
+ * and its wrapper syscall_wrapper() for cases of redirection of raw SYSCALL instructions).
  *
- * This file contains the entry point of system call table in library OS.
+ * The below entry point implementation first saves the CPU context of the current application
+ * thread on the thread's stack, then calls the corresponding LibOS syscall-emulation function, and
+ * then restores the context and passes control back to the application. The context consists of
+ * GPRs, FP control word (fpcw) and the SSE/AVX/... control word (mxcsr).
+ *
+ * Note that LibOS may clobber all FP/SSE/AVX/... (extended) state except the control words. We rely
+ * on the fact that applications do *not* assume that this extended state is preserved across system
+ * calls. Indeed, the extended state (bar control words) is explicitly described as *not* preserved
+ * by the System V ABI, and though syscall ABI is not the same as System V ABI, we assume that no
+ * sane application issues syscalls in a non-System-V compliant manner. See System V ABI docs
+ * (https://uclibc.org/docs/psABI-x86_64.pdf), "Register Usage" for more information.
  */
 
 #include "asm-offsets.h"
@@ -62,7 +73,10 @@ syscalldb:
         cmp $0, %rbx
         je isundef
 
+        # set pointer to shim_regs and save FP Control Word & MXCSR into current thread's TCB
         movq %rbp, %gs:(SHIM_TCB_OFFSET + TCB_REGS)
+        fnstcw %gs:(SHIM_TCB_OFFSET + TCB_FPCW)
+        stmxcsr %gs:(SHIM_TCB_OFFSET + TCB_MXCSR)
 
         /* Translating x86_64 kernel calling convention to user-space
          * calling convention */
@@ -70,7 +84,10 @@ syscalldb:
         andq $~0xF, %rsp  # Required by System V AMD64 ABI.
         call *%rbx
 
+        # invalidate pointer to shim_regs and restore FP Control Word & MXCSR from TCB
         movq $0, %gs:(SHIM_TCB_OFFSET + TCB_REGS)
+        fldcw %gs:(SHIM_TCB_OFFSET + TCB_FPCW)
+        ldmxcsr %gs:(SHIM_TCB_OFFSET + TCB_MXCSR)
 
 ret:
         movq %rbp, %rsp