Relax constraints for reduction vectorization

[sanjoy]

Summary

I'm somewhat unsure here, but gating vectorizing reductions on all
fastmath flags seems unnecessary; reassoc should be sufficient.

[sanjoy]

ping!

[sdesmalen]

lib/Analysis/IVDescriptors.cpp

550

Is there a reason for the condition I->hasAllowContract()? As far as I can tell only hasAllowReassoc() seems required and thus hasAllowContract() would be overly restrictive.

As long as hasAllowReassoc() is true, I think the resulting reduction operations should have the same properties as the original instruction. For example, if an instruction has hasNoNaNs() == hasNoInfs() == false, the vectorised reduction retain those properties under reassocation.

When trying out this patch, the reduction block seems to assume fast instead of just reassoc, e.g.

middle.block:
%bin.rdx = fadd fast <4 x float> %9, %8

You'll probably want to fix that and extend the test to ensure the flags are retained.

[sanjoy]

Propagate only correct FP fast-math flags.

I'd appreciate a careful review here since I'm not sure all of the ways we
codegen reductions in the loop vectorizer.

[sanjoy]

Rebase on trunk.

[sdesmalen]

Thanks for these changes to your patch @sanjoy! The patch is looking in good shape, just a few remaining nits mostly.

include/llvm/Analysis/IVDescriptors.h

244 ↗

(On Diff #189272)

nit: recurrenct -> recurrent

lib/Analysis/IVDescriptors.cpp

254

Is it worth adding a comment describing that FMF will be an intersection of the FastMathFlags from all the reduction operations (and thus needs to start with the full set of flags)?

299

nit: unnecessary curly braces

558

nit: since 'CanVectorizeReduction()' is only used once, it probably makes more sense to just expand it here and remove the function.

lib/CodeGen/ExpandReductions.cpp

124 ↗

(On Diff #189272)

Rather than choosing 'getFast() and going through 'IsOrdered' (which is set when the full 'fast' property is set), do we want to take the FastMathFlags from the IntrinsicInst directly?

124 ↗

(On Diff #189272)

Note that there is a test missing for this change.

lib/Transforms/Utils/LoopUtils.cpp

675 ↗

(On Diff #189272)

nit: unnecessary curly braces.

lib/Transforms/Vectorize/LoopVectorize.cpp

322 ↗

(On Diff #189272)

nit: unnecessary curly braces.

329 ↗

(On Diff #189272)

nit: unnecessary curly braces.

test/Transforms/LoopVectorize/reduction-fastmath.ll

49

Do you want to add a check for the 'fast' attribute on the reductions in the middle.block, similar to what you've done for reduction_sum_float_only_reassoc_and_contract ?

[sanjoy]

• Address review comments

lib/CodeGen/ExpandReductions.cpp

124 ↗

(On Diff #189272)

Rather than choosing 'getFast() and going through 'IsOrdered' (which is set when the full 'fast' property is set), do we want to take the FastMathFlags from the IntrinsicInst directly?

My rationale was that I wanted to keep this as obviously NFC as possible -- previously we would implicitly tag the instructions as fast, and this change just makes it explicit.

Note that there is a test missing for this change.

This change is not supposed to change any behavior (assuming no bugs :) ). Do you want me to add a test case to check existing behavior? CodeGen/Generic/expand-experimental-reductions.ll already tests for fast flags in the expansion.

[sdesmalen]

Thanks for updating your patch @sanjoy ! LGTM (with a side-note on a follow up patch to take more specific flags in expandReductions).
You may want to update the commit message before you commit :)

lib/CodeGen/ExpandReductions.cpp

124 ↗

(On Diff #189272)

My rationale was that I wanted to keep this as obviously NFC as possible -- previously we would implicitly tag the instructions as fast, and this change just makes it explicit.

In a way that feels a bit artificial since the patch is already not NFC and the purpose of this patch is to be more specific about passing the exact flags. However, if you think it makes more sense to do this change in a separate patch, that's fine with me.

[sanjoy]

lib/CodeGen/ExpandReductions.cpp

124 ↗

(On Diff #189272)

As I said above, we already have tests for this behavior (that we propagate fast-math flags into getShuffleReduction.

lib/Analysis/IVDescriptors.cpp

if (!isIntegerRecurrenceKind(Kind))
       return false;
     if (isArithmeticRecurrenceKind(Kind))
       Start = lookThroughAnd(Phi, RecurrenceType, VisitedInsts, CastInsts);
   }
 
   Worklist.push_back(Start);
   VisitedInsts.insert(Start);
 
   // A value in the reduction can be used:
   //  - By the reduction:
   //      - Reduction operation:
   //        - One use of reduction value (safe).
   //        - Multiple use of reduction value (not safe).
   //      - PHI:
   //        - All uses of the PHI must be the reduction (safe).
   //        - Otherwise, not safe.
   //  - By instructions outside of the loop (safe).
while (!Worklist.empty()) {
 
     // Any reduction instruction must be of one of the allowed kinds. We ignore
     // the starting value (the Phi or an AND instruction if the Phi has been
     // type-promoted).
     if (Cur != Start) {
       ReduxDesc = isRecurrenceInstr(Cur, Kind, ReduxDesc, HasFunNoNaNAttr);
       if (!ReduxDesc.isRecurrence())
         return false;
     }
 
     bool IsASelect = isa<SelectInst>(Cur);
 
     // A conditional reduction operation must only have 2 or less uses in
     // VisitedInsts.
     if (IsASelect && (Kind == RK_FloatAdd || Kind == RK_FloatMult) &&
         hasMultipleUsesOf(Cur, VisitedInsts, 2))
       return false;
if ((m_FAdd(m_Value(Op1), m_Value(Op2)).match(I1)  ||
     return InstDesc(Kind == RK_FloatAdd, SI);
 
   if (m_FMul(m_Value(Op1), m_Value(Op2)).match(I1) && (I1->isFast()))
     return InstDesc(Kind == RK_FloatMult, SI);
 
   return InstDesc(false, I);
 }
 
+static bool CanVectorizeReduction(Instruction *I) {
+  return I->hasAllowReassoc() && I->hasAllowContract();
+}
+
 RecurrenceDescriptor::InstDesc
 RecurrenceDescriptor::isRecurrenceInstr(Instruction *I, RecurrenceKind Kind,
                                         InstDesc &Prev, bool HasFunNoNaNAttr) {
   bool FP = I->getType()->isFloatingPointTy();
   Instruction *UAI = Prev.getUnsafeAlgebraInst();
-  if (!UAI && FP && !I->isFast())
+  if (!UAI && FP && !CanVectorizeReduction(I))
     UAI = I; // Found an unsafe (unvectorizable) algebra instruction.
 
   switch (I->getOpcode()) {
   default:
     return InstDesc(false, I);
   case Instruction::PHI:
     return InstDesc(I, Prev.getMinMaxKind(), Prev.getUnsafeAlgebraInst());
   case Instruction::Sub:

test/Transforms/LoopVectorize/reduction-fastmath.ll

+; RUN: opt -S -loop-vectorize < %s | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define float @reduction_sum_float_ieee(i32 %n, float* %array) {
+; CHECK-LABEL: define float @reduction_sum_float_ieee(
+entry:
+  %entry.cond = icmp ne i32 0, 4096
+  br i1 %entry.cond, label %loop, label %loop.exit
+
+loop:
+  %idx = phi i32 [ 0, %entry ], [ %idx.inc, %loop ]
+  %sum = phi float [ 0.000000e+00, %entry ], [ %sum.inc, %loop ]
+  %address = getelementptr float, float* %array, i32 %idx
+  %value = load float, float* %address
+  %sum.inc = fadd float %sum, %value
+  %idx.inc = add i32 %idx, 1
+  %be.cond = icmp ne i32 %idx.inc, 4096
+  br i1 %be.cond, label %loop, label %loop.exit
+
+loop.exit:
+  %sum.lcssa = phi float [ %sum.inc, %loop ], [ 0.000000e+00, %entry ]
+; CHECK-NOT: %wide.load = load <4 x float>, <4 x float>*
+; CHECK: ret float %sum.lcssa
+  ret float %sum.lcssa
+}
+
+define float @reduction_sum_float_fastmath(i32 %n, float* %array) {
+; CHECK-LABEL: define float @reduction_sum_float_fastmath(
+entry:
+  %entry.cond = icmp ne i32 0, 4096
+  br i1 %entry.cond, label %loop, label %loop.exit
+
+loop:
+  %idx = phi i32 [ 0, %entry ], [ %idx.inc, %loop ]
+  %sum = phi float [ 0.000000e+00, %entry ], [ %sum.inc, %loop ]
+  %address = getelementptr float, float* %array, i32 %idx
+  %value = load float, float* %address
+  %sum.inc = fadd fast float %sum, %value
+  %idx.inc = add i32 %idx, 1
+  %be.cond = icmp ne i32 %idx.inc, 4096
+  br i1 %be.cond, label %loop, label %loop.exit
+
+loop.exit:
+  %sum.lcssa = phi float [ %sum.inc, %loop ], [ 0.000000e+00, %entry ]
+; CHECK: %wide.load = load <4 x float>, <4 x float>*
+; CHECK: ret float %sum.lcssa
+  ret float %sum.lcssa
+}
+
+define float @reduction_sum_float_partial_fastmath(i32 %n, float* %array) {
+; CHECK-LABEL: define float @reduction_sum_float_partial_fastmath(
+entry:
+  %entry.cond = icmp ne i32 0, 4096
+  br i1 %entry.cond, label %loop, label %loop.exit
+
+loop:
+  %idx = phi i32 [ 0, %entry ], [ %idx.inc, %loop ]
+  %sum = phi float [ 0.000000e+00, %entry ], [ %sum.inc, %loop ]
+  %address = getelementptr float, float* %array, i32 %idx
+  %value = load float, float* %address
+  %sum.inc = fadd reassoc contract float %sum, %value
+  %idx.inc = add i32 %idx, 1
+  %be.cond = icmp ne i32 %idx.inc, 4096
+  br i1 %be.cond, label %loop, label %loop.exit
+
+loop.exit:
+  %sum.lcssa = phi float [ %sum.inc, %loop ], [ 0.000000e+00, %entry ]
+; CHECK: %wide.load = load <4 x float>, <4 x float>*
+; CHECK: ret float %sum.lcssa
+  ret float %sum.lcssa
+}