[LoopVectorize] Propagate fast-math flags for inloop reductions

[RosieSumpter]

Summary

This patch updates VPReductionRecipe::execute so that the fast-math
flags associated with the underlying instruction of the VPReductionRecipe are
propagated through to the reductions which are created.

Depends on D112547

[RosieSumpter]

Note: this patch builds on the NFC patch D112547

[david-arm]

Thanks for this patch @RosieSumpter - seems like a nice fix! I wonder if it's also worth changing createOrderedReduction to be similar to createTargetReduction and set the flags in there too? Also, you've fixed up the VF=1 cases too.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9752

I completely forgot we can actually just use a guard here, i.e.

IRBuilderBase::FastMathFlagGuard FMFGuard(State.Builder);
State.Builder.setFastMathFlags(RdxDesc->getFastMathFlags());

then you don't need the extra line at the bottom anymore to reset the flags.

9753

Maybe we should do this for all cases here, i.e. ordered and non-ordered?

9755

I just realised that you might be able to replace the whole if block with:

State.Builder.setFastMathFlags(RdxDesc->getFastMathFlags());

here.

9779

I just realised you're actually also fixing the case where VF=1. Is it worth having a test to ensure flags are propagated for VF=1 as well?

[RosieSumpter]

• Use IRBuilderBase::FastMathFlagGuard
• Propagate fast-math flags for both ordered and non-ordered reductions
• Updated tests

[kmclaughlin]

Thank you @RosieSumpter, this patch looks good to me and I think all of the comments have been addressed.

llvm/test/Transforms/LoopVectorize/AArch64/strict-fadd.ll

808

Should fast here be replaced with nnan?

[RosieSumpter]

llvm/test/Transforms/LoopVectorize/AArch64/strict-fadd.ll

808

Yes it should! I will change that before committing.

[david-arm]

llvm/test/Transforms/LoopVectorize/reduction-inloop-pred.ll

1020 ↗

(On Diff #383031)

Nice! Instcombine has folded the call and the fadd together into one!

[paulwalker-arm]

Just adding an extra datapoint to do with as you wish.

llvm/test/Transforms/LoopVectorize/AArch64/strict-fadd.ll

808

Generally speaking it is better for CHECK-NOT lines to be as small/simple as possible so that you do not end up creating a passing test just because some minor detail has changed.

In this instance the CHECK-NOT does not care about the instruction's flags but rather wants to ensure there is never a call to llvm.vector.reduce.fadd. In this regard I think just having CHECK-UNORDERED-NOT: @llvm.vector.reduce.fadd would yield a more resilient test.

That said, given all the expected output is explicitly checked for I'm not sure what value these CHECK-NOT lines provide.

[RosieSumpter]

llvm/test/Transforms/LoopVectorize/AArch64/strict-fadd.ll

808

Thanks @paulwalker-arm, I will change the CHECK-NOT lines. Although I see what you mean about them not having value here, maybe it's best to remove them.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

State.ILV->vectorizeInterleaveGroup(IG, definedValues(), State, getAddr(),
                                       getStoredValues(), getMask());
 }
 
 void VPReductionRecipe::execute(VPTransformState &State) {
   assert(!State.Instance && "Reduction being replicated.");
   Value *PrevInChain = State.get(getChainOp(), 0);
   RecurKind Kind = RdxDesc->getRecurrenceKind();
   bool IsOrdered = State.ILV->useOrderedReductions(*RdxDesc);
+  FastMathFlags Current_FMF = State.Builder.getFastMathFlags();
+  if (IsOrdered) {
+    // Propagate the fast-math flags carried by the underlying instruction.
+    if (auto *FPMO = dyn_cast<FPMathOperator>(getUnderlyingInstr())) {
+      FastMathFlags FMF = FPMO->getFastMathFlags();
+      State.Builder.setFastMathFlags(FMF);
+    }
+  }
   for (unsigned Part = 0; Part < State.UF; ++Part) {
     Value *NewVecOp = State.get(getVecOp(), Part);
     if (VPValue *Cond = getCondOp()) {
       Value *NewCond = State.get(Cond, Part);
       VectorType *VecTy = cast<VectorType>(NewVecOp->getType());
       Value *Iden = RdxDesc->getRecurrenceIdentity(
           Kind, VecTy->getElementType(), RdxDesc->getFastMathFlags());
       Value *IdenVec =
           State.Builder.CreateVectorSplat(VecTy->getElementCount(), Iden);
       Value *Select = State.Builder.CreateSelect(NewCond, NewVecOp, IdenVec);
       NewVecOp = Select;
     }
     Value *NewRed;
     Value *NextInChain;
     if (IsOrdered) {
       if (State.VF.isVector())
         NewRed = createOrderedReduction(State.Builder, *RdxDesc, NewVecOp,
                                         PrevInChain);
       else
         NewRed = State.Builder.CreateBinOp(
             (Instruction::BinaryOps)RdxDesc->getOpcode(Kind), PrevInChain,
             NewVecOp);
       PrevInChain = NewRed;
     } else {
       PrevInChain = State.get(getChainOp(), Part);
       NewRed = createTargetReduction(State.Builder, TTI, *RdxDesc, NewVecOp);
     }
     if (RecurrenceDescriptor::isMinMaxRecurrenceKind(Kind)) {
       NextInChain =
           createMinMaxOp(State.Builder, RdxDesc->getRecurrenceKind(),
                          NewRed, PrevInChain);
     } else if (IsOrdered)
       NextInChain = NewRed;
     else
       NextInChain = State.Builder.CreateBinOp(
           (Instruction::BinaryOps)RdxDesc->getOpcode(Kind), NewRed,
           PrevInChain);
     State.set(this, NextInChain, Part);
   }
+  // Reset fast-math flags.
+  State.Builder.setFastMathFlags(Current_FMF);
 }
 
 void VPReplicateRecipe::execute(VPTransformState &State) {
   if (State.Instance) { // Generate a single instance.
     assert(!State.VF.isScalable() && "Can't scalarize a scalable vector");
     State.ILV->scalarizeInstruction(getUnderlyingInstr(), this, *this,
                                     *State.Instance, IsPredicated, State);
     // Insert scalar instance packing it into a vector.

llvm/test/Transforms/LoopVectorize/AArch64/strict-fadd.ll

 loop:
   %red = phi double [ 0.0, %entry ], [ %red.next, %loop ]
   %for = phi double [ 0.0, %entry ], [ %for.next, %loop ]
   %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop ]
   %red.next = fadd double %for, %red
   %for.next = sitofp i32 %iv to double
   %iv.next = add nsw i32 %iv, 1
   %ec = icmp eq i32 %iv.next, 0
   br i1 %ec, label %exit, label %loop, !llvm.loop !13
 
 exit:
   %res = phi double [ %red.next, %loop ]
   ret double %res
 }
 
+; Test case where fadd has a fast-math flag.
+define float @fadd_strict_fmf(float* noalias nocapture readonly %a, i64 %n) {
+; CHECK-ORDERED-LABEL: @fadd_strict_fmf
+; CHECK-ORDERED: vector.body:
+; CHECK-ORDERED: [[VEC_PHI:%.*]] = phi float [ 0.000000e+00, %vector.ph ], [ [[RDX:%.*]], %vector.body ]
+; CHECK-ORDERED: [[LOAD_VEC:%.*]] = load <8 x float>, <8 x float>*
+; CHECK-ORDERED: [[RDX]] = call nnan float @llvm.vector.reduce.fadd.v8f32(float [[VEC_PHI]], <8 x float> [[LOAD_VEC]])
+; CHECK-ORDERED: for.end:
+; CHECK-ORDERED: [[RES:%.*]] = phi float [ [[SCALAR:%.*]], %for.body ], [ [[RDX]], %middle.block ]
+; CHECK-ORDERED: ret float [[RES]]
+
+; CHECK-UNORDERED-LABEL: @fadd_strict_fmf
+; CHECK-UNORDERED: vector.body:
+; CHECK-UNORDERED: [[VEC_PHI:%.*]] = phi <8 x float> [ <float 0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00>, %vector.ph ], [ [[FADD_VEC:%.*]], %vector.body ]
+; CHECK-UNORDERED: [[LOAD_VEC:%.*]] = load <8 x float>, <8 x float>*
+; CHECK-UNORDERED: [[FADD_VEC]] = fadd nnan <8 x float> [[LOAD_VEC]], [[VEC_PHI]]
+; CHECK-UNORDERED-NOT: call nnan float @llvm.vector.reduce.fadd 
+; CHECK-UNORDERED: middle.block:
+; CHECK-UNORDERED: [[RDX:%.*]] = call nnan float @llvm.vector.reduce.fadd.v8f32(float -0.000000e+00, <8 x float> [[FADD_VEC]])
+; CHECK-UNORDERED: for.body:
+; CHECK-UNORDERED: [[LOAD:%.*]] = load float, float*
+; CHECK-UNORDERED: [[FADD:%.*]] = fadd nnan float [[LOAD]], {{.*}}
+; CHECK-UNORDERED: for.end:
+; CHECK-UNORDERED: [[RES:%.*]] = phi float [ [[FADD]], %for.body ], [ [[RDX]], %middle.block ]
+; CHECK-UNORDERED: ret float [[RES]]
+
+; CHECK-NOT-VECTORIZED-LABEL: @fadd_strict_fmf
+; CHECK-NOT-VECTORIZED-NOT: vector.body
+
+entry:
+  br label %for.body
+
+for.body:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+  %sum.07 = phi float [ 0.000000e+00, %entry ], [ %add, %for.body ]
+  %arrayidx = getelementptr inbounds float, float* %a, i64 %iv
+  %0 = load float, float* %arrayidx, align 4
+  %add = fadd nnan float %0, %sum.07
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond.not = icmp eq i64 %iv.next, %n
+  br i1 %exitcond.not, label %for.end, label %for.body, !llvm.loop !0
+
+for.end:
+  ret float %add
+}
+
 !0 = distinct !{!0, !5, !9, !11}
 !1 = distinct !{!1, !5, !10, !11}
 !2 = distinct !{!2, !6, !9, !11}
 !3 = distinct !{!3, !7, !9, !11, !12}
 !4 = distinct !{!4, !8, !10, !11}
 !5 = !{!"llvm.loop.vectorize.width", i32 8}
 !6 = !{!"llvm.loop.vectorize.width", i32 4}
 !7 = !{!"llvm.loop.vectorize.width", i32 2}
 !8 = !{!"llvm.loop.vectorize.width", i32 1}
 !9 = !{!"llvm.loop.interleave.count", i32 1}
 !10 = !{!"llvm.loop.interleave.count", i32 4}
 !11 = !{!"llvm.loop.vectorize.enable", i1 true}
 !12 = !{!"llvm.loop.vectorize.predicate.enable", i1 true}
 !13 = distinct !{!13, !6, !9, !11}