Stop widening moves from blocking move optimization

Summary:
Follow-up to keeping 32 -> 64 bit zero-extends as `MovZX`.  Two places
still assumed that width of zero-extend was a plain `Move`:

- `emitBoxPrimitive` hand-rolled the unsigned 32-bit case as a 32-bit
  `Move`, the last spot outside codegen that spelled a zero-extend as
  something other than what it is.  It now emits `MovZX`, matching the
  signed branch directly above it.  The encoding is unchanged, `MovZX`
  from a 32-bit source lowers to exactly that 32-bit move.

- `optimizeMoveSequence` caches register-to-spill-slot copies so a later
  read of the slot can use the register instead, and it dropped the whole
  cache on any instruction that was not a Move, Push or Pop.  A widening
  move only writes its register output, so it now invalidates just that
  location.  Before this, a `MovZX` sitting in the middle of a run of
  argument copies would throw away everything learned before it.

`optimizeMoveInstrs` is deliberately left alone.  It drops a `Move` whose
source and destination are the same location, which is not valid for a
widening move, the part of the destination the source does not cover
still has to be written.  The 32-bit return value in `emitBoxPrimitive`
is exactly that case, hence the comment added there.

Reviewed By: yoney

Differential Revision: D116031677

fbshipit-source-id: 8b5234cd3575c2a987d4def76e4874d5ac0a3b9e

Author: alexmalyshev

cinderx/Jit/lir/generator.cpp

@@ -1196,7 +1196,6 @@ void GenerateBoxedReturnWrapperBlocks(
         OutPhyReg{arg0, DataType::k64bit},
         PhyReg{ret32, DataType::k32bit});
   } else if (return_type <= TCUInt32) {
-    // Unsigned 32->64: a 32-bit move zero-extends into the full register.
     // Skippable when the registers already match, except on aarch64 pre-3.14:
     // there arg0 == ret64, but PyLong_FromSize_t reads all 64 bits and AAPCS64
     // leaves the upper half of a 32-bit return unspecified, so we must still
@@ -1205,9 +1204,9 @@ void GenerateBoxedReturnWrapperBlocks(
         PY_VERSION_HEX < 0x030E0000 && kBuildArch == Arch::kAarch64;
     if (arg0 != ret64 || needs_extend) {
       box_block->allocateInstr(
-          Opcode::kMove,
+          Opcode::kMovZX,
           nullptr,
-          OutPhyReg{arg0, DataType::k32bit},
+          OutPhyReg{arg0, DataType::k64bit},
           PhyReg{ret32, DataType::k32bit});
     }
   } else if (return_type <= (TCInt64 | TCUInt64) && arg0 != ret64) {

cinderx/Jit/lir/postalloc.cpp

@@ -636,6 +636,9 @@ RewriteResult rewriteBranchInstrs(Function* function) {
 RewriteResult optimizeMoveInstrs(instr_iter_t instr_iter) {
   auto instr = instr_iter->get();
   auto instr_opcode = instr->opcode();
+  // Deliberately not MovZX/MovSX: a widening move still has to write the part
+  // of the destination that the source does not cover, even when the two name
+  // the same register.
   if (instr_opcode != Opcode::kMove) {
     return kUnchanged;
   }
@@ -1560,19 +1563,16 @@ RewriteResult optimizeMoveSequence(BasicBlock* basicblock) {
       }
     };
 
-    if (instr->isMove() || instr->isPush() || instr->isPop()) {
-      if (instr->isMove()) {
-        Operand* out = instr->output();
-        Operand* in = instr->getInput(0);
-        if (out->isStack() && in->isReg()) {
-          registerMemoryMoves.addRegisterToMemoryMove(
-              in->getPhyRegister(), out->getStackSlot(), instr_iter);
-        } else {
-          invalidateOperand(out);
-        }
-      } else if (instr->isPop()) {
-        auto opnd = instr->output();
-        invalidateOperand(opnd);
+    if (instr->isMove() || instr->isPush() || instr->isPop() ||
+        instr->isMovZX() || instr->isMovSX()) {
+      Operand* out = instr->output();
+      if (instr->isMove() && out->isStack() && instr->getInput(0)->isReg()) {
+        registerMemoryMoves.addRegisterToMemoryMove(
+            instr->getInput(0)->getPhyRegister(),
+            out->getStackSlot(),
+            instr_iter);
+      } else {
+        invalidateOperand(out);
       }
     } else {
       // TODO: for now, we always clear the cache when we hit an instruction

cinderx/RuntimeTests/lir_postalloc_test.cpp

@@ -230,7 +230,6 @@ BB %1 - preds: %0 - section: .coldtext
   ASSERT_TRUE(verifyPostRegAllocInvariants(parsed_func.get(), std::cout));
 }
 
-#if defined(CINDER_AARCH64)
 // Helper to collect instructions from a block into a vector for easy indexing.
 static std::vector<Instruction*> collectInstrs(BasicBlock& bb) {
   std::vector<Instruction*> result;
@@ -240,6 +239,49 @@ static std::vector<Instruction*> collectInstrs(BasicBlock& bb) {
   return result;
 }
 
+// optimizeMoveSequence forwards a spill slot back to the register it was
+// copied from.  A widening move in between only writes its own output, so it
+// must not throw away the rest of what the pass knows.
+TEST_F(LIRPostAllocRewriteTest, MoveSequenceLooksPastWideningMoves) {
+  constexpr PhyLocation kSpilled = ARGUMENT_REGS[0];
+  constexpr PhyLocation kWidenOut = ARGUMENT_REGS[1];
+  constexpr PhyLocation kReloaded = ARGUMENT_REGS[2];
+  constexpr PhyLocation kWidenIn = ARGUMENT_REGS[3];
+  constexpr PhyLocation kSlot{-16, 64};
+
+  Function func;
+  auto* bb = func.allocateBasicBlock();
+
+  bb->allocateInstr(
+      Opcode::kMove,
+      nullptr,
+      OutStk{kSlot, DataType::k64bit},
+      PhyReg{kSpilled, DataType::k64bit});
+  bb->allocateInstr(
+      Opcode::kMovZX,
+      nullptr,
+      OutPhyReg{kWidenOut, DataType::k64bit},
+      PhyReg{kWidenIn, DataType::k32bit});
+  bb->allocateInstr(
+      Opcode::kMove,
+      nullptr,
+      OutPhyReg{kReloaded, DataType::k64bit},
+      Stk{kSlot, DataType::k64bit});
+
+  Environ env;
+  PostRegAllocRewrite rewrite(&func, &env);
+  rewrite.run();
+
+  auto instrs = collectInstrs(*bb);
+  ASSERT_EQ(instrs.size(), 3);
+
+  // The reload reads the register the spill came from rather than the slot.
+  EXPECT_TRUE(instrs[2]->isMove());
+  EXPECT_TRUE(instrs[2]->getInput(0)->isReg());
+  EXPECT_EQ(instrs[2]->getInput(0)->getPhyRegister(), kSpilled);
+}
+
+#if defined(CINDER_AARCH64)
 // kAdd with one register input and one stack input should insert a Move from
 // stack to GP scratch register before the Add, then rewrite the Add's stack
 // input to use the scratch register.