[VPlan] VPlan version of InterleavedAccessInfo.

[fhahn]

Summary

This patch turns InterleaveGroup into a template with the instruction type
being a template parameter. It also adds a VPInterleavedAccessInfo class, which
only contains a mapping from VPInstructions to their respective InterleaveGroup.
As we do not have access to scalar evolution in VPlan, we can re-use
convert InterleavedAccessInfo to VPInterleavedAccess info.

[hsaito]

I like the templatization approach taken here. Thank you!. Changes look fine to me, but will wait for a few others to take a look.

[rkruppe]

lib/Transforms/Vectorize/VPlan.cpp

701

It doesn't seem like this memory is freed anywhere.

Reading D49491 it seems to me that the InterleaveGroups allocated here are only needed while the VPInterleavedAccessInfo object is alive, so it could free them in its destructor (or even better, store unique_ptrs that take care of that automatically).

[fhahn]

Delete allocated InterleaveGroups in destructor.

lib/Transforms/Vectorize/VPlan.cpp

701

Yep thanks! They should be freed in VPInterleavedAccessInfo's destructor.

[dcaballe]

Just a few comments.

Thanks!
Diego

lib/Transforms/Vectorize/VPlan.cpp

686

This algorithm assumes that there is only a region (top region) in the whole H-CFG but this is going to change pretty soon. Could you please change the implementation to recursively go inside other potential nested regions?

693

auto * for casts and dyn_cast (585, 591, 592)?

lib/Transforms/Vectorize/VPlan.h

1468–1496

Just curious. How difficult would it be to templatize the IAI methods needed here?

[fhahn]

lib/Transforms/Vectorize/VPlan.cpp

686

Will do. I just have to think a bit more how everything will fit together with outerloops, as the regular InterleavedAccessInfo only works on a single loop.

lib/Transforms/Vectorize/VPlan.h

1468–1496

I am not entirely sure, but most of the regular IAI is related to building the interleave groups, which is not required here because we re-using the regular one.

[fhahn]

Rebased, updated the use auto * for cast<>.

[fhahn]

Updated to traverse regions.

[rengolin]

So, IIUC, the way you only get the interleave info on instructions and map to VPlan is because we don't yet have scalar evolution in VPlan, so we need to do that in Instruction and then map to VPInstruction.

In the future, once we can (perhaps using the same template technique) have SCEV on VPlan, we won't need the Old2New map anymore, and then the analysis will be independent of the underlying type.

Di I get that right?

[fhahn]

In D49489#1296800, @rengolin wrote:

So, IIUC, the way you only get the interleave info on instructions and map to VPlan is because we don't yet have scalar evolution in VPlan, so we need to do that in Instruction and then map to VPInstruction.

In the future, once we can (perhaps using the same template technique) have SCEV on VPlan, we won't need the Old2New map anymore, and then the analysis will be independent of the underlying type.

Di I get that right?

Exactly, using the underlying plain IR values is used for the initial implementation, until we have something like SCEV for VPlan and can do the whole analysis on top of VPlan.

[rengolin]

Perfect, LGTM. Thanks!

include/llvm/Analysis/VectorUtils.h

 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/IR/IRBuilder.h"
 
 namespace llvm {
 
 template <typename T> class ArrayRef;
 class DemandedBits;
 class GetElementPtrInst;
-class InterleaveGroup; 
+template <typename InstTy> class InterleaveGroup;
 class Loop;
 class ScalarEvolution;
 class TargetTransformInfo;
 class Type;
 class Value;
 
 namespace Intrinsic {
 enum ID : unsigned;

 /// and \p VF 4, that has only its first member present is:
 ///
 ///   <1,0,0,1,0,0,1,0,0,1,0,0>
 ///
 /// Note: The result is a mask of 0's and 1's, as opposed to the other
 /// create[*]Mask() utilities which create a shuffle mask (mask that
 /// consists of indices).
 Constant *createBitMaskForGaps(IRBuilder<> &Builder, unsigned VF,
-                               const InterleaveGroup &Group);
+                               const InterleaveGroup<Instruction> &Group);
 
 /// Create a mask with replicated elements.
 ///
 /// This function creates a shuffle mask for replicating each of the \p VF 
 /// elements in a vector \p ReplicationFactor times. It can be used to
 /// transform a mask of \p VF elements into a mask of
 /// \p VF * \p ReplicationFactor elements used by a predicated
 /// interleaved-group of loads/stores whose Interleaved-factor ==

 ///          A[i]   = a;                         // Member of index 0
 ///          A[i+1] = b;                         // Member of index 1
 ///          A[i+2] = c;                         // Member of index 2
 ///          A[i+3] = d;                         // Member of index 3
 ///        }
 ///
 /// Note: the interleaved load group could have gaps (missing members), but
 /// the interleaved store group doesn't allow gaps.
-class InterleaveGroup {
+template <typename InstTy> class InterleaveGroup {
 public:
-  InterleaveGroup(Instruction *Instr, int Stride, unsigned Align)
+  InterleaveGroup(unsigned Factor, bool Reverse, unsigned Align)
+      : Factor(Factor), Reverse(Reverse), Align(Align), InsertPos(nullptr) {}
+
+  InterleaveGroup(InstTy *Instr, int Stride, unsigned Align)
       : Align(Align), InsertPos(Instr) {
     assert(Align && "The alignment should be non-zero");
 
     Factor = std::abs(Stride);
     assert(Factor > 1 && "Invalid interleave factor");
 
     Reverse = Stride < 0;
     Members[0] = Instr;
   }
 
   bool isReverse() const { return Reverse; }
   unsigned getFactor() const { return Factor; }
   unsigned getAlignment() const { return Align; }
   unsigned getNumMembers() const { return Members.size(); }
 
   /// Try to insert a new member \p Instr with index \p Index and
   /// alignment \p NewAlign. The index is related to the leader and it could be
   /// negative if it is the new leader.
   ///
   /// \returns false if the instruction doesn't belong to the group.
-  bool insertMember(Instruction *Instr, int Index, unsigned NewAlign) {
+  bool insertMember(InstTy *Instr, int Index, unsigned NewAlign) {
     assert(NewAlign && "The new member's alignment should be non-zero");
 
     int Key = Index + SmallestKey;
 
     // Skip if there is already a member with the same index.
     if (Members.find(Key) != Members.end())
       return false;
 
bool insertMember(InstTy *Instr, int Index, unsigned NewAlign) {
     Align = std::min(Align, NewAlign);
     Members[Key] = Instr;
     return true;
   }
 
   /// Get the member with the given index \p Index
   ///
   /// \returns nullptr if contains no such member.
-  Instruction *getMember(unsigned Index) const {
+  InstTy *getMember(unsigned Index) const {
     int Key = SmallestKey + Index;
     auto Member = Members.find(Key);
     if (Member == Members.end())
       return nullptr;
 
     return Member->second;
   }
 
   /// Get the index for the given member. Unlike the key in the member
   /// map, the index starts from 0.
-  unsigned getIndex(Instruction *Instr) const {
-    for (auto I : Members)
+  unsigned getIndex(const InstTy *Instr) const {
+    for (auto I : Members) {
       if (I.second == Instr)
         return I.first - SmallestKey;
+    }
 
     llvm_unreachable("InterleaveGroup contains no such member");
   }
 
-  Instruction *getInsertPos() const { return InsertPos; }
-  void setInsertPos(Instruction *Inst) { InsertPos = Inst; }
+  InstTy *getInsertPos() const { return InsertPos; }
+  void setInsertPos(InstTy *Inst) { InsertPos = Inst; }
 
   /// Add metadata (e.g. alias info) from the instructions in this group to \p
   /// NewInst.
   ///
   /// FIXME: this function currently does not add noalias metadata a'la
   /// addNewMedata.  To do that we need to compute the intersection of the
   /// noalias info from all members.
-  void addMetadata(Instruction *NewInst) const {
-    SmallVector<Value *, 4> VL;
-    std::transform(Members.begin(), Members.end(), std::back_inserter(VL),
-                   [](std::pair<int, Instruction *> p) { return p.second; });
-    propagateMetadata(NewInst, VL);
-  }
+  void addMetadata(InstTy *NewInst) const;
 
   /// Returns true if this Group requires a scalar iteration to handle gaps.
   bool requiresScalarEpilogue() const {
     // If the last member of the Group exists, then a scalar epilog is not
     // needed for this group.
     if (getMember(getFactor() - 1))
       return false;
 
     // We have a group with gaps. It therefore cannot be a group of stores,
     // and it can't be a reversed access, because such groups get invalidated.
     assert(!getMember(0)->mayWriteToMemory() &&
            "Group should have been invalidated");
     assert(!isReverse() && "Group should have been invalidated");
 
     // This is a group of loads, with gaps, and without a last-member
     return true;
   }
 
 private:
   unsigned Factor; // Interleave Factor.
   bool Reverse;
   unsigned Align;
-  DenseMap<int, Instruction *> Members;
+  DenseMap<int, InstTy *> Members;
   int SmallestKey = 0;
   int LargestKey = 0;
 
   // To avoid breaking dependences, vectorized instructions of an interleave
   // group should be inserted at either the first load or the last store in
   // program order.
   //
   // E.g. %even = load i32             // Insert Position
   //      %add = add i32 %even         // Use of %even
   //      %odd = load i32
   //
   //      store i32 %even
   //      %odd = add i32               // Def of %odd
   //      store i32 %odd               // Insert Position
-  Instruction *InsertPos;
+  InstTy *InsertPos;
 };
 
 /// Drive the analysis of interleaved memory accesses in the loop.
 ///
 /// Use this class to analyze interleaved accesses only when we can vectorize
 /// a loop. Otherwise it's meaningless to do analysis as the vectorization
 /// on interleaved accesses is unsafe.
 ///
public:
   /// \p EnableMaskedInterleavedGroup is true.
   void analyzeInterleaving(bool EnableMaskedInterleavedGroup);
 
   /// Invalidate groups, e.g., in case all blocks in loop will be predicated
   /// contrary to original assumption. Although we currently prevent group
   /// formation for predicated accesses, we may be able to relax this limitation
   /// in the future once we handle more complicated blocks.
   void reset() {
-    SmallPtrSet<InterleaveGroup *, 4> DelSet;
+    SmallPtrSet<InterleaveGroup<Instruction> *, 4> DelSet;
     // Avoid releasing a pointer twice.
     for (auto &I : InterleaveGroupMap)
       DelSet.insert(I.second);
     for (auto *Ptr : DelSet)
       delete Ptr;
     InterleaveGroupMap.clear();
     RequiresScalarEpilogue = false;
   }
 
 
   /// Check if \p Instr belongs to any interleave group.
   bool isInterleaved(Instruction *Instr) const {
     return InterleaveGroupMap.find(Instr) != InterleaveGroupMap.end();
   }
 
   /// Get the interleave group that \p Instr belongs to.
   ///
   /// \returns nullptr if doesn't have such group.
-  InterleaveGroup *getInterleaveGroup(Instruction *Instr) const {
-    auto Group = InterleaveGroupMap.find(Instr);
-    if (Group == InterleaveGroupMap.end())
-      return nullptr;
-    return Group->second;
+  InterleaveGroup<Instruction> *
+  getInterleaveGroup(const Instruction *Instr) const {
+    if (InterleaveGroupMap.count(Instr))
+      return InterleaveGroupMap.find(Instr)->second;
+  }
+
+  iterator_range<SmallPtrSetIterator<llvm::InterleaveGroup<Instruction> *>>
+  getInterleaveGroups() {
+    return make_range(InterleaveGroups.begin(), InterleaveGroups.end());
   }
 
   /// Returns true if an interleaved group that may access memory
   /// out-of-bounds requires a scalar epilogue iteration for correctness.
   bool requiresScalarEpilogue() const { return RequiresScalarEpilogue; }
 
   /// Invalidate groups that require a scalar epilogue (due to gaps). This can
   /// happen when optimizing for size forbids a scalar epilogue, and the gap
private:
   const LoopAccessInfo *LAI;
 
   /// True if the loop may contain non-reversed interleaved groups with
   /// out-of-bounds accesses. We ensure we don't speculatively access memory
   /// out-of-bounds by executing at least one scalar epilogue iteration.
   bool RequiresScalarEpilogue = false;
 
   /// Holds the relationships between the members and the interleave group.
-  DenseMap<Instruction *, InterleaveGroup *> InterleaveGroupMap;
+  DenseMap<Instruction *, InterleaveGroup<Instruction> *> InterleaveGroupMap;
+
+  SmallPtrSet<InterleaveGroup<Instruction> *, 4> InterleaveGroups;
 
   /// Holds dependences among the memory accesses in the loop. It maps a source
   /// access to a set of dependent sink accesses.
   DenseMap<Instruction *, SmallPtrSet<Instruction *, 2>> Dependences;
 
   /// The descriptor for a strided memory access.
   struct StrideDescriptor {
     StrideDescriptor() = default;
private:
 
   /// A type for holding instructions and their stride descriptors.
   using StrideEntry = std::pair<Instruction *, StrideDescriptor>;
 
   /// Create a new interleave group with the given instruction \p Instr,
   /// stride \p Stride and alignment \p Align.
   ///
   /// \returns the newly created interleave group.
-  InterleaveGroup *createInterleaveGroup(Instruction *Instr, int Stride,
-                                         unsigned Align) {
-    assert(!isInterleaved(Instr) && "Already in an interleaved access group");
-    InterleaveGroupMap[Instr] = new InterleaveGroup(Instr, Stride, Align);
+  InterleaveGroup<Instruction> *
+  createInterleaveGroup(Instruction *Instr, int Stride, unsigned Align) {
+    assert(!InterleaveGroupMap.count(Instr) &&
+           "Already in an interleaved access group");
+    InterleaveGroupMap[Instr] =
+        new InterleaveGroup<Instruction>(Instr, Stride, Align);
+    InterleaveGroups.insert(InterleaveGroupMap[Instr]);
     return InterleaveGroupMap[Instr];
   }
 
   /// Release the group and remove all the relationships.
-  void releaseGroup(InterleaveGroup *Group) {
+  void releaseGroup(InterleaveGroup<Instruction> *Group) {
     for (unsigned i = 0; i < Group->getFactor(); i++)
       if (Instruction *Member = Group->getMember(i))
         InterleaveGroupMap.erase(Member);
 
+    InterleaveGroups.erase(Group);
     delete Group;
   }
 
   /// Collect all the accesses with a constant stride in program order.
   void collectConstStrideAccesses(
       MapVector<Instruction *, StrideDescriptor> &AccessStrideInfo,
       const ValueToValueMap &Strides);
 

lib/Analysis/VectorUtils.cpp

for (auto Kind :
     }
 
     Inst->setMetadata(Kind, MD);
   }
 
   return Inst;
 }
 
-Constant *llvm::createBitMaskForGaps(IRBuilder<> &Builder, unsigned VF,
-                                           const InterleaveGroup &Group) {
+Constant *
+llvm::createBitMaskForGaps(IRBuilder<> &Builder, unsigned VF,
+                           const InterleaveGroup<Instruction> &Group) {
   // All 1's means mask is not needed.
   if (Group.getNumMembers() == Group.getFactor())
     return nullptr;
 
   // TODO: support reversed access.
   assert(!Group.isReverse() && "Reversed group not supported.");
 
   SmallVector<Constant *, 16> Mask;
void InterleavedAccessInfo::analyzeInterleaving(
 
   if (AccessStrideInfo.empty())
     return;
 
   // Collect the dependences in the loop.
   collectDependences();
 
   // Holds all interleaved store groups temporarily.
-  SmallSetVector<InterleaveGroup *, 4> StoreGroups;
+  SmallSetVector<InterleaveGroup<Instruction> *, 4> StoreGroups;
   // Holds all interleaved load groups temporarily.
-  SmallSetVector<InterleaveGroup *, 4> LoadGroups;
+  SmallSetVector<InterleaveGroup<Instruction> *, 4> LoadGroups;
 
   // Search in bottom-up program order for pairs of accesses (A and B) that can
   // form interleaved load or store groups. In the algorithm below, access A
   // precedes access B in program order. We initialize a group for B in the
   // outer loop of the algorithm, and then in the inner loop, we attempt to
   // insert each A into B's group if:
   //
   //  1. A and B have the same stride,
   //  2. A and B have the same memory object size, and
   //  3. A belongs in B's group according to its distance from B.
   //
   // Special care is taken to ensure group formation will not break any
   // dependences.
   for (auto BI = AccessStrideInfo.rbegin(), E = AccessStrideInfo.rend();
        BI != E; ++BI) {
     Instruction *B = BI->first;
     StrideDescriptor DesB = BI->second;
 
     // Initialize a group for B if it has an allowable stride. Even if we don't
     // create a group for B, we continue with the bottom-up algorithm to ensure
     // we don't break any of B's dependences.
-    InterleaveGroup *Group = nullptr;
+    InterleaveGroup<Instruction> *Group = nullptr;
     if (isStrided(DesB.Stride) && 
         (!isPredicated(B->getParent()) || EnablePredicatedInterleavedMemAccesses)) {
       Group = getInterleaveGroup(B);
       if (!Group) {
         LLVM_DEBUG(dbgs() << "LV: Creating an interleave group with:" << *B
                           << '\n');
         Group = createInterleaveGroup(B, DesB.Stride, DesB.Align);
       }
for (auto AI = std::next(BI); AI != E; ++AI) {
       // the boundaries of the (2, 4) group.
       if (!canReorderMemAccessesForInterleavedGroups(&*AI, &*BI)) {
         // If a dependence exists and A is already in a group, we know that A
         // must be a store since A precedes B and WAR dependences are allowed.
         // Thus, A would be sunk below B. We release A's group to prevent this
         // illegal code motion. A will then be free to form another group with
         // instructions that precede it.
         if (isInterleaved(A)) {
-          InterleaveGroup *StoreGroup = getInterleaveGroup(A);
+          InterleaveGroup<Instruction> *StoreGroup = getInterleaveGroup(A);
           StoreGroups.remove(StoreGroup);
           releaseGroup(StoreGroup);
         }
 
         // If a dependence exists and A is not already in a group (or it was
         // and we just released it), B might be hoisted above A (if B is a
         // load) or another store might be sunk below A (if B is a store). In
         // either case, we can't add additional instructions to B's group. B
for (auto AI = std::next(BI); AI != E; ++AI) {
         // Set the first load in program order as the insert position.
         if (A->mayReadFromMemory())
           Group->setInsertPos(A);
       }
     } // Iteration over A accesses.
   }   // Iteration over B accesses.
 
   // Remove interleaved store groups with gaps.
-  for (InterleaveGroup *Group : StoreGroups)
+  for (auto *Group : StoreGroups)
     if (Group->getNumMembers() != Group->getFactor()) {
       LLVM_DEBUG(
           dbgs() << "LV: Invalidate candidate interleaved store group due "
                     "to gaps.\n");
       releaseGroup(Group);
     }
   // Remove interleaved groups with gaps (currently only loads) whose memory
   // accesses may wrap around. We have to revisit the getPtrStride analysis,
   // this time with ShouldCheckWrap=true, since collectConstStrideAccesses does
   // not check wrapping (see documentation there).
   // FORNOW we use Assume=false;
   // TODO: Change to Assume=true but making sure we don't exceed the threshold
   // of runtime SCEV assumptions checks (thereby potentially failing to
   // vectorize altogether).
   // Additional optional optimizations:
   // TODO: If we are peeling the loop and we know that the first pointer doesn't
   // wrap then we can deduce that all pointers in the group don't wrap.
   // This means that we can forcefully peel the loop in order to only have to
   // check the first pointer for no-wrap. When we'll change to use Assume=true
   // we'll only need at most one runtime check per interleaved group.
-  for (InterleaveGroup *Group : LoadGroups) {
+  for (auto *Group : LoadGroups) {
     // Case 1: A full group. Can Skip the checks; For full groups, if the wide
     // load would wrap around the address space we would do a memory access at
     // nullptr even without the transformation.
     if (Group->getNumMembers() == Group->getFactor())
       continue;
 
     // Case 2: If first and last members of the group don't wrap this implies
     // that all the pointers in the group don't wrap.

 
 void InterleavedAccessInfo::invalidateGroupsRequiringScalarEpilogue() {
   // If no group had triggered the requirement to create an epilogue loop,
   // there is nothing to do.
   if (!requiresScalarEpilogue())
     return;
 
   // Avoid releasing a Group twice.
-  SmallPtrSet<InterleaveGroup *, 4> DelSet;
+  SmallPtrSet<InterleaveGroup<Instruction> *, 4> DelSet;
   for (auto &I : InterleaveGroupMap) {
-    InterleaveGroup *Group = I.second;
+    InterleaveGroup<Instruction> *Group = I.second;
     if (Group->requiresScalarEpilogue())
       DelSet.insert(Group);
   }
   for (auto *Ptr : DelSet) {
     LLVM_DEBUG(
         dbgs()
         << "LV: Invalidate candidate interleaved group due to gaps that "
            "require a scalar epilogue (not allowed under optsize) and cannot "
            "be masked (not enabled). \n");
     releaseGroup(Ptr);
   }
 
   RequiresScalarEpilogue = false;
 }
+
+template <>
+void InterleaveGroup<Instruction>::addMetadata(Instruction *NewInst) const {
+  SmallVector<Value *, 4> VL;
+  std::transform(Members.begin(), Members.end(), std::back_inserter(VL),
+                 [](std::pair<int, Instruction *> p) { return p.second; });
+  propagateMetadata(NewInst, VL);
+}
+
+template <typename InstT>
+void InterleaveGroup<InstT>::addMetadata(InstT *NewInst) const {
+  llvm_unreachable("addMetadata can only be used for Instruction");
+}

lib/Transforms/Vectorize/LoopVectorize.cpp

public:
   void setWideningDecision(Instruction *I, unsigned VF, InstWidening W,
                            unsigned Cost) {
     assert(VF >= 2 && "Expected VF >=2");
     WideningDecisions[std::make_pair(I, VF)] = std::make_pair(W, Cost);
   }
 
   /// Save vectorization decision \p W and \p Cost taken by the cost model for
   /// interleaving group \p Grp and vector width \p VF.
-  void setWideningDecision(const InterleaveGroup *Grp, unsigned VF,
+  void setWideningDecision(const InterleaveGroup<Instruction> *Grp, unsigned VF,
                            InstWidening W, unsigned Cost) {
     assert(VF >= 2 && "Expected VF >=2");
     /// Broadcast this decicion to all instructions inside the group.
     /// But the cost will be assigned to one instruction only.
     for (unsigned i = 0; i < Grp->getFactor(); ++i) {
       if (auto *I = Grp->getMember(i)) {
         if (Grp->getInsertPos() == I)
           WideningDecisions[std::make_pair(I, VF)] = std::make_pair(W, Cost);
public:
   bool interleavedAccessCanBeWidened(Instruction *I, unsigned VF = 1);
 
   /// Check if \p Instr belongs to any interleaved access group.
   bool isAccessInterleaved(Instruction *Instr) {
     return InterleaveInfo.isInterleaved(Instr);
   }
 
   /// Get the interleaved access group that \p Instr belongs to.
-  const InterleaveGroup *getInterleavedAccessGroup(Instruction *Instr) {
+  const InterleaveGroup<Instruction> *
+  getInterleavedAccessGroup(Instruction *Instr) {
     return InterleaveInfo.getInterleaveGroup(Instr);
   }
 
   /// Returns true if an interleaved group requires a scalar iteration
   /// to handle accesses with gaps, and there is nothing preventing us from
   /// creating a scalar epilogue.
   bool requiresScalarEpilogue() const {
     return IsScalarEpilogueAllowed && InterleaveInfo.requiresScalarEpilogue();

 // To:
 //   %R_G.vec = shuffle %R.vec, %G.vec, <0, 1, 2, ..., 7>
 //   %B_U.vec = shuffle %B.vec, undef, <0, 1, 2, 3, u, u, u, u>
 //   %interleaved.vec = shuffle %R_G.vec, %B_U.vec,
 //        <0, 4, 8, 1, 5, 9, 2, 6, 10, 3, 7, 11>    ; Interleave R,G,B elements
 //   store <12 x i32> %interleaved.vec              ; Write 4 tuples of R,G,B
 void InnerLoopVectorizer::vectorizeInterleaveGroup(Instruction *Instr,
                                                    VectorParts *BlockInMask) {
-  const InterleaveGroup *Group = Cost->getInterleavedAccessGroup(Instr);
+  const InterleaveGroup<Instruction> *Group =
+      Cost->getInterleavedAccessGroup(Instr);
   assert(Group && "Fail to get an interleaved access group.");
 
   // Skip if current instruction is not the insert position.
   if (Instr != Group->getInsertPos())
     return;
 
   const DataLayout &DL = Instr->getModule()->getDataLayout();
   Value *Ptr = getLoadStorePointerOperand(Instr);
VPValue *VPRecipeBuilder::createBlockInMask(BasicBlock *BB, VPlanPtr &Plan) {
   }
 
   return BlockMaskCache[BB] = BlockMask;
 }
 
 VPInterleaveRecipe *VPRecipeBuilder::tryToInterleaveMemory(Instruction *I,
                                                            VFRange &Range,
                                                            VPlanPtr &Plan) {
-  const InterleaveGroup *IG = CM.getInterleavedAccessGroup(I);
+  const InterleaveGroup<Instruction> *IG = CM.getInterleavedAccessGroup(I);
   if (!IG)
     return nullptr;
 
   // Now check if IG is relevant for VF's in the given range.
   auto isIGMember = [&](Instruction *I) -> std::function<bool(unsigned)> {
     return [=](unsigned VF) -> bool {
       return (VF >= 2 && // Query is illegal for VF == 1
               CM.getWideningDecision(I, VF) ==
for (Instruction &I : BB->instructionsWithoutDebug()) {
       // First filter out irrelevant instructions, to ensure no recipes are
       // built for them.
       if (isa<BranchInst>(Instr) ||
           DeadInstructions.find(Instr) != DeadInstructions.end())
         continue;
 
       // I is a member of an InterleaveGroup for Range.Start. If it's an adjunct
       // member of the IG, do not construct any Recipe for it.
-      const InterleaveGroup *IG = CM.getInterleavedAccessGroup(Instr);
+      const InterleaveGroup<Instruction> *IG =
+          CM.getInterleavedAccessGroup(Instr);
       if (IG && Instr != IG->getInsertPos() &&
           Range.Start >= 2 && // Query is illegal for VF == 1
           CM.getWideningDecision(Instr, Range.Start) ==
               LoopVectorizationCostModel::CM_Interleave) {
         auto SinkCandidate = SinkAfterInverse.find(Instr);
         if (SinkCandidate != SinkAfterInverse.end())
           Ingredients.push_back(SinkCandidate->second);
         continue;

lib/Transforms/Vectorize/VPlan.h

 #include "llvm/ADT/GraphTraits.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/ADT/ilist.h"
 #include "llvm/ADT/ilist_node.h"
+#include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/IRBuilder.h"
 #include <algorithm>
 #include <cassert>
 #include <cstddef>
 #include <map>
 #include <string>
 
 namespace llvm {
 
 class LoopVectorizationLegality;
 class LoopVectorizationCostModel;
 class BasicBlock;
 class DominatorTree;
 class InnerLoopVectorizer;
-class InterleaveGroup;
+template <class T> class InterleaveGroup;
+class LoopInfo;
 class raw_ostream;
 class Value;
 class VPBasicBlock;
 class VPRegionBlock;
 class VPlan;
 
 /// A range of powers-of-2 vectorization factors with fixed start and
 /// adjustable end. The range includes start and excludes end, e.g.,:
public:
   /// Print the recipe.
   void print(raw_ostream &O, const Twine &Indent) const override;
 };
 
 /// VPInterleaveRecipe is a recipe for transforming an interleave group of load
 /// or stores into one wide load/store and shuffles.
 class VPInterleaveRecipe : public VPRecipeBase {
 private:
-  const InterleaveGroup *IG;
+  const InterleaveGroup<Instruction> *IG;
   std::unique_ptr<VPUser> User;
 
 public:
-  VPInterleaveRecipe(const InterleaveGroup *IG, VPValue *Mask)
+  VPInterleaveRecipe(const InterleaveGroup<Instruction> *IG, VPValue *Mask)
       : VPRecipeBase(VPInterleaveSC), IG(IG) {
     if (Mask) // Create a VPInstruction to register as a user of the mask.
       User.reset(new VPUser({Mask}));
   }
   ~VPInterleaveRecipe() override = default;
 
   /// Method to support type inquiry through isa, cast, and dyn_cast.
   static inline bool classof(const VPRecipeBase *V) {
     return V->getVPRecipeID() == VPRecipeBase::VPInterleaveSC;
   }
 
   /// Generate the wide load or store, and shuffles.
   void execute(VPTransformState &State) override;
 
   /// Print the recipe.
   void print(raw_ostream &O, const Twine &Indent) const override;
 
-  const InterleaveGroup *getInterleaveGroup() { return IG; }
+  const InterleaveGroup<Instruction> *getInterleaveGroup() { return IG; }
 };
 
 /// VPReplicateRecipe replicates a given instruction producing multiple scalar
 /// copies of the original scalar type, one per lane, instead of producing a
 /// single copy of widened type for all lanes. If the instruction is known to be
 /// uniform only one copy, per lane zero, will be generated.
 class VPReplicateRecipe : public VPRecipeBase {
 private:
public:
   /// Disconnect VPBlockBases \p From and \p To bi-directionally. Remove \p To
   /// from the successors of \p From and \p From from the predecessors of \p To.
   static void disconnectBlocks(VPBlockBase *From, VPBlockBase *To) {
     assert(To && "Successor to disconnect is null.");
     From->removeSuccessor(To);
     To->removePredecessor(From);
   }
 };
 
+class VPInterleavedAccessInfo {
+private:
+  DenseMap<VPInstruction *, InterleaveGroup<VPInstruction> *>
+      InterleaveGroupMap;
+
+public:
+  VPInterleavedAccessInfo(VPlan &Plan, InterleavedAccessInfo &IAI);
+
+  ~VPInterleavedAccessInfo() {
+    SmallPtrSet<InterleaveGroup<VPInstruction> *, 4> DelSet;
+    // Avoid releasing a pointer twice.
+    for (auto &I : InterleaveGroupMap)
+      DelSet.insert(I.second);
+    for (auto *Ptr : DelSet)
+      delete Ptr;
+  }
+
+  /// Get the interleave group that \p Instr belongs to.
+  ///
+  /// \returns nullptr if doesn't have such group.
+  InterleaveGroup<VPInstruction> *
+  getInterleaveGroup(VPInstruction *Instr) const {
+    if (InterleaveGroupMap.count(Instr))
+      return InterleaveGroupMap.find(Instr)->second;
+    return nullptr;
+  }
+};
+
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_VECTORIZE_VPLAN_H

lib/Transforms/Vectorize/VPlan.cpp

void VPWidenMemoryInstructionRecipe::print(raw_ostream &O,
   if (User) {
     O << ", ";
     User->getOperand(0)->printAsOperand(O);
   }
   O << "\\l\"";
 }
 
 template void DomTreeBuilder::Calculate<VPDominatorTree>(VPDominatorTree &DT);
+
+VPInterleavedAccessInfo::VPInterleavedAccessInfo(VPlan &Plan,
+                                                 InterleavedAccessInfo &IAI) {
+  DenseMap<InterleaveGroup<Instruction> *, InterleaveGroup<VPInstruction> *>
+      Old2New;
+
+  auto *TopRegion = cast<VPRegionBlock>(Plan.getEntry());
+  ReversePostOrderTraversal<VPBlockBase *> RPOT(TopRegion->getEntry());
+  for (VPBlockBase *Base : RPOT) {
+    VPBasicBlock *VPBB = Base->getEntryBasicBlock();
+    for (auto I = VPBB->begin(), E = VPBB->end(); I != E; I++) {
+      assert(isa<VPInstruction>(&*I) && "Can only handle VPInstructions");
+      auto *VPInst = cast<VPInstruction>(&*I);
+      auto *Inst = cast<Instruction>(VPInst->getUnderlyingValue());
+      auto *IG = IAI.getInterleaveGroup(Inst);
+      if (!IG)
+        continue;
+
+      auto NewIGIter = Old2New.find(IG);
+      if (NewIGIter == Old2New.end())
+        Old2New[IG] = new InterleaveGroup<VPInstruction>(
+            IG->getFactor(), IG->isReverse(), IG->getAlignment());
+
+      if (Inst == IG->getInsertPos())
+        Old2New[IG]->setInsertPos(VPInst);
+
+      InterleaveGroupMap[VPInst] = Old2New[IG];
+      InterleaveGroupMap[VPInst]->insertMember(
+          VPInst, IG->getIndex(Inst),
+          IG->isReverse() ? (-1) * int(IG->getFactor()) : IG->getFactor());
+    }
+  }
+}

lib/Transforms/Vectorize/VPlanValue.h

 // This is the base class of the VPlan Def/Use graph, used for modeling the data
 // flow into, within and out of the VPlan. VPValues can stand for live-ins
 // coming from the input IR, instructions which VPlan will generate if executed
 // and live-outs which the VPlan will need to fix accordingly.
 class VPValue {
   friend class VPBuilder;
   friend class VPlanHCFGTransforms;
   friend class VPBasicBlock;
+  friend class VPInterleavedAccessInfo;
 
 private:
   const unsigned char SubclassID; ///< Subclass identifier (for isa/dyn_cast).
 
   SmallVector<VPUser *, 1> Users;
 
 protected:
   // Hold the underlying Value, if any, attached to this VPValue.