Run DCE after a LoopFlatten test to reduce spurious output [nfc]
[llvm-project.git] / llvm / lib / CodeGen / InterleavedAccessPass.cpp
blob6b3848531569c0b0019b19090d849e24e8b7ba4a
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.
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.
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.
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> poison, <0, 2, 4, 6>
26 // %v1 = shuffle <8 x i32> %wide.vec, <8 x i32> poison, <1, 3, 5, 7>
28 // It could be transformed into a ld2 intrinsic in AArch64 backend or a vld2
29 // intrinsic in ARM backend.
31 // In X86, this can be further optimized into a set of target
32 // specific loads followed by an optimized sequence of shuffles.
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
39 // It could be transformed into a st3 intrinsic in AArch64 backend or a vst3
40 // intrinsic in ARM backend.
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.
45 //===----------------------------------------------------------------------===//
47 #include "llvm/ADT/ArrayRef.h"
48 #include "llvm/ADT/DenseMap.h"
49 #include "llvm/ADT/SetVector.h"
50 #include "llvm/ADT/SmallVector.h"
51 #include "llvm/CodeGen/TargetLowering.h"
52 #include "llvm/CodeGen/TargetPassConfig.h"
53 #include "llvm/CodeGen/TargetSubtargetInfo.h"
54 #include "llvm/IR/Constants.h"
55 #include "llvm/IR/Dominators.h"
56 #include "llvm/IR/Function.h"
57 #include "llvm/IR/IRBuilder.h"
58 #include "llvm/IR/InstIterator.h"
59 #include "llvm/IR/Instruction.h"
60 #include "llvm/IR/Instructions.h"
61 #include "llvm/IR/IntrinsicInst.h"
62 #include "llvm/InitializePasses.h"
63 #include "llvm/Pass.h"
64 #include "llvm/Support/Casting.h"
65 #include "llvm/Support/CommandLine.h"
66 #include "llvm/Support/Debug.h"
67 #include "llvm/Support/MathExtras.h"
68 #include "llvm/Support/raw_ostream.h"
69 #include "llvm/Target/TargetMachine.h"
70 #include "llvm/Transforms/Utils/Local.h"
71 #include <cassert>
72 #include <utility>
74 using namespace llvm;
76 #define DEBUG_TYPE "interleaved-access"
78 static cl::opt<bool> LowerInterleavedAccesses(
79 "lower-interleaved-accesses",
80 cl::desc("Enable lowering interleaved accesses to intrinsics"),
81 cl::init(true), cl::Hidden);
83 namespace {
85 class InterleavedAccess : public FunctionPass {
86 public:
87 static char ID;
89 InterleavedAccess() : FunctionPass(ID) {
90 initializeInterleavedAccessPass(*PassRegistry::getPassRegistry());
93 StringRef getPassName() const override { return "Interleaved Access Pass"; }
95 bool runOnFunction(Function &F) override;
97 void getAnalysisUsage(AnalysisUsage &AU) const override {
98 AU.addRequired<DominatorTreeWrapperPass>();
99 AU.setPreservesCFG();
102 private:
103 DominatorTree *DT = nullptr;
104 const TargetLowering *TLI = nullptr;
106 /// The maximum supported interleave factor.
107 unsigned MaxFactor = 0u;
109 /// Transform an interleaved load into target specific intrinsics.
110 bool lowerInterleavedLoad(LoadInst *LI,
111 SmallVector<Instruction *, 32> &DeadInsts);
113 /// Transform an interleaved store into target specific intrinsics.
114 bool lowerInterleavedStore(StoreInst *SI,
115 SmallVector<Instruction *, 32> &DeadInsts);
117 /// Transform a load and a deinterleave intrinsic into target specific
118 /// instructions.
119 bool lowerDeinterleaveIntrinsic(IntrinsicInst *II,
120 SmallVector<Instruction *, 32> &DeadInsts);
122 /// Transform an interleave intrinsic and a store into target specific
123 /// instructions.
124 bool lowerInterleaveIntrinsic(IntrinsicInst *II,
125 SmallVector<Instruction *, 32> &DeadInsts);
127 /// Returns true if the uses of an interleaved load by the
128 /// extractelement instructions in \p Extracts can be replaced by uses of the
129 /// shufflevector instructions in \p Shuffles instead. If so, the necessary
130 /// replacements are also performed.
131 bool tryReplaceExtracts(ArrayRef<ExtractElementInst *> Extracts,
132 ArrayRef<ShuffleVectorInst *> Shuffles);
134 /// Given a number of shuffles of the form shuffle(binop(x,y)), convert them
135 /// to binop(shuffle(x), shuffle(y)) to allow the formation of an
136 /// interleaving load. Any newly created shuffles that operate on \p LI will
137 /// be added to \p Shuffles. Returns true, if any changes to the IR have been
138 /// made.
139 bool replaceBinOpShuffles(ArrayRef<ShuffleVectorInst *> BinOpShuffles,
140 SmallVectorImpl<ShuffleVectorInst *> &Shuffles,
141 LoadInst *LI);
144 } // end anonymous namespace.
146 char InterleavedAccess::ID = 0;
148 INITIALIZE_PASS_BEGIN(InterleavedAccess, DEBUG_TYPE,
149 "Lower interleaved memory accesses to target specific intrinsics", false,
150 false)
151 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
152 INITIALIZE_PASS_END(InterleavedAccess, DEBUG_TYPE,
153 "Lower interleaved memory accesses to target specific intrinsics", false,
154 false)
156 FunctionPass *llvm::createInterleavedAccessPass() {
157 return new InterleavedAccess();
160 /// Check if the mask is a DE-interleave mask of the given factor
161 /// \p Factor like:
162 /// <Index, Index+Factor, ..., Index+(NumElts-1)*Factor>
163 static bool isDeInterleaveMaskOfFactor(ArrayRef<int> Mask, unsigned Factor,
164 unsigned &Index) {
165 // Check all potential start indices from 0 to (Factor - 1).
166 for (Index = 0; Index < Factor; Index++) {
167 unsigned i = 0;
169 // Check that elements are in ascending order by Factor. Ignore undef
170 // elements.
171 for (; i < Mask.size(); i++)
172 if (Mask[i] >= 0 && static_cast<unsigned>(Mask[i]) != Index + i * Factor)
173 break;
175 if (i == Mask.size())
176 return true;
179 return false;
182 /// Check if the mask is a DE-interleave mask for an interleaved load.
184 /// E.g. DE-interleave masks (Factor = 2) could be:
185 /// <0, 2, 4, 6> (mask of index 0 to extract even elements)
186 /// <1, 3, 5, 7> (mask of index 1 to extract odd elements)
187 static bool isDeInterleaveMask(ArrayRef<int> Mask, unsigned &Factor,
188 unsigned &Index, unsigned MaxFactor,
189 unsigned NumLoadElements) {
190 if (Mask.size() < 2)
191 return false;
193 // Check potential Factors.
194 for (Factor = 2; Factor <= MaxFactor; Factor++) {
195 // Make sure we don't produce a load wider than the input load.
196 if (Mask.size() * Factor > NumLoadElements)
197 return false;
198 if (isDeInterleaveMaskOfFactor(Mask, Factor, Index))
199 return true;
202 return false;
205 /// Check if the mask can be used in an interleaved store.
207 /// It checks for a more general pattern than the RE-interleave mask.
208 /// I.e. <x, y, ... z, x+1, y+1, ...z+1, x+2, y+2, ...z+2, ...>
209 /// E.g. For a Factor of 2 (LaneLen=4): <4, 32, 5, 33, 6, 34, 7, 35>
210 /// E.g. For a Factor of 3 (LaneLen=4): <4, 32, 16, 5, 33, 17, 6, 34, 18, 7, 35, 19>
211 /// E.g. For a Factor of 4 (LaneLen=2): <8, 2, 12, 4, 9, 3, 13, 5>
213 /// The particular case of an RE-interleave mask is:
214 /// I.e. <0, LaneLen, ... , LaneLen*(Factor - 1), 1, LaneLen + 1, ...>
215 /// E.g. For a Factor of 2 (LaneLen=4): <0, 4, 1, 5, 2, 6, 3, 7>
216 static bool isReInterleaveMask(ShuffleVectorInst *SVI, unsigned &Factor,
217 unsigned MaxFactor) {
218 unsigned NumElts = SVI->getShuffleMask().size();
219 if (NumElts < 4)
220 return false;
222 // Check potential Factors.
223 for (Factor = 2; Factor <= MaxFactor; Factor++) {
224 if (SVI->isInterleave(Factor))
225 return true;
228 return false;
231 bool InterleavedAccess::lowerInterleavedLoad(
232 LoadInst *LI, SmallVector<Instruction *, 32> &DeadInsts) {
233 if (!LI->isSimple() || isa<ScalableVectorType>(LI->getType()))
234 return false;
236 // Check if all users of this load are shufflevectors. If we encounter any
237 // users that are extractelement instructions or binary operators, we save
238 // them to later check if they can be modified to extract from one of the
239 // shufflevectors instead of the load.
241 SmallVector<ShuffleVectorInst *, 4> Shuffles;
242 SmallVector<ExtractElementInst *, 4> Extracts;
243 // BinOpShuffles need to be handled a single time in case both operands of the
244 // binop are the same load.
245 SmallSetVector<ShuffleVectorInst *, 4> BinOpShuffles;
247 for (auto *User : LI->users()) {
248 auto *Extract = dyn_cast<ExtractElementInst>(User);
249 if (Extract && isa<ConstantInt>(Extract->getIndexOperand())) {
250 Extracts.push_back(Extract);
251 continue;
253 if (auto *BI = dyn_cast<BinaryOperator>(User)) {
254 if (all_of(BI->users(), [](auto *U) {
255 auto *SVI = dyn_cast<ShuffleVectorInst>(U);
256 return SVI && isa<UndefValue>(SVI->getOperand(1));
257 })) {
258 for (auto *SVI : BI->users())
259 BinOpShuffles.insert(cast<ShuffleVectorInst>(SVI));
260 continue;
263 auto *SVI = dyn_cast<ShuffleVectorInst>(User);
264 if (!SVI || !isa<UndefValue>(SVI->getOperand(1)))
265 return false;
267 Shuffles.push_back(SVI);
270 if (Shuffles.empty() && BinOpShuffles.empty())
271 return false;
273 unsigned Factor, Index;
275 unsigned NumLoadElements =
276 cast<FixedVectorType>(LI->getType())->getNumElements();
277 auto *FirstSVI = Shuffles.size() > 0 ? Shuffles[0] : BinOpShuffles[0];
278 // Check if the first shufflevector is DE-interleave shuffle.
279 if (!isDeInterleaveMask(FirstSVI->getShuffleMask(), Factor, Index, MaxFactor,
280 NumLoadElements))
281 return false;
283 // Holds the corresponding index for each DE-interleave shuffle.
284 SmallVector<unsigned, 4> Indices;
286 Type *VecTy = FirstSVI->getType();
288 // Check if other shufflevectors are also DE-interleaved of the same type
289 // and factor as the first shufflevector.
290 for (auto *Shuffle : Shuffles) {
291 if (Shuffle->getType() != VecTy)
292 return false;
293 if (!isDeInterleaveMaskOfFactor(Shuffle->getShuffleMask(), Factor,
294 Index))
295 return false;
297 assert(Shuffle->getShuffleMask().size() <= NumLoadElements);
298 Indices.push_back(Index);
300 for (auto *Shuffle : BinOpShuffles) {
301 if (Shuffle->getType() != VecTy)
302 return false;
303 if (!isDeInterleaveMaskOfFactor(Shuffle->getShuffleMask(), Factor,
304 Index))
305 return false;
307 assert(Shuffle->getShuffleMask().size() <= NumLoadElements);
309 if (cast<Instruction>(Shuffle->getOperand(0))->getOperand(0) == LI)
310 Indices.push_back(Index);
311 if (cast<Instruction>(Shuffle->getOperand(0))->getOperand(1) == LI)
312 Indices.push_back(Index);
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;
320 bool BinOpShuffleChanged =
321 replaceBinOpShuffles(BinOpShuffles.getArrayRef(), Shuffles, LI);
323 LLVM_DEBUG(dbgs() << "IA: Found an interleaved load: " << *LI << "\n");
325 // Try to create target specific intrinsics to replace the load and shuffles.
326 if (!TLI->lowerInterleavedLoad(LI, Shuffles, Indices, Factor)) {
327 // If Extracts is not empty, tryReplaceExtracts made changes earlier.
328 return !Extracts.empty() || BinOpShuffleChanged;
331 append_range(DeadInsts, Shuffles);
333 DeadInsts.push_back(LI);
334 return true;
337 bool InterleavedAccess::replaceBinOpShuffles(
338 ArrayRef<ShuffleVectorInst *> BinOpShuffles,
339 SmallVectorImpl<ShuffleVectorInst *> &Shuffles, LoadInst *LI) {
340 for (auto *SVI : BinOpShuffles) {
341 BinaryOperator *BI = cast<BinaryOperator>(SVI->getOperand(0));
342 Type *BIOp0Ty = BI->getOperand(0)->getType();
343 ArrayRef<int> Mask = SVI->getShuffleMask();
344 assert(all_of(Mask, [&](int Idx) {
345 return Idx < (int)cast<FixedVectorType>(BIOp0Ty)->getNumElements();
346 }));
348 auto *NewSVI1 =
349 new ShuffleVectorInst(BI->getOperand(0), PoisonValue::get(BIOp0Ty),
350 Mask, SVI->getName(), SVI);
351 auto *NewSVI2 = new ShuffleVectorInst(
352 BI->getOperand(1), PoisonValue::get(BI->getOperand(1)->getType()), Mask,
353 SVI->getName(), SVI);
354 BinaryOperator *NewBI = BinaryOperator::CreateWithCopiedFlags(
355 BI->getOpcode(), NewSVI1, NewSVI2, BI, BI->getName(), SVI);
356 SVI->replaceAllUsesWith(NewBI);
357 LLVM_DEBUG(dbgs() << " Replaced: " << *BI << "\n And : " << *SVI
358 << "\n With : " << *NewSVI1 << "\n And : "
359 << *NewSVI2 << "\n And : " << *NewBI << "\n");
360 RecursivelyDeleteTriviallyDeadInstructions(SVI);
361 if (NewSVI1->getOperand(0) == LI)
362 Shuffles.push_back(NewSVI1);
363 if (NewSVI2->getOperand(0) == LI)
364 Shuffles.push_back(NewSVI2);
367 return !BinOpShuffles.empty();
370 bool InterleavedAccess::tryReplaceExtracts(
371 ArrayRef<ExtractElementInst *> Extracts,
372 ArrayRef<ShuffleVectorInst *> Shuffles) {
373 // If there aren't any extractelement instructions to modify, there's nothing
374 // to do.
375 if (Extracts.empty())
376 return true;
378 // Maps extractelement instructions to vector-index pairs. The extractlement
379 // instructions will be modified to use the new vector and index operands.
380 DenseMap<ExtractElementInst *, std::pair<Value *, int>> ReplacementMap;
382 for (auto *Extract : Extracts) {
383 // The vector index that is extracted.
384 auto *IndexOperand = cast<ConstantInt>(Extract->getIndexOperand());
385 auto Index = IndexOperand->getSExtValue();
387 // Look for a suitable shufflevector instruction. The goal is to modify the
388 // extractelement instruction (which uses an interleaved load) to use one
389 // of the shufflevector instructions instead of the load.
390 for (auto *Shuffle : Shuffles) {
391 // If the shufflevector instruction doesn't dominate the extract, we
392 // can't create a use of it.
393 if (!DT->dominates(Shuffle, Extract))
394 continue;
396 // Inspect the indices of the shufflevector instruction. If the shuffle
397 // selects the same index that is extracted, we can modify the
398 // extractelement instruction.
399 SmallVector<int, 4> Indices;
400 Shuffle->getShuffleMask(Indices);
401 for (unsigned I = 0; I < Indices.size(); ++I)
402 if (Indices[I] == Index) {
403 assert(Extract->getOperand(0) == Shuffle->getOperand(0) &&
404 "Vector operations do not match");
405 ReplacementMap[Extract] = std::make_pair(Shuffle, I);
406 break;
409 // If we found a suitable shufflevector instruction, stop looking.
410 if (ReplacementMap.count(Extract))
411 break;
414 // If we did not find a suitable shufflevector instruction, the
415 // extractelement instruction cannot be modified, so we must give up.
416 if (!ReplacementMap.count(Extract))
417 return false;
420 // Finally, perform the replacements.
421 IRBuilder<> Builder(Extracts[0]->getContext());
422 for (auto &Replacement : ReplacementMap) {
423 auto *Extract = Replacement.first;
424 auto *Vector = Replacement.second.first;
425 auto Index = Replacement.second.second;
426 Builder.SetInsertPoint(Extract);
427 Extract->replaceAllUsesWith(Builder.CreateExtractElement(Vector, Index));
428 Extract->eraseFromParent();
431 return true;
434 bool InterleavedAccess::lowerInterleavedStore(
435 StoreInst *SI, SmallVector<Instruction *, 32> &DeadInsts) {
436 if (!SI->isSimple())
437 return false;
439 auto *SVI = dyn_cast<ShuffleVectorInst>(SI->getValueOperand());
440 if (!SVI || !SVI->hasOneUse() || isa<ScalableVectorType>(SVI->getType()))
441 return false;
443 // Check if the shufflevector is RE-interleave shuffle.
444 unsigned Factor;
445 if (!isReInterleaveMask(SVI, Factor, MaxFactor))
446 return false;
448 LLVM_DEBUG(dbgs() << "IA: Found an interleaved store: " << *SI << "\n");
450 // Try to create target specific intrinsics to replace the store and shuffle.
451 if (!TLI->lowerInterleavedStore(SI, SVI, Factor))
452 return false;
454 // Already have a new target specific interleaved store. Erase the old store.
455 DeadInsts.push_back(SI);
456 DeadInsts.push_back(SVI);
457 return true;
460 bool InterleavedAccess::lowerDeinterleaveIntrinsic(
461 IntrinsicInst *DI, SmallVector<Instruction *, 32> &DeadInsts) {
462 LoadInst *LI = dyn_cast<LoadInst>(DI->getOperand(0));
464 if (!LI || !LI->hasOneUse() || !LI->isSimple())
465 return false;
467 LLVM_DEBUG(dbgs() << "IA: Found a deinterleave intrinsic: " << *DI << "\n");
469 // Try and match this with target specific intrinsics.
470 if (!TLI->lowerDeinterleaveIntrinsicToLoad(DI, LI))
471 return false;
473 // We now have a target-specific load, so delete the old one.
474 DeadInsts.push_back(DI);
475 DeadInsts.push_back(LI);
476 return true;
479 bool InterleavedAccess::lowerInterleaveIntrinsic(
480 IntrinsicInst *II, SmallVector<Instruction *, 32> &DeadInsts) {
481 if (!II->hasOneUse())
482 return false;
484 StoreInst *SI = dyn_cast<StoreInst>(*(II->users().begin()));
486 if (!SI || !SI->isSimple())
487 return false;
489 LLVM_DEBUG(dbgs() << "IA: Found an interleave intrinsic: " << *II << "\n");
491 // Try and match this with target specific intrinsics.
492 if (!TLI->lowerInterleaveIntrinsicToStore(II, SI))
493 return false;
495 // We now have a target-specific store, so delete the old one.
496 DeadInsts.push_back(SI);
497 DeadInsts.push_back(II);
498 return true;
501 bool InterleavedAccess::runOnFunction(Function &F) {
502 auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
503 if (!TPC || !LowerInterleavedAccesses)
504 return false;
506 LLVM_DEBUG(dbgs() << "*** " << getPassName() << ": " << F.getName() << "\n");
508 DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
509 auto &TM = TPC->getTM<TargetMachine>();
510 TLI = TM.getSubtargetImpl(F)->getTargetLowering();
511 MaxFactor = TLI->getMaxSupportedInterleaveFactor();
513 // Holds dead instructions that will be erased later.
514 SmallVector<Instruction *, 32> DeadInsts;
515 bool Changed = false;
517 for (auto &I : instructions(F)) {
518 if (auto *LI = dyn_cast<LoadInst>(&I))
519 Changed |= lowerInterleavedLoad(LI, DeadInsts);
521 if (auto *SI = dyn_cast<StoreInst>(&I))
522 Changed |= lowerInterleavedStore(SI, DeadInsts);
524 if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
525 // At present, we only have intrinsics to represent (de)interleaving
526 // with a factor of 2.
527 if (II->getIntrinsicID() == Intrinsic::experimental_vector_deinterleave2)
528 Changed |= lowerDeinterleaveIntrinsic(II, DeadInsts);
529 if (II->getIntrinsicID() == Intrinsic::experimental_vector_interleave2)
530 Changed |= lowerInterleaveIntrinsic(II, DeadInsts);
534 for (auto *I : DeadInsts)
535 I->eraseFromParent();
537 return Changed;