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 @@ -1076,8 +1076,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,127 @@ } }; +/// Class to record LLVM IR flag for a recipe. +class VPOpWithIRFlags { + enum class OperationType { + OverflowingBinOp, + PossiblyExactOp, + FPMathOp, + Other + }; + OperationType OpType; + + 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: + VPOpWithIRFlags() { + OpType = OperationType::Other; + HasNUW = 0; + HasNSW = 0; + IsExact = 0; + AllowReassoc = 0; + NoNaNs = 0; + NoInfs = 0; + NoSignedZeros = 0; + AllowReciprocal = 0; + AllowContract = 0; + ApproxFunc = 0; + } + + VPOpWithIRFlags(Instruction &I) : VPOpWithIRFlags() { + if (auto *Op = dyn_cast(&I)) { + HasNUW = Op->hasNoUnsignedWrap(); + HasNSW = Op->hasNoSignedWrap(); + OpType = OperationType::OverflowingBinOp; + } else if (auto *Op = dyn_cast(&I)) { + IsExact = Op->isExact(); + OpType = OperationType::PossiblyExactOp; + } 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(); + OpType = OperationType::FPMathOp; + } + } + + /// Drop all poison-generating flags. + void dropPoisonGeneratingFlags() { + HasNUW = false; + HasNSW = false; + IsExact = false; + NoNaNs = false; + NoInfs = false; + } + + bool hasNoUnsignedWrap() const { + assert(OpType == OperationType::OverflowingBinOp && + "NUW flag not available"); + return HasNUW; + } + bool hasNoSignedWrap() const { + assert(OpType == OperationType::OverflowingBinOp && + "NSW flag not available"); + return HasNSW; + } + bool isExact() const { + assert(OpType == OperationType::PossiblyExactOp && + "Exact flag not available"); + return IsExact; + } + + FastMathFlags getFastMathFlags() const { + assert(OpType == OperationType::FPMathOp && "FMFs not available"); + 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 (OpType == OperationType::PossiblyExactOp && isExact()) + I->setIsExact(); + else if (OpType == OperationType::OverflowingBinOp) { + if (hasNoUnsignedWrap()) + I->setHasNoUnsignedWrap(); + if (hasNoSignedWrap()) + I->setHasNoSignedWrap(); + } else if (OpType == OperationType::FPMathOp && 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 VPOpWithIRFlags { + public: template VPWidenRecipe(Instruction &I, iterator_range Operands) - : VPRecipeBase(VPDef::VPWidenSC, Operands), VPValue(this, &I) {} + : VPRecipeBase(VPDef::VPWidenSC, Operands), VPValue(this, &I), + VPOpWithIRFlags(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);