Support the umwait instruction.

Author: khuey

CMakeLists.txt

@@ -1462,6 +1462,7 @@ set(BASIC_TESTS
   tun
   two_signals_with_mask
   ulimit_low
+  x86/umwait
   uname
   unexpected_exit
   unexpected_exit_execve

src/record_signal.cc

@@ -87,6 +87,8 @@ static bool try_handle_trapped_instruction(RecordTask* t, siginfo_t* si) {
     case SpecialInstOpcode::X86_RDTSCP:
     case SpecialInstOpcode::X86_TPAUSE:
     case SpecialInstOpcode::X86_TPAUSE_REX:
+    case SpecialInstOpcode::X86_UMWAIT:
+    case SpecialInstOpcode::X86_UMWAIT_REX:
       if (t->tsc_mode == PR_TSC_SIGSEGV) {
         return false;
       }
@@ -139,7 +141,9 @@ static bool try_handle_trapped_instruction(RecordTask* t, siginfo_t* si) {
                        cpuid_data.edx);
     LOG(debug) << " trapped for cpuid: " << HEX(eax) << ":" << HEX(ecx);
   } else if (special_instruction.opcode == SpecialInstOpcode::X86_TPAUSE ||
-             special_instruction.opcode == SpecialInstOpcode::X86_TPAUSE_REX) {
+             special_instruction.opcode == SpecialInstOpcode::X86_TPAUSE_REX ||
+             special_instruction.opcode == SpecialInstOpcode::X86_UMWAIT ||
+             special_instruction.opcode == SpecialInstOpcode::X86_UMWAIT_REX) {
     // The Intel SDM states:
     //
     // > Prior to executing the TPAUSE instruction, an operating system may
@@ -155,16 +159,16 @@ static bool try_handle_trapped_instruction(RecordTask* t, siginfo_t* si) {
     // > - IA32_UMWAIT_CONTROL[0] — C0.2 is not allowed by the OS. Value of “1”
     // >   means all C0.2 requests revert to C0.1.
     // >
-    // > If the processor that executed a TPAUSE instruction wakes due to the
-    // > expiration of the operating system time-limit, the instructions sets
+    // > If the processor that executed a TPAUSE/UMWAIT instruction wakes due to
+    // > the expiration of the operating system time-limit, the instructions sets
     // > RFLAGS.CF; otherwise, that flag is cleared.
     //
     // As the "operating system", we choose a time-limit of zero. The point of
-    // TPAUSE is to make spin-wait loops more efficient without yielding to the
-    // operating system scheduler and triggering a full context switch. But
-    // because we serialize all execution to a single physical core, no progress
-    // will ever be made without yielding, and of course we've already paid the
-    // cost of a context switch and then some.
+    // TPAUSE/UMWAIT is to make spin-wait loops more efficient without yielding
+    // to the operating system scheduler and triggering a full context switch.
+    // But because we serialize all execution to a single physical core, no
+    // progress will ever be made without yielding, and of course we've already
+    // paid the cost of a context switch and then some.
     r.set_flags((r.flags() & ~X86_ALL_ARITH_FLAGS) | X86_CF_FLAG);
   }
 

src/test/x86/umwait.c

@@ -0,0 +1,66 @@
+/* -*- Mode: C; tab-width: 8; c-basic-offset: 2; indent-tabs-mode: nil; -*- */
+
+#include "util.h"
+
+enum cpuid_requests {
+  CPUID_GETEXTENDEDFEATURES = 0x07,
+};
+
+static void cpuid(int code, int subrequest, unsigned int* a, unsigned int* c,
+                  unsigned int* d) {
+  asm volatile("cpuid"
+               : "=a"(*a), "=c"(*c), "=d"(*d)
+               : "a"(code), "c"(subrequest)
+               : "ebx");
+}
+
+int main(void) {
+  uint32_t eax, ecx, edx, state = 0;
+#ifdef __x86_64__
+  register uint32_t r15d __asm__("r15") = state;
+#endif
+  uint64_t tsc;
+
+  cpuid(CPUID_GETEXTENDEDFEATURES, 0, &eax, &ecx, &edx);
+  if (!(ecx & (1 << 5))) {
+    // WAITPKG not present.
+    atomic_puts("umwait not supported on this system, skipping");
+    atomic_puts("EXIT-SUCCESS");
+    return 0;
+  }
+
+  tsc = __rdtsc() + 1000;
+  eax = tsc;
+  edx = tsc >> 32;
+
+  asm volatile(
+    ".byte 0xf2, 0x0f, 0xae, 0xf7;"  // umwait %edi
+    "lahf;"
+    : "+a"(eax), "+d"(edx)
+    : "D"(state)
+    : "cc", "memory");
+
+  // All the arith flags should be cleared except CF, which should be set.
+  test_assert((uint8_t)(eax >> 8) == 0x3);
+
+#ifdef __x86_64__
+  tsc = __rdtsc() + 1000;
+  eax = tsc;
+  edx = tsc >> 32;
+
+  // Test umwait with a REX prefix since the instruction length
+  // is different.
+  asm volatile(
+    ".byte 0xf2, 0x41, 0x0f, 0xae, 0xf7;"  // umwait %r15d
+    "lahf;"
+    : "+a"(eax), "+d"(edx)
+    : "r"(r15d)
+    : "cc", "memory");
+
+  // All the arith flags should be cleared except CF, which should be set.
+  test_assert((uint8_t)(eax >> 8) == 0x3);
+#endif
+
+  atomic_puts("EXIT-SUCCESS");
+  return 0;
+}

src/util.cc

@@ -1927,6 +1927,8 @@ const uint8_t rdtsc_insn[2] = { 0x0f, 0x31 };
 static const uint8_t rdtscp_insn[] = { 0x0f, 0x01, 0xf9 };
 static const uint8_t tpause_insn[] = { 0x66, 0x0f, 0xae, 0xf0 };
 static const uint8_t tpause_rex_insn[] = { 0x66, 0x41, 0x0f, 0xae, 0xf0 };
+static const uint8_t umwait_insn[] = { 0xf2, 0x0f, 0xae, 0xf0 };
+static const uint8_t umwait_rex_insn[] = { 0xf2, 0x41, 0x0f, 0xae, 0xf0 };
 static const uint8_t cpuid_insn[] = { 0x0f, 0xa2 };
 static const uint8_t int3_insn[] = { 0xcc };
 static const uint8_t pushf_insn[] = { 0x9c };
@@ -1973,13 +1975,19 @@ SpecialInst special_instruction_at(Task* t, remote_code_ptr ip) {
       if (!memcmp(insn, tpause_rex_insn, sizeof(tpause_rex_insn))) {
         return {SpecialInstOpcode::X86_TPAUSE_REX};
       }
+      if (!memcmp(insn, umwait_rex_insn, sizeof(umwait_rex_insn))) {
+        return {SpecialInstOpcode::X86_UMWAIT_REX};
+      }
     }
     if (len >= sizeof(tpause_insn)) {
       // Mask off the register selector.
       insn[3] &= ~0x07;
       if (!memcmp(insn, tpause_insn, sizeof(tpause_insn))) {
         return {SpecialInstOpcode::X86_TPAUSE};
       }
+      if (!memcmp(insn, umwait_insn, sizeof(umwait_insn))) {
+        return {SpecialInstOpcode::X86_UMWAIT};
+      }
     }
   } else if (t->arch() == aarch64) {
     uint8_t insn[4];
@@ -2020,6 +2028,10 @@ size_t special_instruction_len(SpecialInstOpcode insn) {
     return sizeof(tpause_insn);
   } else if (insn == SpecialInstOpcode::X86_TPAUSE_REX) {
     return sizeof(tpause_rex_insn);
+  } else if (insn == SpecialInstOpcode::X86_UMWAIT) {
+    return sizeof(umwait_insn);
+  } else if (insn == SpecialInstOpcode::X86_UMWAIT_REX) {
+    return sizeof(umwait_rex_insn);
   } else if (insn == SpecialInstOpcode::ARM_MRS_CNTFRQ_EL0 ||
              insn == SpecialInstOpcode::ARM_MRS_CNTVCT_EL0 ||
              insn == SpecialInstOpcode::ARM_MRS_CNTVCTSS_EL0) {

src/util.h

@@ -495,6 +495,8 @@ enum class SpecialInstOpcode {
   X86_PUSHF16,
   X86_TPAUSE,
   X86_TPAUSE_REX,
+  X86_UMWAIT,
+  X86_UMWAIT_REX,
 };
 
 struct SpecialInst {