[LV] Don't sink into replication regions

[dmgreen]

Summary

The new test case here contains a first order recurrences and an instruction that is replicated. The first order recurrence forces an instruction to be sunk _into_, as opposed to after the replication region. That causes several things to go wrong including registering vector instructions multiple times and failing to create dominance relations correctly.

Instead we should be sinking to after the replication region, which is what this patch makes sure happens.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8610

can this just be Sink->moveBefore(NextBlock->getFirstInsertionPt())? With that, Sink->removeFromParent() shouldn't be needed?

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

650

It's probably worth to check the IR for the whole vector body for the sink-after case.

664

Can you also add another instruction that will get predicate? I think currently things will work out as expected, because we predicate each instruction separately, but it could lead to problems if we predicate multiple instructions together and cannot sink to the right position. But this will get re-worked as part of the VPlan transition, so I think a test case to catch this should be sufficient.

[dmgreen]

Thanks.

Add a VPRecipeBase::moveBefore and use it and getFirstNonPhi.
Added an extra test with predicated stores and removed some undefs.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8610

Oh right, moveBefore was still missing. Thanks for adding it. I'm not sure if there's a benefit of having the variant that takes a basic block & an iterator? For consistency with the existing functions, I think it would be slightly better to just have the version that moves before a VPRecipeBase (same as moveAfter)? Those versions are easier & more convenient to use, IMO.

llvm/lib/Transforms/Vectorize/VPlan.h

670 ↗

(On Diff #314090)

Could you also add a unit test to VPlanTEst.cpp?

[dmgreen]

Added a unit test, that caught that the Parent was not set correctly.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8610

It was added because the iterator might be end(), in which case it's not possible to tell which block it originally referred to. (The block might be empty, that's why it can't use moveAfter). It is modelled after this equivalent method, from BasicBlock/Instruction: https://github.com/llvm/llvm-project/blob/master/llvm/include/llvm/IR/Instruction.h#L142

[fhahn]

In D93629#2481794, @dmgreen wrote:

Added a unit test, that caught that the Parent was not set correctly.

Great thanks!

LGTM

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8610

oh right, because there are no terminator instructions in a VPlan BB. That's a bit unfortunate but not much we can do in that case. Thanks

[dmgreen]

Thanks Florian.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8602

Thanks for fixing!
Note that in general, sinking an instruction more than originally asked for, might place it after its use; in this case, if its use were to belong to the same replication region - potentially possible in the future.
Is it possible for NextBlock to be another replicating Region rather than a BasicBlock?
Going forward, such legality-based scheduling motions should probably be applied earlier, to an initial Region-free VF=UF=1 VPlan.

[dmgreen]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8602

Is it possible for NextBlock to be another replicating Region rather than a BasicBlock?

No, Not currently. There is always an empty VPBasicBlock between one replication region and the next. That is what sink_into_replication_region_multiple is testing, along with the two replication regions not being combined into a single block. We can't (currently) sink past the wrong instruction, as far as I understand, because the replication region is not combined into anything else and we are only sinking past that region, not anything else.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

VPlanPtr LoopVectorizationPlanner::buildVPlanWithVPRecipes(
   // Transform initial VPlan: Apply previously taken decisions, in order, to
   // bring the VPlan to its final state.
   // ---------------------------------------------------------------------------
 
   // Apply Sink-After legal constraints.
   for (auto &Entry : SinkAfter) {
     VPRecipeBase *Sink = RecipeBuilder.getRecipe(Entry.first);
     VPRecipeBase *Target = RecipeBuilder.getRecipe(Entry.second);
+    // If the target is in a replication region, make sure to move Sink to the
+    // block after it, not into the replication region itself.
+    if (auto *Region =
+            dyn_cast_or_null<VPRegionBlock>(Target->getParent()->getParent())) {
+      if (Region->isReplicator()) {
+        assert(Region->getNumSuccessors() == 1 && "Expected SESE region!");
+        VPBasicBlock *NextBlock =
+            cast<VPBasicBlock>(Region->getSuccessors().front());
+        Sink->removeFromParent();
+        NextBlock->insert(Sink, NextBlock->begin());
+        continue;
+      }
+    }
     Sink->moveAfter(Target);
   }
 
   // Interleave memory: for each Interleave Group we marked earlier as relevant
   // for this VPlan, replace the Recipes widening its memory instructions with a
   // single VPInterleaveRecipe at its insertion point.
   for (auto IG : InterleaveGroups) {
     auto *Recipe = cast<VPWidenMemoryInstructionRecipe>(

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

for.cond:
   %iv.next = add i16 %iv, 1
   %rec.2.prev = zext i16 %iv.next to i32
   %rec.1.prev = add i16 %iv.next, 5
   br i1 %cmp, label %for.end, label %for.cond
 
 for.end:
   ret void
 }
+
+; CHECK-LABEL: sink_into_replication_region
+; CHECK: <4 x i32>
+define i32 @sink_into_replication_region() {
+bb:
+  br label %bb2
+
+  bb1:                                              ; preds = %bb2
+  %tmp = phi i32 [ %tmp6, %bb2 ]
+  ret i32 %tmp
+
+  bb2:                                              ; preds = %bb2, %bb
+  %tmp3 = phi i32 [ %tmp8, %bb2 ], [ undef, %bb ]
+  %tmp4 = phi i32 [ %tmp7, %bb2 ], [ undef, %bb ]
+  %tmp5 = phi i32 [ %tmp6, %bb2 ], [ undef, %bb ]
+  %tmp6 = add i32 %tmp5, %tmp4
+  %tmp7 = udiv i32 219220132, %tmp3
+  %tmp8 = add nsw i32 %tmp3, -1
+  %tmp9 = icmp slt i32 %tmp3, 2
+  br i1 %tmp9, label %bb1, label %bb2, !prof !2
+}
+
+!2 = !{!"branch_weights", i32 1, i32 1}