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

[lukel97]

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 LiveIntervals 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.

[llvmbot]

@llvm/pr-subscribers-backend-risc-v

Author: Luke Lau (lukel97)

Changes

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 #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.

---

Full diff: https://github.com/llvm/llvm-project/pull/98286.diff

2 Files Affected:

• (modified) llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp (+4)
• (modified) llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.mir (+26)

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

[preames]

Calling MI->uses() after MI->eraseFromParent() is definitely unsafe.

[preames]

I realized when taking a deeper look that you're duplicating functionality here. This same case is supposed to be handled in the "if (DefReg.isVirtual() && LIS) {" above. I think what happens is that you've got an order of operations bug. You need to actually modify the instruction first, then perform all of the LIS update. (That is, move that block done to just before the OldVReg update, and maybe move the setDesc call up.

[lukel97]

DefReg is a different register, it's the output VL, this is for the input AVL. But now that you mention it, we were also duplicating the liveintervals update for the AVL too. We already handle it for the normal case, but the bug was ultimately caused by this case here where we continue early for two consecutive vsetvlis with no pseudos in between:

https://github.com/llvm/llvm-project/pull/98286/files#diff-a186cfa1b7676c0a95a09935007221c321cfc52ce54b85d15064963d413725a3R1657-R1661

I've removed that redundant case so we're only updating the AVL in one place.

[preames]

Sorry for the confusion, you're right about DefReg vs OldVLReg. I think you're also right about the early continue being the trigger, but looking at that code, it looks like we're also missing the ADDI prune in that case as well. Maybe pull out a lambda for prepareToDelete which does all the LIS fixup, ADDI deletion, and addition to the ToDelete list?

[preames]

LGTM to the current patch.

I talked with Luke offline, and was convinced this is correct as structured. I'm not a fan of the style - doing the updates in two places, and out of order - but I'm explicitly okay with this landing, and I will post an NFC tomorrow.