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 @@ -1099,8 +1099,12 @@ // load/store. Collect recipe if its underlying instruction has // poison-generating flags. Instruction *Instr = CurRec->getUnderlyingInstr(); - if (Instr && Instr->hasPoisonGeneratingFlags()) - State.MayGeneratePoisonRecipes.insert(CurRec); + if (Instr && Instr->hasPoisonGeneratingFlags()) { + if (auto *VPW = dyn_cast(CurRec)) + VPW->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,85 @@ } }; +/// Class to record LLVM IR flag for a recipe. +class VPIRFlags { + unsigned HasNUW : 1; + unsigned HasNSW : 1; + unsigned IsExact : 1; + unsigned AllowReassoc : 1; + unsigned NoNaNs : 1; + unsigned NoInfs : 1; + unsigned NoSignedZeros : 1; + unsigned AllowReciprocal : 1; + unsigned AllowContract : 1; + unsigned ApproxFunc : 1; + +public: + VPIRFlags(Instruction &I) { + if (auto *Op = dyn_cast(&I)) { + HasNUW = Op->hasNoUnsignedWrap(); + HasNSW = Op->hasNoSignedWrap(); + } else if (auto *Op = dyn_cast(&I)) { + IsExact = Op->isExact(); + } else if (auto *Op = dyn_cast(&I)) { + FastMathFlags FMF = Op->getFastMathFlags(); + AllowReassoc = FMF.allowReassoc(); + NoNaNs = FMF.noNaNs(); + NoInfs = FMF.noInfs(); + NoSignedZeros = FMF.noSignedZeros(); + AllowReciprocal = FMF.allowReciprocal(); + AllowContract = FMF.allowContract(); + ApproxFunc = FMF.approxFunc(); + } + } + + /// Drop all poison-generating flags. + void dropPoisonGeneratingFlags() { + HasNUW = false; + HasNSW = false; + IsExact = false; + NoNaNs = false; + NoInfs = false; + } + + bool hasNoUnsignedWrap() const { return HasNUW; } + bool hasNoSignedWrap() const { return HasNSW; } + bool isExact() const { return IsExact; } + FastMathFlags getFastMathFlags() const { + FastMathFlags FMFs; + FMFs.setAllowReassoc(AllowReassoc); + FMFs.setNoNaNs(NoNaNs); + FMFs.setNoInfs(NoInfs); + FMFs.setNoSignedZeros(NoSignedZeros); + FMFs.setAllowReciprocal(AllowReciprocal); + FMFs.setAllowContract(AllowContract); + FMFs.setApproxFunc(ApproxFunc); + return FMFs; + } + + /// Add the enabled IR flags to \p I. + void transferFlags(Instruction *I) const { + if (isExact()) + I->setIsExact(); + if (hasNoUnsignedWrap()) + I->setHasNoUnsignedWrap(); + if (hasNoSignedWrap()) + I->setHasNoSignedWrap(); + if (getFastMathFlags().any()) + I->setFastMathFlags(getFastMathFlags()); + } +}; + /// 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 VPRecipeBase, public VPValue, public VPIRFlags { + public: template VPWidenRecipe(Instruction &I, iterator_range Operands) - : VPRecipeBase(VPDef::VPWidenSC, Operands), VPValue(this, &I) {} + : VPRecipeBase(VPDef::VPWidenSC, Operands), VPValue(this, &I), + VPIRFlags(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 @@ -630,17 +630,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)) + transferFlags(VecOp); // Use this vector value for all users of the original instruction. State.set(this, V, Part);