[llvm] [RISCV] Shrink uses of coalesced vsetvlis (PR #98286)

[Luke Lau via llvm-commits]

https://github.com/lukel97 updated https://github.com/llvm/llvm-project/pull/98286

>From bd01c59463bb23f6284e77b3542f1fc9fd23336c Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Wed, 10 Jul 2024 16:38:43 +0800
Subject: [PATCH 1/3] [RISCV] Shrink uses of coalesced vsetvlis

When we coalesce and delete a vsetvli, we should shrink the LiveInterval of its AVL register now that there is one less use.

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  |  4 +++
 .../test/CodeGen/RISCV/rvv/vsetvli-insert.mir | 26 +++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp b/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
index 67e1b76cd304f..06f09685a8119 100644
--- a/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
+++ b/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
@@ -1727,6 +1727,10 @@ void RISCVInsertVSETVLI::coalesceVSETVLIs(MachineBasicBlock &MBB) const {
     if (LIS)
       LIS->RemoveMachineInstrFromMaps(*MI);
     MI->eraseFromParent();
+    if (LIS)
+      for (MachineOperand &MO : MI->uses())
+        if (MO.isReg() && MO.getReg().isVirtual())
+          LIS->shrinkToUses(&LIS->getInterval(MO.getReg()));
   }
 }
 
diff --git a/llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.mir b/llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.mir
index 817bb3a905985..deff36835a84e 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

>From 3d7ed5f9580e027acea835df6780abccb22c0774 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Thu, 11 Jul 2024 00:15:46 +0800
Subject: [PATCH 2/3] Store the old register instead of iterating through uses
 after MI is erased

---
 llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp b/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
index 06f09685a8119..9d4afa3ebca56 100644
--- a/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
+++ b/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
@@ -1726,11 +1726,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)
-      for (MachineOperand &MO : MI->uses())
-        if (MO.isReg() && MO.getReg().isVirtual())
-          LIS->shrinkToUses(&LIS->getInterval(MO.getReg()));
+    if (LIS && OldAVLReg && OldAVLReg.isVirtual())
+      LIS->shrinkToUses(&LIS->getInterval(OldAVLReg));
   }
 }
 

>From 42ae7460218d2645af88eac705d43fd618bf5857 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Thu, 11 Jul 2024 23:13:48 +0800
Subject: [PATCH 3/3] Remove redundant shrink

Also fix comment: MI originally used OldVLReg, not NextMI
---
 llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp b/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
index 9d4afa3ebca56..b926ac7cffb3a 100644
--- a/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
+++ b/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
@@ -1692,12 +1692,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));