[LoopVectorize] Propagate fast-math flags for VPInstruction

[RosieSumpter]

Summary

In-loop vector reductions which use the llvm.fmuladd intrinsic involve
the creation of two recipes; a VPReductionRecipe for the fadd and a
VPInstruction for the fmul. If the call to llvm.fmuladd has fast-math flags
these should be propagated through to the fmul instruction, so an
interface setFastMathFlags has been added to the VPInstruction class to
enable this.

Depends on D111555

[dmgreen]

What happens when the instruction is printed inside VPInstruction::print? Can (should?) we get it to display the fast math flags?

[fhahn]

In D113125#3106989, @dmgreen wrote:

What happens when the instruction is printed inside VPInstruction::print? Can (should?) we get it to display the fast math flags?

Yeah, I think the flags should be handled in ::print and it would be good to also add a test in llvm/test/Transforms/LoopVectorize/vplan-printing.ll

llvm/lib/Transforms/Vectorize/VPlan.h

876

can we assert that this is only used for opcodes that can take fast-math flags? On the IR side this is done by checking for FPMathOperator, but there's no helper to check if an opcode is a FP operation. Perhaps it would be sufficient to check for the relevant opcodes?

Also, for now we only need to set fast-math flags at construction time I think, so could this be set directly in the constructor?

[RosieSumpter]

• Added a print method to the FastMathFlags class and used it in VPInstruction::print
• Added a test to vplan-printing.ll
• Assert that the VPInstruction opcode is for a floating-point operation when setting fast-math flags

[RosieSumpter]

Thanks for the comments @dmgreen and @fhahn. RE printing the fast-math flags, it seems the only way is to check for each flag and print the correct word, as was being done in WriteOptimizationInfo in llvm/lib/IR/AsmWriter.cpp. Since this is now occurring in two places, I've added a print method to FastMathFlags. Let me know if you think this is the right way to go.

llvm/lib/Transforms/Vectorize/VPlan.h

876

Hey @fhahn, thanks for this, I've added an assert to check for all the fp opcodes in setFastMathFlags. I haven't set the fast-math flags directly in the constructor for now since the setFastMathFlags approach makes the flag setting explicit, and mirrors the Instruction class. As always though, I'm happy to change it if you still think setting the flags in the constructor is the best way.

[sdesmalen]

llvm/lib/IR/AsmWriter.cpp

1318

Hi @RosieSumpter, could you maybe also define an overloaded operator<< for FastMathFlags, so that you can write Out << FPO->getFastMathFlags()?

[RosieSumpter]

• Defined overloaded operator << for FastMathFlags and used it in AsmWriter.cpp

[sdesmalen]

llvm/lib/Transforms/Vectorize/VPlan.cpp

765

nit: O << FMF ?

llvm/test/Transforms/LoopVectorize/vplan-printing.ll

241

nit: Should VPReductionRecipe::print also print the fast-math flags?

[RosieSumpter]

• Changed FMF.print(O) to O << FMF
• Added fast-math flags to VPReductionRecipe::print

[sdesmalen]

Thanks for the changes @RosieSumpter, LGTM!

[fhahn]

llvm/lib/Transforms/Vectorize/VPlan.cpp

1215

That change is not really related to the patch title/description. I think it would be good to either adjust the title or commit it separately.

llvm/lib/Transforms/Vectorize/VPlan.h

876

Thanks for adding the assert! For VPInstructions we do not need the full flexibility provided by setFastMathFlags, as for now we only setting them at construction time. This was the main reason to suggest doing it directly in the constructor. But I guess in the future this may be helpful in case we need to drop fast-math flags.

[RosieSumpter]

llvm/lib/Transforms/Vectorize/VPlan.cpp

1215

Hi @fhahn, agreed - I split off this change into a separate commit before merging.

llvm/include/llvm/IR/Operator.h

public:
   void setFast(bool B = true) { B ? set() : clear(); }
 
   void operator&=(const FastMathFlags &OtherFlags) {
     Flags &= OtherFlags.Flags;
   }
   void operator|=(const FastMathFlags &OtherFlags) {
     Flags |= OtherFlags.Flags;
   }
+
+  /// Print fast-math flags to \p O.
+  void print(raw_ostream &O) const;
 };
 
 /// Utility class for floating point operations which can have
 /// information about relaxed accuracy requirements attached to them.
 class FPMathOperator : public Operator {
 private:
   friend class Instruction;
 

llvm/lib/IR/AsmWriter.cpp

static void WriteAsOperandInternal(raw_ostream &Out, const Value *V,
                                    AsmWriterContext &WriterCtx);
 
 static void WriteAsOperandInternal(raw_ostream &Out, const Metadata *MD,
                                    AsmWriterContext &WriterCtx,
                                    bool FromValue = false);
 
 static void WriteOptimizationInfo(raw_ostream &Out, const User *U) {
   if (const FPMathOperator *FPO = dyn_cast<const FPMathOperator>(U)) {
-    // 'Fast' is an abbreviation for all fast-math-flags.
-    if (FPO->isFast())
-      Out << " fast";
-    else {
-      if (FPO->hasAllowReassoc())
-        Out << " reassoc";
-      if (FPO->hasNoNaNs())
-        Out << " nnan";
-      if (FPO->hasNoInfs())
-        Out << " ninf";
-      if (FPO->hasNoSignedZeros())
-        Out << " nsz";
-      if (FPO->hasAllowReciprocal())
-        Out << " arcp";
-      if (FPO->hasAllowContract())
-        Out << " contract";
-      if (FPO->hasApproxFunc())
-        Out << " afn";
-    }
+    FastMathFlags FMF = FPO->getFastMathFlags();
+    FMF.print(Out);
   }
 
   if (const OverflowingBinaryOperator *OBO =
         dyn_cast<OverflowingBinaryOperator>(U)) {
     if (OBO->hasNoUnsignedWrap())
       Out << " nuw";
     if (OBO->hasNoSignedWrap())
       Out << " nsw";

llvm/lib/IR/Operator.cpp

for (gep_type_iterator GTI = gep_type_begin(this), GTE = gep_type_end(this);
     // increment the offset by IndexedSize.
     if (!IndexedSize.isZero()) {
       VariableOffsets.insert({V, APInt(BitWidth, 0)});
       VariableOffsets[V] += IndexedSize;
     }
   }
   return true;
 }
+
+void FastMathFlags::print(raw_ostream &O) const {
+  if (all())
+    O << " fast";
+  else {
+    if (allowReassoc())
+      O << " reassoc";
+    if (noNaNs())
+      O << " nnan";
+    if (noInfs())
+      O << " ninf";
+    if (noSignedZeros())
+      O << " nsz";
+    if (allowReciprocal())
+      O << " arcp";
+    if (allowContract())
+      O << " contract";
+    if (approxFunc())
+      O << " afn";
+  }
+}
 } // namespace llvm

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

for (Instruction *R : ReductionOperations) {
       if (IsFMulAdd) {
         // If the instruction is a call to the llvm.fmuladd intrinsic then we
         // need to create an fmul recipe to use as the vector operand for the
         // fadd reduction.
         SmallVector<VPValue *, 2> FMulOps = {
             VecOp, Plan->getVPValue(R->getOperand(1))};
         VPInstruction *FMulRecipe =
             new VPInstruction(Instruction::FMul, FMulOps);
+        FMulRecipe->setFastMathFlags(R->getFastMathFlags());
         WidenRecipe->getParent()->insert(FMulRecipe,
                                          WidenRecipe->getIterator());
         VecOp = FMulRecipe;
       }
       VPReductionRecipe *RedRecipe =
           new VPReductionRecipe(&RdxDesc, R, ChainOp, VecOp, CondOp, TTI);
       WidenRecipe->getVPSingleValue()->replaceAllUsesWith(RedRecipe);
       Plan->removeVPValueFor(R);

llvm/lib/Transforms/Vectorize/VPlan.h

enum {
     SLPLoad,
     SLPStore,
     ActiveLaneMask,
   };
 
 private:
   typedef unsigned char OpcodeTy;
   OpcodeTy Opcode;
+  FastMathFlags FMF;
 
   /// Utility method serving execute(): generates a single instance of the
   /// modeled instruction.
   void generateInstruction(VPTransformState &State, unsigned Part);
 
 protected:
   void setUnderlyingInstr(Instruction *I) { setUnderlyingValue(I); }
 
bool hasResult() const {
     case Instruction::Unreachable:
     case Instruction::Fence:
     case Instruction::AtomicRMW:
       return false;
     default:
       return true;
     }
   }
+
+  /// Set the fast-math flags.
+  void setFastMathFlags(FastMathFlags FMFNew);
 };
 
 /// VPWidenRecipe is a recipe for producing a copy of vector type its
 /// ingredient. This recipe covers most of the traditional vectorization cases
 /// where each ingredient transforms into a vectorized version of itself.
 class VPWidenRecipe : public VPRecipeBase, public VPValue {
 public:
   template <typename IterT>

llvm/lib/Transforms/Vectorize/VPlan.cpp

void VPInstruction::generateInstruction(VPTransformState &State,
   }
   default:
     llvm_unreachable("Unsupported opcode for instruction");
   }
 }
 
 void VPInstruction::execute(VPTransformState &State) {
   assert(!State.Instance && "VPInstruction executing an Instance");
+  IRBuilderBase::FastMathFlagGuard FMFGuard(State.Builder);
+  State.Builder.setFastMathFlags(FMF);
   for (unsigned Part = 0; Part < State.UF; ++Part)
     generateInstruction(State, Part);
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 void VPInstruction::dump() const {
   VPSlotTracker SlotTracker(getParent()->getPlan());
   print(dbgs(), "", SlotTracker);
case VPInstruction::ActiveLaneMask:
     break;
   case VPInstruction::FirstOrderRecurrenceSplice:
     O << "first-order splice";
     break;
   default:
     O << Instruction::getOpcodeName(getOpcode());
   }
 
+  FMF.print(O);
+
   for (const VPValue *Operand : operands()) {
     O << " ";
     Operand->printAsOperand(O, SlotTracker);
   }
 }
 #endif
 
+void VPInstruction::setFastMathFlags(FastMathFlags FMFNew) {
+  // Make sure the VPInstruction is a floating-point operation.
+  assert((Opcode == Instruction::FAdd || Opcode == Instruction::FMul ||
+          Opcode == Instruction::FNeg || Opcode == Instruction::FSub ||
+          Opcode == Instruction::FDiv || Opcode == Instruction::FRem ||
+          Opcode == Instruction::FCmp) &&
+         "this op can't take fast-math flags");
+  FMF = FMFNew;
+}
+
 /// Generate the code inside the body of the vectorized loop. Assumes a single
 /// LoopVectorBody basic-block was created for this. Introduce additional
 /// basic-blocks as needed, and fill them all.
 void VPlan::execute(VPTransformState *State) {
   // -1. Check if the backedge taken count is needed, and if so build it.
   if (BackedgeTakenCount && BackedgeTakenCount->getNumUsers()) {
     Value *TC = State->TripCount;
     IRBuilder<> Builder(State->CFG.PrevBB->getTerminator());

 void VPReductionRecipe::print(raw_ostream &O, const Twine &Indent,
                               VPSlotTracker &SlotTracker) const {
   O << Indent << "REDUCE ";
   printAsOperand(O, SlotTracker);
   O << " = ";
   getChainOp()->printAsOperand(O, SlotTracker);
   O << " + reduce." << Instruction::getOpcodeName(RdxDesc->getOpcode())
     << " (";
   getVecOp()->printAsOperand(O, SlotTracker);
   if (getCondOp()) {
     O << ", ";
     getCondOp()->printAsOperand(O, SlotTracker);
   }
   O << ")";
 }
 
 void VPReplicateRecipe::print(raw_ostream &O, const Twine &Indent,

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

 ; CHECK-ORDERED: [[WIDE_LOAD:%.*]] = load <vscale x 8 x float>, <vscale x 8 x float>*
 ; CHECK-ORDERED: [[WIDE_LOAD1:%.*]] = load <vscale x 8 x float>, <vscale x 8 x float>*
 ; CHECK-ORDERED: [[WIDE_LOAD2:%.*]] = load <vscale x 8 x float>, <vscale x 8 x float>*
 ; CHECK-ORDERED: [[WIDE_LOAD3:%.*]] = load <vscale x 8 x float>, <vscale x 8 x float>*
 ; CHECK-ORDERED: [[WIDE_LOAD4:%.*]] = load <vscale x 8 x float>, <vscale x 8 x float>*
 ; CHECK-ORDERED: [[WIDE_LOAD5:%.*]] = load <vscale x 8 x float>, <vscale x 8 x float>*
 ; CHECK-ORDERED: [[WIDE_LOAD6:%.*]] = load <vscale x 8 x float>, <vscale x 8 x float>*
 ; CHECK-ORDERED: [[WIDE_LOAD7:%.*]] = load <vscale x 8 x float>, <vscale x 8 x float>*
-; CHECK-ORDERED: [[FMUL:%.*]] = fmul <vscale x 8 x float> [[WIDE_LOAD]], [[WIDE_LOAD4]]
-; CHECK-ORDERED: [[FMUL1:%.*]] = fmul <vscale x 8 x float> [[WIDE_LOAD1]], [[WIDE_LOAD5]]
-; CHECK-ORDERED: [[FMUL2:%.*]] = fmul <vscale x 8 x float> [[WIDE_LOAD2]], [[WIDE_LOAD6]]
-; CHECK-ORDERED: [[FMUL3:%.*]] = fmul <vscale x 8 x float> [[WIDE_LOAD3]], [[WIDE_LOAD7]]
+; CHECK-ORDERED: [[FMUL:%.*]] = fmul nnan <vscale x 8 x float> [[WIDE_LOAD]], [[WIDE_LOAD4]]
+; CHECK-ORDERED: [[FMUL1:%.*]] = fmul nnan <vscale x 8 x float> [[WIDE_LOAD1]], [[WIDE_LOAD5]]
+; CHECK-ORDERED: [[FMUL2:%.*]] = fmul nnan <vscale x 8 x float> [[WIDE_LOAD2]], [[WIDE_LOAD6]]
+; CHECK-ORDERED: [[FMUL3:%.*]] = fmul nnan <vscale x 8 x float> [[WIDE_LOAD3]], [[WIDE_LOAD7]]
 ; CHECK-ORDERED: [[RDX:%.*]] = call nnan float @llvm.vector.reduce.fadd.nxv8f32(float [[VEC_PHI]], <vscale x 8 x float> [[FMUL]])
 ; CHECK-ORDERED: [[RDX1:%.*]] = call nnan float @llvm.vector.reduce.fadd.nxv8f32(float [[RDX]], <vscale x 8 x float> [[FMUL1]])
 ; CHECK-ORDERED: [[RDX2:%.*]] = call nnan float @llvm.vector.reduce.fadd.nxv8f32(float [[RDX1]], <vscale x 8 x float> [[FMUL2]])
 ; CHECK-ORDERED: [[RDX3:%.*]] = call nnan float @llvm.vector.reduce.fadd.nxv8f32(float [[RDX2]], <vscale x 8 x float> [[FMUL3]])
 ; CHECK-ORDERED: for.end
 ; CHECK-ORDERED: [[RES:%.*]] = phi float [ [[SCALAR:%.*]], %for.body ], [ [[RDX3]], %middle.block ]
 ; CHECK-ORDERED: ret float [[RES]]
 

llvm/test/Transforms/LoopVectorize/vplan-printing.ll

for.body:
   %iv.next = add nuw nsw i64 %iv, 4
   %cmp = icmp slt i64 %iv.next, 1024
   br i1 %cmp, label %for.body, label %for.end
 
 for.end:
   ret void
 }
 
+define float @print_fmuladd_strict(float* %a, float* %b, i64 %n) {
+; CHECK-LABEL: Checking a loop in "print_fmuladd_strict"
+; CHECK:      VPlan 'Initial VPlan for VF={4},UF>=1' {
+; CHECK-NEXT: for.body:
+; CHECK-NEXT:   WIDEN-INDUCTION %iv = phi 0, %iv.next
+; CHECK-NEXT:   WIDEN-REDUCTION-PHI ir<%sum.07> = phi ir<0.000000e+00>, ir<%muladd>
+; CHECK-NEXT:   CLONE ir<%arrayidx> = getelementptr ir<%a>, ir<%iv>
+; CHECK-NEXT:   WIDEN ir<%0> = load ir<%arrayidx>
+; CHECK-NEXT:   CLONE ir<%arrayidx2> = getelementptr ir<%b>, ir<%iv>
+; CHECK-NEXT:   WIDEN ir<%1> = load ir<%arrayidx2>
+; CHECK-NEXT:   EMIT vp<%6> = fmul nnan ninf nsz ir<%0> ir<%1>
+; CHECK-NEXT:   REDUCE ir<%muladd> = ir<%sum.07> + reduce.fadd (vp<%6>)
+; CHECK-NEXT:   No successors
+; CHECK-NEXT: }
+
+entry:
+  br label %for.body
+
+for.body:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+  %sum.07 = phi float [ 0.000000e+00, %entry ], [ %muladd, %for.body ]
+  %arrayidx = getelementptr inbounds float, float* %a, i64 %iv
+  %0 = load float, float* %arrayidx, align 4
+  %arrayidx2 = getelementptr inbounds float, float* %b, i64 %iv
+  %1 = load float, float* %arrayidx2, align 4
+  %muladd = tail call nnan ninf nsz float @llvm.fmuladd.f32(float %0, float %1, float %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
+
+for.end:
+  ret float %muladd
+}
+
 declare float @llvm.sqrt.f32(float) nounwind readnone
+declare float @llvm.fmuladd.f32(float, float, float)