[VPlan] Record IR flags on VPWidenRecipe directly (NFC).

[fhahn]

Summary

This patch introduces a VPIRFlags class to record various IR flags. This
allows de-coupling of IR flags from the underlying instructions. The
main benefit is that it allows dropping of IR flags from recipes
directly, without the need to go through
State::MayGeneratePoisonRecipes. The plan is to remove
MayGeneratePoisonRecipes once all relevant recipes are transitioned.

It also allows dropping IR flags during VPlan-to-VPlan transforms, which
will be used in a follow-up patch to implement truncateToMinimalBitwidths
as VPlan-to-VPlan transform.

[Ayal]

llvm/lib/Transforms/Vectorize/VPlan.h

954

Record which of the three operators is recorded, to ensure appropriate flags are accessed?
Initialize all to off?

972

Ensuring the generated code will be treated as if via Instruction::dropPoisonGeneratingFlags().
Can instead keep a flag in VPIRFlags indicating that such flags should be dropped, instead of replicating that logic here.

992

Better than having VPIRFlags hold FastMathFlags as an internal struct/class/union?

997

This replaces Instruction::copyIRFlags(const Value *V, bool IncludeWrapFlags = true) which sets the flags of "this" according to those of V, whether on or off, but does so in the opposite direction - setting flags of I according to those of "this", as in:

if (isa<OverflowingBinaryOperator>(I)) {
  I->setHasNoSignedWrap(hasNoSignedWrap());
  I->setHasNoUnsignedWrap(hasNoUnsignedWrap());
} else if (isa<PossiblyExactOperator>(I)) {
  I->setIsExact(isExact());
} else if (isa<FPMathOperator>(I)) {
  I->copyFastMathFlags(getFastMathFlags());
}

where each clause assumes or checks that VPIRFlags also records the same operator.
But it does so by only turning-on in I every flag that is turned-on in VPIRFlags - assuming are all initialized off in both VPIRFlags and I?

The current approach could use each turn-on flag to indicate which operator is recorded, i.e., could do

if (isExact())
  I->setIsExact();
else if (getFastMathFlags().any())
  I->setFastMathFlags(getFastMathFlags());
else {
  if (hasNoUnsignedWrap())
    I->setHasNoUnsignedWrap();
  if (hasNoSignedWrap())
    I->setHasNoSignedWrap();
}

(just to clarify understanding.)

(Note that GEPs should record an inBounds flag.)

1071–1072

Should the recipe "have a" VPIRFlags rather than "be a" VPIRFlags?

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

642

Retain comment somewhere?

[fhahn]

Address latest comments, thanks!

[fhahn]

llvm/lib/Transforms/Vectorize/VPlan.h

954

Done. It looks like there's no easy way to check if an opcode on its own is a overflowing/possibly exact/FP op, so I added a helper enum class.

972

I am not sure how this would work when we don't have an Instruction available from which to drop.

We could record a new flag when calling dropPoisonGeneratingFlags on a recipe, but this would then also require to call another function to call Instruction::dropPoisonGeneratingFlags() during code-gen. We still need to look up the flags to set during code-gen, so it seems cleaner to properly unset them here.

Also when thinking about printing the flags as part of the VPlan printing (which isn't the case yet)

992

This is mostly to keep keep things tightly packed, but I can change that.

997

Updated using the recorded type of the op, thanks!

1071–1072

The main benefit of inheriting is to make the common transferFlags and dropPoisonGeneratingFlags available. But the name VPIRFlags is not ideal for inheriting.. I changed it to VPOpWithIRFlags, but can also change the VPWidenRecipe (and others transitioned in future patches) to hold an instance instead.

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

642

I am not sure where a good place would be. I am not sure it is necessary, because now it it is sufficient to directly drop the flags in collectPoisonGeneratingRecipes.

[fhahn]

Change to VPRecipeWithIRFlags inherting from VPRecipeBase and add classof implementation to allow for more generic invalidation (dyn_cast<VPRecipeWithIRFlags>) which simplifies follow-on patches that extend invalidation to additional recipes.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1076–1079

Update comment - recipes that support flags can drop them now rather than be collected to have them dropped later. This suggests that recipeBase should indicate if it supports flags or not? In the long run, any recipe for an underlying having poison generating flags should be capable of dropping them, right?

llvm/lib/Transforms/Vectorize/VPlan.h

954

An alternative may be to keep the OpType next to its flags, retaining a simple bit vector mask, which could be reset by a single zeroing of a short/int, as in:

unsigned HasWrapFlags : 1;
unsigned HasNUW : 1;
unsigned HasNSW : 1;

unsigned HasFastMathFlags : 1;
unsigned AllowReassoc : 1;
unsigned NoNaNs : 1;
unsigned NoInfs : 1;
unsigned NoSignedZeros : 1;
unsigned AllowReciprocal : 1;
unsigned AllowContract : 1;
unsigned ApproxFunc : 1;

unsigned HasInBoundFlag : 1;
unsigned HasInBound : 1;

where "Other" corresponds to all "Has*Flags" (or everything) being false.

972

Agreed: better drop poison generating flags eagerly rather than record a dropPoison flag and (continuing to) drop them lazily.

A main motivation for trying to reuse existing FMF struct and dropPoisonGeneratingFlags() API rather than reimplement them is to help keep the two in sync - which deserves a comment.

TL;DR:
The Instruction whose flags needs to drop will be available when generated, as the concluding part of transfer/copy flags, as done now.

Admittedly the flags still need to be recorded and looked up.

If a dropPoison flag is recorded, it should effect isExact() and hasNo*Wrap() for them to be consistent, so may as well flip their bit, or else inline them into transferFlags() - but they may be needed in the future, including by VPlan printing.

992

Keeping things tightly packed looks great.

A union OTOH would emphasize that these are mutually exclusive alternatives, and also allow initializing all to zero at once.

The FMF flags are currently used only as a whole - to record and apply all using setFastMathFlags(getFastMathFlags()). I.e., access to individual flags is not needed nor provided. If/when such access is needed, it could and should be provided via getFastMathFlags(), as that provides their full API, but currently requires to first construct all FMFs, redundantly.
WDYT?

1007

Can drop only relevant flags according to OpType.

1044

This conceptually "sets" the IR flags of I rather than only adding those enabled? The two are identical iff I has (initially) all its flags unset. Cf. copyIRFlags().

1054

May as well set each flag of I directly here rather than first having getFastMathFlags() set each in FMFs and then copying all to I. Unless these flags are held in an FMF to be set directly.

1071–1072

Class essentially holds IR flags, as in FMF (and more), so original name sounds fine as does having VPWidenRecipe hold an instance.
Worth reasoning about how this aspect would relate to the class hierarchy of recipes now or as more are encountered?
BTW, CmpInst::Predicate "flags" may also be supported in a similar way, cf. D50823

[fhahn]

Address latest comments, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1076–1079

Added a comment.

With the current system, recipes which have IR flags can be casted to VPRecipeWithIRFlags, similar how IR instructions/operators can be casted to OverflowingBinaryOperator & co. D149082 migrates VPWidenGEPRecipe and D150027 migrates VPRecplicateRecipe and removes MayGeneratePoisonRecipes

llvm/lib/Transforms/Vectorize/VPlan.h

954

Updated to simply include it in the bitfield, thanks!

972

I added a note to dropPoisonGeneratingFlags that it needs to be kept in sync with Instruction::dropPoisonGeneratingFlags.

I think it is very unlikely that one of the existing flags will get re-classified as poison-generating. If a new poison-generating flag gets added to the IR, we will never set this flags unless it is added here. And then it should also be added to dropPoisonGeneratingFlags, so keeping things in sync should not be too much trouble (once all recipes migrated).

Printing is added in D150029. Please let me know if you still would prefer to separately track whether poison-generating flags need to be dropped.

992

Updated to use a union, thanks!

Also updated to set fast-math flags individually. I can also add individual getters.

1007

Updated, thanks!

1044

Changed the name to setFlags to make it clearer.

1054

updated to set them directly here.

1071–1072

The current modeling is similar to OverflwoingBinaryOperator/PossiblyExactOperator in IR, providing a convenient way to check if an object may have flags, allowing for checking via isa/dyn_cast.

IMO CmpInst::Predicate is slightly different to the wrapping flags, as it is more like an (essential) argument, similar to designation type for casts.

[Ayal]

Thanks! Looks good to me, adding last minor nits.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1079

1082–1083

llvm/lib/Transforms/Vectorize/VPlan.h

940

949

Can reuse FastMathFlags, but that would complicate initialization and turn the union from char to int?

967

Better as unsigned char?

976

(0 >> false)
Ah, the purpose of having Has*Flags as indicator bits was to embed them with all other bits so that all could be initialized with one zeroing assignment, and they could direct relate to the bits they cover by placing them next to each other. Now an enum may (again) be preferred, referring to named structs, and possibly switched upon.

1034

These getters can be dropped after inlining them into their only use below in setFlags(), similar to how the latter accesses FMFs flags - directly w/o getters.

1036

This getFastMathFlags() is now useless and can be dropped.
May be needed in the future, at which point it should also be tested.

1049

[fhahn]

Address latest comments, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1079

Adjusted, thanks!

1082–1083

Adjusted. thanks!

llvm/lib/Transforms/Vectorize/VPlan.h

940

Adjusted, thanks!

949

Yeah I think that would make things more complicated unfortunately and unnecessarily increase the size for now.

967

Updated, thanks!

976

Switched back to the enum and updated code to use switch.

1034

Yep, removed them.

1036

Yep, removed!

1049

updated, thanks!

[Ayal]

Ship it!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

while (!Worklist.empty()) {
       if (isa<VPWidenMemoryInstructionRecipe>(CurRec) ||
           isa<VPInterleaveRecipe>(CurRec) ||
           isa<VPScalarIVStepsRecipe>(CurRec) ||
           isa<VPCanonicalIVPHIRecipe>(CurRec) ||
           isa<VPActiveLaneMaskPHIRecipe>(CurRec))
         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())
+      if (Instr && Instr->hasPoisonGeneratingFlags()) {
+        if (auto *OpWithFlags = dyn_cast<VPRecipeWithIRFlags>(CurRec))
+          OpWithFlags->dropPoisonGeneratingFlags();
+        else
           State.MayGeneratePoisonRecipes.insert(CurRec);
+      }
 
       // Add new definitions to the worklist.
       for (VPValue *operand : CurRec->operands())
         if (VPRecipeBase *OpDef = operand->getDefiningRecipe())
           Worklist.push_back(OpDef);
     }
   });
 

llvm/lib/Transforms/Vectorize/VPlan.h

 #include "llvm/ADT/Twine.h"
 #include "llvm/ADT/ilist.h"
 #include "llvm/ADT/ilist_node.h"
 #include "llvm/Analysis/IVDescriptors.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/IR/FMF.h"
+#include "llvm/IR/Operator.h"
 #include <algorithm>
 #include <cassert>
 #include <cstddef>
 #include <string>
 
 namespace llvm {
 
 class BasicBlock;
bool onlyFirstLaneUsed(const VPValue *Op) const override {
     case VPInstruction::CanonicalIVIncrementForPartNUW:
     case VPInstruction::BranchOnCount:
       return true;
     };
     llvm_unreachable("switch should return");
   }
 };
 
+/// 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 <typename IterT>
+  VPRecipeWithIRFlags(const unsigned char SC, iterator_range<IterT> Operands)
+      : VPRecipeBase(SC, Operands) {
+    HasWrapFlags = 0;
+    HasExact = 0;
+    HasFastMathFlags = 0;
+    AllFlags = 0;
+  }
+
+  template <typename IterT>
+  VPRecipeWithIRFlags(const unsigned char SC, iterator_range<IterT> Operands,
+                      Instruction &I)
+      : VPRecipeWithIRFlags(SC, Operands) {
+    if (auto *Op = dyn_cast<OverflowingBinaryOperator>(&I)) {
+      HasWrapFlags = 1;
+      WrapFlags.HasNUW = Op->hasNoUnsignedWrap();
+      WrapFlags.HasNSW = Op->hasNoSignedWrap();
+    } else if (auto *Op = dyn_cast<PossiblyExactOperator>(&I)) {
+      HasExact = 1;
+      ExactFlags.IsExact = Op->isExact();
+    } else if (auto *Op = dyn_cast<FPMathOperator>(&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 <typename IterT>
   VPWidenRecipe(Instruction &I, iterator_range<IterT> Operands)
-      : VPRecipeBase(VPDef::VPWidenSC, Operands), VPValue(this, &I) {}
+      : VPRecipeWithIRFlags(VPDef::VPWidenSC, Operands, I), VPValue(this, &I) {}
 
   ~VPWidenRecipe() override = default;
 
   VP_CLASSOF_IMPL(VPDef::VPWidenSC)
 
   /// Produce widened copies of all Ingredients.
   void execute(VPTransformState &State) override;
 

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

case Instruction::Xor: {
 
     for (unsigned Part = 0; Part < State.UF; ++Part) {
       SmallVector<Value *, 2> Ops;
       for (VPValue *VPOp : operands())
         Ops.push_back(State.get(VPOp, Part));
 
       Value *V = Builder.CreateNAryOp(I.getOpcode(), Ops);
 
-      if (auto *VecOp = dyn_cast<Instruction>(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<Instruction>(V))
+        setFlags(VecOp);
 
       // Use this vector value for all users of the original instruction.
       State.set(this, V, Part);
       State.addMetadata(V, &I);
     }
 
     break;
   }