From ff8a03a7acca35eeba15bf4f1af8afa009cd04f3 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke@igalia.com>
Date: Fri, 12 Jul 2024 09:40:20 +0800
Subject: [PATCH] [RISCV] Fix coalesced vsetvli's AVL LiveInterval not always
 being shrunk (#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
https://github.com/llvm/llvm-project/pull/97264#issuecomment-2218036877.
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  | 11 ++++----
 .../test/CodeGen/RISCV/rvv/vsetvli-insert.mir | 26 +++++++++++++++++++
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp b/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
index 1f7d322be4d2a2..c74e7ac929c6b0 100644
--- a/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
+++ b/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));
   }
 }
 
diff --git a/llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.mir b/llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.mir
index 817bb3a905985f..deff36835a84ef 100644
--- a/llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.mir
+++ b/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