[AArch64][ARM] Match interleaved memory accesses into ldN/stN/vldN/vstN intrinsics.

[• HaoLiu]

Summary

Hi,

I refactored the patch in D10335 according to the comments. The main concern is to share as much code as possible.

This patch adds a new pass called InterleavedAccessPass in the lib/CodeGen. It constains the code about identifying interleaved accesses. As we can not share the code about creating target specific intrinsics, I put such code in the target backends.

Review please.

Thanks,
-Hao

[sbaranga]

lib/Target/ARM/ARMTargetMachine.cpp

41 ↗

(On Diff #27923)

Would it be better to only have one switch in the interleave pass instead of having a separate switch in each backend?

The pass could return when executing runOnFunction if the option is not enabled.

[• HaoLiu]

Updated a new patch refactored according to Silviu's comment.

Review please.

Thanks,
-Hao

lib/Target/ARM/ARMTargetMachine.cpp

41 ↗

(On Diff #27923)

Reasonable.

[rengolin]

lib/CodeGen/InterleavedAccessPass.cpp

113 ↗

(On Diff #28006)

If the mask index can't be negative, why use ArrayRef<int>?

134 ↗

(On Diff #28006)

Checking for all factors "up to" in isDeInterleaveMaskOfFactor() is redundant with this line.

Though, I see that you're using it in other functions that may need that functionality.

Not sure how to split this, but it looks inefficient...

[rengolin]

lib/CodeGen/InterleavedAccessPass.cpp

222 ↗

(On Diff #28006)

CGP?

[• HaoLiu]

Updated a patch according to Renato's comments.

Review please.

Thanks,
-Hao

lib/CodeGen/InterleavedAccessPass.cpp

113 ↗

(On Diff #28006)

It could be negative. When a mask is undef, it is -1.
Here we only compare non-negative masks and ignore undef masks.

134 ↗

(On Diff #28006)

I merged isDeInterleaveMask() and isDeInterleaveMaskOfFactor() into one function isDeInterleaveMask().

222 ↗

(On Diff #28006)

Fixed.

[• HaoLiu]

ping...

Thanks,
-Hao

[sbaranga]

LGTM from me, but I think Michael also wanted to have a look and perhaps Renato has further comments?

[rengolin]

lib/CodeGen/InterleavedAccessPass.cpp

134 ↗

(On Diff #28006)

Hum, I'm still seeing isDeInterleaveMaskOfFactor in the latest patch...

[mzolotukhin]

Hi Hao,

The code generally looks fine, but I have a question regarding lowerInterleavedStore (please see inline).

Thanks,
Michael

lib/CodeGen/InterleavedAccessPass.cpp

32–33 ↗

(On Diff #28203)

How would IR look for 4 vectors? Will we have a shuffle of shuffles?

56–57 ↗

(On Diff #28203)

Do these names comply with the coding standards?

240–242 ↗

(On Diff #28203)

Will it work for Factor != 2? If not, and other factors aren't supported for now, please add an explicit assert and TODO for it. If yes, should we also check the other shuffles?

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

415 ↗

(On Diff #28203)

Nitpick: I'd prefer comparing with 2 and 4, instead of 1 and 5. I.e.

if (Factor >= 2 && Factor <= 4)

Also, could we somehow reuse MIN_FACTOR and MAX_FACTOR from InterleavedAccessPass.cpp here? Having the same constants in different places will lead to bugs in future.

lib/Transforms/Vectorize/LoopVectorize.cpp

142 ↗

(On Diff #28203)

This change doesn't belong here and anyway needs a separate discussion.

[• HaoLiu]

Hi,

I refactored the patch according to Michael's comments. As well as inline comments to answer questions from Renato and Michael.

Review please.

Thanks,
-Hao

lib/CodeGen/InterleavedAccessPass.cpp

32–33 ↗

(On Diff #28203)

Will have a shuffle.
E.g. An interleaved store of factor 4.

%i.vec = shuffle <8 x i32> %v0, <8 x i32> %v2, <0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15>
store <16 x i32> %i.vec, <16 x i32>* %ptr

%v0 and %v2 could be concatenated from other small vectors, like:

%v0 = shuffle <4 x i32> %A, <4 x i32> %B, <0, 1, 2, 3, 4, 5, 6, 7>
%v1 = shuffle <4 x i32> %C, <4 x i32> %D, <0, 1, 2, 3, 4, 5, 6, 7>

but we only need to check the last shuffle with the RE-interleaved mask.

240–242 ↗

(On Diff #28203)

Yes, it will work for other factor.
As the previous example of factor 4, we only need to check the last shuffle with RE-interleaved mask.

134 ↗

(On Diff #28006)

Sorry. I misunderstood and misleaded. I merged isReInterleaveMask() and isReInterleaveMaskOfFactor().

I cannot merge isDeInterleaveMask() and isDeInterleaveMaskOfFactor(), which are both used in lowerInterleavedLoad(). The former is used to check and find an interleave factor. The later is only used to check whether the given mask is the DE-interleaved of the given factor.

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

415 ↗

(On Diff #28203)

I refactored to add a hook called getMaxSupportedInterleaveFactor(), which is used to share the maximum factor supported by a target.
No need to get the minimum factor, which is always 2.

lib/Transforms/Vectorize/LoopVectorize.cpp

142 ↗

(On Diff #28203)

Yes. Forgot to clean.

[rengolin]

lib/CodeGen/InterleavedAccessPass.cpp

134 ↗

(On Diff #28006)

Right, I thought it was weird that you had merged them. :)

[rengolin]

Thanks Hao, LGTM.

[mzolotukhin]

Hi Hao,

Thanks for the explanation!

The code looks good to me now. I'd also recommend to check-in shared part and the target implementations in separate commits - by the logic everything should work fine even without them, right? (and that's what will happen on other targets)

Best regards,
Michael

test/CodeGen/AArch64/aarch64-interleaved-accesses.ll

13–15 ↗

(On Diff #28439)

s/delat/delta/

test/CodeGen/ARM/arm-interleaved-accesses.ll

13–15 ↗

(On Diff #28439)

s/delat/delta/ ?

[• HaoLiu]

Hi

I committed this patch separately in r240751, r240754 and r240755. Thanks again for your review comments.

As interleaved access is disabled by default in LoopVectorize, currently this pass is also disabled by default. The next step is to enable it by default. My dear colleague (probably Silviu) will do benchmarking work and enable it by default later.

Thanks,
-Hao

[MatzeB]

Is there a reason this was implemented as an IR pass and not at the SelectionDAG level?

[rengolin]

In D10533#660846, @MatzeB wrote:

Is there a reason this was implemented as an IR pass and not at the SelectionDAG level?

Wow, this brings back memories, but unfortunately not the reason for an additional pass. I think it was something to do with being after the register allocation, though I may be misquoting someone.

llvm/trunk/include/llvm/CodeGen/Passes.h

/// MachineDominanaceFrontier - This pass is a machine dominators analysis pass.
   extern char &StackMapLivenessID;
 
   /// createJumpInstrTables - This pass creates jump-instruction tables.
   ModulePass *createJumpInstrTablesPass();
 
   /// createForwardControlFlowIntegrityPass - This pass adds control-flow
   /// integrity.
   ModulePass *createForwardControlFlowIntegrityPass();
+
+  /// InterleavedAccess Pass - This pass identifies and matches interleaved
+  /// memory accesses to target specific intrinsics.
+  ///
+  FunctionPass *createInterleavedAccessPass(const TargetMachine *TM);
 } // End llvm namespace
 
 /// Target machine pass initializer for passes with dependencies. Use with
 /// INITIALIZE_TM_PASS_END.
 #define INITIALIZE_TM_PASS_BEGIN INITIALIZE_PASS_BEGIN
 
 /// Target machine pass initializer for passes with dependencies. Use with
 /// INITIALIZE_TM_PASS_BEGIN.

llvm/trunk/include/llvm/Target/TargetLowering.h

virtual bool hasPairedLoad(Type * /*LoadedType*/,
     return false;
   }
 
   virtual bool hasPairedLoad(EVT /*LoadedType*/,
                              unsigned & /*RequiredAligment*/) const {
     return false;
   }
 
+  /// \brief Get the maximum supported factor for interleaved memory accesses.
+  /// Default to be the minimum interleave factor: 2.
+  virtual unsigned getMaxSupportedInterleaveFactor() const { return 2; }
+
+  /// \brief Lower an interleaved load to target specific intrinsics. Return
+  /// true on success.
+  ///
+  /// \p LI is the vector load instruction.
+  /// \p Shuffles is the shufflevector list to DE-interleave the loaded vector.
+  /// \p Indices is the corresponding indices for each shufflevector.
+  /// \p Factor is the interleave factor.
+  virtual bool lowerInterleavedLoad(LoadInst *LI,
+                                    ArrayRef<ShuffleVectorInst *> Shuffles,
+                                    ArrayRef<unsigned> Indices,
+                                    unsigned Factor) const {
+    return false;
+  }
+
+  /// \brief Lower an interleaved store to target specific intrinsics. Return
+  /// true on success.
+  ///
+  /// \p SI is the vector store instruction.
+  /// \p SVI is the shufflevector to RE-interleave the stored vector.
+  /// \p Factor is the interleave factor.
+  virtual bool lowerInterleavedStore(StoreInst *SI, ShuffleVectorInst *SVI,
+                                     unsigned Factor) const {
+    return false;
+  }
+
   /// Return true if zero-extending the specific node Val to type VT2 is free
   /// (either because it's implicitly zero-extended such as ARM ldrb / ldrh or
   /// because it's folded such as X86 zero-extending loads).
   virtual bool isZExtFree(SDValue Val, EVT VT2) const {
     return isZExtFree(Val.getValueType(), VT2);
   }
 
   /// Return true if an fpext operation is free (for instance, because

llvm/trunk/lib/CodeGen/CMakeLists.txt

add_llvm_library(LLVMCodeGen
   GCMetadataPrinter.cpp
   GCRootLowering.cpp
   GCStrategy.cpp
   GlobalMerge.cpp
   IfConversion.cpp
   ImplicitNullChecks.cpp
   InlineSpiller.cpp
   InterferenceCache.cpp
+  InterleavedAccessPass.cpp
   IntrinsicLowering.cpp
   LLVMTargetMachine.cpp
   LatencyPriorityQueue.cpp
   LexicalScopes.cpp
   LiveDebugVariables.cpp
   LiveInterval.cpp
   LiveIntervalAnalysis.cpp
   LiveIntervalUnion.cpp

llvm/trunk/lib/CodeGen/InterleavedAccessPass.cpp

+//=----------------------- InterleavedAccessPass.cpp -----------------------==//
+//
+// The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the Interleaved Access pass, which identifies
+// interleaved memory accesses and transforms into target specific intrinsics.
+//
+// An interleaved load reads data from memory into several vectors, with
+// DE-interleaving the data on a factor. An interleaved store writes several
+// vectors to memory with RE-interleaving the data on a factor.
+//
+// As interleaved accesses are hard to be identified in CodeGen (mainly because
+// the VECTOR_SHUFFLE DAG node is quite different from the shufflevector IR),
+// we identify and transform them to intrinsics in this pass. So the intrinsics
+// can be easily matched into target specific instructions later in CodeGen.
+//
+// E.g. An interleaved load (Factor = 2):
+//        %wide.vec = load <8 x i32>, <8 x i32>* %ptr
+//        %v0 = shuffle <8 x i32> %wide.vec, <8 x i32> undef, <0, 2, 4, 6>
+//        %v1 = shuffle <8 x i32> %wide.vec, <8 x i32> undef, <1, 3, 5, 7>
+//
+// It could be transformed into a ld2 intrinsic in AArch64 backend or a vld2
+// intrinsic in ARM backend.
+//
+// E.g. An interleaved store (Factor = 3):
+//        %i.vec = shuffle <8 x i32> %v0, <8 x i32> %v1,
+//                                    <0, 4, 8, 1, 5, 9, 2, 6, 10, 3, 7, 11>
+//        store <12 x i32> %i.vec, <12 x i32>* %ptr
+//
+// It could be transformed into a st3 intrinsic in AArch64 backend or a vst3
+// intrinsic in ARM backend.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/MathExtras.h"
+#include "llvm/Target/TargetLowering.h"
+#include "llvm/Target/TargetSubtargetInfo.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "interleaved-access"
+
+static cl::opt<bool> LowerInterleavedAccesses(
+    "lower-interleaved-accesses",
+    cl::desc("Enable lowering interleaved accesses to intrinsics"),
+    cl::init(false), cl::Hidden);
+
+static unsigned MaxFactor; // The maximum supported interleave factor.
+
+namespace llvm {
+static void initializeInterleavedAccessPass(PassRegistry &);
+}
+
+namespace {
+
+class InterleavedAccess : public FunctionPass {
+
+public:
+  static char ID;
+  InterleavedAccess(const TargetMachine *TM = nullptr)
+      : FunctionPass(ID), TM(TM), TLI(nullptr) {
+    initializeInterleavedAccessPass(*PassRegistry::getPassRegistry());
+  }
+
+  const char *getPassName() const override { return "Interleaved Access Pass"; }
+
+  bool runOnFunction(Function &F) override;
+
+private:
+  const TargetMachine *TM;
+  const TargetLowering *TLI;
+
+  /// \brief Transform an interleaved load into target specific intrinsics.
+  bool lowerInterleavedLoad(LoadInst *LI,
+                            SmallVector<Instruction *, 32> &DeadInsts);
+
+  /// \brief Transform an interleaved store into target specific intrinsics.
+  bool lowerInterleavedStore(StoreInst *SI,
+                             SmallVector<Instruction *, 32> &DeadInsts);
+};
+} // end anonymous namespace.
+
+char InterleavedAccess::ID = 0;
+INITIALIZE_TM_PASS(InterleavedAccess, "interleaved-access",
+    "Lower interleaved memory accesses to target specific intrinsics",
+    false, false)
+
+FunctionPass *llvm::createInterleavedAccessPass(const TargetMachine *TM) {
+  return new InterleavedAccess(TM);
+}
+
+/// \brief Check if the mask is a DE-interleave mask of the given factor
+/// \p Factor like:
+///     <Index, Index+Factor, ..., Index+(NumElts-1)*Factor>
+static bool isDeInterleaveMaskOfFactor(ArrayRef<int> Mask, unsigned Factor,
+                                       unsigned &Index) {
+  // Check all potential start indices from 0 to (Factor - 1).
+  for (Index = 0; Index < Factor; Index++) {
+    unsigned i = 0;
+
+    // Check that elements are in ascending order by Factor. Ignore undef
+    // elements.
+    for (; i < Mask.size(); i++)
+      if (Mask[i] >= 0 && static_cast<unsigned>(Mask[i]) != Index + i * Factor)
+        break;
+
+    if (i == Mask.size())
+      return true;
+  }
+
+  return false;
+}
+
+/// \brief Check if the mask is a DE-interleave mask for an interleaved load.
+///
+/// E.g. DE-interleave masks (Factor = 2) could be:
+///     <0, 2, 4, 6>    (mask of index 0 to extract even elements)
+///     <1, 3, 5, 7>    (mask of index 1 to extract odd elements)
+static bool isDeInterleaveMask(ArrayRef<int> Mask, unsigned &Factor,
+                               unsigned &Index) {
+  if (Mask.size() < 2)
+    return false;
+
+  // Check potential Factors.
+  for (Factor = 2; Factor <= MaxFactor; Factor++)
+    if (isDeInterleaveMaskOfFactor(Mask, Factor, Index))
+      return true;
+
+  return false;
+}
+
+/// \brief Check if the mask is RE-interleave mask for an interleaved store.
+///
+/// I.e. <0, NumSubElts, ... , NumSubElts*(Factor - 1), 1, NumSubElts + 1, ...>
+///
+/// E.g. The RE-interleave mask (Factor = 2) could be:
+///     <0, 4, 1, 5, 2, 6, 3, 7>
+static bool isReInterleaveMask(ArrayRef<int> Mask, unsigned &Factor) {
+  unsigned NumElts = Mask.size();
+  if (NumElts < 4)
+    return false;
+
+  // Check potential Factors.
+  for (Factor = 2; Factor <= MaxFactor; Factor++) {
+    if (NumElts % Factor)
+      continue;
+
+    unsigned NumSubElts = NumElts / Factor;
+    if (!isPowerOf2_32(NumSubElts))
+      continue;
+
+    // Check whether each element matchs the RE-interleaved rule. Ignore undef
+    // elements.
+    unsigned i = 0;
+    for (; i < NumElts; i++)
+      if (Mask[i] >= 0 &&
+          static_cast<unsigned>(Mask[i]) !=
+              (i % Factor) * NumSubElts + i / Factor)
+        break;
+
+    // Find a RE-interleaved mask of current factor.
+    if (i == NumElts)
+      return true;
+  }
+
+  return false;
+}
+
+bool InterleavedAccess::lowerInterleavedLoad(
+    LoadInst *LI, SmallVector<Instruction *, 32> &DeadInsts) {
+  if (!LI->isSimple())
+    return false;
+
+  SmallVector<ShuffleVectorInst *, 4> Shuffles;
+
+  // Check if all users of this load are shufflevectors.
+  for (auto UI = LI->user_begin(), E = LI->user_end(); UI != E; UI++) {
+    ShuffleVectorInst *SVI = dyn_cast<ShuffleVectorInst>(*UI);
+    if (!SVI || !isa<UndefValue>(SVI->getOperand(1)))
+      return false;
+
+    Shuffles.push_back(SVI);
+  }
+
+  if (Shuffles.empty())
+    return false;
+
+  unsigned Factor, Index;
+
+  // Check if the first shufflevector is DE-interleave shuffle.
+  if (!isDeInterleaveMask(Shuffles[0]->getShuffleMask(), Factor, Index))
+    return false;
+
+  // Holds the corresponding index for each DE-interleave shuffle.
+  SmallVector<unsigned, 4> Indices;
+  Indices.push_back(Index);
+
+  Type *VecTy = Shuffles[0]->getType();
+
+  // Check if other shufflevectors are also DE-interleaved of the same type
+  // and factor as the first shufflevector.
+  for (unsigned i = 1; i < Shuffles.size(); i++) {
+    if (Shuffles[i]->getType() != VecTy)
+      return false;
+
+    if (!isDeInterleaveMaskOfFactor(Shuffles[i]->getShuffleMask(), Factor,
+                                    Index))
+      return false;
+
+    Indices.push_back(Index);
+  }
+
+  DEBUG(dbgs() << "IA: Found an interleaved load: " << *LI << "\n");
+
+  // Try to create target specific intrinsics to replace the load and shuffles.
+  if (!TLI->lowerInterleavedLoad(LI, Shuffles, Indices, Factor))
+    return false;
+
+  for (auto SVI : Shuffles)
+    DeadInsts.push_back(SVI);
+
+  DeadInsts.push_back(LI);
+  return true;
+}
+
+bool InterleavedAccess::lowerInterleavedStore(
+    StoreInst *SI, SmallVector<Instruction *, 32> &DeadInsts) {
+  if (!SI->isSimple())
+    return false;
+
+  ShuffleVectorInst *SVI = dyn_cast<ShuffleVectorInst>(SI->getValueOperand());
+  if (!SVI || !SVI->hasOneUse())
+    return false;
+
+  // Check if the shufflevector is RE-interleave shuffle.
+  unsigned Factor;
+  if (!isReInterleaveMask(SVI->getShuffleMask(), Factor))
+    return false;
+
+  DEBUG(dbgs() << "IA: Found an interleaved store: " << *SI << "\n");
+
+  // Try to create target specific intrinsics to replace the store and shuffle.
+  if (!TLI->lowerInterleavedStore(SI, SVI, Factor))
+    return false;
+
+  // Already have a new target specific interleaved store. Erase the old store.
+  DeadInsts.push_back(SI);
+  DeadInsts.push_back(SVI);
+  return true;
+}
+
+bool InterleavedAccess::runOnFunction(Function &F) {
+  if (!TM || !LowerInterleavedAccesses)
+    return false;
+
+  DEBUG(dbgs() << "*** " << getPassName() << ": " << F.getName() << "\n");
+
+  TLI = TM->getSubtargetImpl(F)->getTargetLowering();
+  MaxFactor = TLI->getMaxSupportedInterleaveFactor();
+
+  // Holds dead instructions that will be erased later.
+  SmallVector<Instruction *, 32> DeadInsts;
+  bool Changed = false;
+
+  for (auto &I : inst_range(F)) {
+    if (LoadInst *LI = dyn_cast<LoadInst>(&I))
+      Changed |= lowerInterleavedLoad(LI, DeadInsts);
+
+    if (StoreInst *SI = dyn_cast<StoreInst>(&I))
+      Changed |= lowerInterleavedStore(SI, DeadInsts);
+  }
+
+  for (auto I : DeadInsts)
+    I->eraseFromParent();
+
+  return Changed;
+}