Visit instruction arguments in reverse in register allocator when calculating last-use data

Summary:
(This is another attempt at fixing the issue papered over in D104499077.)

optimizeMoveSequence() caches register-to-spill moves and rewrites later stack operands to use the cached register. The bug is that a rewritten operand can be marked `lastUse` even when the same spill slot is read again later in the block. D104499077 makes this easy to hit by reusing one stripped value twice for the same runtime call, but the optimizer bug exists independently of that change.

  [spill] = Move RDI
  RSI = Move [spill]        ; rewritten to RSI = Move RDI
  RDX = Move [spill]        ; still reads the spill later
  Call ...

If the first reload carries `lastUse`, `optimizeMoveSequence()` deletes `[spill] = Move RDI` as dead immediately after rewriting `RSI = Move [spill]`. The later reload still reads `[spill]`, so the call runs with stale stack data.

The cause involves the same value being read from the same spill slot by multiple instructions in the same basic block. This is aa situation created by D102409172's tag stripping:

* D102409172 adds a strip_cache — when the same tagged value needs stripping for two call arguments, one `And` instruction is emitted and both call arguments become `LinkedOperand`s pointing to the same vreg output.
* After regalloc, if that vreg is spilled, both operands become stack-slot reads from the same physical location.
* After `rewriteCallInstrs` extracts them into separate `Move`s, you get two `Move`s reading the same spill slot.
* However, before the fix, the allocator iterates inputs in forward order while processing instructions in reverse. For `Call func, [spill], [spill]`: the first input is encountered first, def is not in the live set, so it gets `lastUse`. The second input finds def already live — no `lastUse`. So, after call rewriting, `lastUse` is on the chronologically first `Move` — exactly the wrong one for spill deletion.

The bug specifically triggers in the self-reload variant: when the spilled register happens to match one of the argument registers, producing `Move RDI, [spill]` (self-reload) followed by `Move RSI, [spill]`.

To deal with this we change the register allocator to visit instruction inputs in reverse (apparently) consistently with visiting instructions in reverse.  Then `rewriteCallInstrs()` can correctly apply `lastUse` to the chronologically last `Move`.

Reviewed By: yoney

Differential Revision: D105856708

fbshipit-source-id: 60af6cbc087576f2c18796104c390ca538884ae0

Author: jbower-fb

cinderx/Jit/lir/postalloc.cpp

@@ -1480,10 +1480,6 @@ RewriteResult optimizeMoveSequence(BasicBlock* basicblock) {
       auto out_reg = instr->output()->isReg()
           ? instr->output()->getPhyRegister()
           : PhyLocation::REG_INVALID;
-      // for moves only we can generate A = Move A, which will get optimized out
-      if (instr->isMove()) {
-        out_reg = PhyLocation::REG_INVALID;
-      }
       instr->foreachInputOperand([&](OperandBase* operand) {
         if (!operand->isStack()) {
           return;
@@ -1498,9 +1494,6 @@ RewriteResult optimizeMoveSequence(BasicBlock* basicblock) {
         auto opnd = static_cast<Operand*>(operand);
         auto data_type = opnd->dataType();
         auto old_opnd = fmt::to_string(*opnd);
-        bool replacement_is_self_move = instr->isMove() &&
-            instr->output()->isReg() &&
-            instr->output()->getPhyRegister() == reg;
         opnd->setPhyRegister(reg);
         JIT_CHECK(
             bitSize(data_type) == bitSize(opnd->dataType()),
@@ -1511,13 +1504,8 @@ RewriteResult optimizeMoveSequence(BasicBlock* basicblock) {
             *instr);
         changed = kChanged;
 
-        // if the stack location operand can be replaced by the register it came
-        // from and this is the last use of the operand, we can remove the move
-        // instruction moving from the register to the stack location.
-        // Keep the spill alive when the replacement would turn this consumer
-        // into a self-move: later loads from the same stack slot may still
-        // need the spilled value after this move is optimized away.
-        if (opnd->isLastUse() && !replacement_is_self_move) {
+        // If this is the last use of the stack operand, remove the spill.
+        if (opnd->isLastUse()) {
           auto opt_iter = registerMemoryMoves.getInstrFromMemory(stack_slot);
           JIT_CHECK(opt_iter.has_value(), "There must be a def instruction.");
           basicblock->instructions().erase(*opt_iter);

cinderx/Jit/lir/regalloc.cpp

@@ -468,9 +468,15 @@ void LinearScanAllocator::calculateLiveIntervals() {
         visit_indirect(output_opnd);
       }
 
-      // inputs
-      for (size_t i = 0; i < instr->getNumInputs(); i++) {
-        const OperandBase* opnd = instr->getInput(i);
+      // Inputs are iterated in reverse so that when two operands reference
+      // the same vreg, lastUse lands on the last operand by index. After
+      // call rewriting splits operands into separate Move instructions, the
+      // last operand becomes the chronologically last Move, which is the
+      // correct one to carry lastUse for spill-deletion decisions.
+      size_t num_inputs = instr->getNumInputs();
+      for (size_t i = 0; i < num_inputs; ++i) {
+        size_t idx = num_inputs - i - 1;
+        const OperandBase* opnd = instr->getInput(idx);
         if (!opnd->isVreg() && !opnd->isInd()) {
           continue;
         }
@@ -480,7 +486,7 @@ void LinearScanAllocator::calculateLiveIntervals() {
           continue;
         }
 
-        register_input(opnd, instr->getInputPhyRegUse(i));
+        register_input(opnd, instr->getInputPhyRegUse(idx));
       }
 
       if (instr_opcode == Instruction::kCall ||

cinderx/RuntimeTests/backend_test.cpp

@@ -768,19 +768,34 @@ TEST_F(BackendTest, MoveSequenceOptTest) {
   [RBP - 16]:Object = Move RAX:Object
   [RBP - 24]:Object = Move RSI:Object
         RDI:Object = Move RAX:Object
+        RSI:Object = Move [RBP - 24]:Object
         RDX:Object = Move RCX:Object
                      Call Object
+
+  [RBP - 32] is deleted: lastUse with no later stack reads.
+  RSI = Move [RBP - 24] is a self-reload (RSI spilled and loaded back to RSI).
+  It is left intact because reg == out_reg skips the rewrite.
   */
-  ASSERT_EQ(bb->getNumInstrs(), 5);
+  ASSERT_EQ(bb->getNumInstrs(), 6);
   auto& instrs = bb->instructions();
 
   auto iter = instrs.begin();
 
-  ASSERT_EQ((*(iter++))->opcode(), Instruction::kMove);
-  ASSERT_EQ((*(iter++))->opcode(), Instruction::kMove);
-  ASSERT_EQ((*(iter++))->opcode(), Instruction::kMove);
-  ASSERT_EQ((*(iter++))->opcode(), Instruction::kMove);
-  ASSERT_EQ((*(iter++))->opcode(), Instruction::kCall);
+  auto* spill0 = (*(iter++)).get();
+  auto* spill1 = (*(iter++)).get();
+  auto* arg0 = (*(iter++)).get();
+  auto* arg1 = (*(iter++)).get();
+  auto* arg2 = (*(iter++)).get();
+  auto* call_instr = (*(iter++)).get();
+
+  ASSERT_EQ(spill0->opcode(), Instruction::kMove);
+  ASSERT_EQ(spill1->opcode(), Instruction::kMove);
+  ASSERT_EQ(arg0->opcode(), Instruction::kMove);
+  ASSERT_EQ(arg0->getInput(0)->type(), OperandBase::kReg);
+  ASSERT_EQ(arg1->opcode(), Instruction::kMove);
+  ASSERT_EQ(arg1->getInput(0)->type(), OperandBase::kStack);
+  ASSERT_EQ(arg2->opcode(), Instruction::kMove);
+  ASSERT_EQ(call_instr->opcode(), Instruction::kCall);
 }
 
 TEST_F(BackendTest, MoveSequenceOpt2Test) {
@@ -820,47 +835,47 @@ TEST_F(BackendTest, MoveSequenceOpt2Test) {
   ASSERT_EQ((*iter)->getInput(1)->type(), OperandBase::kStack);
 }
 
-TEST_F(BackendTest, MoveSequenceOptKeepsSharedSpillAliveAcrossSelfMove) {
+TEST_F(BackendTest, MoveSequenceOptLeavesSelfReloadsIntact) {
   auto lirfunc = std::make_unique<Function>();
   auto bb = lirfunc->allocateBasicBlock();
   auto epilogue = lirfunc->allocateBasicBlock();
 
   const PhyLocation kSharedSlot{-16, 64};
-  const PhyLocation kSelfMoveReg = ARGUMENT_REGS[0];
+  const PhyLocation kReloadReg = ARGUMENT_REGS[0];
   constexpr uint64_t kExpected = 4;
 
-  // Seed the spill slot with a known stale value. If the optimizer later
-  // deletes the real spill by mistake, the final reload reads this zero and the
-  // test fails deterministically instead of depending on uninitialized stack
-  // contents.
+  // Set up the previously failing case:
+  //
+  //   [RBP - 16] = Move RSI
+  //          RSI = Move [RBP - 16]   ; writes RSI, does not consume cached RSI
+  //          RAX = Move [RBP - 16]   ; later stack read still needs the spill
+  //
+  // A bad rewrite would turn the middle instruction into `RSI = Move RSI` and
+  // then conclude the spill is dead. This test checks that we keep both the
+  // spill store and the explicit self-reload in the block.
+  // Make a deleted spill observable instead of reading arbitrary stack data.
   bb->allocateInstr(
       Instruction::kMove,
       nullptr,
       OutStk{kSharedSlot, OperandBase::k64bit},
-      Imm{0, OperandBase::k64bit});
+      Imm{kExpected - kExpected, OperandBase::k64bit});
   bb->allocateInstr(
       Instruction::kMove,
       nullptr,
-      OutPhyReg{kSelfMoveReg, OperandBase::k64bit},
+      OutPhyReg{kReloadReg, OperandBase::k64bit},
       Imm{kExpected, OperandBase::k64bit});
   bb->allocateInstr(
       Instruction::kMove,
       nullptr,
       OutStk{kSharedSlot, OperandBase::k64bit},
-      PhyReg{kSelfMoveReg, OperandBase::k64bit});
-  // Reload the same spilled value back into the same physical register. Once
-  // the stack input is rewritten to use the remembered register directly, this
-  // instruction becomes `Move reg, reg` and is dropped later as a no-op.
-  auto self_move = bb->allocateInstr(
+      PhyReg{kReloadReg, OperandBase::k64bit});
+
+  auto self_reload = bb->allocateInstr(
       Instruction::kMove,
       nullptr,
-      OutPhyReg{kSelfMoveReg, OperandBase::k64bit},
+      OutPhyReg{kReloadReg, OperandBase::k64bit},
       Stk{kSharedSlot, OperandBase::k64bit});
-  // Mark the stack operand as a last use to model the exact case that used to
-  // trigger spill deletion inside optimizeMoveSequence().
-  self_move->getInput(0)->setLastUse();
-  // A second reload still needs the shared spill slot after the self-move has
-  // been simplified away. That is what the old optimization missed.
+  self_reload->getInput(0)->setLastUse();
   bb->allocateInstr(
       Instruction::kMove,
       nullptr,
@@ -871,30 +886,26 @@ TEST_F(BackendTest, MoveSequenceOptKeepsSharedSpillAliveAcrossSelfMove) {
 
   auto func = reinterpret_cast<uint64_t (*)()>(SimpleCompile(lirfunc.get()));
 
-  bool saw_preserved_spill = false;
-  for (auto iter = bb->instructions().begin(); iter != bb->instructions().end();
-       ++iter) {
-    auto* instr = iter->get();
+  bool saw_spill = false;
+  bool saw_self_reload = false;
+  for (auto& instr : bb->instructions()) {
     if (!instr->isMove()) {
       continue;
     }
     auto* out = instr->output();
     auto* in = instr->getInput(0);
     if (out->isStack() && out->getStackSlot().loc == kSharedSlot.loc &&
-        in->isReg() && in->getPhyRegister() == kSelfMoveReg) {
-      saw_preserved_spill = true;
-      break;
+        in->isReg() && in->getPhyRegister() == kReloadReg) {
+      saw_spill = true;
+    }
+    if (out->isReg() && out->getPhyRegister() == kReloadReg && in->isStack() &&
+        in->getStackSlot().loc == kSharedSlot.loc) {
+      saw_self_reload = true;
     }
   }
 
-  // Check both the internal LIR shape and the final machine-code behavior.
-  // Before the fix, optimizeMoveSequence() treated the self-move's stack input
-  // as the last consumer of the spill, erased the defining store, and relied on
-  // the register value staying live. That was wrong because the later reload
-  // still legitimately reads the stack slot. Keeping the spill is safe: it only
-  // skips the deletion in the self-move case, where the move itself is about to
-  // disappear and therefore cannot serve as the final use of the spilled value.
-  EXPECT_TRUE(saw_preserved_spill);
+  EXPECT_TRUE(saw_spill);
+  EXPECT_TRUE(saw_self_reload);
   EXPECT_EQ(func(), kExpected);
 }