[RISCV] Check VL dominates and potentially move in tryReduceVL (#106753)

Similar to what we do in foldVMV_V_V with the passthru, if we end up
changing the Src's VL in tryReduceVL we need to make sure it dominates.

Fixes #106735

Author: lukel97

llvm/lib/Target/RISCV/RISCVVectorPeephole.cpp

@@ -70,6 +70,7 @@ class RISCVVectorPeephole : public MachineFunctionPass {
 
   bool isAllOnesMask(const MachineInstr *MaskDef) const;
   std::optional<unsigned> getConstant(const MachineOperand &VL) const;
+  bool ensureDominates(const MachineOperand &Use, MachineInstr &Src) const;
 
   /// Maps uses of V0 to the corresponding def of V0.
   DenseMap<const MachineInstr *, const MachineInstr *> V0Defs;
@@ -165,6 +166,9 @@ bool RISCVVectorPeephole::tryToReduceVL(MachineInstr &MI) const {
   if (VL.isIdenticalTo(SrcVL) || !isVLKnownLE(VL, SrcVL))
     return false;
 
+  if (!ensureDominates(VL, *Src))
+    return false;
+
   if (VL.isImm())
     SrcVL.ChangeToImmediate(VL.getImm());
   else if (VL.isReg())
@@ -456,6 +460,26 @@ static bool dominates(MachineBasicBlock::const_iterator A,
   return &*I == A;
 }
 
+/// If the register in \p MO doesn't dominate \p Src, try to move \p Src so it
+/// does. Returns false if doesn't dominate and we can't move. \p MO must be in
+/// the same basic block as \Src.
+bool RISCVVectorPeephole::ensureDominates(const MachineOperand &MO,
+                                          MachineInstr &Src) const {
+  assert(MO.getParent()->getParent() == Src.getParent());
+  if (!MO.isReg() || MO.getReg() == RISCV::NoRegister)
+    return true;
+
+  MachineInstr *Def = MRI->getVRegDef(MO.getReg());
+  if (Def->getParent() == Src.getParent() && !dominates(Def, Src)) {
+    if (!isSafeToMove(Src, *Def->getNextNode()))
+      return false;
+    // FIXME: Update V0Defs
+    Src.moveBefore(Def->getNextNode());
+  }
+
+  return true;
+}
+
 /// If a PseudoVMV_V_V is the only user of its input, fold its passthru and VL
 /// into it.
 ///
@@ -501,15 +525,8 @@ bool RISCVVectorPeephole::foldVMV_V_V(MachineInstr &MI) {
     return false;
 
   // If the new passthru doesn't dominate Src, try to move Src so it does.
-  if (Passthru.getReg() != RISCV::NoRegister) {
-    MachineInstr *PassthruDef = MRI->getVRegDef(Passthru.getReg());
-    if (PassthruDef->getParent() == Src->getParent() &&
-        !dominates(PassthruDef, Src)) {
-      if (!isSafeToMove(*Src, *PassthruDef->getNextNode()))
-        return false;
-      Src->moveBefore(PassthruDef->getNextNode());
-    }
-  }
+  if (!ensureDominates(Passthru, *Src))
+    return false;
 
   if (SrcPassthru.getReg() != Passthru.getReg()) {
     SrcPassthru.setReg(Passthru.getReg());

llvm/test/CodeGen/RISCV/rvv/reduce-vl-peephole.ll

@@ -0,0 +1,19 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -mtriple=riscv64 -mattr=+v -verify-machineinstrs | FileCheck %s
+
+define void @avl_not_dominated(<vscale x 2 x i32> %v, ptr %p) {
+; CHECK-LABEL: avl_not_dominated:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetivli zero, 1, e32, m1, ta, ma
+; CHECK-NEXT:    vmv.x.s a1, v8
+; CHECK-NEXT:    slli a1, a1, 32
+; CHECK-NEXT:    srli a1, a1, 32
+; CHECK-NEXT:    vsetvli zero, a1, e32, m1, ta, ma
+; CHECK-NEXT:    vadd.vi v8, v8, 1
+; CHECK-NEXT:    vse32.v v8, (a0)
+; CHECK-NEXT:    ret
+  %w = add <vscale x 2 x i32> %v, splat (i32 1)
+  %evl = extractelement <vscale x 2 x i32> %v, i32 0
+  call void @llvm.vp.store(<vscale x 2 x i32> %w, ptr %p, <vscale x 2 x i1> splat(i1 true), i32 %evl)
+  ret void
+}

llvm/test/CodeGen/RISCV/rvv/reduce-vl-peephole.mir

@@ -0,0 +1,15 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
+# RUN: llc %s -o - -mtriple=riscv64 -mattr=+v -run-pass=riscv-vector-peephole \
+# RUN:     -verify-machineinstrs | FileCheck %s
+---
+name: avl_not_dominated
+body: |
+  bb.0:
+    ; CHECK-LABEL: name: avl_not_dominated
+    ; CHECK: %evl:gprnox0 = ADDI $x0, 1
+    ; CHECK-NEXT: %x:vr = PseudoVADD_VV_M1 $noreg, $noreg, $noreg, %evl, 5 /* e32 */, 0 /* tu, mu */
+    ; CHECK-NEXT: PseudoVSE32_V_M1 %x, $noreg, %evl, 5 /* e32 */
+    %x:vr = PseudoVADD_VV_M1 $noreg, $noreg, $noreg, -1, 5 /* e32 */, 0 /* tu, mu */
+    %evl:gprnox0 = ADDI $x0, 1
+    PseudoVSE32_V_M1 %x:vr, $noreg, %evl, 5 /* e32 */
+...