[LoopVectorize] Fix incorrect order of invariant stores when there are multiple reductions.

[igor.kirillov]

Summary

When a loop has multiple reductions, each with an intermediate invariant
store, the order in which those reductions are processed is not considered.
This can result in the invariant stores outside the loop not preserving the
original order.
This patch sorts VPReductionPHIRecipes by the order in which they have
stores in the original loop before running
InnerLoopVectorizer::fixReduction function, and it helps to maintain
the correct order of stores.

Fixes https://github.com/llvm/llvm-project/issues/64047

Depends on D157630

[david-arm]

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

582

Is it worth also having a test where you're storing to different pointers, given that they could alias? We should still maintain the order in that case too.

[igor.kirillov]

Added tests with different store pointers

[igor.kirillov]

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

582

Pre-committed two tests

[david-arm]

LGTM! Thanks for the fix @igor.kirillov. I think this patch is good - we're not reordering any PHI values themselves in the vector loop, but rather the patch-up code in the middle block.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

void InnerLoopVectorizer::fixCrossIterationPHIs(VPTransformState &State) {
   // In order to support recurrences we need to be able to vectorize Phi nodes.
   // Phi nodes have cycles, so we need to vectorize them in two stages. This is
   // stage #2: We now need to fix the recurrences by adding incoming edges to
   // the currently empty PHI nodes. At this point every instruction in the
   // original loop is widened to a vector form so we can use them to construct
   // the incoming edges.
   VPBasicBlock *Header =
       State.Plan->getVectorLoopRegion()->getEntryBasicBlock();
+
+  // Gather all VPReductionPHIRecipe and sort them so that Intermediate stores
+  // sank outside of the loop would keep the same order as they had in the
+  // original loop.
+  SmallVector<VPReductionPHIRecipe *> ReductionPHIList;
   for (VPRecipeBase &R : Header->phis()) {
     if (auto *ReductionPhi = dyn_cast<VPReductionPHIRecipe>(&R))
-      fixReduction(ReductionPhi, State);
-    else if (auto *FOR = dyn_cast<VPFirstOrderRecurrencePHIRecipe>(&R))
+      ReductionPHIList.emplace_back(ReductionPhi);
+  }
+  stable_sort(ReductionPHIList, [this](const VPReductionPHIRecipe *R1,
+                                       const VPReductionPHIRecipe *R2) {
+    auto *IS1 = R1->getRecurrenceDescriptor().IntermediateStore;
+    auto *IS2 = R2->getRecurrenceDescriptor().IntermediateStore;
+
+    // If neither of the recipes has an intermediate store, keep the order the
+    // same.
+    if (!IS1 && !IS2)
+      return false;
+
+    // If only one of the recipes has an intermediate store, then move it
+    // towards the beginning of the list.
+    if (IS1 && !IS2)
+      return true;
+
+    if (!IS1 && IS2)
+      return false;
+
+    // If both recipes have an intermediate store, then the recipe with the
+    // later store should be processed earlier. So it should go to the beginning
+    // of the list.
+    return DT->dominates(IS2, IS1);
+  });
+
+  for (VPReductionPHIRecipe *ReductionPhi : ReductionPHIList)
+
+  for (VPRecipeBase &R : Header->phis()) {
+    if (auto *FOR = dyn_cast<VPFirstOrderRecurrencePHIRecipe>(&R))
       fixFixedOrderRecurrence(FOR, State);
   }
 }
 
 void InnerLoopVectorizer::fixFixedOrderRecurrence(
     VPFirstOrderRecurrencePHIRecipe *PhiR, VPTransformState &State) {
   // This is the second phase of vectorizing first-order recurrences. An
   // overview of the transformation is described below. Suppose we have the

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

for.body:                                         ; preds = %for.body, %entry
   %exitcond = icmp eq i64 %iv.next, 1000
   br i1 %exitcond, label %exit, label %for.body
 
 exit:                                             ; preds = %for.body
   %add.lcssa = phi i32 [ %add, %for.body ]
   ret i32 %add.lcssa
 }
 
-; FIXME: This tests currently shows incorrect behavior and it will fixed in the following patch
 ; Make sure that if there are several reductions in the loop, the order of invariant stores sank outside of the loop is preserved
 define void @reduc_add_mul_store(ptr %dst, ptr readonly %src) {
 ; CHECK-LABEL: define void @reduc_add_mul_store
 ; CHECK:       middle.block:
-; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @llvm.vector.reduce.mul.v4i32(<4 x i32> [[TMP1:%.*]])
+; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> [[TMP1:%.*]])
 ; CHECK-NEXT:    store i32 [[TMP2]], ptr %dst, align 4
-; CHECK-NEXT:    [[TMP4:%.*]] = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> [[TMP3:%.*]])
+; CHECK-NEXT:    [[TMP4:%.*]] = call i32 @llvm.vector.reduce.mul.v4i32(<4 x i32> [[TMP3:%.*]])
 ; CHECK-NEXT:    store i32 [[TMP4]], ptr %dst, align 4
 ;
 entry:
   br label %for.body
 
 for.body:
   %sum = phi i32 [ 0, %entry ], [ %sum.next, %for.body ]
   %mul = phi i32 [ 1, %entry ], [ %mul.next, %for.body ]
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
   %gep.src = getelementptr inbounds i32, ptr %src, i64 %iv
   %0 = load i32, ptr %gep.src, align 4
   %sum.next = add nsw i32 %sum, %0
   store i32 %sum.next, ptr %dst, align 4
   %mul.next = mul nsw i32 %mul, %0
   store i32 %mul.next, ptr %dst, align 4
   %iv.next = add nuw nsw i64 %iv, 1
   %exitcond = icmp eq i64 %iv.next, 1000
   br i1 %exitcond, label %exit, label %for.body
 
 exit:
   ret void