[RISCV] Allow vsetvlis with same register AVL in doLocalPostpass

[lukel97]

If we want to delete a vsetvli in the backwards pass, we need to be able to replace the AVL in MI with PrevMI's AVL.
We currently check to make sure that AVL isn't demanded so that we can clobber it, which works fine for immediate values.

However for registers we need to make sure that we can the instruction defining the AVL also dominates PrevMI.
Rather than go all in and use MachineDominatorTree or similar, this handles a simpler case by just checking that the two AVLs are the same virtual register.

Based off the draft patch in
#75544 (comment).

[llvmbot]

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

Author: Luke Lau (lukel97)

Changes

One of the requirements to be able to delete a vsetvli in the backwards pass is
that the preceding vsetvli must have the same AVL. This handles the case where
the AVLs are registers by using MachineRegisterInfo to check if there are any
definitions between the two vsetvlis.

The Dominates helper was taken from MachineDominatorTree and scans through the
instructions in the block which is less than ideal. But it's only called
whenever the two registers are the same, which should be rare.

This also replaces the equally-zero check with the existing hasEquallyZeroAVL
function, which is needed to handle the case where the AVLs are the same.

Based off the draft patch in
#75544 (comment).

---

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

3 Files Affected:

• (modified) llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp (+42-18)
• (modified) llvm/test/CodeGen/RISCV/rvv/fixed-vectors-insert.ll (+2-4)
• (modified) llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.ll (+1-2)

diff --git a/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp b/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
index 3400b24e0abb01..b124778a4de21c 100644
--- a/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
+++ b/llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp
@@ -1459,20 +1459,6 @@ static void doUnion(DemandedFields &A, DemandedFields B) {
   A.MaskPolicy |= B.MaskPolicy;
 }
 
-static bool isNonZeroAVL(const MachineOperand &MO,
-                         const MachineRegisterInfo &MRI) {
-  if (MO.isReg()) {
-    if (MO.getReg() == RISCV::X0)
-      return true;
-    if (MachineInstr *MI = MRI.getVRegDef(MO.getReg());
-        MI && isNonZeroLoadImmediate(*MI))
-      return true;
-    return false;
-  }
-  assert(MO.isImm());
-  return 0 != MO.getImm();
-}
-
 // Return true if we can mutate PrevMI to match MI without changing any the
 // fields which would be observed.
 static bool canMutatePriorConfig(const MachineInstr &PrevMI,
@@ -1491,16 +1477,54 @@ static bool canMutatePriorConfig(const MachineInstr &PrevMI,
     if (Used.VLZeroness) {
       if (isVLPreservingConfig(PrevMI))
         return false;
-      if (!isNonZeroAVL(MI.getOperand(1), MRI) ||
-          !isNonZeroAVL(PrevMI.getOperand(1), MRI))
+      if (!getInfoForVSETVLI(PrevMI).hasEquallyZeroAVL(getInfoForVSETVLI(MI),
+                                                       MRI))
         return false;
     }
 
-    // TODO: Track whether the register is defined between
-    // PrevMI and MI.
     if (MI.getOperand(1).isReg() &&
         RISCV::X0 != MI.getOperand(1).getReg())
       return false;
+
+    // Taken from MachineDominatorTree::dominates
+    auto Dominates = [](const MachineInstr &A, const MachineInstr &B) {
+      assert(A.getParent() == B.getParent());
+      // Loop through the basic block until we find A or B.
+      MachineBasicBlock::const_iterator I = A.getParent()->begin();
+      for (; I != A && I != B; ++I)
+        /*empty*/;
+      return I == A;
+    };
+
+    // Given A and B are in the same block and A comes before (dominates) B,
+    // return whether or not Reg is defined between A and B.
+    auto IsDefinedBetween = [&MRI, &Dominates](const Register Reg,
+                                               const MachineInstr &A,
+                                               const MachineInstr &B) {
+      assert(Dominates(A, B));
+      for (const auto &Def : MRI.def_instructions(Reg)) {
+        if (Def.getParent() != A.getParent())
+          continue;
+        // If B defines Reg, assume it early clobbers for now.
+        if (&Def == &B)
+          return true;
+        if (Dominates(Def, A) && !Dominates(Def, B))
+          return true;
+      }
+
+      // Reg isn't defined between PrevMI and MI.
+      return false;
+    };
+
+    auto &AVL = MI.getOperand(1);
+    auto &PrevAVL = PrevMI.getOperand(1);
+
+    if (AVL.isReg() && AVL.getReg() != RISCV::X0) {
+      bool AreSameAVL = PrevAVL.isReg() && AVL.getReg() == PrevAVL.getReg() &&
+                        !IsDefinedBetween(AVL.getReg(), PrevMI, MI);
+      if (!AreSameAVL)
+        return false;
+    }
   }
 
   if (!PrevMI.getOperand(2).isImm() || !MI.getOperand(2).isImm())
diff --git a/llvm/test/CodeGen/RISCV/rvv/fixed-vectors-insert.ll b/llvm/test/CodeGen/RISCV/rvv/fixed-vectors-insert.ll
index 57760070603b2e..4954827876c19a 100644
--- a/llvm/test/CodeGen/RISCV/rvv/fixed-vectors-insert.ll
+++ b/llvm/test/CodeGen/RISCV/rvv/fixed-vectors-insert.ll
@@ -63,9 +63,8 @@ define <32 x i32> @insertelt_v32i32_31(<32 x i32> %a, i32 %y) {
 ; CHECK-LABEL: insertelt_v32i32_31:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    li a1, 32
-; CHECK-NEXT:    vsetvli zero, a1, e32, m1, ta, ma
-; CHECK-NEXT:    vmv.s.x v16, a0
 ; CHECK-NEXT:    vsetvli zero, a1, e32, m8, ta, ma
+; CHECK-NEXT:    vmv.s.x v16, a0
 ; CHECK-NEXT:    vslideup.vi v8, v16, 31
 ; CHECK-NEXT:    ret
   %b = insertelement <32 x i32> %a, i32 %y, i32 31
@@ -101,9 +100,8 @@ define <64 x i32> @insertelt_v64i32_63(<64 x i32> %a, i32 %y) {
 ; CHECK-LABEL: insertelt_v64i32_63:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    li a1, 32
-; CHECK-NEXT:    vsetvli zero, a1, e32, m1, ta, ma
-; CHECK-NEXT:    vmv.s.x v24, a0
 ; CHECK-NEXT:    vsetvli zero, a1, e32, m8, ta, ma
+; CHECK-NEXT:    vmv.s.x v24, a0
 ; CHECK-NEXT:    vslideup.vi v16, v24, 31
 ; CHECK-NEXT:    ret
   %b = insertelement <64 x i32> %a, i32 %y, i32 63
diff --git a/llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.ll b/llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.ll
index e15c5a3323cbe2..7c95d81306655b 100644
--- a/llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.ll
+++ b/llvm/test/CodeGen/RISCV/rvv/vsetvli-insert.ll
@@ -643,9 +643,8 @@ define <vscale x 2 x float> @fp_reduction_vfmv_s_f(float %0, <vscale x 8 x float
 define dso_local <vscale x 2 x i32> @int_reduction_vmv_s_x(i32 signext %0, <vscale x 8 x i32> %1, i64 %2) {
 ; CHECK-LABEL: int_reduction_vmv_s_x:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    vsetvli zero, a1, e32, m1, ta, ma
-; CHECK-NEXT:    vmv.s.x v12, a0
 ; CHECK-NEXT:    vsetvli zero, a1, e32, m4, ta, ma
+; CHECK-NEXT:    vmv.s.x v12, a0
 ; CHECK-NEXT:    vredsum.vs v8, v8, v12
 ; CHECK-NEXT:    ret
   %4 = tail call <vscale x 8 x i32> @llvm.riscv.vmv.s.x.nxv8i32.i64(<vscale x 8 x i32> poison, i32 %0, i64 %2)

[lukel97]

Ping

[preames]

LGTM

[lukel97]

Just noticed, the commit message doesn't seem to have the PR body in it. Not sure what happened: I just clicked the squash and merge button.