Merge pull request #3107 from lioncash/pext

Arm64/ALUOps: Remove spills in PEXT

FEXCore/Source/Interface/Core/JIT/Arm64/ALUOps.cpp

@@ -654,6 +654,7 @@ DEF_OP(PDep) {
 DEF_OP(PExt) {
   auto Op = IROp->C<IR::IROp_PExt>();
   const auto OpSize = IROp->Size;
+  const auto OpSizeBitsM1 = (OpSize * 8) - 1;
 
   LOGMAN_THROW_AA_FMT(OpSize == 4 || OpSize == 8, "Unsupported {} size: {}", __func__, OpSize);
   const auto EmitSize = OpSize == 8 ? ARMEmitter::Size::i64Bit : ARMEmitter::Size::i32Bit;
@@ -662,11 +663,9 @@ DEF_OP(PExt) {
   const auto Mask = GetReg(Op->Mask.ID());
   const auto Dest = GetReg(Node);
 
-  const auto MaskReg    = TMP1;
-  const auto BitReg     = TMP2;
-  const auto SubMaskReg = TMP3;
-  const auto Offset     = TMP4;
-  const auto ZeroReg    = ARMEmitter::Reg::zr;
+  const auto MaskReg  = TMP1;
+  const auto BitReg   = TMP2;
+  const auto ValueReg = TMP3;
 
   ARMEmitter::ForwardLabel EarlyExit;
   ARMEmitter::BackwardLabel NextBit;
@@ -674,35 +673,22 @@ DEF_OP(PExt) {
 
   cbz(EmitSize, Mask, &EarlyExit);
   mov(EmitSize, MaskReg, Mask);
-  mov(EmitSize, Offset, ZeroReg);
-
-  // We sadly need to spill a reg for this for the time being
-  // TODO: Remove when scratch registers can be allocated
-  //       explicitly.
-  SpillStaticRegs(TMP2, false, 1U << Mask.Idx());
-  mov(EmitSize, Mask, ZeroReg);
+  mov(EmitSize, ValueReg, Input);
+  mov(EmitSize, Dest, ARMEmitter::Reg::zr);
 
   // Main loop
   Bind(&NextBit);
-  rbit(EmitSize, BitReg, MaskReg);
-  clz(EmitSize, BitReg, BitReg);
-  sub(EmitSize, SubMaskReg, MaskReg, 1);
-  ands(EmitSize, MaskReg.R(), SubMaskReg.R(), MaskReg.R());
-  lsrv(EmitSize, BitReg, Input, BitReg);
-  and_(EmitSize, BitReg, BitReg, 1);
-  lslv(EmitSize, BitReg, BitReg, Offset);
-  add(EmitSize, Offset, Offset, 1);
-  orr(EmitSize, Mask, BitReg, Mask);
-  b(ARMEmitter::Condition::CC_NE, &NextBit);
-  mov(EmitSize, Dest, Mask);
-  // Restore our mask register before leaving
-  // TODO: Also remove along with above TODO.
-  FillStaticRegs(false, 1U << Mask.Idx());
-  b(&Done);
+  cbz(EmitSize, MaskReg, &Done);
+  clz(EmitSize, BitReg, MaskReg);
+  lslv(EmitSize, ValueReg, ValueReg, BitReg);
+  lslv(EmitSize, MaskReg, MaskReg, BitReg);
+  extr(EmitSize, Dest, Dest, ValueReg, OpSizeBitsM1);
+  bfc(EmitSize, MaskReg, OpSizeBitsM1, 1);
+  b(&NextBit);
 
   // Early exit
   Bind(&EarlyExit);
-  mov(EmitSize, Dest, ZeroReg);
+  mov(EmitSize, Dest, ARMEmitter::Reg::zr);
 
   // All done with nothing to do.
   Bind(&Done);

unittests/ASM/VEX/pext.asm

@@ -6,7 +6,11 @@
       "RCX": "0x00012567",
       "RDX": "0x1234567812345678",
       "RSI": "0xFF00FF00FF00FF00",
-      "RDI": "0x12561256"
+      "RDI": "0x12561256",
+      "R8":  "0x1234567812345678",
+      "R10": "0x12345678",
+      "R11": "0x12345678",
+      "R12": "0x00005678"
   },
   "HostFeatures": ["BMI2"]
 }
@@ -17,9 +21,29 @@ mov eax, 0x12345678
 mov ebx, 0xFF00FFF0
 pext ecx, eax, ebx
 
+; 32-bit full mask
+mov r10d,  0x12345678
+mov r9d, 0xFFFFFFFF
+pext r10d, r10d, r9d
+
+; 32-bit half mask
+mov r12d, 0x12345678
+mov r9d, 0x0000FFFF
+pext r12d, r12d, r9d
+
 ; 64-bit
 mov rdx, 0x1234567812345678
 mov rsi, 0xFF00FF00FF00FF00
 pext rdi, rdx, rsi
 
+; 64-bit full mask
+mov r8, 0x1234567812345678
+mov r9, 0xFFFFFFFFFFFFFFFF
+pext r8, r8, r9
+
+; 64-bit half mask
+mov r11, 0x1234567812345678
+mov r9,  0x00000000FFFFFFFF
+pext r11, r11, r9
+
 hlt

unittests/InstructionCountCI/VEX_map2.json

@@ -3794,7 +3794,7 @@
       ]
     },
     "pext eax, ebx, ecx": {
-      "ExpectedInstructionCount": 19,
+      "ExpectedInstructionCount": 14,
       "Optimal": "No",
       "Comment": [
         "Map 2 0b10 0xf5 32-bit"
@@ -3802,50 +3802,38 @@
       "ExpectedArm64ASM": [
         "mov w20, w7",
         "mov w21, w5",
-        "cbz w21, #+0x40",
+        "cbz w21, #+0x2c",
         "mov w0, w21",
-        "mov w3, wzr",
-        "mov w21, wzr",
-        "rbit w1, w0",
-        "clz w1, w1",
-        "sub w2, w0, #0x1 (1)",
-        "ands w0, w2, w0",
-        "lsr w1, w20, w1",
-        "and w1, w1, #0x1",
-        "lsl w1, w1, w3",
-        "add w3, w3, #0x1 (1)",
-        "orr w21, w1, w21",
-        "b.ne #-0x24",
-        "mov w4, w21",
-        "b #+0x8",
+        "mov w2, w20",
+        "mov w4, wzr",
+        "cbz w0, #+0x20",
+        "clz w1, w0",
+        "lsl w2, w2, w1",
+        "lsl w0, w0, w1",
+        "extr w4, w4, w2, #31",
+        "bfc w0, #31, #1",
+        "b #-0x18",
         "mov w4, wzr"
       ]
     },
     "pext rax, rbx, rcx": {
-      "ExpectedInstructionCount": 19,
+      "ExpectedInstructionCount": 12,
       "Optimal": "No",
       "Comment": [
         "Map 2 0b10 0xf5 64-bit"
       ],
       "ExpectedArm64ASM": [
-        "cbz x5, #+0x48",
+        "cbz x5, #+0x2c",
         "mov x0, x5",
-        "mov x3, xzr",
-        "str x5, [x28, #16]",
-        "mov x5, xzr",
-        "rbit x1, x0",
-        "clz x1, x1",
-        "sub x2, x0, #0x1 (1)",
-        "ands x0, x2, x0",
-        "lsr x1, x7, x1",
-        "and x1, x1, #0x1",
-        "lsl x1, x1, x3",
-        "add x3, x3, #0x1 (1)",
-        "orr x5, x1, x5",
-        "b.ne #-0x24",
-        "mov x4, x5",
-        "ldr x5, [x28, #16]",
-        "b #+0x8",
+        "mov x2, x7",
+        "mov x4, xzr",
+        "cbz x0, #+0x20",
+        "clz x1, x0",
+        "lsl x2, x2, x1",
+        "lsl x0, x0, x1",
+        "extr x4, x4, x2, #63",
+        "bfc x0, #63, #1",
+        "b #-0x18",
         "mov x4, xzr"
       ]
     },