[VPlan] Disallow sinking side-effecting first order recurrence users

[ebrevnov]

Summary

We started seeing new failure after D142886. Looks like it enabled new cases and we hit an assert:
assert(Current->getNumDefinedValues() == 1 &&

"only recipes with a single defined value expected");

When we do instruction sinking for the first order recurrence we hit an assert if instruction doesn't have single def. In case instruction doesn't produce any new def there is no new users and nothing to sink.

[fhahn]

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

682

I *think* this would not be safe in all cases; e.g. if we have a store here than we would need to sink a store, which may not be legal if there are other loads in the loop and we would sink past them (would be good to add a test case)

I *think* to start with, it would be good to prevent sinking instructions with side-effects in TryPushSinkCandidate.

There's a patch series that moves memory legality checking to VPlan (D142895 and following) which support s sinking some memory operations.

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

112

I *think* it would also be good to check actual codegen for the test in llvm/test/Transforms/LoopVectorize/first-order-recurrence-chains.ll.

151

nit: using loop as block name would be more descriptive.

152

nit: It would be easier to read if the induction phi would be named iv, and the recurrence phis something like for.X.

156

does this need addrspace(1)? Also, would be good to pass in the base pointer as arg, to avoid UB in test.

[ebrevnov]

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

682

My understanding that legality of sinking is checked in isFixedOrderRecurrence. It's checked for every use regardless of number of defs. Thus we should fail to recognize FOR in your case (because it rejects any side-effecting instruction currently). Am I missing something?

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

112

Do you suggest to update CHECKs for all cases inside this file?

151

ok. will fix.

152

ok. will fix.

156

ok. will fix.

[ebrevnov]

ping

[fhahn]

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

682

This is not longer completely true, whether the users of recurrences can be re-ordered or not is checked based on VPlan, which is more powerful. Note that the test case hasn't been vectorized before re-ordering based on VPlan (see https://llvm.godbolt.org/z/8d9PM353T). Until we check for aliasing based on VPlan (patches mentioned in earlier comments), we should bail out when trying to sink recipes that may have side effects.

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

112

No, adding a separate test here llvm/test/Transforms/LoopVectorize/first-order-recurrence-chains.ll, but it is probably not needed, checking the VPlan only is simpler.

[ebrevnov]

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

682

This is not longer completely true, ...

I understand the intent to do legality checking on VPlan, but at the moment it doesn't seem to be fully the case. In particular, there is a check that sink candidate doesn't have side effects in RecurrenceDescriptor::isFixedOrderRecurrence:

if (SinkCandidate->getParent() != PhiBB ||
    SinkCandidate->mayHaveSideEffects() ||
    SinkCandidate->mayReadFromMemory() || SinkCandidate->isTerminator())
  return false;

That means we can't meet side effecting user of FOR in sinkRecurrenceUsersAfterPrevious. So I don't understand why we need another check here.

Maybe, what you mean is to move the check from RecurrenceDescriptor::isFixedOrderRecurrence to sinkRecurrenceUsersAfterPrevious? Could you clarify.

PS: It looks like comment for 'adjustFixedOrderRecurrences' is a bit outdated since it still assumes that legality is checked before.

/// The current implementation assumes all users can be sunk after the previous
/// value, which is enforced by earlier legality checks.

[ebrevnov]

gentle ping...

[fhahn]

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

682

That means we can't meet side effecting user of FOR in sinkRecurrenceUsersAfterPrevious. So I don't understand why we need another check here.

This is true in most cases, but your test case uncovered a scenario where this is not the case. Before sinkRecurrenceUsersAfterPrevious was introduced to handle the re-ordering, the code was unable to reorder the code in the test case due to limitations in the old implementation.

Now we can re-order, but this means we need to check for side effects to catch cases where the original order of instructions doesn't require sinking an instruction with side effects, but does so after sinking operands for a different FOR.

[fhahn]

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

682

Maybe, what you mean is to move the check from RecurrenceDescriptor::isFixedOrderRecurrence to sinkRecurrenceUsersAfterPrevious? Could you clarify.

I think for now we need the checks in both places still, but the checks in RecurrenceDescriptor::isFixedOrderRecurrence will be removed as part of D142895 and following.

[ebrevnov]

Prevent sinking side-effecting instructions

[ebrevnov]

Florian, I've addressed yours concerns. Please take a look now.

[fhahn]

LGTM with inline comments, thanks! first-order-recurrence-multiply-recurrences.ll needs updating, but the change there is fine and expected. We will be able to vectorize this again once FOR splice doesn't need to be considered as having side-effects.

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

668

nit: indentation seems off here, please make sure the patch is formatted properly before landing.

680

For now, I think with the fix above the original assert should be kept and handling for recipes with zero defined values should not be needed. If not, at least the assert will help to surface a test case.

[ebrevnov]

Florian, thank you for review.

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

668

Fixed.

680

Ok, restored assert back.

[fhahn]

Just realized the commit title is not in line with the implementation :(

[ebrevnov]

In D151204#4553881, @fhahn wrote:

Just realized the commit title is not in line with the implementation :(

Indeed, fixed. Thanks!

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

auto TryToPushSinkCandidate = [&](VPRecipeBase *SinkCandidate) {
     if (SinkCandidate == Previous)
       return false;
 
     if (isa<VPHeaderPHIRecipe>(SinkCandidate) ||
         !Seen.insert(SinkCandidate).second ||
         properlyDominates(Previous, SinkCandidate, VPDT))
       return true;
 
+    if (SinkCandidate->mayHaveSideEffects())
+      return false;
+
     WorkList.push_back(SinkCandidate);
     return true;
   };
 
   // Recursively sink users of FOR after Previous.
   WorkList.push_back(FOR);
   for (unsigned I = 0; I != WorkList.size(); ++I) {
     VPRecipeBase *Current = WorkList[I];
     assert(Current->getNumDefinedValues() == 1 &&
            "only recipes with a single defined value expected");
+
     for (VPUser *User : Current->getVPSingleValue()->users()) {
       if (auto *R = dyn_cast<VPRecipeBase>(User))
         if (!TryToPushSinkCandidate(R))
           return false;
     }
   }
 
   // Keep recipes to sink ordered by dominance so earlier instructions are
   // processed first.
   sort(WorkList, [&VPDT](const VPRecipeBase *A, const VPRecipeBase *B) {

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

loop:
   %add.2 = add i16 %add.1, %for.3
   store i16 %add.2, ptr %gep.ptr
   %exitcond.not = icmp eq i64 %iv.next, 1000
   br i1 %exitcond.not, label %exit, label %loop
 
 exit:
   ret void
 }
+
+; 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). Due to
+; this dependency all uses of the former FOR (for.y) should be sunk after
+; incoming value from the previous iteration (for.x.prev) of te latter FOR (for.y).
+; That means side-effecting user (store i64 %for.y.i64, ptr %gep) of the latter
+; FOR (for.y) should be moved which is not currently supported.
+define i32 @test_chained_first_order_recurrences_4(ptr %base) {
+; CHECK-LABEL: 'test_chained_first_order_recurrences_4'
+; CHECK: No VPlan could be built for
+
+entry:
+  br label %loop
+
+ret:
+  ret i32 0
+
+loop:
+  %iv = phi i64 [ %iv.next, %loop ], [ 0, %entry ]
+  %for.x = phi i64 [ %for.x.next, %loop ], [ 0, %entry ]
+  %for.y = phi i32 [ %for.x.prev, %loop ], [ 0, %entry ]
+  %iv.next = add i64 %iv, 1
+  %gep = getelementptr i64, ptr %base, i64 %iv
+  %for.x.prev = trunc i64 %for.x to i32
+  %for.y.i64 = sext i32 %for.y to i64
+  store i64 %for.y.i64, ptr %gep
+  %for.x.next = mul i64 0, 0
+  %icmp = icmp ugt i64 %iv, 4096
+  br i1 %icmp, label %ret, label %loop
+}

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

 
 exit:
   ret void
 }
 
 define void @test_pr54233_for_depend_on_each_other(ptr noalias %a, ptr noalias %b) {
 ; CHECK-LABEL: @test_pr54233_for_depend_on_each_other(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    br i1 false, label [[SCALAR_PH:%.*]], label [[VECTOR_PH:%.*]]
-; CHECK:       vector.ph:
-; CHECK-NEXT:    br label [[VECTOR_BODY:%.*]]
-; CHECK:       vector.body:
-; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
-; CHECK-NEXT:    [[VECTOR_RECUR:%.*]] = phi <4 x i32> [ <i32 poison, i32 poison, i32 poison, i32 0>, [[VECTOR_PH]] ], [ [[TMP4:%.*]], [[VECTOR_BODY]] ]
-; CHECK-NEXT:    [[VECTOR_RECUR1:%.*]] = phi <4 x i32> [ <i32 poison, i32 poison, i32 poison, i32 0>, [[VECTOR_PH]] ], [ [[BROADCAST_SPLAT:%.*]], [[VECTOR_BODY]] ]
-; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[INDEX]], 0
-; CHECK-NEXT:    [[TMP1:%.*]] = load i32, ptr [[B:%.*]], align 4
-; CHECK-NEXT:    [[BROADCAST_SPLATINSERT:%.*]] = insertelement <4 x i32> poison, i32 [[TMP1]], i64 0
-; CHECK-NEXT:    [[BROADCAST_SPLAT]] = shufflevector <4 x i32> [[BROADCAST_SPLATINSERT]], <4 x i32> poison, <4 x i32> zeroinitializer
-; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <4 x i32> [[VECTOR_RECUR1]], <4 x i32> [[BROADCAST_SPLAT]], <4 x i32> <i32 3, i32 4, i32 5, i32 6>
-; CHECK-NEXT:    [[TMP3:%.*]] = or <4 x i32> [[TMP2]], <i32 10, i32 10, i32 10, i32 10>
-; CHECK-NEXT:    [[TMP4]] = xor <4 x i32> <i32 12, i32 12, i32 12, i32 12>, [[TMP2]]
-; CHECK-NEXT:    [[TMP5:%.*]] = shufflevector <4 x i32> [[VECTOR_RECUR]], <4 x i32> [[TMP4]], <4 x i32> <i32 3, i32 4, i32 5, i32 6>
-; CHECK-NEXT:    [[TMP6:%.*]] = shl <4 x i32> [[TMP2]], [[TMP5]]
-; CHECK-NEXT:    [[TMP7:%.*]] = xor <4 x i32> [[TMP6]], <i32 255, i32 255, i32 255, i32 255>
-; CHECK-NEXT:    [[TMP8:%.*]] = and <4 x i32> [[TMP7]], [[TMP3]]
-; CHECK-NEXT:    [[TMP9:%.*]] = getelementptr inbounds i32, ptr [[A:%.*]], i64 [[TMP0]]
-; CHECK-NEXT:    [[TMP10:%.*]] = getelementptr inbounds i32, ptr [[TMP9]], i32 0
-; CHECK-NEXT:    store <4 x i32> [[TMP8]], ptr [[TMP10]], align 4
-; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 4
-; CHECK-NEXT:    [[TMP11:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1000
-; CHECK-NEXT:    br i1 [[TMP11]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP4:![0-9]+]]
-; CHECK:       middle.block:
-; CHECK-NEXT:    [[CMP_N:%.*]] = icmp eq i64 1001, 1000
-; CHECK-NEXT:    [[VECTOR_RECUR_EXTRACT:%.*]] = extractelement <4 x i32> [[TMP4]], i32 3
-; CHECK-NEXT:    [[VECTOR_RECUR_EXTRACT2:%.*]] = extractelement <4 x i32> [[BROADCAST_SPLAT]], i32 3
-; CHECK-NEXT:    br i1 [[CMP_N]], label [[EXIT:%.*]], label [[SCALAR_PH]]
-; CHECK:       scalar.ph:
-; CHECK-NEXT:    [[SCALAR_RECUR_INIT3:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[VECTOR_RECUR_EXTRACT2]], [[MIDDLE_BLOCK]] ]
-; CHECK-NEXT:    [[SCALAR_RECUR_INIT:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[VECTOR_RECUR_EXTRACT]], [[MIDDLE_BLOCK]] ]
-; CHECK-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i64 [ 1000, [[MIDDLE_BLOCK]] ], [ 0, [[ENTRY]] ]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
-; CHECK-NEXT:    [[SCALAR_RECUR:%.*]] = phi i32 [ [[SCALAR_RECUR_INIT]], [[SCALAR_PH]] ], [ [[FOR_1_NEXT:%.*]], [[LOOP]] ]
-; CHECK-NEXT:    [[SCALAR_RECUR4:%.*]] = phi i32 [ [[SCALAR_RECUR_INIT3]], [[SCALAR_PH]] ], [ [[FOR_2_NEXT:%.*]], [[LOOP]] ]
-; CHECK-NEXT:    [[OR:%.*]] = or i32 [[SCALAR_RECUR4]], 10
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[SCALAR_RECUR4]], [[SCALAR_RECUR]]
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[FOR_1:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[FOR_1_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[FOR_2:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[FOR_2_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[OR:%.*]] = or i32 [[FOR_2]], 10
+; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[FOR_2]], [[FOR_1]]
 ; CHECK-NEXT:    [[XOR:%.*]] = xor i32 [[SHL]], 255
 ; CHECK-NEXT:    [[AND:%.*]] = and i32 [[XOR]], [[OR]]
-; CHECK-NEXT:    [[FOR_1_NEXT]] = xor i32 12, [[SCALAR_RECUR4]]
-; CHECK-NEXT:    [[FOR_2_NEXT]] = load i32, ptr [[B]], align 4
-; CHECK-NEXT:    [[A_GEP:%.*]] = getelementptr inbounds i32, ptr [[A]], i64 [[IV]]
+; CHECK-NEXT:    [[FOR_1_NEXT]] = xor i32 12, [[FOR_2]]
+; CHECK-NEXT:    [[FOR_2_NEXT]] = load i32, ptr [[B:%.*]], align 4
+; CHECK-NEXT:    [[A_GEP:%.*]] = getelementptr inbounds i32, ptr [[A:%.*]], i64 [[IV]]
 ; CHECK-NEXT:    store i32 [[AND]], ptr [[A_GEP]], align 4
 ; CHECK-NEXT:    [[IV_NEXT]] = add nuw i64 [[IV]], 1
 ; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp eq i64 [[IV]], 1000
-; CHECK-NEXT:    br i1 [[EXITCOND]], label [[EXIT]], label [[LOOP]], !llvm.loop [[LOOP5:![0-9]+]]
+; CHECK-NEXT:    br i1 [[EXITCOND]], label [[EXIT:%.*]], label [[LOOP]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %loop
 
 loop:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]