contrib/llvm-project/llvm/lib/CodeGen/InterleavedAccessPass.cpp

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