[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 ↗

(On Diff #385036)

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 ↗

(On Diff #385770)

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

1203

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

1203

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

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) { FMF = 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());
   }
 
   for (const VPValue *Operand : operands()) {
     O << " ";
     Operand->printAsOperand(O, SlotTracker);
   }
 }
 #endif
 
 /// Generate the code inside the body of the vectorized loop. Assumes a single
 /// LoopVectorBody basic-block was created for this. Introduce additional

 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]]