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

[Philip Reames via llvm-commits]

https://github.com/preames updated https://github.com/llvm/llvm-project/pull/98647

>From 5f403a9fe03ee1d1781917b3eb6cf417ac016a8f Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Fri, 12 Jul 2024 07:33:29 -0700
Subject: [PATCH 1/2] [RISCV] Prune dead LI in vsetvli coalescing with
 trivially dead vsetvli

This is a follow up to ff8a03a7.  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.
---
 llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp  | 40 +++++++++++--------
 .../test/CodeGen/RISCV/rvv/vsetvli-insert.mir |  1 -
 2 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp b/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
index c74e7ac929c6b..76aed9b4083de 100644
--- a/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
+++ b/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
@@ -1633,6 +1633,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)) {
@@ -1685,22 +1703,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());
@@ -1720,8 +1725,9 @@ 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);
+
   }
 }
 
diff --git a/llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.mir b/llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.mir
index deff36835a84e..a9cf741b1cb2a 100644
--- a/llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.mir
+++ b/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

>From 23d8efc5abe33df764d21cdc191a801afdc6031b Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Mon, 15 Jul 2024 08:35:01 -0700
Subject: [PATCH 2/2] Update llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp

Co-authored-by: Luke Lau <luke_lau at icloud.com>
---
 llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp b/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
index 76aed9b4083de..e8c22ff97a5a1 100644
--- a/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
+++ b/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
@@ -1727,7 +1727,6 @@ void RISCVInsertVSETVLI::coalesceVSETVLIs(MachineBasicBlock &MBB) const {
     MI->eraseFromParent();
     if (OldAVLReg && OldAVLReg.isVirtual())
       afterDroppedAVLUse(OldAVLReg);
-
   }
 }