[RISCV] Allow vsetvlis with same register AVL in doLocalPostpass

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
llvm#75544 (comment).

Author: lukel97

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,52 @@ 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);
+    bool AreSameAVL = AVL.isReg() && PrevAVL.isReg() &&
+                      AVL.getReg() == PrevAVL.getReg() &&
+                      !IsDefinedBetween(AVL.getReg(), PrevMI, MI);
+    if (AVL.isReg() && AVL.getReg() != RISCV::X0 && !AreSameAVL)
+      return false;
   }
 
   if (!PrevMI.getOperand(2).isImm() || !MI.getOperand(2).isImm())

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

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)