diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp --- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp +++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp @@ -1073,11 +1073,17 @@ continue; // This recipe contributes to the address computation of a widen - // load/store. Collect recipe if its underlying instruction has - // poison-generating flags. + // load/store. If the underlying instruction has poison-generating flags, + // either drop them directly if the recipe already models the flags or + // collect them in a set. + // TODO: Migrate all recipes to hold their own flags. Instruction *Instr = CurRec->getUnderlyingInstr(); - if (Instr && Instr->hasPoisonGeneratingFlags()) - State.MayGeneratePoisonRecipes.insert(CurRec); + if (Instr && Instr->hasPoisonGeneratingFlags()) { + if (auto *OpWithFlags = dyn_cast(CurRec)) + OpWithFlags->dropPoisonGeneratingFlags(); + else + State.MayGeneratePoisonRecipes.insert(CurRec); + } // Add new definitions to the worklist. for (VPValue *operand : CurRec->operands()) diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h --- a/llvm/lib/Transforms/Vectorize/VPlan.h +++ b/llvm/lib/Transforms/Vectorize/VPlan.h @@ -37,6 +37,7 @@ #include "llvm/Analysis/VectorUtils.h" #include "llvm/IR/DebugLoc.h" #include "llvm/IR/FMF.h" +#include "llvm/IR/Operator.h" #include #include #include @@ -936,14 +937,143 @@ } }; +/// Class to record LLVM IR flag for a recipe. +class VPRecipeWithIRFlags : public VPRecipeBase { + struct WrapFlagsTy { + char HasNUW : 1; + char HasNSW : 1; + }; + struct ExactFlagsTy { + char IsExact : 1; + }; + struct FastMathFlagsTy { + char AllowReassoc : 1; + char NoNaNs : 1; + char NoInfs : 1; + char NoSignedZeros : 1; + char AllowReciprocal : 1; + char AllowContract : 1; + char ApproxFunc : 1; + }; + + char HasWrapFlags : 1; + char HasExact : 1; + char HasFastMathFlags : 1; + + union { + WrapFlagsTy WrapFlags; + ExactFlagsTy ExactFlags; + FastMathFlagsTy FMFs; + char AllFlags; + }; + +public: + template + VPRecipeWithIRFlags(const unsigned char SC, iterator_range Operands) + : VPRecipeBase(SC, Operands) { + HasWrapFlags = 0; + HasExact = 0; + HasFastMathFlags = 0; + AllFlags = 0; + } + + template + VPRecipeWithIRFlags(const unsigned char SC, iterator_range Operands, + Instruction &I) + : VPRecipeWithIRFlags(SC, Operands) { + if (auto *Op = dyn_cast(&I)) { + HasWrapFlags = 1; + WrapFlags.HasNUW = Op->hasNoUnsignedWrap(); + WrapFlags.HasNSW = Op->hasNoSignedWrap(); + } else if (auto *Op = dyn_cast(&I)) { + HasExact = 1; + ExactFlags.IsExact = Op->isExact(); + } else if (auto *Op = dyn_cast(&I)) { + HasFastMathFlags = 1; + FastMathFlags FMF = Op->getFastMathFlags(); + FMFs.AllowReassoc = FMF.allowReassoc(); + FMFs.NoNaNs = FMF.noNaNs(); + FMFs.NoInfs = FMF.noInfs(); + FMFs.NoSignedZeros = FMF.noSignedZeros(); + FMFs.AllowReciprocal = FMF.allowReciprocal(); + FMFs.AllowContract = FMF.allowContract(); + FMFs.ApproxFunc = FMF.approxFunc(); + } + } + + static inline bool classof(const VPRecipeBase *R) { + return R->getVPDefID() == VPRecipeBase::VPWidenSC; + } + + /// Drop all poison-generating flags. + void dropPoisonGeneratingFlags() { + // NOTE: This needs to be kept in-sync with + // Instruction::dropPoisonGeneratingFlags. + if (HasWrapFlags) { + WrapFlags.HasNUW = false; + WrapFlags.HasNSW = false; + } else if (HasExact) { + ExactFlags.IsExact = false; + } else if (HasFastMathFlags) { + FMFs.NoNaNs = false; + FMFs.NoInfs = false; + } + } + + bool hasNoUnsignedWrap() const { + assert(HasWrapFlags && "NUW flag not available"); + return WrapFlags.HasNUW; + } + bool hasNoSignedWrap() const { + assert(HasWrapFlags && "NSW flag not available"); + return WrapFlags.HasNSW; + } + bool isExact() const { + assert(HasExact && "Exact flag not available"); + return ExactFlags.IsExact; + } + + FastMathFlags getFastMathFlags() const { + assert(HasFastMathFlags && "FMFs not available"); + FastMathFlags Res; + Res.setAllowReassoc(FMFs.AllowReassoc); + Res.setNoNaNs(FMFs.NoNaNs); + Res.setNoInfs(FMFs.NoInfs); + Res.setNoSignedZeros(FMFs.NoSignedZeros); + Res.setAllowReciprocal(FMFs.AllowReciprocal); + Res.setAllowContract(FMFs.AllowContract); + Res.setApproxFunc(FMFs.ApproxFunc); + return Res; + } + + /// Set the enabled IR flags for \p I. + void setFlags(Instruction *I) const { + if (HasExact) + I->setIsExact(isExact()); + else if (HasWrapFlags) { + I->setHasNoUnsignedWrap(hasNoUnsignedWrap()); + I->setHasNoSignedWrap(hasNoSignedWrap()); + } else if (HasFastMathFlags) { + I->setHasAllowReassoc(FMFs.AllowReassoc); + I->setHasNoNaNs(FMFs.NoNaNs); + I->setHasNoInfs(FMFs.NoInfs); + I->setHasNoSignedZeros(FMFs.NoSignedZeros); + I->setHasAllowReciprocal(FMFs.AllowReciprocal); + I->setHasAllowContract(FMFs.AllowContract); + I->setHasApproxFunc(FMFs.ApproxFunc); + } + } +}; + /// 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 { +class VPWidenRecipe : public VPRecipeWithIRFlags, public VPValue { + public: template VPWidenRecipe(Instruction &I, iterator_range Operands) - : VPRecipeBase(VPDef::VPWidenSC, Operands), VPValue(this, &I) {} + : VPRecipeWithIRFlags(VPDef::VPWidenSC, Operands, I), VPValue(this, &I) {} ~VPWidenRecipe() override = default; diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp --- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp +++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp @@ -632,17 +632,8 @@ Value *V = Builder.CreateNAryOp(I.getOpcode(), Ops); - if (auto *VecOp = dyn_cast(V)) { - VecOp->copyIRFlags(&I); - - // If the instruction is vectorized and was in a basic block that needed - // predication, we can't propagate poison-generating flags (nuw/nsw, - // exact, etc.). The control flow has been linearized and the - // instruction is no longer guarded by the predicate, which could make - // the flag properties to no longer hold. - if (State.MayGeneratePoisonRecipes.contains(this)) - VecOp->dropPoisonGeneratingFlags(); - } + if (auto *VecOp = dyn_cast(V)) + setFlags(VecOp); // Use this vector value for all users of the original instruction. State.set(this, V, Part);