[RISCV] Prune dead LI in vsetvli coalescing with trivially dead vsetvli (#98647)

This is a follow up to ff8a03a. On the review for that, I'd suggested a
stylistic rework, but after discussion we decided to move forward with
the fix as it was. This change is a small part of that suggested rework.
Once I sat down and wrote the code, I think I've convinced myself of an
entirely different approach (tbd), but for the moment, let's use a
lambda to share code so that we can pickup a missed optimization, and
reduce some duplication.

---------

Co-authored-by: Luke Lau <[email protected]>

Author: preames

llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp

@@ -1644,6 +1644,24 @@ void RISCVInsertVSETVLI::coalesceVSETVLIs(MachineBasicBlock &MBB) const {
   Used.demandVL();
   Used.demandVTYPE();
   SmallVector<MachineInstr*> ToDelete;
+
+  // Update LIS and cleanup dead AVLs given a value which has
+  // has had one use (as an AVL) removed.
+  auto afterDroppedAVLUse = [&](Register OldVLReg) {
+    if (LIS)
+      LIS->shrinkToUses(&LIS->getInterval(OldVLReg));
+
+    MachineInstr *VLOpDef = MRI->getUniqueVRegDef(OldVLReg);
+    if (VLOpDef && TII->isAddImmediate(*VLOpDef, OldVLReg) &&
+        MRI->use_nodbg_empty(OldVLReg)) {
+      if (LIS) {
+        LIS->removeInterval(OldVLReg);
+        LIS->RemoveMachineInstrFromMaps(*VLOpDef);
+      }
+      VLOpDef->eraseFromParent();
+    }
+  };
+
   for (MachineInstr &MI : make_range(MBB.rbegin(), MBB.rend())) {
 
     if (!isVectorConfigInstr(MI)) {
@@ -1696,22 +1714,9 @@ void RISCVInsertVSETVLI::coalesceVSETVLIs(MachineBasicBlock &MBB) const {
             MI.getOperand(1).ChangeToImmediate(NextMI->getOperand(1).getImm());
           else
             MI.getOperand(1).ChangeToRegister(NextMI->getOperand(1).getReg(), false);
+          if (OldVLReg && OldVLReg.isVirtual())
+            afterDroppedAVLUse(OldVLReg);
 
-          if (OldVLReg && OldVLReg.isVirtual()) {
-            // MI no longer uses OldVLReg so shrink its LiveInterval.
-            if (LIS)
-              LIS->shrinkToUses(&LIS->getInterval(OldVLReg));
-
-            MachineInstr *VLOpDef = MRI->getUniqueVRegDef(OldVLReg);
-            if (VLOpDef && TII->isAddImmediate(*VLOpDef, OldVLReg) &&
-                MRI->use_nodbg_empty(OldVLReg)) {
-              if (LIS) {
-                LIS->removeInterval(OldVLReg);
-                LIS->RemoveMachineInstrFromMaps(*VLOpDef);
-              }
-              VLOpDef->eraseFromParent();
-            }
-          }
           MI.setDesc(NextMI->getDesc());
         }
         MI.getOperand(2).setImm(NextMI->getOperand(2).getImm());
@@ -1731,8 +1736,8 @@ void RISCVInsertVSETVLI::coalesceVSETVLIs(MachineBasicBlock &MBB) const {
     if (MI->getOperand(1).isReg())
       OldAVLReg = MI->getOperand(1).getReg();
     MI->eraseFromParent();
-    if (LIS && OldAVLReg && OldAVLReg.isVirtual())
-      LIS->shrinkToUses(&LIS->getInterval(OldAVLReg));
+    if (OldAVLReg && OldAVLReg.isVirtual())
+      afterDroppedAVLUse(OldAVLReg);
   }
 }
 

llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.mir

@@ -589,7 +589,6 @@ body: |
     ; CHECK-LABEL: name: coalesce_shrink_removed_vsetvlis_uses
     ; CHECK: liveins: $x10, $v8
     ; CHECK-NEXT: {{  $}}
-    ; CHECK-NEXT: %avl1:gprnox0 = ADDI $x0, 1
     ; CHECK-NEXT: %avl2:gprnox0 = ADDI $x0, 2
     ; CHECK-NEXT: dead $x0 = PseudoVSETVLI %avl2, 209 /* e32, m2, ta, ma */, implicit-def $vl, implicit-def $vtype
     ; CHECK-NEXT: %x:gpr = COPY $x10