[LV] Mark first-order recurrences as allowed exit

[Ayal]

Summary

First-order recurrences require special treatment when they are live-out; such treatment is provided by fixFirstOrderRecurrence(), so they should be included in AllowedExit set. (Should probably have been included originally in D16197.)

Fixes PR45526: AllowedExit set is used by prepareToFoldTailByMasking() to check whether the treatment for live-outs also holds when folding the tail, which is currently the case for reductions but not (yet) for first-order recurrences or other values.

[skatkov]

I'm not sure how correct this fix is but I verified that it fixes my original problem.

[fhahn]

LGTM, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

671

nit: it seems a bit counter-intuitive to add to something called AllowedExit to reject loops with users of Phi outside the loop. There might be potential for re-naming/refactoring to make it clearer in future. In the meantime it might make sense to add a comment.

llvm/test/Transforms/LoopVectorize/optsize.ll

129

This seems like a very subtle way to check vectorization is not happening. Given that the function is small, it might be better to check the IR for the whole function (or just the control-flow)

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

671

Not sure what comment may be worth adding where. AllowedExit was introduced originally to record all exits that can be vectorized, and first-order-recurring phi's should have been included in it when support for vectorizing them was first introduced (in D16197, as noted in the summary).
AllowedExit is examined by prepareToFoldTail because only a subset of generally vectorizable exits are also supported when vectorizing with foldTail. Perhaps rename AllowedExit to AllowedGeneralExit ?
Above TODO was added to revisit the use of AllowedExit in general.

llvm/test/Transforms/LoopVectorize/optsize.ll

129

ok, replacing this CHECK-NOT with a positive check that the IR remains intact, which is small enough in this case.
Other tests check non-vectorization by CHECK_NOT's of specific vector types or of vector.body.

[Ayal]

Addressed comment regarding test; rebased.

[fhahn]

LGTM, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

671

Above TODO was added to revisit the use of AllowedExit in general.

Right, I think the existing TODO is fine.

[Ayal]

In D78210#1985815, @skatkov wrote:

I'm not sure how correct this fix is but I verified that it fixes my original problem.

Thanks @skatkov; patch committed, PR45526 closed.

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

for (Instruction &I : *BB) {
           addInductionPhi(Phi, ID, AllowedExit);
           if (ID.hasUnsafeAlgebra() && !HasFunNoNaNAttr)
             Requirements->addUnsafeAlgebraInst(ID.getUnsafeAlgebraInst());
           continue;
         }
 
         if (RecurrenceDescriptor::isFirstOrderRecurrence(Phi, TheLoop,
                                                          SinkAfter, DT)) {
+          AllowedExit.insert(Phi);
           FirstOrderRecurrences.insert(Phi);
           continue;
         }
 
         // As a last resort, coerce the PHI to a AddRec expression
         // and re-try classifying it a an induction PHI.
         if (InductionDescriptor::isInductionPHI(Phi, TheLoop, PSE, ID, true)) {
           addInductionPhi(Phi, ID, AllowedExit);

llvm/test/Transforms/LoopVectorize/optsize.ll

for.body29:
   %idxprom34 = zext i16 %add33 to i32
   %arrayidx35 = getelementptr [2592 x i16], [2592 x i16] * @cm_array, i32 0, i32 %idxprom34
   store i16 0, i16 * %arrayidx35, align 1
   %inc37 = add i16 %i24.0170, 1
   %cmp26 = icmp ult i16 %inc37, 756
   br i1 %cmp26, label %for.body29, label %for.cond.cleanup28
 }
 
+; PR45526: don't vectorize with fold-tail if first-order-recurrence is live-out.
+;
+define i32 @pr45526() optsize {
+;
+; CHECK-LABEL: @pr45526
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   br label %loop
+; CHECK-EMPTY:
+; CHECK-NEXT: loop:
+; CHECK-NEXT:   %piv = phi i32 [ 0, %entry ], [ %pivPlus1, %loop ]
+; CHECK-NEXT:   %for = phi i32 [ 5, %entry ], [ %pivPlus1, %loop ]
+; CHECK-NEXT:   %pivPlus1 = add nuw nsw i32 %piv, 1
+; CHECK-NEXT:   %cond = icmp ult i32 %piv, 510
+; CHECK-NEXT:   br i1 %cond, label %loop, label %exit
+; CHECK-EMPTY:
+; CHECK-NEXT: exit:
+; CHECK-NEXT:   %for.lcssa = phi i32 [ %for, %loop ]
+; CHECK-NEXT:   ret i32 %for.lcssa
+;
+entry:
+  br label %loop
+
+loop:
+  %piv = phi i32 [ 0, %entry ], [ %pivPlus1, %loop ]
+  %for = phi i32 [ 5, %entry ], [ %pivPlus1, %loop ]
+  %pivPlus1 = add nuw nsw i32 %piv, 1
+  %cond = icmp ult i32 %piv, 510
+  br i1 %cond, label %loop, label %exit
+
+exit:
+  ret i32 %for
+}
+
 !llvm.module.flags = !{!0}
 !0 = !{i32 1, !"ProfileSummary", !1}
 !1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
 !2 = !{!"ProfileFormat", !"InstrProf"}
 !3 = !{!"TotalCount", i64 10000}
 !4 = !{!"MaxCount", i64 10}
 !5 = !{!"MaxInternalCount", i64 1}
 !6 = !{!"MaxFunctionCount", i64 1000}
 !7 = !{!"NumCounts", i64 3}
 !8 = !{!"NumFunctions", i64 3}
 !9 = !{!"DetailedSummary", !10}
 !10 = !{!11, !12, !13}
 !11 = !{i32 10000, i64 100, i32 1}
 !12 = !{i32 999000, i64 100, i32 1}
 !13 = !{i32 999999, i64 1, i32 2}
 !14 = !{!"function_entry_count", i64 0}