[llvm-branch-commits] [llvm] [CodeGen] Fix incorrect rematerialization rollback order (PR #197576)

[Matt Arsenault via llvm-branch-commits]

================
@@ -762,56 +758,75 @@ Printable Rematerializer::printUser(const MachineInstr *MI,
   });
 }
 
-Rollbacker::RollbackInfo::RollbackInfo(const Rematerializer &Remater,
-                                       RegisterIdx RegIdx) {
-  const Rematerializer::Reg &Reg = Remater.getReg(RegIdx);
-  DefReg = Reg.getDefReg();
-  DefRegion = Reg.DefRegion;
-  Dependencies = Reg.Dependencies;
-
-  InsertPos = std::next(Reg.DefMI->getIterator());
-  if (InsertPos != Reg.DefMI->getParent()->end())
-    NextRegIdx = Remater.getDefRegIdx(*InsertPos);
-  else
-    NextRegIdx = Rematerializer::NoReg;
-}
-
 void Rollbacker::rematerializerNoteRegCreated(const Rematerializer &Remater,
                                               RegisterIdx RegIdx) {
   if (RollingBack)
     return;
+  assert(Remater.isRematerializedRegister(RegIdx) && "only remats are created");
   Rematerializations[Remater.getOriginOf(RegIdx)].insert(RegIdx);
 }
 
-void Rollbacker::rematerializerNoteRegDeleted(const Rematerializer &Remater,
-                                              RegisterIdx RegIdx) {
-  if (RollingBack || Remater.isRematerializedRegister(RegIdx))
+void Rollbacker::rematerializerNoteRegWillBeDeleted(
+    const Rematerializer &Remater, RegisterIdx RegIdx) {
+  if (RollingBack)
+    return;
+
+  // Find a valid re-creation position after the register's definition.
+  MachineInstr *DefMI = Remater.getReg(RegIdx).DefMI;
+  MachineBasicBlock *ParentMBB = DefMI->getParent();
+  MachineBasicBlock::iterator ValidPos = std::next(DefMI->getIterator());
+  while (ValidPos != ParentMBB->end() && isRollbackableMI(*ValidPos, Remater))
+    ValidPos = std::next(ValidPos);
+
+  if (Remater.isRematerializedRegister(RegIdx)) {
+    // Rematerializations will not be re-created. Previously deleted registers
+    // that reference this register's defining instruction as their re-creation
+    // position should instead be re-created at a valid position after the
+    // deleted MI.
+    invalidatePosition(DefMI, ValidPos);
     return;
-  DeadRegs.try_emplace(RegIdx, Remater, RegIdx);
+  }
+
+  // Original registers can be re-created. Add a re-creation position for the
+  // definition of the rematerializable register.
+  DeadRegs.push_back(DeadReg(RegIdx, Remater));
+  const InsertBeforePos InsertPos = makePos(ValidPos, ParentMBB);
+  PosToIdx[InsertPos].insert(Positions.size());
+  Positions.push_back(InsertPos);
 }
 
 void Rollbacker::rollback(Rematerializer &Remater) {
   RollingBack = true;
 
-  // Re-create deleted registers.
-  for (auto &[RegIdx, Info] : DeadRegs) {
-    assert(!Remater.getReg(RegIdx).isAlive() && "register should be dead");
-
-    // The MI that was originally just after the MI defining the register we
-    // are trying to re-create may have been deleted. In such cases, we can
-    // re-create at that MI's own insert position (and apply the same logic
-    // recursively).
-    MachineBasicBlock::iterator InsertPos = Info.InsertPos;
-    RegisterIdx NextRegIdx = Info.NextRegIdx;
-    while (NextRegIdx != Rematerializer::NoReg) {
-      const auto *NextRegRollback = DeadRegs.find(NextRegIdx);
-      if (NextRegRollback == DeadRegs.end())
-        break;
-      InsertPos = NextRegRollback->second.InsertPos;
-      NextRegIdx = NextRegRollback->second.NextRegIdx;
+  // As we re-create registers, map deleted definitions to re-created ones. This
+  // allows to replace invalid re-creation positions that reference deleted
+  // definitions to valid new positions while restoring original MI order.
+  DenseMap<MachineInstr *, MachineInstr *> Replacements;
+  unsigned PositionIndex = Positions.size();
+
+  // Re-create deleted registers in reverse order of deletion. Related registers
+  // are deleted in reverse def-use order so this ensures we re-create registers
+  // in def-use order. This also ensures that re-creation positions that became
+  // invalid due to later MI deletions can be corrected as we go.
+  for (const DeadReg &Reg : reverse(DeadRegs)) {
+    assert(!Remater.getReg(Reg.Idx).isAlive() && "register should be dead");
+
+    // Determine re-creation position for the register's definition.
+    MachineBasicBlock::iterator InsertPosition;
+    const auto [Ptr, IsMBB] = Positions[--PositionIndex];
+    if (IsMBB) {
+      InsertPosition = static_cast<MachineBasicBlock *>(Ptr)->end();
----------------
arsenm wrote:

Should not have these casts 

https://github.com/llvm/llvm-project/pull/197576