[MIR][MachineInstr] Update MachineInstr::eraseFromParent() to return an iterator (#179787)

Unlike LLVM IR `Instruction::eraseFromParent()`,
`MachineInstr::eraseFromParent()` is void and does not return the
iterator following the erased instruction. Returning an iterator can be
very helpful for example when we are erasing MachineInstrs while
iterating, as it provides a convenient way to get a valid iterator.

This patch updates `MachineInstr::eraseFromParent()` to return a
`MachineBlock::iterator` (which is a
`MachineInstrBundleIterator<MachineInstr>`). If the erased instruction
is the head of a bundle, then the returned iterator points to the next
bundle (see unittest).

llvm/include/llvm/CodeGen/MachineInstr.h

@@ -22,6 +22,7 @@
 #include "llvm/ADT/ilist_node.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Analysis/MemoryLocation.h"
+#include "llvm/CodeGen/MachineInstrBundleIterator.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
@@ -1329,7 +1330,10 @@ public:
   /// If this instruction is the header of a bundle, the whole bundle is erased.
   /// This function can not be used for instructions inside a bundle, use
   /// eraseFromBundle() to erase individual bundled instructions.
-  LLVM_ABI void eraseFromParent();
+  /// \returns the iterator following the erased instruction. If this is the
+  /// header of a bundle it returns the iterator following the erased bundle
+  /// iterator.
+  LLVM_ABI MachineInstrBundleIterator<MachineInstr> eraseFromParent();
 
   /// Unlink 'this' from its basic block and delete it.
   ///

llvm/lib/CodeGen/MachineInstr.cpp

@@ -797,9 +797,9 @@ MachineInstr *MachineInstr::removeFromBundle() {
   return getParent()->remove_instr(this);
 }
 
-void MachineInstr::eraseFromParent() {
+MachineBasicBlock::iterator MachineInstr::eraseFromParent() {
   assert(getParent() && "Not embedded in a basic block!");
-  getParent()->erase(this);
+  return getParent()->erase(this);
 }
 
 void MachineInstr::eraseFromBundle() {

llvm/lib/Target/AMDGPU/SIFixSGPRCopies.cpp

@@ -926,8 +926,7 @@ bool SIFixSGPRCopies::lowerSpecialCase(MachineInstr &MI,
         llvm_unreachable("failed to constrain register");
     } else if (tryMoveVGPRConstToSGPR(MI.getOperand(1), DstReg, MI.getParent(),
                                       MI, MI.getDebugLoc())) {
-      I = std::next(I);
-      MI.eraseFromParent();
+      I = MI.eraseFromParent();
     }
     return true;
   }

llvm/lib/Target/AMDGPU/SIMemoryLegalizer.cpp

@@ -2481,8 +2481,7 @@ bool SIMemoryLegalizer::run(MachineFunction &MF) {
               MO.setIsInternalRead(false);
         }
 
-        MI->eraseFromParent();
-        MI = II->getIterator();
+        MI = MI->eraseFromParent();
       }
 
       if (MI->getDesc().TSFlags & SIInstrFlags::maybeAtomic) {

llvm/unittests/CodeGen/MachineInstrTest.cpp

@@ -613,4 +613,62 @@ TEST(MachineInstrTest, SpliceOperands) {
   EXPECT_EQ(MI->getNumOperands(), 10U);
 }
 
+// Checks the iterator returned by MacineInstr::eraseFromParent().
+TEST(MachineInstr, EraseFromParentReturnedIterator) {
+  LLVMContext Ctx;
+  Module Mod("Module", Ctx);
+  auto MF = createMachineFunction(Ctx, Mod);
+  auto MBB = MF->CreateMachineBasicBlock();
+
+  MCInstrDesc MCID = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
+  MachineInstr *MI1 = MF->CreateMachineInstr(MCID, DebugLoc());
+  MBB->insert(MBB->end(), MI1);
+  MachineInstr *MI2 = MF->CreateMachineInstr(MCID, DebugLoc());
+  MBB->insert(MBB->end(), MI2);
+
+  MachineBasicBlock::iterator It1 = MI1->eraseFromParent();
+  EXPECT_EQ(It1, MI2->getIterator());
+  MachineBasicBlock::iterator It2 = MI2->eraseFromParent();
+  EXPECT_EQ(It2, MBB->end());
+}
+
+// Checks the iterator returned by MacineInstr::eraseFromParent() when
+// instructions are in bundles.
+TEST(MachineInstr, EraseFromParentReturnedIteratorBundle) {
+  LLVMContext Ctx;
+  Module Mod("Module", Ctx);
+  auto MF = createMachineFunction(Ctx, Mod);
+  auto MBB = MF->CreateMachineBasicBlock();
+  MCInstrDesc MCID = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
+
+  // Bundle1 {
+  //   MI1B1
+  //   MI2B1
+  // }
+  // Bundle2 {
+  //   MI1B2
+  //   MI2B2
+  // }
+  MachineInstr *MI1B1 = MF->CreateMachineInstr(MCID, DebugLoc());
+  MBB->insert(MBB->end(), MI1B1);
+  MachineInstr *MI2B1 = MF->CreateMachineInstr(MCID, DebugLoc());
+  MBB->insert(MBB->end(), MI2B1);
+  MI2B1->bundleWithPred();
+
+  MachineInstr *MI1B2 = MF->CreateMachineInstr(MCID, DebugLoc());
+  MBB->insert(MBB->end(), MI1B2);
+  MachineInstr *MI2B2 = MF->CreateMachineInstr(MCID, DebugLoc());
+  MBB->insert(MBB->end(), MI2B2);
+  MI2B2->bundleWithPred();
+
+  // MI1B1->eraseFromParent() erases the whole Bundle1.
+  // The returned iterator matches the head of Bundle2.
+  MachineBasicBlock::iterator It1 = MI1B1->eraseFromParent();
+  EXPECT_EQ(It1, MI1B2->getIterator());
+
+  // Erasing MI1B2 erases the whole Bundle2.
+  MachineBasicBlock::iterator It2 = MI1B2->eraseFromParent();
+  EXPECT_EQ(It2, MBB->end());
+}
+
 } // end namespace