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 relevant recipes to hold their own flags. Instruction *Instr = CurRec->getUnderlyingInstr(); - if (Instr && Instr->hasPoisonGeneratingFlags()) - State.MayGeneratePoisonRecipes.insert(CurRec); + if (Instr && Instr->hasPoisonGeneratingFlags()) { + if (auto *RecWithFlags = dyn_cast(CurRec)) + RecWithFlags->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,131 @@ } }; +/// Class to record LLVM IR flag for a recipe along with it. +class VPRecipeWithIRFlags : public VPRecipeBase { + enum class OperationType : unsigned char { + OverflowingBinOp, + PossiblyExactOp, + FPMathOp, + Other + }; + 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; + }; + + OperationType OpType; + + union { + WrapFlagsTy WrapFlags; + ExactFlagsTy ExactFlags; + FastMathFlagsTy FMFs; + unsigned char AllFlags; + }; + +public: + template + VPRecipeWithIRFlags(const unsigned char SC, iterator_range Operands) + : VPRecipeBase(SC, Operands) { + OpType = OperationType::Other; + AllFlags = 0; + } + + template + VPRecipeWithIRFlags(const unsigned char SC, iterator_range Operands, + Instruction &I) + : VPRecipeWithIRFlags(SC, Operands) { + if (auto *Op = dyn_cast(&I)) { + OpType = OperationType::OverflowingBinOp; + WrapFlags.HasNUW = Op->hasNoUnsignedWrap(); + WrapFlags.HasNSW = Op->hasNoSignedWrap(); + } else if (auto *Op = dyn_cast(&I)) { + OpType = OperationType::PossiblyExactOp; + ExactFlags.IsExact = Op->isExact(); + } else if (auto *Op = dyn_cast(&I)) { + OpType = OperationType::FPMathOp; + 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. + switch (OpType) { + case OperationType::OverflowingBinOp: + WrapFlags.HasNUW = false; + WrapFlags.HasNSW = false; + break; + case OperationType::PossiblyExactOp: + ExactFlags.IsExact = false; + break; + case OperationType::FPMathOp: + FMFs.NoNaNs = false; + FMFs.NoInfs = false; + break; + case OperationType::Other: + break; + } + } + + /// Set the IR flags for \p I. + void setFlags(Instruction *I) const { + switch (OpType) { + case OperationType::OverflowingBinOp: + I->setHasNoUnsignedWrap(WrapFlags.HasNUW); + I->setHasNoSignedWrap(WrapFlags.HasNSW); + break; + case OperationType::PossiblyExactOp: + I->setIsExact(ExactFlags.IsExact); + break; + case OperationType::FPMathOp: + 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); + break; + case OperationType::Other: + break; + } + } +}; + /// 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);