From 5f3343985b20e0ae27368a1d3fd15896c4ed4fa5 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo@fhahn.com>
Date: Tue, 9 May 2023 12:33:14 +0100
Subject: [PATCH] [VPlan] Use VPRecipeWithIRFlags for VPWidenGEPRecipe (NFCI).

Extend VPRecipeWithIRFlags to also include InBounds and use for VPWidenGEPRecipe.

The last remaining recipe that needs updating for
MayGeneratePoisonRecipes is VPReplicateRecipe.

Depends on D149081.

Reviewed By: Ayal

Differential Revision: https://reviews.llvm.org/D149082
---
 llvm/lib/Transforms/Vectorize/VPlan.h          | 22 +++++++++++++++++++---
 llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp | 13 ++++---------
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 4ddf45a..95b8371 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -943,6 +943,7 @@ class VPRecipeWithIRFlags : public VPRecipeBase {
     OverflowingBinOp,
     PossiblyExactOp,
     FPMathOp,
+    GEPOp,
     Other
   };
   struct WrapFlagsTy {
@@ -952,6 +953,9 @@ class VPRecipeWithIRFlags : public VPRecipeBase {
   struct ExactFlagsTy {
     char IsExact : 1;
   };
+  struct GEPFlagsTy {
+    char IsInBounds : 1;
+  };
   struct FastMathFlagsTy {
     char AllowReassoc : 1;
     char NoNaNs : 1;
@@ -967,6 +971,7 @@ class VPRecipeWithIRFlags : public VPRecipeBase {
   union {
     WrapFlagsTy WrapFlags;
     ExactFlagsTy ExactFlags;
+    GEPFlagsTy GEPFlags;
     FastMathFlagsTy FMFs;
     unsigned char AllFlags;
   };
@@ -990,6 +995,9 @@ public:
     } else if (auto *Op = dyn_cast<PossiblyExactOperator>(&I)) {
       OpType = OperationType::PossiblyExactOp;
       ExactFlags.IsExact = Op->isExact();
+    } else if (auto *GEP = dyn_cast<GetElementPtrInst>(&I)) {
+      OpType = OperationType::GEPOp;
+      GEPFlags.IsInBounds = GEP->isInBounds();
     } else if (auto *Op = dyn_cast<FPMathOperator>(&I)) {
       OpType = OperationType::FPMathOp;
       FastMathFlags FMF = Op->getFastMathFlags();
@@ -1004,7 +1012,8 @@ public:
   }
 
   static inline bool classof(const VPRecipeBase *R) {
-    return R->getVPDefID() == VPRecipeBase::VPWidenSC;
+    return R->getVPDefID() == VPRecipeBase::VPWidenSC ||
+           R->getVPDefID() == VPRecipeBase::VPWidenGEPSC;
   }
 
   /// Drop all poison-generating flags.
@@ -1019,6 +1028,9 @@ public:
     case OperationType::PossiblyExactOp:
       ExactFlags.IsExact = false;
       break;
+    case OperationType::GEPOp:
+      GEPFlags.IsInBounds = false;
+      break;
     case OperationType::FPMathOp:
       FMFs.NoNaNs = false;
       FMFs.NoInfs = false;
@@ -1038,6 +1050,9 @@ public:
     case OperationType::PossiblyExactOp:
       I->setIsExact(ExactFlags.IsExact);
       break;
+    case OperationType::GEPOp:
+      cast<GetElementPtrInst>(I)->setIsInBounds(GEPFlags.IsInBounds);
+      break;
     case OperationType::FPMathOp:
       I->setHasAllowReassoc(FMFs.AllowReassoc);
       I->setHasNoNaNs(FMFs.NoNaNs);
@@ -1177,7 +1192,7 @@ struct VPWidenSelectRecipe : public VPRecipeBase, public VPValue {
 };
 
 /// A recipe for handling GEP instructions.
-class VPWidenGEPRecipe : public VPRecipeBase, public VPValue {
+class VPWidenGEPRecipe : public VPRecipeWithIRFlags, public VPValue {
   bool isPointerLoopInvariant() const {
     return getOperand(0)->isDefinedOutsideVectorRegions();
   }
@@ -1195,7 +1210,8 @@ class VPWidenGEPRecipe : public VPRecipeBase, public VPValue {
 public:
   template <typename IterT>
   VPWidenGEPRecipe(GetElementPtrInst *GEP, iterator_range<IterT> Operands)
-      : VPRecipeBase(VPDef::VPWidenGEPSC, Operands), VPValue(this, GEP) {}
+      : VPRecipeWithIRFlags(VPDef::VPWidenGEPSC, Operands, *GEP),
+        VPValue(this, GEP) {}
 
   ~VPWidenGEPRecipe() override = default;
 
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index 20e811b..596afc9 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -800,6 +800,7 @@ void VPWidenGEPRecipe::execute(VPTransformState &State) {
     //       collectLoopScalars() and teach getVectorValue() to broadcast
     //       the lane-zero scalar value.
     auto *Clone = State.Builder.Insert(GEP->clone());
+    setFlags(Clone);
     for (unsigned Part = 0; Part < State.UF; ++Part) {
       Value *EntryPart = State.Builder.CreateVectorSplat(State.VF, Clone);
       State.set(this, EntryPart, Part);
@@ -831,18 +832,12 @@ void VPWidenGEPRecipe::execute(VPTransformState &State) {
           Indices.push_back(State.get(Operand, Part));
       }
 
-      // If the GEP instruction is vectorized and was in a basic block that
-      // needed predication, we can't propagate the poison-generating 'inbounds'
-      // flag. The control flow has been linearized and the GEP is no longer
-      // guarded by the predicate, which could make the 'inbounds' properties to
-      // no longer hold.
-      bool IsInBounds =
-          GEP->isInBounds() && State.MayGeneratePoisonRecipes.count(this) == 0;
-
       // Create the new GEP. Note that this GEP may be a scalar if VF == 1,
       // but it should be a vector, otherwise.
       auto *NewGEP = State.Builder.CreateGEP(GEP->getSourceElementType(), Ptr,
-                                             Indices, "", IsInBounds);
+                                             Indices, "");
+      if (auto *I = dyn_cast<GetElementPtrInst>(NewGEP))
+        setFlags(I);
       assert((State.VF.isScalar() || NewGEP->getType()->isVectorTy()) &&
              "NewGEP is not a pointer vector");
       State.set(this, NewGEP, Part);
-- 
2.7.4