!fixup address comments, thanks!

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -781,78 +781,83 @@ static bool hoistPreviousBeforeFORUsers(VPFirstOrderRecurrencePHIRecipe *FOR,
     return false;
 
   // Collect recipes that need hoisting.
-  SmallVector<VPRecipeBase *> WorkList;
+  SmallVector<VPRecipeBase *> HoistCandidates;
   SmallPtrSet<VPRecipeBase *, 8> Seen;
-  VPBasicBlock *HoistBlock = FOR->getParent();
-  auto HoistPoint = HoistBlock->getFirstNonPhi();
-  auto TryToPushHoistCandidate = [&](VPRecipeBase *HoistCandidate) {
-    // If we reach FOR, it means the original Previous depends on some other
-    // recurrence that in turn depends on FOR. If that is the case, we would
-    // also need to hoist recipes involving the other FOR, which may break
-    // dependencies.
-    if (HoistCandidate == FOR)
-      return false;
-
-    // Hoist candidate outside any region, no need to hoist.
-    if (!HoistCandidate->getParent()->getParent())
-      return true;
-
-    // Hoist candidate is a header phi or already visited, no need to hoist.
-    if (isa<VPHeaderPHIRecipe>(HoistCandidate) ||
-        !Seen.insert(HoistCandidate).second)
-      return true;
+  VPRecipeBase *HoistPoint = nullptr;
+  // Find the closest hoist point by looking at all users of FOR and selecting
+  // the recipe dominating all other users.
+  for (VPUser *U : FOR->users()) {
+    auto *R = dyn_cast<VPRecipeBase>(U);
+    if (!R)
+      continue;
+    if (!HoistPoint || VPDT.properlyDominates(R, HoistPoint))
+      HoistPoint = R;
+  }
 
-    // If we reached a recipe that dominates all users of FOR, we don't need to
+  auto NeedsHoisting = [HoistPoint, &VPDT,
+                        &Seen](VPValue *HoistCandidateV) -> VPRecipeBase * {
+    VPRecipeBase *HoistCandidate = HoistCandidateV->getDefiningRecipe();
+    if (!HoistCandidate)
+      return nullptr;
+    assert((!HoistCandidate->getParent()->getParent() ||
+            HoistCandidate->getParent()->getParent() ==
+                HoistCandidate->getParent()->getEnclosingLoopRegion()) &&
+           "CFG in VPlan should still be flattened, without replicate regions");
+    // Hoist candidate has already beend visited, no need to hoist.
+    if (!Seen.insert(HoistCandidate).second)
+      return nullptr;
+
+    // Candidate is outside loop region or a header phi, dominates FOR users w/o
+    // hoisting.
+    if (!HoistCandidate->getParent()->getEnclosingLoopRegion() ||
+        isa<VPHeaderPHIRecipe>(HoistCandidate))
+      return nullptr;
+
+    // If we reached a recipe that dominates HoistPoint, we don't need to
     // hoist the recipe.
-    if (all_of(FOR->users(), [&VPDT, HoistCandidate](VPUser *U) {
-          return VPDT.properlyDominates(HoistCandidate, cast<VPRecipeBase>(U));
-        })) {
-      if (VPDT.properlyDominates(&*HoistPoint, HoistCandidate)) {
-        // This HoistCandidate domiantes all users of FOR and is closer to them
-        // than the previous HoistPoint.
-        HoistPoint = std::next(HoistCandidate->getIterator());
-        HoistBlock = HoistCandidate->getParent();
-      }
-      return true;
-    }
-
-    // Don't move candiates with sideeffects, as we do not yet analyze recipes
-    // between candidate and hoist destination yet.
-    if (HoistCandidate->mayHaveSideEffects())
-      return false;
-
-    WorkList.push_back(HoistCandidate);
-    return true;
+    if (VPDT.properlyDominates(HoistCandidate, HoistPoint))
+      return nullptr;
+    return HoistCandidate;
+  };
+  auto CanHoist = [&](VPRecipeBase *HoistCandidate) {
+    // Avoid hoisting candidates with side-effects, as we do not yet analyze
+    // associated dependencies.
+    return !HoistCandidate->mayHaveSideEffects();
   };
 
   // Recursively try to hoist Previous and its operands before all users of FOR.
-  // Update HoistPoint to the closest recipe that dominates all users of FOR.
-  if (!TryToPushHoistCandidate(Previous))
-    return false;
+  if (NeedsHoisting(Previous->getVPSingleValue()))
+    HoistCandidates.push_back(Previous);
 
-  for (unsigned I = 0; I != WorkList.size(); ++I) {
-    VPRecipeBase *Current = WorkList[I];
+  for (unsigned I = 0; I != HoistCandidates.size(); ++I) {
+    VPRecipeBase *Current = HoistCandidates[I];
     assert(Current->getNumDefinedValues() == 1 &&
            "only recipes with a single defined value expected");
+    if (!CanHoist(Current))
+      return false;
 
-    for (VPValue *Op : Current->operands())
-      if (auto *R = Op->getDefiningRecipe())
-        if (!TryToPushHoistCandidate(R))
-          return false;
+    for (VPValue *Op : Current->operands()) {
+      // If we reach FOR, it means the original Previous depends on some other
+      // recurrence that in turn depends on FOR. If that is the case, we would
+      // also need to hoist recipes involving the other FOR, which may break
+      // dependencies.
+      if (Op == FOR)
+        return false;
+
+      if (auto *R = NeedsHoisting(Op))
+        HoistCandidates.push_back(R);
+    }
   }
 
   // Keep recipes to hoist ordered by dominance so earlier instructions are
   // processed first.
-  sort(WorkList, [&VPDT](const VPRecipeBase *A, const VPRecipeBase *B) {
+  sort(HoistCandidates, [&VPDT](const VPRecipeBase *A, const VPRecipeBase *B) {
     return VPDT.properlyDominates(A, B);
   });
 
-  for (VPRecipeBase *HoistCandidate : WorkList) {
-    if (HoistPoint == HoistCandidate->getIterator()) {
-      HoistPoint = std::next(HoistCandidate->getIterator());
-      continue;
-    }
-    HoistCandidate->moveBefore(*HoistBlock, HoistPoint);
+  for (VPRecipeBase *HoistCandidate : HoistCandidates) {
+    HoistCandidate->moveBefore(*HoistPoint->getParent(),
+                               HoistPoint->getIterator());
   }
 
   return true;
@@ -881,9 +886,9 @@ bool VPlanTransforms::adjustFixedOrderRecurrences(VPlan &Plan,
       Previous = PrevPhi->getBackedgeValue()->getDefiningRecipe();
     }
 
-    if (!sinkRecurrenceUsersAfterPrevious(FOR, Previous, VPDT))
-      if (!hoistPreviousBeforeFORUsers(FOR, Previous, VPDT))
-        return false;
+    if (!sinkRecurrenceUsersAfterPrevious(FOR, Previous, VPDT) &&
+        !hoistPreviousBeforeFORUsers(FOR, Previous, VPDT))
+      return false;
 
     // Introduce a recipe to combine the incoming and previous values of a
     // fixed-order recurrence.

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -36,13 +36,13 @@ struct VPlanTransforms {
                                 GetIntOrFpInductionDescriptor,
                             ScalarEvolution &SE, const TargetLibraryInfo &TLI);
 
-  /// Try to move users of fixed-order recurrences after the recipe defining
-  /// their previous value, either by sinking them or hoisting the recipe
+  /// Try to have all users of fixed-order recurrences appear after the recipe
+  /// defining their previous value, by either sinking or hoisting the recipe
   /// defining their previous value (and its operands). Then introduce
   /// FirstOrderRecurrenceSplice VPInstructions to combine the value from the
   /// recurrence phis and previous values. The current implementation assumes
-  /// all users can be sunk after the previous value, which is enforced by
-  /// earlier legality checks.
+  /// all users can be sunk after the previous value, or the previous value can
+  /// be hoisted before all users, which is enforced by earlier legality checks.
   /// \returns true if all users of fixed-order recurrences could be re-arranged
   /// as needed or false if it is not possible. In the latter case, \p Plan is
   /// not valid.

llvm/test/Transforms/LoopVectorize/X86/fixed-order-recurrence.ll

@@ -280,6 +280,18 @@ exit:
 }
 
 ; Test for https://github.com/llvm/llvm-project/issues/106523.
+; %for.2 requires no code motion, as its previous (%or) precedes its (first)
+; user (store). Furthermore, its user cannot sink, being a store.
+;
+; %for.1 requires code motion, as its previous (%trunc) follows its (first)
+; user (%or). Sinking %or past %trunc seems possible, as %or has no uses
+; (except for feeding %for.2; worth strengthening VPlan's dce?). However, %or
+; is both the user of %for.1 and the previous of %for.2, and we refrain from
+; sinking instructions that act as previous because they (may) serve points to
+; sink after.
+
+; Instead, %for.1 can be reconciled by hoisting its previous above its user
+; %or, as this user %trunc depends only on %iv.
 define void @for_iv_trunc_optimized(ptr %dst) {
 ; CHECK-LABEL: @for_iv_trunc_optimized(
 ; CHECK-NEXT:  bb:
@@ -294,8 +306,8 @@ define void @for_iv_trunc_optimized(ptr %dst) {
 ; CHECK-NEXT:    [[STEP_ADD]] = add <4 x i32> [[VEC_IND]], <i32 4, i32 4, i32 4, i32 4>
 ; CHECK-NEXT:    [[TMP0:%.*]] = shufflevector <4 x i32> [[VECTOR_RECUR]], <4 x i32> [[VEC_IND]], <4 x i32> <i32 3, i32 4, i32 5, i32 6>
 ; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i32> [[VEC_IND]], <4 x i32> [[STEP_ADD]], <4 x i32> <i32 3, i32 4, i32 5, i32 6>
-; CHECK-NEXT:    [[TMP2:%.*]] = or <4 x i32> [[TMP0]], zeroinitializer
-; CHECK-NEXT:    [[TMP3]] = or <4 x i32> [[TMP1]], zeroinitializer
+; CHECK-NEXT:    [[TMP2:%.*]] = or <4 x i32> [[TMP0]], <i32 3, i32 3, i32 3, i32 3>
+; CHECK-NEXT:    [[TMP3]] = or <4 x i32> [[TMP1]], <i32 3, i32 3, i32 3, i32 3>
 ; CHECK-NEXT:    [[TMP5:%.*]] = shufflevector <4 x i32> [[TMP2]], <4 x i32> [[TMP3]], <4 x i32> <i32 3, i32 4, i32 5, i32 6>
 ; CHECK-NEXT:    [[TMP6:%.*]] = extractelement <4 x i32> [[TMP5]], i32 3
 ; CHECK-NEXT:    store i32 [[TMP6]], ptr [[DST:%.*]], align 4
@@ -316,7 +328,7 @@ define void @for_iv_trunc_optimized(ptr %dst) {
 ; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ [[ADD:%.*]], [[LOOP]] ], [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ]
 ; CHECK-NEXT:    [[FOR_1:%.*]] = phi i32 [ [[TRUNC:%.*]], [[LOOP]] ], [ [[SCALAR_RECUR_INIT]], [[SCALAR_PH]] ]
 ; CHECK-NEXT:    [[FOR_2:%.*]] = phi i32 [ [[OR:%.*]], [[LOOP]] ], [ [[SCALAR_RECUR_INIT4]], [[SCALAR_PH]] ]
-; CHECK-NEXT:    [[OR]] = or i32 [[FOR_1]], 0
+; CHECK-NEXT:    [[OR]] = or i32 [[FOR_1]], 3
 ; CHECK-NEXT:    [[ADD]] = add i64 [[IV]], 1
 ; CHECK-NEXT:    store i32 [[FOR_2]], ptr [[DST]], align 4
 ; CHECK-NEXT:    [[ICMP:%.*]] = icmp ult i64 [[IV]], 337
@@ -329,10 +341,10 @@ bb:
   br label %loop
 
 loop:
-  %iv = phi i64 [ %add, %loop ], [ 1, %bb ]
-  %for.1 = phi i32 [ %trunc, %loop ], [ 1, %bb ]
-  %for.2 = phi i32 [ %or, %loop ], [ 0, %bb ]
-  %or = or i32 %for.1, 0
+  %iv = phi i64 [ 1, %bb ], [ %add, %loop ]
+  %for.1 = phi i32 [ 1, %bb ], [ %trunc, %loop ]
+  %for.2 = phi i32 [ 0, %bb ], [ %or, %loop ]
+  %or = or i32 %for.1, 3
   %add = add i64 %iv, 1
   store i32 %for.2, ptr %dst, align 4
   %icmp = icmp ult i64 %iv, 337

llvm/test/Transforms/LoopVectorize/first-order-recurrence-chains-vplan.ll

@@ -149,7 +149,7 @@ exit:
 ; This test has two FORs (for.x and for.y) where incoming value from the previous
 ; iteration (for.x.prev) of one FOR (for.y) depends on another FOR (for.x).
 ; Sinking would require moving a recipe with side effects (store). Instead,
-; for.x.prev can be hoisted.
+; for.x.next can be hoisted.
 define i32 @test_chained_first_order_recurrences_4(ptr %base, i64 %x) {
 ; CHECK-LABEL: 'test_chained_first_order_recurrences_4'
 ; CHECK:      VPlan 'Initial VPlan for VF={4},UF>=1' {
@@ -167,10 +167,10 @@ define i32 @test_chained_first_order_recurrences_4(ptr %base, i64 %x) {
 ; CHECK-NEXT:     FIRST-ORDER-RECURRENCE-PHI ir<%for.x> = phi ir<0>, ir<%for.x.next>
 ; CHECK-NEXT:     FIRST-ORDER-RECURRENCE-PHI ir<%for.y> = phi ir<0>, ir<%for.x.prev>
 ; CHECK-NEXT:     vp<[[SCALAR_STEPS:%.+]]> = SCALAR-STEPS vp<[[CAN_IV]]>, ir<1>
+; CHECK-NEXT:     CLONE ir<%gep> = getelementptr ir<%base>, vp<[[SCALAR_STEPS]]>
 ; CHECK-NEXT:     EMIT vp<[[SPLICE_X:%.]]> = first-order splice ir<%for.x>, ir<%for.x.next>
 ; CHECK-NEXT:     WIDEN-CAST ir<%for.x.prev> = trunc  vp<[[SPLICE_X]]> to i32
 ; CHECK-NEXT:     EMIT vp<[[SPLICE_Y:%.+]]> = first-order splice ir<%for.y>, ir<%for.x.prev>
-; CHECK-NEXT:     CLONE ir<%gep> = getelementptr ir<%base>, vp<[[SCALAR_STEPS]]>
 ; CHECK-NEXT:     WIDEN-CAST ir<%for.y.i64> = sext  vp<[[SPLICE_Y]]> to i64
 ; CHECK-NEXT:     vp<[[VEC_PTR:%.+]]> = vector-pointer ir<%gep>
 ; CHECK-NEXT:     WIDEN store vp<[[VEC_PTR]]>, ir<%for.y.i64>

llvm/test/Transforms/LoopVectorize/first-order-recurrence-multiply-recurrences.ll

@@ -382,7 +382,10 @@ exit:
   ret void
 }
 
-define void @hoist_previous_value_and_operands(ptr %dst, i64 %mask) {
+; Similar to the test cases for https://github.com/llvm/llvm-project/issues/106523.
+; The previous truncation (%trunc) gets vectorized (rather than folded into an
+; IV) and hoisted along with its AND operand above the user 'or'.
+define void @hoist_previous_value_and_operand(ptr %dst, i64 %mask) {
 ; CHECK-LABEL: @hoist_previous_value_and_operands(
 ; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    br i1 false, label [[SCALAR_PH:%.*]], label [[VECTOR_PH:%.*]]
@@ -397,13 +400,13 @@ define void @hoist_previous_value_and_operands(ptr %dst, i64 %mask) {
 ; CHECK-NEXT:    [[VECTOR_RECUR1:%.*]] = phi <4 x i32> [ <i32 poison, i32 poison, i32 poison, i32 0>, [[VECTOR_PH]] ], [ [[TMP4:%.*]], [[VECTOR_BODY]] ]
 ; CHECK-NEXT:    [[OFFSET_IDX:%.*]] = add i64 1, [[INDEX]]
 ; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[OFFSET_IDX]], 0
+; CHECK-NEXT:    [[TMP6:%.*]] = getelementptr inbounds i32, ptr [[DST:%.*]], i64 [[TMP0]]
+; CHECK-NEXT:    [[TMP7:%.*]] = getelementptr inbounds i32, ptr [[TMP6]], i32 0
 ; CHECK-NEXT:    [[TMP1:%.*]] = and <4 x i64> [[VEC_IND]], [[BROADCAST_SPLAT]]
 ; CHECK-NEXT:    [[TMP2]] = trunc <4 x i64> [[TMP1]] to <4 x i32>
 ; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <4 x i32> [[VECTOR_RECUR]], <4 x i32> [[TMP2]], <4 x i32> <i32 3, i32 4, i32 5, i32 6>
 ; CHECK-NEXT:    [[TMP4]] = or <4 x i32> [[TMP3]], zeroinitializer
 ; CHECK-NEXT:    [[TMP5:%.*]] = shufflevector <4 x i32> [[VECTOR_RECUR1]], <4 x i32> [[TMP4]], <4 x i32> <i32 3, i32 4, i32 5, i32 6>
-; CHECK-NEXT:    [[TMP6:%.*]] = getelementptr inbounds i32, ptr [[DST:%.*]], i64 [[TMP0]]
-; CHECK-NEXT:    [[TMP7:%.*]] = getelementptr inbounds i32, ptr [[TMP6]], i32 0
 ; CHECK-NEXT:    store <4 x i32> [[TMP5]], ptr [[TMP7]], align 4
 ; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 4
 ; CHECK-NEXT:    [[VEC_IND_NEXT]] = add <4 x i64> [[VEC_IND]], <i64 4, i64 4, i64 4, i64 4>
@@ -437,10 +440,10 @@ bb:
   br label %loop
 
 loop:
-  %iv = phi i64 [ %add, %loop ], [ 1, %bb ]
-  %for.1 = phi i32 [ %trunc, %loop ], [ 1, %bb ]
-  %for.2 = phi i32 [ %or, %loop ], [ 0, %bb ]
-  %or = or i32 %for.1, 0
+  %iv = phi i64 [ 1, %bb ], [ %add, %loop ]
+  %for.1 = phi i32 [ 1, %bb ], [ %trunc, %loop ]
+  %for.2 = phi i32 [ 0, %bb ], [ %or, %loop ]
+  %or = or i32 %for.1, 3
   %add = add i64 %iv, 1
   %gep = getelementptr inbounds i32, ptr %dst, i64 %iv
   store i32 %for.2, ptr %gep, align 4