[RISCV] Fix coalesced vsetvli's AVL LiveInterval not always being shr…

…unk (#98286)

Most of the time when we coalesce and delete a vsetvli, we shrink the
LiveInterval of its AVL register now that there is one less use. However
there's one edge case we were missing where if we have two vsetvlis with
no users of vl or vtype in between, we coalesced a vsetvli without
shrinking it's AVL.

This fixes it by shrinking the LiveInterval whenever we delete a
vsetvli, and also makes the LiveIntervals consistent in-situ by not
removing the use before shrinking.

This fixes a -verify-machineinstrs assertion in an MIR test case I found
while investigating
#97264 (comment).
I couldn't recreate this at the LLVM IR level, seemingly because
RISCVInsertVSETVLI will just avoid inserting extra vsetvlis that don't
need coalesced.

llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp

@@ -1686,12 +1686,8 @@ void RISCVInsertVSETVLI::coalesceVSETVLIs(MachineBasicBlock &MBB) const {
           else
             MI.getOperand(1).ChangeToRegister(NextMI->getOperand(1).getReg(), false);
 
-          // Clear NextMI's AVL early so we're not counting it as a use.
-          if (NextMI->getOperand(1).isReg())
-            NextMI->getOperand(1).setReg(RISCV::NoRegister);
-
           if (OldVLReg && OldVLReg.isVirtual()) {
-            // NextMI no longer uses OldVLReg so shrink its LiveInterval.
+            // MI no longer uses OldVLReg so shrink its LiveInterval.
             if (LIS)
               LIS->shrinkToUses(&LIS->getInterval(OldVLReg));
 
@@ -1720,7 +1716,12 @@ void RISCVInsertVSETVLI::coalesceVSETVLIs(MachineBasicBlock &MBB) const {
   for (auto *MI : ToDelete) {
     if (LIS)
       LIS->RemoveMachineInstrFromMaps(*MI);
+    Register OldAVLReg;
+    if (MI->getOperand(1).isReg())
+      OldAVLReg = MI->getOperand(1).getReg();
     MI->eraseFromParent();
+    if (LIS && OldAVLReg && OldAVLReg.isVirtual())
+      LIS->shrinkToUses(&LIS->getInterval(OldAVLReg));
   }
 }
 

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

@@ -92,6 +92,10 @@
     ret void
   }
 
+  define void @coalesce_shrink_removed_vsetvlis_uses() {
+    ret void
+  }
+
   declare <vscale x 1 x i64> @llvm.riscv.vadd.nxv1i64.nxv1i64.i64(<vscale x 1 x i64>, <vscale x 1 x i64>, <vscale x 1 x i64>, i64) #1
 
   declare <vscale x 1 x i64> @llvm.riscv.vle.nxv1i64.i64(<vscale x 1 x i64>, ptr nocapture, i64) #4
@@ -576,3 +580,25 @@ body:             |
     $v0 = PseudoVADD_VV_M1 $noreg, $noreg, $noreg, 3, 6, 0
     PseudoRET
 ...
+---
+name: coalesce_shrink_removed_vsetvlis_uses
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $x10, $v8
+    ; 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
+    ; CHECK-NEXT: renamable $v8 = PseudoVMV_S_X undef renamable $v8, %x, 1, 5 /* e32 */, implicit $vl, implicit $vtype
+    ; CHECK-NEXT: PseudoRET implicit $v8
+    %avl1:gprnox0 = ADDI $x0, 1
+    dead $x0 = PseudoVSETVLI %avl1:gprnox0, 209, implicit-def dead $vl, implicit-def dead $vtype
+    %avl2:gprnox0 = ADDI $x0, 2
+    dead $x0 = PseudoVSETVLI %avl2:gprnox0, 209, implicit-def dead $vl, implicit-def dead $vtype
+    %x:gpr = COPY $x10
+    renamable $v8 = PseudoVMV_S_X undef renamable $v8, killed renamable %x, 1, 5
+    PseudoRET implicit $v8