contrib/llvm/lib/CodeGen/InterleavedAccessPass.cpp

   1 //===--------------------- InterleavedAccessPass.cpp ----------------------===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 // This file implements the Interleaved Access pass, which identifies
  11 // interleaved memory accesses and transforms them into target specific
  12 // intrinsics.
  13 //
  14 // An interleaved load reads data from memory into several vectors, with
  15 // DE-interleaving the data on a factor. An interleaved store writes several
  16 // vectors to memory with RE-interleaving the data on a factor.
  17 //
  18 // As interleaved accesses are difficult to identified in CodeGen (mainly
  19 // because the VECTOR_SHUFFLE DAG node is quite different from the shufflevector
  20 // IR), we identify and transform them to intrinsics in this pass so the
  21 // intrinsics can be easily matched into target specific instructions later in
  22 // CodeGen.
  23 //
  24 // E.g. An interleaved load (Factor = 2):
  25 //        %wide.vec = load <8 x i32>, <8 x i32>* %ptr
  26 //        %v0 = shuffle <8 x i32> %wide.vec, <8 x i32> undef, <0, 2, 4, 6>
  27 //        %v1 = shuffle <8 x i32> %wide.vec, <8 x i32> undef, <1, 3, 5, 7>
  28 //
  29 // It could be transformed into a ld2 intrinsic in AArch64 backend or a vld2
  30 // intrinsic in ARM backend.
  31 //
  32 // In X86, this can be further optimized into a set of target
  33 // specific loads followed by an optimized sequence of shuffles.
  34 //
  35 // E.g. An interleaved store (Factor = 3):
  36 //        %i.vec = shuffle <8 x i32> %v0, <8 x i32> %v1,
  37 //                                    <0, 4, 8, 1, 5, 9, 2, 6, 10, 3, 7, 11>
  38 //        store <12 x i32> %i.vec, <12 x i32>* %ptr
  39 //
  40 // It could be transformed into a st3 intrinsic in AArch64 backend or a vst3
  41 // intrinsic in ARM backend.
  42 //
  43 // Similarly, a set of interleaved stores can be transformed into an optimized
  44 // sequence of shuffles followed by a set of target specific stores for X86.
  45 //===----------------------------------------------------------------------===//
  46
  47 #include "llvm/CodeGen/Passes.h"
  48 #include "llvm/CodeGen/TargetPassConfig.h"
  49 #include "llvm/IR/Dominators.h"
  50 #include "llvm/IR/InstIterator.h"
  51 #include "llvm/Support/Debug.h"
  52 #include "llvm/Support/MathExtras.h"
  53 #include "llvm/Support/raw_ostream.h"
  54 #include "llvm/Target/TargetLowering.h"
  55 #include "llvm/Target/TargetSubtargetInfo.h"
  56
  57 using namespace llvm;
  58
  59 #define DEBUG_TYPE "interleaved-access"
  60
  61 static cl::opt<bool> LowerInterleavedAccesses(
  62     "lower-interleaved-accesses",
  63     cl::desc("Enable lowering interleaved accesses to intrinsics"),
  64     cl::init(true), cl::Hidden);
  65
  66 namespace {
  67
  68 class InterleavedAccess : public FunctionPass {
  69
  70 public:
  71   static char ID;
  72   InterleavedAccess() : FunctionPass(ID), DT(nullptr), TLI(nullptr) {
  73     initializeInterleavedAccessPass(*PassRegistry::getPassRegistry());
  74   }
  75
  76   StringRef getPassName() const override { return "Interleaved Access Pass"; }
  77
  78   bool runOnFunction(Function &F) override;
  79
  80   void getAnalysisUsage(AnalysisUsage &AU) const override {
  81     AU.addRequired<DominatorTreeWrapperPass>();
  82     AU.addPreserved<DominatorTreeWrapperPass>();
  83   }
  84
  85 private:
  86   DominatorTree *DT;
  87   const TargetLowering *TLI;
  88
  89   /// The maximum supported interleave factor.
  90   unsigned MaxFactor;
  91
  92   /// \brief Transform an interleaved load into target specific intrinsics.
  93   bool lowerInterleavedLoad(LoadInst *LI,
  94                             SmallVector<Instruction *, 32> &DeadInsts);
  95
  96   /// \brief Transform an interleaved store into target specific intrinsics.
  97   bool lowerInterleavedStore(StoreInst *SI,
  98                              SmallVector<Instruction *, 32> &DeadInsts);
  99
 100   /// \brief Returns true if the uses of an interleaved load by the
 101   /// extractelement instructions in \p Extracts can be replaced by uses of the
 102   /// shufflevector instructions in \p Shuffles instead. If so, the necessary
 103   /// replacements are also performed.
 104   bool tryReplaceExtracts(ArrayRef<ExtractElementInst *> Extracts,
 105                           ArrayRef<ShuffleVectorInst *> Shuffles);
 106 };
 107 } // end anonymous namespace.
 108
 109 char InterleavedAccess::ID = 0;
 110 INITIALIZE_PASS_BEGIN(
 111     InterleavedAccess, "interleaved-access",
 112     "Lower interleaved memory accesses to target specific intrinsics", false,
 113     false)
 114 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 115 INITIALIZE_PASS_END(
 116     InterleavedAccess, "interleaved-access",
 117     "Lower interleaved memory accesses to target specific intrinsics", false,
 118     false)
 119
 120 FunctionPass *llvm::createInterleavedAccessPass() {
 121   return new InterleavedAccess();
 122 }
 123
 124 /// \brief Check if the mask is a DE-interleave mask of the given factor
 125 /// \p Factor like:
 126 ///     <Index, Index+Factor, ..., Index+(NumElts-1)*Factor>
 127 static bool isDeInterleaveMaskOfFactor(ArrayRef<int> Mask, unsigned Factor,
 128                                        unsigned &Index) {
 129   // Check all potential start indices from 0 to (Factor - 1).
 130   for (Index = 0; Index < Factor; Index++) {
 131     unsigned i = 0;
 132
 133     // Check that elements are in ascending order by Factor. Ignore undef
 134     // elements.
 135     for (; i < Mask.size(); i++)
 136       if (Mask[i] >= 0 && static_cast<unsigned>(Mask[i]) != Index + i * Factor)
 137         break;
 138
 139     if (i == Mask.size())
 140       return true;
 141   }
 142
 143   return false;
 144 }
 145
 146 /// \brief Check if the mask is a DE-interleave mask for an interleaved load.
 147 ///
 148 /// E.g. DE-interleave masks (Factor = 2) could be:
 149 ///     <0, 2, 4, 6>    (mask of index 0 to extract even elements)
 150 ///     <1, 3, 5, 7>    (mask of index 1 to extract odd elements)
 151 static bool isDeInterleaveMask(ArrayRef<int> Mask, unsigned &Factor,
 152                                unsigned &Index, unsigned MaxFactor) {
 153   if (Mask.size() < 2)
 154     return false;
 155
 156   // Check potential Factors.
 157   for (Factor = 2; Factor <= MaxFactor; Factor++)
 158     if (isDeInterleaveMaskOfFactor(Mask, Factor, Index))
 159       return true;
 160
 161   return false;
 162 }
 163
 164 /// \brief Check if the mask can be used in an interleaved store.
 165 //
 166 /// It checks for a more general pattern than the RE-interleave mask.
 167 /// I.e. <x, y, ... z, x+1, y+1, ...z+1, x+2, y+2, ...z+2, ...>
 168 /// E.g. For a Factor of 2 (LaneLen=4): <4, 32, 5, 33, 6, 34, 7, 35>
 169 /// E.g. For a Factor of 3 (LaneLen=4): <4, 32, 16, 5, 33, 17, 6, 34, 18, 7, 35, 19>
 170 /// E.g. For a Factor of 4 (LaneLen=2): <8, 2, 12, 4, 9, 3, 13, 5>
 171 ///
 172 /// The particular case of an RE-interleave mask is:
 173 /// I.e. <0, LaneLen, ... , LaneLen*(Factor - 1), 1, LaneLen + 1, ...>
 174 /// E.g. For a Factor of 2 (LaneLen=4): <0, 4, 1, 5, 2, 6, 3, 7>
 175 static bool isReInterleaveMask(ArrayRef<int> Mask, unsigned &Factor,
 176                                unsigned MaxFactor, unsigned OpNumElts) {
 177   unsigned NumElts = Mask.size();
 178   if (NumElts < 4)
 179     return false;
 180
 181   // Check potential Factors.
 182   for (Factor = 2; Factor <= MaxFactor; Factor++) {
 183     if (NumElts % Factor)
 184       continue;
 185
 186     unsigned LaneLen = NumElts / Factor;
 187     if (!isPowerOf2_32(LaneLen))
 188       continue;
 189
 190     // Check whether each element matches the general interleaved rule.
 191     // Ignore undef elements, as long as the defined elements match the rule.
 192     // Outer loop processes all factors (x, y, z in the above example)
 193     unsigned I = 0, J;
 194     for (; I < Factor; I++) {
 195       unsigned SavedLaneValue;
 196       unsigned SavedNoUndefs = 0;
 197
 198       // Inner loop processes consecutive accesses (x, x+1... in the example)
 199       for (J = 0; J < LaneLen - 1; J++) {
 200         // Lane computes x's position in the Mask
 201         unsigned Lane = J * Factor + I;
 202         unsigned NextLane = Lane + Factor;
 203         int LaneValue = Mask[Lane];
 204         int NextLaneValue = Mask[NextLane];
 205
 206         // If both are defined, values must be sequential
 207         if (LaneValue >= 0 && NextLaneValue >= 0 &&
 208             LaneValue + 1 != NextLaneValue)
 209           break;
 210
 211         // If the next value is undef, save the current one as reference
 212         if (LaneValue >= 0 && NextLaneValue < 0) {
 213           SavedLaneValue = LaneValue;
 214           SavedNoUndefs = 1;
 215         }
 216
 217         // Undefs are allowed, but defined elements must still be consecutive:
 218         // i.e.: x,..., undef,..., x + 2,..., undef,..., undef,..., x + 5, ....
 219         // Verify this by storing the last non-undef followed by an undef
 220         // Check that following non-undef masks are incremented with the
 221         // corresponding distance.
 222         if (SavedNoUndefs > 0 && LaneValue < 0) {
 223           SavedNoUndefs++;
 224           if (NextLaneValue >= 0 &&
 225               SavedLaneValue + SavedNoUndefs != (unsigned)NextLaneValue)
 226             break;
 227         }
 228       }
 229
 230       if (J < LaneLen - 1)
 231         break;
 232
 233       int StartMask = 0;
 234       if (Mask[I] >= 0) {
 235         // Check that the start of the I range (J=0) is greater than 0
 236         StartMask = Mask[I];
 237       } else if (Mask[(LaneLen - 1) * Factor + I] >= 0) {
 238         // StartMask defined by the last value in lane
 239         StartMask = Mask[(LaneLen - 1) * Factor + I] - J;
 240       } else if (SavedNoUndefs > 0) {
 241         // StartMask defined by some non-zero value in the j loop
 242         StartMask = SavedLaneValue - (LaneLen - 1 - SavedNoUndefs);
 243       }
 244       // else StartMask remains set to 0, i.e. all elements are undefs
 245
 246       if (StartMask < 0)
 247         break;
 248       // We must stay within the vectors; This case can happen with undefs.
 249       if (StartMask + LaneLen > OpNumElts*2)
 250         break;
 251     }
 252
 253     // Found an interleaved mask of current factor.
 254     if (I == Factor)
 255       return true;
 256   }
 257
 258   return false;
 259 }
 260
 261 bool InterleavedAccess::lowerInterleavedLoad(
 262     LoadInst *LI, SmallVector<Instruction *, 32> &DeadInsts) {
 263   if (!LI->isSimple())
 264     return false;
 265
 266   SmallVector<ShuffleVectorInst *, 4> Shuffles;
 267   SmallVector<ExtractElementInst *, 4> Extracts;
 268
 269   // Check if all users of this load are shufflevectors. If we encounter any
 270   // users that are extractelement instructions, we save them to later check if
 271   // they can be modifed to extract from one of the shufflevectors instead of
 272   // the load.
 273   for (auto UI = LI->user_begin(), E = LI->user_end(); UI != E; UI++) {
 274     auto *Extract = dyn_cast<ExtractElementInst>(*UI);
 275     if (Extract && isa<ConstantInt>(Extract->getIndexOperand())) {
 276       Extracts.push_back(Extract);
 277       continue;
 278     }
 279     ShuffleVectorInst *SVI = dyn_cast<ShuffleVectorInst>(*UI);
 280     if (!SVI || !isa<UndefValue>(SVI->getOperand(1)))
 281       return false;
 282
 283     Shuffles.push_back(SVI);
 284   }
 285
 286   if (Shuffles.empty())
 287     return false;
 288
 289   unsigned Factor, Index;
 290
 291   // Check if the first shufflevector is DE-interleave shuffle.
 292   if (!isDeInterleaveMask(Shuffles[0]->getShuffleMask(), Factor, Index,
 293                           MaxFactor))
 294     return false;
 295
 296   // Holds the corresponding index for each DE-interleave shuffle.
 297   SmallVector<unsigned, 4> Indices;
 298   Indices.push_back(Index);
 299
 300   Type *VecTy = Shuffles[0]->getType();
 301
 302   // Check if other shufflevectors are also DE-interleaved of the same type
 303   // and factor as the first shufflevector.
 304   for (unsigned i = 1; i < Shuffles.size(); i++) {
 305     if (Shuffles[i]->getType() != VecTy)
 306       return false;
 307
 308     if (!isDeInterleaveMaskOfFactor(Shuffles[i]->getShuffleMask(), Factor,
 309                                     Index))
 310       return false;
 311
 312     Indices.push_back(Index);
 313   }
 314
 315   // Try and modify users of the load that are extractelement instructions to
 316   // use the shufflevector instructions instead of the load.
 317   if (!tryReplaceExtracts(Extracts, Shuffles))
 318     return false;
 319
 320   DEBUG(dbgs() << "IA: Found an interleaved load: " << *LI << "\n");
 321
 322   // Try to create target specific intrinsics to replace the load and shuffles.
 323   if (!TLI->lowerInterleavedLoad(LI, Shuffles, Indices, Factor))
 324     return false;
 325
 326   for (auto SVI : Shuffles)
 327     DeadInsts.push_back(SVI);
 328
 329   DeadInsts.push_back(LI);
 330   return true;
 331 }
 332
 333 bool InterleavedAccess::tryReplaceExtracts(
 334     ArrayRef<ExtractElementInst *> Extracts,
 335     ArrayRef<ShuffleVectorInst *> Shuffles) {
 336
 337   // If there aren't any extractelement instructions to modify, there's nothing
 338   // to do.
 339   if (Extracts.empty())
 340     return true;
 341
 342   // Maps extractelement instructions to vector-index pairs. The extractlement
 343   // instructions will be modified to use the new vector and index operands.
 344   DenseMap<ExtractElementInst *, std::pair<Value *, int>> ReplacementMap;
 345
 346   for (auto *Extract : Extracts) {
 347
 348     // The vector index that is extracted.
 349     auto *IndexOperand = cast<ConstantInt>(Extract->getIndexOperand());
 350     auto Index = IndexOperand->getSExtValue();
 351
 352     // Look for a suitable shufflevector instruction. The goal is to modify the
 353     // extractelement instruction (which uses an interleaved load) to use one
 354     // of the shufflevector instructions instead of the load.
 355     for (auto *Shuffle : Shuffles) {
 356
 357       // If the shufflevector instruction doesn't dominate the extract, we
 358       // can't create a use of it.
 359       if (!DT->dominates(Shuffle, Extract))
 360         continue;
 361
 362       // Inspect the indices of the shufflevector instruction. If the shuffle
 363       // selects the same index that is extracted, we can modify the
 364       // extractelement instruction.
 365       SmallVector<int, 4> Indices;
 366       Shuffle->getShuffleMask(Indices);
 367       for (unsigned I = 0; I < Indices.size(); ++I)
 368         if (Indices[I] == Index) {
 369           assert(Extract->getOperand(0) == Shuffle->getOperand(0) &&
 370                  "Vector operations do not match");
 371           ReplacementMap[Extract] = std::make_pair(Shuffle, I);
 372           break;
 373         }
 374
 375       // If we found a suitable shufflevector instruction, stop looking.
 376       if (ReplacementMap.count(Extract))
 377         break;
 378     }
 379
 380     // If we did not find a suitable shufflevector instruction, the
 381     // extractelement instruction cannot be modified, so we must give up.
 382     if (!ReplacementMap.count(Extract))
 383       return false;
 384   }
 385
 386   // Finally, perform the replacements.
 387   IRBuilder<> Builder(Extracts[0]->getContext());
 388   for (auto &Replacement : ReplacementMap) {
 389     auto *Extract = Replacement.first;
 390     auto *Vector = Replacement.second.first;
 391     auto Index = Replacement.second.second;
 392     Builder.SetInsertPoint(Extract);
 393     Extract->replaceAllUsesWith(Builder.CreateExtractElement(Vector, Index));
 394     Extract->eraseFromParent();
 395   }
 396
 397   return true;
 398 }
 399
 400 bool InterleavedAccess::lowerInterleavedStore(
 401     StoreInst *SI, SmallVector<Instruction *, 32> &DeadInsts) {
 402   if (!SI->isSimple())
 403     return false;
 404
 405   ShuffleVectorInst *SVI = dyn_cast<ShuffleVectorInst>(SI->getValueOperand());
 406   if (!SVI || !SVI->hasOneUse())
 407     return false;
 408
 409   // Check if the shufflevector is RE-interleave shuffle.
 410   unsigned Factor;
 411   unsigned OpNumElts = SVI->getOperand(0)->getType()->getVectorNumElements();
 412   if (!isReInterleaveMask(SVI->getShuffleMask(), Factor, MaxFactor, OpNumElts))
 413     return false;
 414
 415   DEBUG(dbgs() << "IA: Found an interleaved store: " << *SI << "\n");
 416
 417   // Try to create target specific intrinsics to replace the store and shuffle.
 418   if (!TLI->lowerInterleavedStore(SI, SVI, Factor))
 419     return false;
 420
 421   // Already have a new target specific interleaved store. Erase the old store.
 422   DeadInsts.push_back(SI);
 423   DeadInsts.push_back(SVI);
 424   return true;
 425 }
 426
 427 bool InterleavedAccess::runOnFunction(Function &F) {
 428   auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
 429   if (!TPC || !LowerInterleavedAccesses)
 430     return false;
 431
 432   DEBUG(dbgs() << "*** " << getPassName() << ": " << F.getName() << "\n");
 433
 434   DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
 435   auto &TM = TPC->getTM<TargetMachine>();
 436   TLI = TM.getSubtargetImpl(F)->getTargetLowering();
 437   MaxFactor = TLI->getMaxSupportedInterleaveFactor();
 438
 439   // Holds dead instructions that will be erased later.
 440   SmallVector<Instruction *, 32> DeadInsts;
 441   bool Changed = false;
 442
 443   for (auto &I : instructions(F)) {
 444     if (LoadInst *LI = dyn_cast<LoadInst>(&I))
 445       Changed |= lowerInterleavedLoad(LI, DeadInsts);
 446
 447     if (StoreInst *SI = dyn_cast<StoreInst>(&I))
 448       Changed |= lowerInterleavedStore(SI, DeadInsts);
 449   }
 450
 451   for (auto I : DeadInsts)
 452     I->eraseFromParent();
 453
 454   return Changed;
 455 }