From 706065b0e23fa69f9dd590b4c7e97114aaaf9449 Mon Sep 17 00:00:00 2001
From: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Date: Fri, 29 Mar 2024 10:42:01 -0400
Subject: [PATCH] OpcodeDispatcher: accelerate cmpxchg with flagm

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
---
 .../Interface/Core/JIT/Arm64/ALUOps.cpp       | 26 +++++++++++++++++++
 .../Interface/Core/OpcodeDispatcher.cpp       | 12 ++-------
 FEXCore/Source/Interface/IR/IR.json           |  5 ++++
 .../RedundantFlagCalculationElimination.cpp   |  6 +++++
 4 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/FEXCore/Source/Interface/Core/JIT/Arm64/ALUOps.cpp b/FEXCore/Source/Interface/Core/JIT/Arm64/ALUOps.cpp
index 0055812e87..129d05a4f4 100644
--- a/FEXCore/Source/Interface/Core/JIT/Arm64/ALUOps.cpp
+++ b/FEXCore/Source/Interface/Core/JIT/Arm64/ALUOps.cpp
@@ -5,6 +5,7 @@ tags: backend|arm64
 $end_info$
 */
 
+#include "FEXCore/IR/IR.h"
 #include "Interface/Context/Context.h"
 #include "Interface/Core/ArchHelpers/CodeEmitter/Emitter.h"
 #include "Interface/Core/ArchHelpers/CodeEmitter/Registers.h"
@@ -299,6 +300,31 @@ DEF_OP(SubNZCV) {
   }
 }
 
+DEF_OP(CmpPairZ) {
+  auto Op = IROp->C<IR::IROp_CmpPairZ>();
+  const uint8_t OpSize = IROp->Size;
+
+  const auto EmitSize = OpSize == IR::i64Bit ? ARMEmitter::Size::i64Bit : ARMEmitter::Size::i32Bit;
+
+  // Save NZCV
+  mrs(TMP1, ARMEmitter::SystemRegister::NZCV);
+
+  // Compare, setting Z and clobbering NzCV
+  const auto Src1 = GetRegPair(Op->Src1.ID());
+  const auto Src2 = GetRegPair(Op->Src2.ID());
+  cmp(EmitSize, Src1.first, Src2.first);
+  ccmp(EmitSize, Src1.second, Src2.second, ARMEmitter::StatusFlags::None, ARMEmitter::Condition::CC_EQ);
+
+  // Restore NzCV
+  if (CTX->HostFeatures.SupportsFlagM) {
+    rmif(TMP1, 0, 0xb /* NzCV */);
+  } else {
+    cset(ARMEmitter::Size::i32Bit, TMP2, ARMEmitter::Condition::CC_EQ);
+    bfi(ARMEmitter::Size::i32Bit, TMP1, TMP2, 30 /* lsb: Z */, 1);
+    msr(ARMEmitter::SystemRegister::NZCV, TMP1);
+  }
+}
+
 DEF_OP(CarryInvert) {
   LOGMAN_THROW_A_FMT(CTX->HostFeatures.SupportsFlagM, "Unsupported flagm op");
   cfinv();
diff --git a/FEXCore/Source/Interface/Core/OpcodeDispatcher.cpp b/FEXCore/Source/Interface/Core/OpcodeDispatcher.cpp
index 12091359d6..2214f4c923 100644
--- a/FEXCore/Source/Interface/Core/OpcodeDispatcher.cpp
+++ b/FEXCore/Source/Interface/Core/OpcodeDispatcher.cpp
@@ -4258,16 +4258,8 @@ void OpDispatchBuilder::CMPXCHGPairOp(OpcodeArgs) {
   OrderedNode *Result_Lower = _ExtractElementPair(IR::SizeToOpSize(Size), CASResult, 0);
   OrderedNode *Result_Upper = _ExtractElementPair(IR::SizeToOpSize(Size), CASResult, 1);
 
-  // Set ZF if memory result was expected
-  auto OneConst = _Constant(1);
-  auto ZeroConst = _Constant(0);
-
-  OrderedNode *ZFResult = _Select(FEXCore::IR::COND_EQ,
-    CASResult, Expected,
-    OneConst, ZeroConst);
-
-  // Set ZF
-  SetRFLAG<FEXCore::X86State::RFLAG_ZF_RAW_LOC>(ZFResult);
+  HandleNZCV_RMW();
+  _CmpPairZ(IR::SizeToOpSize(Size), CASResult, Expected);
   CalculateDeferredFlags();
 
   auto UpdateIfNotZF = [this](auto Reg, auto Value) {
diff --git a/FEXCore/Source/Interface/IR/IR.json b/FEXCore/Source/Interface/IR/IR.json
index f3e8f640fe..abf9b8586d 100644
--- a/FEXCore/Source/Interface/IR/IR.json
+++ b/FEXCore/Source/Interface/IR/IR.json
@@ -1094,6 +1094,11 @@
           "Size == FEXCore::IR::OpSize::i32Bit || Size == FEXCore::IR::OpSize::i64Bit"
         ]
       },
+      "CmpPairZ OpSize:#Size, GPRPair:$Src1, GPRPair:$Src2": {
+        "Desc": ["Compares register pairs and sets Z accordingly, preserving N/Z/V.",
+                 "This accelerates cmpxchg."],
+        "HasSideEffects": true
+      },
       "SubNZCV OpSize:#Size, GPR:$Src1, GPR:$Src2": {
         "Desc": ["Set NZCV for the difference of two GPRs. ",
                  "Carry flag uses arm64 definition, inverted x86.",
diff --git a/FEXCore/Source/Interface/IR/Passes/RedundantFlagCalculationElimination.cpp b/FEXCore/Source/Interface/IR/Passes/RedundantFlagCalculationElimination.cpp
index 4a0fde2314..2df6270d8d 100644
--- a/FEXCore/Source/Interface/IR/Passes/RedundantFlagCalculationElimination.cpp
+++ b/FEXCore/Source/Interface/IR/Passes/RedundantFlagCalculationElimination.cpp
@@ -176,6 +176,12 @@ DeadFlagCalculationEliminination::Classify(IROp_Header *IROp)
         .CanEliminate = true,
       };
 
+    case OP_CMPPAIRZ:
+      return {
+        .Write = FLAG_Z,
+        .CanEliminate = true,
+      };
+
     case OP_CARRYINVERT:
       return {
         .Read = FLAG_C,