1 //===- MVELaneInterleaving.cpp - Inverleave for MVE instructions ----------===//
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 pass interleaves around sext/zext/trunc instructions. MVE does not have
10 // a single sext/zext or trunc instruction that takes the bottom half of a
11 // vector and extends to a full width, like NEON has with MOVL. Instead it is
12 // expected that this happens through top/bottom instructions. So the MVE
13 // equivalent VMOVLT/B instructions take either the even or odd elements of the
14 // input and extend them to the larger type, producing a vector with half the
15 // number of elements each of double the bitwidth. As there is no simple
16 // instruction, we often have to turn sext/zext/trunc into a series of lane
17 // moves (or stack loads/stores, which we do not do yet).
18 //
19 // This pass takes vector code that starts at truncs, looks for interconnected
20 // blobs of operations that end with sext/zext (or constants/splats) of the
21 // form:
22 //   %sa = sext v8i16 %a to v8i32
23 //   %sb = sext v8i16 %b to v8i32
24 //   %add = add v8i32 %sa, %sb
25 //   %r = trunc %add to v8i16
26 // And adds shuffles to allow the use of VMOVL/VMOVN instrctions:
27 //   %sha = shuffle v8i16 %a, undef, <0, 2, 4, 6, 1, 3, 5, 7>
28 //   %sa = sext v8i16 %sha to v8i32
29 //   %shb = shuffle v8i16 %b, undef, <0, 2, 4, 6, 1, 3, 5, 7>
30 //   %sb = sext v8i16 %shb to v8i32
31 //   %add = add v8i32 %sa, %sb
32 //   %r = trunc %add to v8i16
33 //   %shr = shuffle v8i16 %r, undef, <0, 4, 1, 5, 2, 6, 3, 7>
34 // Which can then be split and lowered to MVE instructions efficiently:
35 //   %sa_b = VMOVLB.s16 %a
36 //   %sa_t = VMOVLT.s16 %a
37 //   %sb_b = VMOVLB.s16 %b
38 //   %sb_t = VMOVLT.s16 %b
39 //   %add_b = VADD.i32 %sa_b, %sb_b
40 //   %add_t = VADD.i32 %sa_t, %sb_t
41 //   %r = VMOVNT.i16 %add_b, %add_t
42 //
43 //===----------------------------------------------------------------------===//
44 
45 #include "ARM.h"
46 #include "ARMBaseInstrInfo.h"
47 #include "ARMSubtarget.h"
48 #include "llvm/ADT/SetVector.h"
49 #include "llvm/Analysis/TargetTransformInfo.h"
50 #include "llvm/CodeGen/TargetLowering.h"
51 #include "llvm/CodeGen/TargetPassConfig.h"
52 #include "llvm/CodeGen/TargetSubtargetInfo.h"
53 #include "llvm/IR/BasicBlock.h"
54 #include "llvm/IR/Constant.h"
55 #include "llvm/IR/Constants.h"
56 #include "llvm/IR/DerivedTypes.h"
57 #include "llvm/IR/Function.h"
58 #include "llvm/IR/IRBuilder.h"
59 #include "llvm/IR/InstIterator.h"
60 #include "llvm/IR/InstrTypes.h"
61 #include "llvm/IR/Instruction.h"
62 #include "llvm/IR/Instructions.h"
63 #include "llvm/IR/IntrinsicInst.h"
64 #include "llvm/IR/Intrinsics.h"
65 #include "llvm/IR/IntrinsicsARM.h"
66 #include "llvm/IR/PatternMatch.h"
67 #include "llvm/IR/Type.h"
68 #include "llvm/IR/Value.h"
69 #include "llvm/InitializePasses.h"
70 #include "llvm/Pass.h"
71 #include "llvm/Support/Casting.h"
72 #include <algorithm>
73 #include <cassert>
74 
75 using namespace llvm;
76 
77 #define DEBUG_TYPE "mve-laneinterleave"
78 
79 cl::opt<bool> EnableInterleave(
80     "enable-mve-interleave", cl::Hidden, cl::init(true),
81     cl::desc("Enable interleave MVE vector operation lowering"));
82 
83 namespace {
84 
85 class MVELaneInterleaving : public FunctionPass {
86 public:
87   static char ID; // Pass identification, replacement for typeid
88 
89   explicit MVELaneInterleaving() : FunctionPass(ID) {
90     initializeMVELaneInterleavingPass(*PassRegistry::getPassRegistry());
91   }
92 
93   bool runOnFunction(Function &F) override;
94 
95   StringRef getPassName() const override { return "MVE lane interleaving"; }
96 
97   void getAnalysisUsage(AnalysisUsage &AU) const override {
98     AU.setPreservesCFG();
99     AU.addRequired<TargetPassConfig>();
100     FunctionPass::getAnalysisUsage(AU);
101   }
102 };
103 
104 } // end anonymous namespace
105 
106 char MVELaneInterleaving::ID = 0;
107 
108 INITIALIZE_PASS(MVELaneInterleaving, DEBUG_TYPE, "MVE lane interleaving", false,
109                 false)
110 
111 Pass *llvm::createMVELaneInterleavingPass() {
112   return new MVELaneInterleaving();
113 }
114 
115 static bool isProfitableToInterleave(SmallSetVector<Instruction *, 4> &Exts,
116                                      SmallSetVector<Instruction *, 4> &Truncs) {
117   // This is not always beneficial to transform. Exts can be incorporated into
118   // loads, Truncs can be folded into stores.
119   // Truncs are usually the same number of instructions,
120   //  VSTRH.32(A);VSTRH.32(B) vs VSTRH.16(VMOVNT A, B) with interleaving
121   // Exts are unfortunately more instructions in the general case:
122   //  A=VLDRH.32; B=VLDRH.32;
123   // vs with interleaving:
124   //  T=VLDRH.16; A=VMOVNB T; B=VMOVNT T
125   // But those VMOVL may be folded into a VMULL.
126 
127   // But expensive extends/truncs are always good to remove. FPExts always
128   // involve extra VCVT's so are always considered to be beneficial to convert.
129   for (auto *E : Exts) {
130     if (isa<FPExtInst>(E) || !isa<LoadInst>(E->getOperand(0))) {
131       LLVM_DEBUG(dbgs() << "Beneficial due to " << *E << "\n");
132       return true;
133     }
134   }
135   for (auto *T : Truncs) {
136     if (T->hasOneUse() && !isa<StoreInst>(*T->user_begin())) {
137       LLVM_DEBUG(dbgs() << "Beneficial due to " << *T << "\n");
138       return true;
139     }
140   }
141 
142   // Otherwise, we know we have a load(ext), see if any of the Extends are a
143   // vmull. This is a simple heuristic and certainly not perfect.
144   for (auto *E : Exts) {
145     if (!E->hasOneUse() ||
146         cast<Instruction>(*E->user_begin())->getOpcode() != Instruction::Mul) {
147       LLVM_DEBUG(dbgs() << "Not beneficial due to " << *E << "\n");
148       return false;
149     }
150   }
151   return true;
152 }
153 
154 static bool tryInterleave(Instruction *Start,
155                           SmallPtrSetImpl<Instruction *> &Visited) {
156   LLVM_DEBUG(dbgs() << "tryInterleave from " << *Start << "\n");
157   auto *VT = cast<FixedVectorType>(Start->getType());
158 
159   if (!isa<Instruction>(Start->getOperand(0)))
160     return false;
161 
162   // Look for connected operations starting from Ext's, terminating at Truncs.
163   std::vector<Instruction *> Worklist;
164   Worklist.push_back(Start);
165   Worklist.push_back(cast<Instruction>(Start->getOperand(0)));
166 
167   SmallSetVector<Instruction *, 4> Truncs;
168   SmallSetVector<Instruction *, 4> Exts;
169   SmallSetVector<Use *, 4> OtherLeafs;
170   SmallSetVector<Instruction *, 4> Ops;
171 
172   while (!Worklist.empty()) {
173     Instruction *I = Worklist.back();
174     Worklist.pop_back();
175 
176     switch (I->getOpcode()) {
177     // Truncs
178     case Instruction::Trunc:
179     case Instruction::FPTrunc:
180       if (Truncs.count(I))
181         continue;
182       Truncs.insert(I);
183       Visited.insert(I);
184       break;
185 
186     // Extend leafs
187     case Instruction::SExt:
188     case Instruction::ZExt:
189     case Instruction::FPExt:
190       if (Exts.count(I))
191         continue;
192       for (auto *Use : I->users())
193         Worklist.push_back(cast<Instruction>(Use));
194       Exts.insert(I);
195       break;
196 
197     case Instruction::Call: {
198       IntrinsicInst *II = dyn_cast<IntrinsicInst>(I);
199       if (!II)
200         return false;
201 
202       switch (II->getIntrinsicID()) {
203       case Intrinsic::abs:
204       case Intrinsic::smin:
205       case Intrinsic::smax:
206       case Intrinsic::umin:
207       case Intrinsic::umax:
208       case Intrinsic::sadd_sat:
209       case Intrinsic::ssub_sat:
210       case Intrinsic::uadd_sat:
211       case Intrinsic::usub_sat:
212       case Intrinsic::minnum:
213       case Intrinsic::maxnum:
214       case Intrinsic::fabs:
215       case Intrinsic::fma:
216       case Intrinsic::ceil:
217       case Intrinsic::floor:
218       case Intrinsic::rint:
219       case Intrinsic::round:
220       case Intrinsic::trunc:
221         break;
222       default:
223         return false;
224       }
225       LLVM_FALLTHROUGH; // Fall through to treating these like an operator below.
226     }
227     // Binary/tertiary ops
228     case Instruction::Add:
229     case Instruction::Sub:
230     case Instruction::Mul:
231     case Instruction::AShr:
232     case Instruction::LShr:
233     case Instruction::Shl:
234     case Instruction::ICmp:
235     case Instruction::FCmp:
236     case Instruction::FAdd:
237     case Instruction::FMul:
238     case Instruction::Select:
239       if (Ops.count(I))
240         continue;
241       Ops.insert(I);
242 
243       for (Use &Op : I->operands()) {
244         if (!isa<FixedVectorType>(Op->getType()))
245           continue;
246         if (isa<Instruction>(Op))
247           Worklist.push_back(cast<Instruction>(&Op));
248         else
249           OtherLeafs.insert(&Op);
250       }
251 
252       for (auto *Use : I->users())
253         Worklist.push_back(cast<Instruction>(Use));
254       break;
255 
256     case Instruction::ShuffleVector:
257       // A shuffle of a splat is a splat.
258       if (cast<ShuffleVectorInst>(I)->isZeroEltSplat())
259         continue;
260       LLVM_FALLTHROUGH;
261 
262     default:
263       LLVM_DEBUG(dbgs() << "  Unhandled instruction: " << *I << "\n");
264       return false;
265     }
266   }
267 
268   if (Exts.empty() && OtherLeafs.empty())
269     return false;
270 
271   LLVM_DEBUG({
272     dbgs() << "Found group:\n  Exts:";
273     for (auto *I : Exts)
274       dbgs() << "  " << *I << "\n";
275     dbgs() << "  Ops:";
276     for (auto *I : Ops)
277       dbgs() << "  " << *I << "\n";
278     dbgs() << "  OtherLeafs:";
279     for (auto *I : OtherLeafs)
280       dbgs() << "  " << *I->get() << " of " << *I->getUser() << "\n";
281     dbgs() << "Truncs:";
282     for (auto *I : Truncs)
283       dbgs() << "  " << *I << "\n";
284   });
285 
286   assert(!Truncs.empty() && "Expected some truncs");
287 
288   // Check types
289   unsigned NumElts = VT->getNumElements();
290   unsigned BaseElts = VT->getScalarSizeInBits() == 16
291                           ? 8
292                           : (VT->getScalarSizeInBits() == 8 ? 16 : 0);
293   if (BaseElts == 0 || NumElts % BaseElts != 0) {
294     LLVM_DEBUG(dbgs() << "  Type is unsupported\n");
295     return false;
296   }
297   if (Start->getOperand(0)->getType()->getScalarSizeInBits() !=
298       VT->getScalarSizeInBits() * 2) {
299     LLVM_DEBUG(dbgs() << "  Type not double sized\n");
300     return false;
301   }
302   for (Instruction *I : Exts)
303     if (I->getOperand(0)->getType() != VT) {
304       LLVM_DEBUG(dbgs() << "  Wrong type on " << *I << "\n");
305       return false;
306     }
307   for (Instruction *I : Truncs)
308     if (I->getType() != VT) {
309       LLVM_DEBUG(dbgs() << "  Wrong type on " << *I << "\n");
310       return false;
311     }
312 
313   // Check that it looks beneficial
314   if (!isProfitableToInterleave(Exts, Truncs))
315     return false;
316 
317   // Create new shuffles around the extends / truncs / other leaves.
318   IRBuilder<> Builder(Start);
319 
320   SmallVector<int, 16> LeafMask;
321   SmallVector<int, 16> TruncMask;
322   // LeafMask : 0, 2, 4, 6, 1, 3, 5, 7   8, 10, 12, 14,  9, 11, 13, 15
323   // TruncMask: 0, 4, 1, 5, 2, 6, 3, 7   8, 12,  9, 13, 10, 14, 11, 15
324   for (unsigned Base = 0; Base < NumElts; Base += BaseElts) {
325     for (unsigned i = 0; i < BaseElts / 2; i++)
326       LeafMask.push_back(Base + i * 2);
327     for (unsigned i = 0; i < BaseElts / 2; i++)
328       LeafMask.push_back(Base + i * 2 + 1);
329   }
330   for (unsigned Base = 0; Base < NumElts; Base += BaseElts) {
331     for (unsigned i = 0; i < BaseElts / 2; i++) {
332       TruncMask.push_back(Base + i);
333       TruncMask.push_back(Base + i + BaseElts / 2);
334     }
335   }
336 
337   for (Instruction *I : Exts) {
338     LLVM_DEBUG(dbgs() << "Replacing ext " << *I << "\n");
339     Builder.SetInsertPoint(I);
340     Value *Shuffle = Builder.CreateShuffleVector(I->getOperand(0), LeafMask);
341     bool FPext = isa<FPExtInst>(I);
342     bool Sext = isa<SExtInst>(I);
343     Value *Ext = FPext ? Builder.CreateFPExt(Shuffle, I->getType())
344                        : Sext ? Builder.CreateSExt(Shuffle, I->getType())
345                               : Builder.CreateZExt(Shuffle, I->getType());
346     I->replaceAllUsesWith(Ext);
347     LLVM_DEBUG(dbgs() << "  with " << *Shuffle << "\n");
348   }
349 
350   for (Use *I : OtherLeafs) {
351     LLVM_DEBUG(dbgs() << "Replacing leaf " << *I << "\n");
352     Builder.SetInsertPoint(cast<Instruction>(I->getUser()));
353     Value *Shuffle = Builder.CreateShuffleVector(I->get(), LeafMask);
354     I->getUser()->setOperand(I->getOperandNo(), Shuffle);
355     LLVM_DEBUG(dbgs() << "  with " << *Shuffle << "\n");
356   }
357 
358   for (Instruction *I : Truncs) {
359     LLVM_DEBUG(dbgs() << "Replacing trunc " << *I << "\n");
360 
361     Builder.SetInsertPoint(I->getParent(), ++I->getIterator());
362     Value *Shuf = Builder.CreateShuffleVector(I, TruncMask);
363     I->replaceAllUsesWith(Shuf);
364     cast<Instruction>(Shuf)->setOperand(0, I);
365 
366     LLVM_DEBUG(dbgs() << "  with " << *Shuf << "\n");
367   }
368 
369   return true;
370 }
371 
372 bool MVELaneInterleaving::runOnFunction(Function &F) {
373   if (!EnableInterleave)
374     return false;
375   auto &TPC = getAnalysis<TargetPassConfig>();
376   auto &TM = TPC.getTM<TargetMachine>();
377   auto *ST = &TM.getSubtarget<ARMSubtarget>(F);
378   if (!ST->hasMVEIntegerOps())
379     return false;
380 
381   bool Changed = false;
382 
383   SmallPtrSet<Instruction *, 16> Visited;
384   for (Instruction &I : reverse(instructions(F))) {
385     if (I.getType()->isVectorTy() &&
386         (isa<TruncInst>(I) || isa<FPTruncInst>(I)) && !Visited.count(&I))
387       Changed |= tryInterleave(&I, Visited);
388   }
389 
390   return Changed;
391 }
392