Legalize AArch64 Select stack inputs in target selection

Summary: Move the AArch64 stack-input legalization for `Select` out of the post-generation fixed-point rewrite pass and into target selection. This preserves the existing behavior of lowering all stack-backed `Select` inputs through virtual registers while leaving the immediate `Select` legalization in postgen, where generated `Move(Imm)` instructions can still participate in later postgen rewrites.

Reviewed By: alexmalyshev

Differential Revision: D109608870

fbshipit-source-id: f06d6c4be7f4a55479aadc5d0d529f3f63689d4a

Author: kddnewton

cinderx/Jit/lir/postgen.cpp

@@ -523,6 +523,7 @@ RewriteResult rewriteLoadSecondCallResult(instr_iter_t instr_iter) {
 //   - Call: late-created by PostRegAllocRewrite via setOpcode()
 //   - Negate/Invert: handled by AArch64 target selection
 //   - Inc/Dec: handled by AArch64 target selection
+//   - Select: handled by AArch64 target selection
 //   - EpilogueEnd: special return-value handling
 //   - Pop: stack output, not input
 bool lowerStackInputToVreg(instr_iter_t instr_iter, size_t idx) {
@@ -570,8 +571,6 @@ RewriteResult rewriteSingleStackInputToVreg(
   switch (instr->opcode()) {
     case Instruction::kPush:
       return rewriteSingleStackInputToVreg(instr_iter, 0);
-    case Instruction::kSelect:
-      return rewriteAllStackInputsToVreg(instr_iter);
     default:
       return kUnchanged;
   }

cinderx/Jit/lir/target_select.cpp

@@ -231,6 +231,18 @@ void legalizeA64UnaryStackInput(BasicBlock* block, instr_iter_t instr_iter) {
   moveA64StackInputToVreg(block, instr_iter, 0);
 }
 
+/* AArch64 Select lowers to register-only csel. */
+void legalizeA64SelectStackInputs(BasicBlock* block, instr_iter_t instr_iter) {
+  Instruction* instr = instr_iter->get();
+  JIT_DCHECK(instr->isSelect(), "Expected Select, got {}", instr->opname());
+
+  for (size_t i = 0; i < instr->getNumInputs(); i++) {
+    if (instr->getInput(i)->isStack()) {
+      moveA64StackInputToVreg(block, instr_iter, i);
+    }
+  }
+}
+
 /* AArch64 Inc/Dec only operate on registers. Rewrite stack updates through a
  * virtual register so register allocation handles the temporary.
  */
@@ -558,6 +570,9 @@ void selectA64Opcodes(Function* func) {
         case Instruction::kInvert:
           legalizeA64UnaryStackInput(block, cur_iter);
           break;
+        case Instruction::kSelect:
+          legalizeA64SelectStackInputs(block, cur_iter);
+          break;
         case Instruction::kInc:
         case Instruction::kDec:
           legalizeA64StackInputForIncDec(block, cur_iter);