Legalize AArch64 unary stack inputs in target selection

Summary: Move the AArch64 stack-input legalization for unary `Negate` and `Invert` out of the post-generation fixed-point rewrite pass and into target selection.

Reviewed By: alexmalyshev

Differential Revision: D109608869

fbshipit-source-id: 39491038e241657dc172badd6876615bdf7438d4

Author: kddnewton

cinderx/Jit/lir/postgen.cpp

@@ -521,6 +521,7 @@ RewriteResult rewriteLoadSecondCallResult(instr_iter_t instr_iter) {
 //   - MovZX/MovSX/MovSXD: specialized sign/zero-extending loads from stack
 //   - Lea: takes the ADDRESS of a stack slot, not the value
 //   - Call: late-created by PostRegAllocRewrite via setOpcode()
+//   - Negate/Invert: handled by AArch64 target selection
 //   - Inc/Dec: handled by AArch64 target selection
 //   - EpilogueEnd: special return-value handling
 //   - Pop: stack output, not input
@@ -567,8 +568,6 @@ RewriteResult rewriteSingleStackInputToVreg(
   }
 
   switch (instr->opcode()) {
-    case Instruction::kNegate:
-    case Instruction::kInvert:
     case Instruction::kPush:
       return rewriteSingleStackInputToVreg(instr_iter, 0);
     case Instruction::kSelect:

cinderx/Jit/lir/target_select.cpp

@@ -200,6 +200,26 @@ void legalizeA64GuardFPInput(BasicBlock* block, instr_iter_t instr_iter) {
       kGuardVarIndex, std::make_unique<Operand>(move, Operand::kLinked));
 }
 
+/* AArch64 unary arithmetic instructions only operate on registers. */
+void legalizeA64UnaryStackInput(BasicBlock* block, instr_iter_t instr_iter) {
+  Instruction* instr = instr_iter->get();
+  JIT_DCHECK(
+      instr->isNegate() || instr->isInvert(),
+      "Expected Negate or Invert, got {}",
+      instr->opname());
+
+  Operand* input = instr->getInput(0);
+  if (!input->isStack()) {
+    return;
+  }
+
+  PhyLocation loc = input->getStackSlot();
+  DataType dt = input->dataType();
+  Instruction* move = block->allocateInstrBefore(
+      instr_iter, Instruction::kMove, OutVReg{dt}, Stk{loc, dt});
+  instr->setInput(0, std::make_unique<Operand>(move, Operand::kLinked));
+}
+
 /* AArch64 Inc/Dec only operate on registers. Rewrite stack updates through a
  * virtual register so register allocation handles the temporary.
  */
@@ -525,6 +545,10 @@ void selectA64Opcodes(Function* func) {
           legalizeA64GuardFPInput(block, cur_iter);
           selectA64Guard(block, cur_iter, use_counts);
           break;
+        case Instruction::kNegate:
+        case Instruction::kInvert:
+          legalizeA64UnaryStackInput(block, cur_iter);
+          break;
         case Instruction::kInc:
         case Instruction::kDec:
           legalizeA64StackInputForIncDec(block, cur_iter);