[RISCV] postpone removal in initundef pass (#71661)

InitUndef pass need replace the implicit def with Undef pseudo, but
current remove method will make noreg2implicit borken.

This patch postpone the removal until all basicblock be processed.

Author: BeMg

llvm/lib/Target/RISCV/RISCVRVVInitUndef.cpp

@@ -42,6 +42,7 @@
 #include "RISCV.h"
 #include "RISCVSubtarget.h"
 #include "llvm/ADT/SmallSet.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/CodeGen/DetectDeadLanes.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 using namespace llvm;
@@ -59,6 +60,8 @@ class RISCVInitUndef : public MachineFunctionPass {
 
   // Newly added vregs, assumed to be fully rewritten
   SmallSet<Register, 8> NewRegs;
+  SmallVector<MachineInstr *, 8> DeadInsts;
+
 public:
   static char ID;
 
@@ -174,7 +177,7 @@ bool RISCVInitUndef::handleImplicitDef(MachineBasicBlock &MBB,
   BuildMI(MBB, Inst, Inst->getDebugLoc(), TII->get(Opcode), NewDest);
 
   if (!HasOtherUse)
-    Inst = MBB.erase(Inst);
+    DeadInsts.push_back(&(*Inst));
 
   for (auto MO : UseMOs) {
     MO->setReg(NewDest);
@@ -298,6 +301,10 @@ bool RISCVInitUndef::runOnMachineFunction(MachineFunction &MF) {
   for (MachineBasicBlock &BB : MF)
     Changed |= processBasicBlock(MF, BB, DLD);
 
+  for (auto *DeadMI : DeadInsts)
+    DeadMI->eraseFromParent();
+  DeadInsts.clear();
+
   return Changed;
 }
 

llvm/test/CodeGen/RISCV/regalloc-last-chance-recoloring-failure.ll

@@ -36,7 +36,7 @@ define void @last_chance_recoloring_failure() {
 ; CHECK-NEXT:    vmclr.m v0
 ; CHECK-NEXT:    li s0, 36
 ; CHECK-NEXT:    vsetvli zero, s0, e16, m4, ta, ma
-; CHECK-NEXT:    vfwadd.vv v16, v8, v8, v0.t
+; CHECK-NEXT:    vfwadd.vv v16, v8, v12, v0.t
 ; CHECK-NEXT:    csrr a0, vlenb
 ; CHECK-NEXT:    slli a0, a0, 3
 ; CHECK-NEXT:    add a0, sp, a0
@@ -45,7 +45,7 @@ define void @last_chance_recoloring_failure() {
 ; CHECK-NEXT:    call func@plt
 ; CHECK-NEXT:    li a0, 32
 ; CHECK-NEXT:    vsetvli zero, a0, e16, m4, ta, ma
-; CHECK-NEXT:    vrgather.vv v16, v8, v8, v0.t
+; CHECK-NEXT:    vrgather.vv v16, v8, v12, v0.t
 ; CHECK-NEXT:    vsetvli zero, s0, e16, m4, ta, ma
 ; CHECK-NEXT:    addi a1, sp, 16
 ; CHECK-NEXT:    csrr a2, vlenb
@@ -105,13 +105,13 @@ define void @last_chance_recoloring_failure() {
 ; SUBREGLIVENESS-NEXT:    vmclr.m v0
 ; SUBREGLIVENESS-NEXT:    li s0, 36
 ; SUBREGLIVENESS-NEXT:    vsetvli zero, s0, e16, m4, ta, ma
-; SUBREGLIVENESS-NEXT:    vfwadd.vv v16, v8, v8, v0.t
+; SUBREGLIVENESS-NEXT:    vfwadd.vv v16, v8, v12, v0.t
 ; SUBREGLIVENESS-NEXT:    addi a0, sp, 16
 ; SUBREGLIVENESS-NEXT:    vs8r.v v16, (a0) # Unknown-size Folded Spill
 ; SUBREGLIVENESS-NEXT:    call func@plt
 ; SUBREGLIVENESS-NEXT:    li a0, 32
 ; SUBREGLIVENESS-NEXT:    vsetvli zero, a0, e16, m4, ta, ma
-; SUBREGLIVENESS-NEXT:    vrgather.vv v16, v8, v8, v0.t
+; SUBREGLIVENESS-NEXT:    vrgather.vv v16, v8, v12, v0.t
 ; SUBREGLIVENESS-NEXT:    vsetvli zero, s0, e16, m4, ta, ma
 ; SUBREGLIVENESS-NEXT:    csrr a1, vlenb
 ; SUBREGLIVENESS-NEXT:    slli a1, a1, 3

llvm/test/CodeGen/RISCV/rvv/handle-noreg-with-implicit-def.mir

@@ -0,0 +1,20 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 3
+# RUN: llc -mtriple=riscv64 -mattr=+v -verify-machineinstrs -run-pass=riscv-init-undef -o - %s | FileCheck %s --check-prefix=MIR
+...
+---
+name:            vrgather_all_undef
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    ; MIR-LABEL: name: vrgather_all_undef
+    ; MIR: [[PseudoRVVInitUndefM1_:%[0-9]+]]:vr = PseudoRVVInitUndefM1
+    ; MIR-NEXT: [[DEF:%[0-9]+]]:vr = IMPLICIT_DEF
+    ; MIR-NEXT: early-clobber %1:vr = PseudoVRGATHER_VI_M1 [[DEF]], killed [[PseudoRVVInitUndefM1_]], 0, 0, 5 /* e32 */, 0 /* tu, mu */
+    ; MIR-NEXT: $v8 = COPY %1
+    ; MIR-NEXT: PseudoRET implicit $v8
+    %2:vr = IMPLICIT_DEF
+    early-clobber %1:vr = PseudoVRGATHER_VI_M1 $noreg, killed undef %2, 0, 0, 5 /* e32 */, 0 /* tu, mu */
+    $v8 = COPY %1
+    PseudoRET implicit $v8
+
+...