FEX: Implement support for CPUID faulting

FEXCore/Source/Interface/Core/ArchHelpers/Arm64Emitter.h

@@ -235,6 +235,8 @@ protected:
 
   void Align16B();
 
+  ARMEmitter::Register GetX86GPRMapping(FEXCore::X86State::X86Reg Reg);
+
 #ifdef VIXL_SIMULATOR
   // Generates a vixl simulator runtime call.
   //

FEXCore/Source/Interface/Core/Dispatcher/Dispatcher.cpp

@@ -417,6 +417,7 @@ void Dispatcher::EmitDispatcher() {
     brk(0);
   }
 
+  ARMEmitter::BackwardLabel GuestSIGSEGVLabelPostSpill;
   {
     // Guest Overflow handler
     // Needs to be distinct from the SignalHandlerReturnAddress
@@ -424,6 +425,8 @@ void Dispatcher::EmitDispatcher() {
 
     SpillStaticRegs(TMP1);
 
+    (void)Bind(&GuestSIGSEGVLabelPostSpill);
+
     // hlt/udf = SIGILL
     // brk = SIGTRAP
     // ??? = SIGSEGV
@@ -534,6 +537,28 @@ void Dispatcher::EmitDispatcher() {
     (void)b(&LoopTop);
   }
 
+  auto EmitL1Lookup = [&](int rip_offset = 0) {
+    // L1 lookup after syscall. Mirrors the implementation in `ExitFunction`.
+    // Most syscall instructions just return to the following RIP, so it is highly likely to be in the L1 in tight loops.
+    ldr(TMP3, STATE_PTR(CpuStateFrame, State.rip));
+    ldp<ARMEmitter::IndexType::OFFSET>(TMP1, TMP2, STATE, offsetof(FEXCore::Core::CpuStateFrame, State.L1Pointer));
+
+    if (rip_offset) {
+      add(ARMEmitter::Size::i64Bit, TMP3, TMP3, rip_offset);
+    }
+
+    // Calculate (tmp1 + ((ripreg & L1_ENTRIES_MASK) << 4)) for the address
+    // L1Mask is pre-shifted.
+    and_(ARMEmitter::Size::i64Bit, TMP2, TMP2, TMP3, ARMEmitter::ShiftType::LSL, FEXCore::ilog2(sizeof(LookupCache::LookupCacheEntry)));
+    add(TMP1, TMP1, TMP2);
+
+    ldp<ARMEmitter::IndexType::OFFSET>(TMP2, TMP1, TMP1, 0);
+
+    sub(TMP1, TMP1, TMP3);
+    (void)cbnz(ARMEmitter::Size::i64Bit, TMP1, &LoopTop);
+    br(TMP2);
+  };
+
   {
     // All dynamic and static registers are spilled coming in to this handler.
     // It's also the end of block and RIP might have changed, so we jump directly to the top of the loop.
@@ -566,23 +591,90 @@ void Dispatcher::EmitDispatcher() {
     // We can safely claim we are no longer in a syscall
     str(ARMEmitter::XReg::zr, STATE, offsetof(FEXCore::Core::CpuStateFrame, InSyscallInfo));
 
-    {
-      // L1 lookup after syscall. Mirrors the implementation in `ExitFunction`.
-      // Most syscall instructions just return to the following RIP, so it is highly likely to be in the L1 in tight loops.
-      ldr(TMP3, STATE_PTR(CpuStateFrame, State.rip));
-      ldp<ARMEmitter::IndexType::OFFSET>(TMP1, TMP2, STATE, offsetof(FEXCore::Core::CpuStateFrame, State.L1Pointer));
+    EmitL1Lookup();
+  }
 
-      // Calculate (tmp1 + ((ripreg & L1_ENTRIES_MASK) << 4)) for the address
-      // L1Mask is pre-shifted.
-      and_(ARMEmitter::Size::i64Bit, TMP2, TMP2, TMP3, ARMEmitter::ShiftType::LSL, FEXCore::ilog2(sizeof(LookupCache::LookupCacheEntry)));
-      add(TMP1, TMP1, TMP2);
+  {
+    // All dynamic and static registers are spilled coming in to this handler.
+    // It's also the end of block and RIP might have changed, so we jump directly to the top of the loop.
+    ThreadDispatchCPUIDHandler = GetCursorAddress<uint64_t>();
 
-      ldp<ARMEmitter::IndexType::OFFSET>(TMP2, TMP1, TMP1, 0);
+    // CPUID behaves like an ISB.
+    isb();
 
-      sub(TMP1, TMP1, TMP3);
-      (void)cbnz(ARMEmitter::Size::i64Bit, TMP1, &LoopTop);
-      br(TMP2);
+    auto RAXReg = GetX86GPRMapping(FEXCore::X86State::REG_RAX);
+    auto RBXReg = GetX86GPRMapping(FEXCore::X86State::REG_RBX);
+    auto RCXReg = GetX86GPRMapping(FEXCore::X86State::REG_RCX);
+    auto RDXReg = GetX86GPRMapping(FEXCore::X86State::REG_RDX);
+
+    ARMEmitter::ForwardLabel SetupCPUIDFault {};
+    ldr(TMP2, STATE, offsetof(FEXCore::Core::CPUState, CoreFlags));
+    (void)tbz(TMP2, FEXCore::ilog2(FEXCore::ToUnderlying(FEXCore::Core::CoreFlagsEnum::CPUID_ENABLED)), &SetupCPUIDFault);
+
+    // TMP1 passes the instruction size. Usually only 2-bytes but prefixes can make it 15.
+    ldr(TMP2, STATE_PTR(CpuStateFrame, State.rip));
+    add(TMP2, TMP2, TMP1);
+    str(TMP2, STATE_PTR(CpuStateFrame, State.rip));
+
+    mov(ARMEmitter::XReg::x1, RAXReg.X());
+    mov(ARMEmitter::XReg::x2, RCXReg.X());
+
+    // x0 = CPUID Handler
+    // x1 = CPUID Function (RAX)
+    // x2 = CPUID Leaf (RCX)
+    ldr(ARMEmitter::XReg::x0, STATE, offsetof(FEXCore::Core::CpuStateFrame, Pointers.CPUIDObj));
+    ldr(ARMEmitter::XReg::x3, STATE, offsetof(FEXCore::Core::CpuStateFrame, Pointers.CPUIDFunction));
+
+    if (!CTX->Config.DisableVixlIndirectCalls) [[unlikely]] {
+      GenerateIndirectRuntimeCall<__uint128_t, void*, uint64_t, uint64_t>(ARMEmitter::Reg::r3);
+    } else {
+      blr(ARMEmitter::Reg::r3);
     }
+
+    // Results are in x0, x1
+    auto Res0 = ARMEmitter::Reg::r0;
+    auto Res1 = ARMEmitter::Reg::r1;
+
+    if (!TMP_ABIARGS) {
+      // SRA Overlaps, need to move.
+      mov(ARMEmitter::Size::i64Bit, TMP1, Res0);
+      mov(ARMEmitter::Size::i64Bit, TMP2, Res1);
+      Res0 = TMP1;
+      Res1 = TMP2;
+    }
+
+    FillStaticRegs();
+
+    // Results need to be moved to 32-bit EAX, ECX, EBX, EDX.
+    mov(ARMEmitter::Size::i32Bit, RAXReg, Res0);
+    mov(ARMEmitter::Size::i32Bit, RCXReg, Res1);
+    ubfx(ARMEmitter::Size::i64Bit, RBXReg, Res0, 32, 32);
+    ubfx(ARMEmitter::Size::i64Bit, RDXReg, Res1, 32, 32);
+
+    EmitL1Lookup();
+
+    // CPUID fault handling
+    (void)Bind(&SetupCPUIDFault);
+
+    constexpr Core::CpuStateFrame::SynchronousFaultDataStruct State = {
+      .FaultToTopAndGeneratedException = 1,
+      .Signal = Core::FAULT_SIGSEGV,
+      .TrapNo = X86State::X86_TRAPNO_GP,
+      .si_code = 0x80,
+      .err_code = 0,
+      .FaultAddress = 0, // Comes from register.
+    };
+
+    uint64_t Constant {};
+    memcpy(&Constant, &State, sizeof(uint64_t));
+
+    LoadConstant(ARMEmitter::Size::i64Bit, TMP1, static_cast<uint64_t>(Constant));
+
+    str(TMP1, STATE, offsetof(FEXCore::Core::CpuStateFrame, SynchronousFaultData));
+    ldr(TMP2, STATE_PTR(CpuStateFrame, State.rip));
+    str(TMP2, STATE, offsetof(FEXCore::Core::CpuStateFrame, SynchronousFaultData) + 8);
+
+    (void)b(&GuestSIGSEGVLabelPostSpill);
   }
 
   {
@@ -2683,6 +2775,7 @@ void Dispatcher::InitThreadPointers(FEXCore::Core::InternalThreadState* Thread)
     Ptrs.ThreadStopHandlerSpillSRA = ThreadStopHandlerAddressSpillSRA;
     Ptrs.ThreadPauseHandlerSpillSRA = ThreadPauseHandlerAddressSpillSRA;
     Ptrs.ThreadDispatchSyscallHandler = ThreadDispatchSyscallHandler;
+    Ptrs.ThreadDispatchCPUIDHandler = ThreadDispatchCPUIDHandler;
     Ptrs.ThreadDispatchRemoveCodeEntry = ThreadDispatchRemoveCodeEntry;
     Ptrs.GuestSignal_SIGILL = GuestSignal_SIGILL;
     Ptrs.GuestSignal_SIGTRAP = GuestSignal_SIGTRAP;

FEXCore/Source/Interface/Core/Dispatcher/Dispatcher.h

@@ -82,6 +82,7 @@ private:
   uint64_t ThreadPauseHandlerAddress {};
   uint64_t ThreadPauseHandlerAddressSpillSRA {};
   uint64_t ThreadDispatchSyscallHandler {};
+  uint64_t ThreadDispatchCPUIDHandler {};
   uint64_t ThreadDispatchRemoveCodeEntry {};
   uint64_t ExitFunctionLinkerAddress {};
   uint64_t SignalHandlerReturnAddress {};

FEXCore/Source/Interface/Core/JIT/BranchOps.cpp

@@ -395,45 +395,12 @@ DEF_OP(ThreadRemoveCodeEntry) {
 DEF_OP(CPUID) {
   auto Op = IROp->C<IR::IROp_CPUID>();
 
-  isb();
-  mov(ARMEmitter::Size::i64Bit, TMP2, GetReg(Op->Function));
-  mov(ARMEmitter::Size::i64Bit, TMP3, GetReg(Op->Leaf));
+  SpillStaticRegs(TMP1);
 
-  PushDynamicRegs(TMP4);
-  SpillStaticRegs(TMP4);
-
-  // x0 = CPUID Handler
-  // x1 = CPUID Function
-  // x2 = CPUID Leaf
-  ldr(ARMEmitter::XReg::x0, STATE, offsetof(FEXCore::Core::CpuStateFrame, Pointers.CPUIDObj));
-  ldr(ARMEmitter::XReg::x3, STATE, offsetof(FEXCore::Core::CpuStateFrame, Pointers.CPUIDFunction));
-
-  if (!TMP_ABIARGS) {
-    mov(ARMEmitter::Size::i64Bit, ARMEmitter::Reg::r1, TMP2);
-    mov(ARMEmitter::Size::i64Bit, ARMEmitter::Reg::r2, TMP3);
-  }
-
-  if (!CTX->Config.DisableVixlIndirectCalls) [[unlikely]] {
-    GenerateIndirectRuntimeCall<__uint128_t, void*, uint64_t, uint64_t>(ARMEmitter::Reg::r3);
-  } else {
-    blr(ARMEmitter::Reg::r3);
-  }
-
-  if (!TMP_ABIARGS) {
-    mov(ARMEmitter::Size::i64Bit, TMP1, ARMEmitter::Reg::r0);
-    mov(ARMEmitter::Size::i64Bit, TMP2, ARMEmitter::Reg::r1);
-  }
-
-  FillStaticRegs();
-
-  PopDynamicRegs();
-
-  // Results are in x0, x1
-  // Results want to be 4xi32 scalars
-  mov(ARMEmitter::Size::i32Bit, GetReg(Op->OutEAX), TMP1);
-  mov(ARMEmitter::Size::i32Bit, GetReg(Op->OutECX), TMP2);
-  ubfx(ARMEmitter::Size::i64Bit, GetReg(Op->OutEBX), TMP1, 32, 32);
-  ubfx(ARMEmitter::Size::i64Bit, GetReg(Op->OutEDX), TMP2, 32, 32);
+  LoadConstant(ARMEmitter::Size::i32Bit, TMP1, Op->InstSize);
+  // Jump to the syscall dispatch handler. We won't return after this.
+  ldr(TMP2, STATE, offsetof(FEXCore::Core::CpuStateFrame, Pointers.ThreadDispatchCPUIDHandler));
+  br(TMP2);
 }
 
 DEF_OP(XGetBV) {

FEXCore/Source/Interface/Core/OpcodeDispatcher.cpp

@@ -1370,20 +1370,17 @@ void OpDispatchBuilder::MOVOffsetOp(OpcodeArgs) {
 }
 
 void OpDispatchBuilder::CPUIDOp(OpcodeArgs) {
-  Ref Src = LoadGPRRegister(X86State::REG_RAX);
-  Ref Leaf = LoadGPRRegister(X86State::REG_RCX);
+  // Calculate flags early.
+  CalculateDeferredFlags();
 
-  Ref RAX = _AllocateGPR(false);
-  Ref RBX = _AllocateGPR(false);
-  Ref RCX = _AllocateGPR(false);
-  Ref RDX = _AllocateGPR(false);
+  // Similar to Syscall. Store the RIP to the start of the instruction.
+  // CPUID operation will modify the RIP itself.
+  const auto GPRSize = GetGPROpSize();
+  auto NewRIP = GetRelocatedPC(Op, -Op->InstSize);
+  _StoreContextGPR(GPRSize, NewRIP, offsetof(FEXCore::Core::CPUState, rip));
 
-  _CPUID(Src, Leaf, RAX, RBX, RCX, RDX);
-
-  StoreGPRRegister(X86State::REG_RAX, RAX);
-  StoreGPRRegister(X86State::REG_RBX, RBX);
-  StoreGPRRegister(X86State::REG_RCX, RCX);
-  StoreGPRRegister(X86State::REG_RDX, RDX);
+  FlushRegisterCache();
+  _CPUID(Op->InstSize);
 }
 
 uint32_t OpDispatchBuilder::GetConstantShift(X86Tables::DecodedOp Op, bool Is1Bit) {

FEXCore/Source/Interface/Core/X86Tables/SecondaryTables.cpp

@@ -204,7 +204,7 @@ constexpr std::array<X86InstInfo, MAX_SECOND_TABLE_SIZE> SecondBaseOps = []() co
 
     {0xA0, 1, X86InstInfo{"",        TYPE_ARCH_DISPATCHER, FLAGS_DEBUG_MEM_ACCESS | FLAGS_NO_OVERLAY, 0, { .Indirect = Secondary_ArchSelect_LUT[ENTRY_A0] }}},
     {0xA1, 1, X86InstInfo{"",        TYPE_ARCH_DISPATCHER, FLAGS_DEBUG_MEM_ACCESS | FLAGS_NO_OVERLAY, 0, { .Indirect = Secondary_ArchSelect_LUT[ENTRY_A1] }}},
-    {0xA2, 1, X86InstInfo{"CPUID",   TYPE_INST,     FLAGS_NO_OVERLAY,                                              0}},
+    {0xA2, 1, X86InstInfo{"CPUID",   TYPE_INST,     FLAGS_NO_OVERLAY | FLAGS_BLOCK_END,                                              0}},
     {0xA3, 1, X86InstInfo{"BT",      TYPE_INST,     FLAGS_DEBUG_MEM_ACCESS | FLAGS_MODRM | FLAGS_SF_MOD_DST | FLAGS_NO_OVERLAY | FLAGS_SUPPORTS_LOCK, 0}},
     {0xA4, 1, X86InstInfo{"SHLD",    TYPE_INST,     FLAGS_MODRM | FLAGS_SF_MOD_DST | FLAGS_NO_OVERLAY,                                              1}},
     {0xA5, 1, X86InstInfo{"SHLD",    TYPE_INST,     FLAGS_MODRM | FLAGS_SF_MOD_DST | FLAGS_NO_OVERLAY,                           0}},

FEXCore/Source/Interface/IR/IR.json

@@ -337,9 +337,8 @@
         "HasSideEffects": true
       },
 
-      "GPR:$EAX, GPR:$EBX, GPR:$ECX, GPR:$EDX = CPUID GPR:$Function, GPR:$Leaf": {
+      "CPUID i8:$InstSize": {
         "Desc": ["Calls in to the CPUID handler function to return emulated CPUID"],
-        "DestSize": "OpSize::i32Bit",
         "HasSideEffects": true
       },
       "GPR:$EAX, GPR:$EDX = XGetBV GPR:$Function": {

FEXCore/Source/Interface/IR/IREmitter.cpp

@@ -21,6 +21,7 @@ static bool IsFragmentExit(FEXCore::IR::IROps Op) {
   switch (Op) {
   case OP_THREADREMOVECODEENTRY:
   case OP_SYSCALL:
+  case OP_CPUID:
   case OP_EXITFUNCTION:
   case OP_BREAK: return true;
   default: return false;

FEXCore/Source/Interface/IR/Passes/RegisterAllocationPass.cpp

@@ -527,27 +527,6 @@ bool ConstrainedRAPass::TryPostRAMerge(Ref LastNode, Ref CodeNode, IROp_Header*
       IREmit->RemovePostRA(CodeNode);
       return false;
     }
-  } else if (IROp->Op == OP_CPUID && PhysicalRegister(IROp->Args[0]) == PhysicalRegister(LastNode) && LastOp->Op == OP_CONSTANT) {
-    // Try to constant fold. As a limitation of merging only 2 instructions, we
-    // can only handle constant functions, not constant leafs. This could be
-    // lifted if we generalized at a (significant) complexity cost.
-    uint64_t ConstantFunction = LastOp->C<IROp_Constant>()->Constant;
-    auto Op = IROp->CW<IR::IROp_CPUID>();
-
-    const auto SupportsConstant = CPUID->DoesFunctionReportConstantData(ConstantFunction);
-    if (SupportsConstant.SupportsConstantFunction == CPUIDEmu::SupportsConstant::CONSTANT &&
-        SupportsConstant.NeedsLeaf != CPUIDEmu::NeedsLeafConstant::NEEDSLEAFCONSTANT) {
-      const auto Result = CPUID->RunFunction(ConstantFunction, 0 /* leaf */);
-
-      IREmit->SetWriteCursorBefore(CodeNode);
-      IREmit->_Fence(IR::FenceType::Inst);
-      IREmit->_Constant(Result.eax).Node->Reg = PhysicalRegister(Op->OutEAX).Raw;
-      IREmit->_Constant(Result.ebx).Node->Reg = PhysicalRegister(Op->OutEBX).Raw;
-      IREmit->_Constant(Result.ecx).Node->Reg = PhysicalRegister(Op->OutECX).Raw;
-      IREmit->_Constant(Result.edx).Node->Reg = PhysicalRegister(Op->OutEDX).Raw;
-      IREmit->RemovePostRA(CodeNode);
-      return false;
-    }
   }
 
   // Merge moves that are immediately consumed.

FEXCore/include/FEXCore/Core/CoreState.h

@@ -71,6 +71,12 @@ static_assert(std::is_standard_layout_v<NonAtomicRefCounter<uint64_t>>, "Needs t
 static_assert(std::is_trivially_copyable_v<NonAtomicRefCounter<uint64_t>>, "needs to be trivially copyable");
 static_assert(sizeof(NonAtomicRefCounter<uint64_t>) == sizeof(uint64_t), "Needs to be correct size");
 
+// Flags about the core state that the JIT can query.
+enum class CoreFlagsEnum : uint64_t {
+  CPUID_ENABLED = (1ULL << 0),
+};
+FEX_DEF_NUM_OPS(CoreFlagsEnum)
+
 struct alignas(64) CPUState {
   // Allows more efficient handling of the register
   // file in the event AVX is not supported.
@@ -107,7 +113,9 @@ struct alignas(64) CPUState {
   uint64_t L1Pointer {};
   uint64_t L1Mask {};
   uint64_t callret_sp {};
-  uint64_t _pad1 {};
+
+  // CPUID is default enabled.
+  CoreFlagsEnum CoreFlags {CoreFlagsEnum::CPUID_ENABLED};
 
   // Cacheline: 1,2,3,4
   // The high 128-bits of AVX registers when not being emulated by SVE256.
@@ -406,6 +414,7 @@ struct JITPointers {
   uint64_t ThreadStopHandlerSpillSRA {};
   uint64_t ThreadPauseHandlerSpillSRA {};
   uint64_t ThreadDispatchSyscallHandler {};
+  uint64_t ThreadDispatchCPUIDHandler {};
   uint64_t ThreadDispatchRemoveCodeEntry {};
   uint64_t GuestSignal_SIGILL {};
   uint64_t GuestSignal_SIGTRAP {};

Source/Tools/LinuxEmulation/LinuxSyscalls/Syscalls/Thread.cpp

@@ -546,11 +546,16 @@ void RegisterThread(FEX::HLE::SyscallHandler* Handler) {
       Result = -EINVAL; // We don't support CET, return EINVAL
       break;
     case 0x1011: // ARCH_GET_CPUID
-      return 1;
-      break;
-    case 0x1012:      // ARCH_SET_CPUID
-      return -ENODEV; // Claim we don't support faulting on CPUID
-      break;
+      return (Frame->State.CoreFlags & FEXCore::Core::CoreFlagsEnum::CPUID_ENABLED) == FEXCore::Core::CoreFlagsEnum::CPUID_ENABLED ? 1 : 0;
+    case 0x1012: // ARCH_SET_CPUID
+      if (addr == 0) {
+        // Disable CPUID (thus enabling CPUID faulting)
+        Frame->State.CoreFlags = Frame->State.CoreFlags & ~FEXCore::Core::CoreFlagsEnum::CPUID_ENABLED;
+      } else {
+        // Enable CPUID (thus disabling CPUID faulting)
+        Frame->State.CoreFlags = Frame->State.CoreFlags | FEXCore::Core::CoreFlagsEnum::CPUID_ENABLED;
+      }
+      return 0;
     default:
       LogMan::Msg::EFmt("Unknown prctl: 0x{:x}", code);
       Result = -EINVAL;