1 //===-- AArch64A57FPLoadBalancing.cpp - Balance FP ops statically on A57---===//
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 // For best-case performance on Cortex-A57, we should try to use a balanced
9 // mix of odd and even D-registers when performing a critical sequence of
10 // independent, non-quadword FP/ASIMD floating-point multiply or
11 // multiply-accumulate operations.
12 //
13 // This pass attempts to detect situations where the register allocation may
14 // adversely affect this load balancing and to change the registers used so as
15 // to better utilize the CPU.
16 //
17 // Ideally we'd just take each multiply or multiply-accumulate in turn and
18 // allocate it alternating even or odd registers. However, multiply-accumulates
19 // are most efficiently performed in the same functional unit as their
20 // accumulation operand. Therefore this pass tries to find maximal sequences
21 // ("Chains") of multiply-accumulates linked via their accumulation operand,
22 // and assign them all the same "color" (oddness/evenness).
23 //
24 // This optimization affects S-register and D-register floating point
25 // multiplies and FMADD/FMAs, as well as vector (floating point only) muls and
26 // FMADD/FMA. Q register instructions (and 128-bit vector instructions) are
27 // not affected.
28 //===----------------------------------------------------------------------===//
29 
30 #include "AArch64.h"
31 #include "AArch64InstrInfo.h"
32 #include "AArch64Subtarget.h"
33 #include "llvm/ADT/BitVector.h"
34 #include "llvm/ADT/EquivalenceClasses.h"
35 #include "llvm/CodeGen/MachineFunction.h"
36 #include "llvm/CodeGen/MachineFunctionPass.h"
37 #include "llvm/CodeGen/MachineInstr.h"
38 #include "llvm/CodeGen/MachineInstrBuilder.h"
39 #include "llvm/CodeGen/MachineRegisterInfo.h"
40 #include "llvm/CodeGen/RegisterClassInfo.h"
41 #include "llvm/CodeGen/RegisterScavenging.h"
42 #include "llvm/Support/CommandLine.h"
43 #include "llvm/Support/Debug.h"
44 #include "llvm/Support/raw_ostream.h"
45 using namespace llvm;
46 
47 #define DEBUG_TYPE "aarch64-a57-fp-load-balancing"
48 
49 // Enforce the algorithm to use the scavenged register even when the original
50 // destination register is the correct color. Used for testing.
51 static cl::opt<bool>
52 TransformAll("aarch64-a57-fp-load-balancing-force-all",
53              cl::desc("Always modify dest registers regardless of color"),
54              cl::init(false), cl::Hidden);
55 
56 // Never use the balance information obtained from chains - return a specific
57 // color always. Used for testing.
58 static cl::opt<unsigned>
59 OverrideBalance("aarch64-a57-fp-load-balancing-override",
60               cl::desc("Ignore balance information, always return "
61                        "(1: Even, 2: Odd)."),
62               cl::init(0), cl::Hidden);
63 
64 //===----------------------------------------------------------------------===//
65 // Helper functions
66 
67 // Is the instruction a type of multiply on 64-bit (or 32-bit) FPRs?
68 static bool isMul(MachineInstr *MI) {
69   switch (MI->getOpcode()) {
70   case AArch64::FMULSrr:
71   case AArch64::FNMULSrr:
72   case AArch64::FMULDrr:
73   case AArch64::FNMULDrr:
74     return true;
75   default:
76     return false;
77   }
78 }
79 
80 // Is the instruction a type of FP multiply-accumulate on 64-bit (or 32-bit) FPRs?
81 static bool isMla(MachineInstr *MI) {
82   switch (MI->getOpcode()) {
83   case AArch64::FMSUBSrrr:
84   case AArch64::FMADDSrrr:
85   case AArch64::FNMSUBSrrr:
86   case AArch64::FNMADDSrrr:
87   case AArch64::FMSUBDrrr:
88   case AArch64::FMADDDrrr:
89   case AArch64::FNMSUBDrrr:
90   case AArch64::FNMADDDrrr:
91     return true;
92   default:
93     return false;
94   }
95 }
96 
97 //===----------------------------------------------------------------------===//
98 
99 namespace {
100 /// A "color", which is either even or odd. Yes, these aren't really colors
101 /// but the algorithm is conceptually doing two-color graph coloring.
102 enum class Color { Even, Odd };
103 #ifndef NDEBUG
104 static const char *ColorNames[2] = { "Even", "Odd" };
105 #endif
106 
107 class Chain;
108 
109 class AArch64A57FPLoadBalancing : public MachineFunctionPass {
110   MachineRegisterInfo *MRI;
111   const TargetRegisterInfo *TRI;
112   RegisterClassInfo RCI;
113 
114 public:
115   static char ID;
116   explicit AArch64A57FPLoadBalancing() : MachineFunctionPass(ID) {
117     initializeAArch64A57FPLoadBalancingPass(*PassRegistry::getPassRegistry());
118   }
119 
120   bool runOnMachineFunction(MachineFunction &F) override;
121 
122   MachineFunctionProperties getRequiredProperties() const override {
123     return MachineFunctionProperties().set(
124         MachineFunctionProperties::Property::NoVRegs);
125   }
126 
127   StringRef getPassName() const override {
128     return "A57 FP Anti-dependency breaker";
129   }
130 
131   void getAnalysisUsage(AnalysisUsage &AU) const override {
132     AU.setPreservesCFG();
133     MachineFunctionPass::getAnalysisUsage(AU);
134   }
135 
136 private:
137   bool runOnBasicBlock(MachineBasicBlock &MBB);
138   bool colorChainSet(std::vector<Chain*> GV, MachineBasicBlock &MBB,
139                      int &Balance);
140   bool colorChain(Chain *G, Color C, MachineBasicBlock &MBB);
141   int scavengeRegister(Chain *G, Color C, MachineBasicBlock &MBB);
142   void scanInstruction(MachineInstr *MI, unsigned Idx,
143                        std::map<unsigned, Chain*> &Active,
144                        std::vector<std::unique_ptr<Chain>> &AllChains);
145   void maybeKillChain(MachineOperand &MO, unsigned Idx,
146                       std::map<unsigned, Chain*> &RegChains);
147   Color getColor(unsigned Register);
148   Chain *getAndEraseNext(Color PreferredColor, std::vector<Chain*> &L);
149 };
150 }
151 
152 char AArch64A57FPLoadBalancing::ID = 0;
153 
154 INITIALIZE_PASS_BEGIN(AArch64A57FPLoadBalancing, DEBUG_TYPE,
155                       "AArch64 A57 FP Load-Balancing", false, false)
156 INITIALIZE_PASS_END(AArch64A57FPLoadBalancing, DEBUG_TYPE,
157                     "AArch64 A57 FP Load-Balancing", false, false)
158 
159 namespace {
160 /// A Chain is a sequence of instructions that are linked together by
161 /// an accumulation operand. For example:
162 ///
163 ///   fmul def d0, ?
164 ///   fmla def d1, ?, ?, killed d0
165 ///   fmla def d2, ?, ?, killed d1
166 ///
167 /// There may be other instructions interleaved in the sequence that
168 /// do not belong to the chain. These other instructions must not use
169 /// the "chain" register at any point.
170 ///
171 /// We currently only support chains where the "chain" operand is killed
172 /// at each link in the chain for simplicity.
173 /// A chain has three important instructions - Start, Last and Kill.
174 ///   * The start instruction is the first instruction in the chain.
175 ///   * Last is the final instruction in the chain.
176 ///   * Kill may or may not be defined. If defined, Kill is the instruction
177 ///     where the outgoing value of the Last instruction is killed.
178 ///     This information is important as if we know the outgoing value is
179 ///     killed with no intervening uses, we can safely change its register.
180 ///
181 /// Without a kill instruction, we must assume the outgoing value escapes
182 /// beyond our model and either must not change its register or must
183 /// create a fixup FMOV to keep the old register value consistent.
184 ///
185 class Chain {
186 public:
187   /// The important (marker) instructions.
188   MachineInstr *StartInst, *LastInst, *KillInst;
189   /// The index, from the start of the basic block, that each marker
190   /// appears. These are stored so we can do quick interval tests.
191   unsigned StartInstIdx, LastInstIdx, KillInstIdx;
192   /// All instructions in the chain.
193   std::set<MachineInstr*> Insts;
194   /// True if KillInst cannot be modified. If this is true,
195   /// we cannot change LastInst's outgoing register.
196   /// This will be true for tied values and regmasks.
197   bool KillIsImmutable;
198   /// The "color" of LastInst. This will be the preferred chain color,
199   /// as changing intermediate nodes is easy but changing the last
200   /// instruction can be more tricky.
201   Color LastColor;
202 
203   Chain(MachineInstr *MI, unsigned Idx, Color C)
204       : StartInst(MI), LastInst(MI), KillInst(nullptr),
205         StartInstIdx(Idx), LastInstIdx(Idx), KillInstIdx(0),
206         LastColor(C) {
207     Insts.insert(MI);
208   }
209 
210   /// Add a new instruction into the chain. The instruction's dest operand
211   /// has the given color.
212   void add(MachineInstr *MI, unsigned Idx, Color C) {
213     LastInst = MI;
214     LastInstIdx = Idx;
215     LastColor = C;
216     assert((KillInstIdx == 0 || LastInstIdx < KillInstIdx) &&
217            "Chain: broken invariant. A Chain can only be killed after its last "
218            "def");
219 
220     Insts.insert(MI);
221   }
222 
223   /// Return true if MI is a member of the chain.
224   bool contains(MachineInstr &MI) { return Insts.count(&MI) > 0; }
225 
226   /// Return the number of instructions in the chain.
227   unsigned size() const {
228     return Insts.size();
229   }
230 
231   /// Inform the chain that its last active register (the dest register of
232   /// LastInst) is killed by MI with no intervening uses or defs.
233   void setKill(MachineInstr *MI, unsigned Idx, bool Immutable) {
234     KillInst = MI;
235     KillInstIdx = Idx;
236     KillIsImmutable = Immutable;
237     assert((KillInstIdx == 0 || LastInstIdx < KillInstIdx) &&
238            "Chain: broken invariant. A Chain can only be killed after its last "
239            "def");
240   }
241 
242   /// Return the first instruction in the chain.
243   MachineInstr *getStart() const { return StartInst; }
244   /// Return the last instruction in the chain.
245   MachineInstr *getLast() const { return LastInst; }
246   /// Return the "kill" instruction (as set with setKill()) or NULL.
247   MachineInstr *getKill() const { return KillInst; }
248   /// Return an instruction that can be used as an iterator for the end
249   /// of the chain. This is the maximum of KillInst (if set) and LastInst.
250   MachineBasicBlock::iterator end() const {
251     return ++MachineBasicBlock::iterator(KillInst ? KillInst : LastInst);
252   }
253   MachineBasicBlock::iterator begin() const { return getStart(); }
254 
255   /// Can the Kill instruction (assuming one exists) be modified?
256   bool isKillImmutable() const { return KillIsImmutable; }
257 
258   /// Return the preferred color of this chain.
259   Color getPreferredColor() {
260     if (OverrideBalance != 0)
261       return OverrideBalance == 1 ? Color::Even : Color::Odd;
262     return LastColor;
263   }
264 
265   /// Return true if this chain (StartInst..KillInst) overlaps with Other.
266   bool rangeOverlapsWith(const Chain &Other) const {
267     unsigned End = KillInst ? KillInstIdx : LastInstIdx;
268     unsigned OtherEnd = Other.KillInst ?
269       Other.KillInstIdx : Other.LastInstIdx;
270 
271     return StartInstIdx <= OtherEnd && Other.StartInstIdx <= End;
272   }
273 
274   /// Return true if this chain starts before Other.
275   bool startsBefore(const Chain *Other) const {
276     return StartInstIdx < Other->StartInstIdx;
277   }
278 
279   /// Return true if the group will require a fixup MOV at the end.
280   bool requiresFixup() const {
281     return (getKill() && isKillImmutable()) || !getKill();
282   }
283 
284   /// Return a simple string representation of the chain.
285   std::string str() const {
286     std::string S;
287     raw_string_ostream OS(S);
288 
289     OS << "{";
290     StartInst->print(OS, /* SkipOpers= */true);
291     OS << " -> ";
292     LastInst->print(OS, /* SkipOpers= */true);
293     if (KillInst) {
294       OS << " (kill @ ";
295       KillInst->print(OS, /* SkipOpers= */true);
296       OS << ")";
297     }
298     OS << "}";
299 
300     return OS.str();
301   }
302 
303 };
304 
305 } // end anonymous namespace
306 
307 //===----------------------------------------------------------------------===//
308 
309 bool AArch64A57FPLoadBalancing::runOnMachineFunction(MachineFunction &F) {
310   if (skipFunction(F.getFunction()))
311     return false;
312 
313   if (!F.getSubtarget<AArch64Subtarget>().balanceFPOps())
314     return false;
315 
316   bool Changed = false;
317   LLVM_DEBUG(dbgs() << "***** AArch64A57FPLoadBalancing *****\n");
318 
319   MRI = &F.getRegInfo();
320   TRI = F.getRegInfo().getTargetRegisterInfo();
321   RCI.runOnMachineFunction(F);
322 
323   for (auto &MBB : F) {
324     Changed |= runOnBasicBlock(MBB);
325   }
326 
327   return Changed;
328 }
329 
330 bool AArch64A57FPLoadBalancing::runOnBasicBlock(MachineBasicBlock &MBB) {
331   bool Changed = false;
332   LLVM_DEBUG(dbgs() << "Running on MBB: " << MBB
333                     << " - scanning instructions...\n");
334 
335   // First, scan the basic block producing a set of chains.
336 
337   // The currently "active" chains - chains that can be added to and haven't
338   // been killed yet. This is keyed by register - all chains can only have one
339   // "link" register between each inst in the chain.
340   std::map<unsigned, Chain*> ActiveChains;
341   std::vector<std::unique_ptr<Chain>> AllChains;
342   unsigned Idx = 0;
343   for (auto &MI : MBB)
344     scanInstruction(&MI, Idx++, ActiveChains, AllChains);
345 
346   LLVM_DEBUG(dbgs() << "Scan complete, " << AllChains.size()
347                     << " chains created.\n");
348 
349   // Group the chains into disjoint sets based on their liveness range. This is
350   // a poor-man's version of graph coloring. Ideally we'd create an interference
351   // graph and perform full-on graph coloring on that, but;
352   //   (a) That's rather heavyweight for only two colors.
353   //   (b) We expect multiple disjoint interference regions - in practice the live
354   //       range of chains is quite small and they are clustered between loads
355   //       and stores.
356   EquivalenceClasses<Chain*> EC;
357   for (auto &I : AllChains)
358     EC.insert(I.get());
359 
360   for (auto &I : AllChains)
361     for (auto &J : AllChains)
362       if (I != J && I->rangeOverlapsWith(*J))
363         EC.unionSets(I.get(), J.get());
364   LLVM_DEBUG(dbgs() << "Created " << EC.getNumClasses() << " disjoint sets.\n");
365 
366   // Now we assume that every member of an equivalence class interferes
367   // with every other member of that class, and with no members of other classes.
368 
369   // Convert the EquivalenceClasses to a simpler set of sets.
370   std::vector<std::vector<Chain*> > V;
371   for (auto I = EC.begin(), E = EC.end(); I != E; ++I) {
372     std::vector<Chain*> Cs(EC.member_begin(I), EC.member_end());
373     if (Cs.empty()) continue;
374     V.push_back(std::move(Cs));
375   }
376 
377   // Now we have a set of sets, order them by start address so
378   // we can iterate over them sequentially.
379   llvm::sort(V,
380              [](const std::vector<Chain *> &A, const std::vector<Chain *> &B) {
381                return A.front()->startsBefore(B.front());
382              });
383 
384   // As we only have two colors, we can track the global (BB-level) balance of
385   // odds versus evens. We aim to keep this near zero to keep both execution
386   // units fed.
387   // Positive means we're even-heavy, negative we're odd-heavy.
388   //
389   // FIXME: If chains have interdependencies, for example:
390   //   mul r0, r1, r2
391   //   mul r3, r0, r1
392   // We do not model this and may color each one differently, assuming we'll
393   // get ILP when we obviously can't. This hasn't been seen to be a problem
394   // in practice so far, so we simplify the algorithm by ignoring it.
395   int Parity = 0;
396 
397   for (auto &I : V)
398     Changed |= colorChainSet(std::move(I), MBB, Parity);
399 
400   return Changed;
401 }
402 
403 Chain *AArch64A57FPLoadBalancing::getAndEraseNext(Color PreferredColor,
404                                                   std::vector<Chain*> &L) {
405   if (L.empty())
406     return nullptr;
407 
408   // We try and get the best candidate from L to color next, given that our
409   // preferred color is "PreferredColor". L is ordered from larger to smaller
410   // chains. It is beneficial to color the large chains before the small chains,
411   // but if we can't find a chain of the maximum length with the preferred color,
412   // we fuzz the size and look for slightly smaller chains before giving up and
413   // returning a chain that must be recolored.
414 
415   // FIXME: Does this need to be configurable?
416   const unsigned SizeFuzz = 1;
417   unsigned MinSize = L.front()->size() - SizeFuzz;
418   for (auto I = L.begin(), E = L.end(); I != E; ++I) {
419     if ((*I)->size() <= MinSize) {
420       // We've gone past the size limit. Return the previous item.
421       Chain *Ch = *--I;
422       L.erase(I);
423       return Ch;
424     }
425 
426     if ((*I)->getPreferredColor() == PreferredColor) {
427       Chain *Ch = *I;
428       L.erase(I);
429       return Ch;
430     }
431   }
432 
433   // Bailout case - just return the first item.
434   Chain *Ch = L.front();
435   L.erase(L.begin());
436   return Ch;
437 }
438 
439 bool AArch64A57FPLoadBalancing::colorChainSet(std::vector<Chain*> GV,
440                                               MachineBasicBlock &MBB,
441                                               int &Parity) {
442   bool Changed = false;
443   LLVM_DEBUG(dbgs() << "colorChainSet(): #sets=" << GV.size() << "\n");
444 
445   // Sort by descending size order so that we allocate the most important
446   // sets first.
447   // Tie-break equivalent sizes by sorting chains requiring fixups before
448   // those without fixups. The logic here is that we should look at the
449   // chains that we cannot change before we look at those we can,
450   // so the parity counter is updated and we know what color we should
451   // change them to!
452   // Final tie-break with instruction order so pass output is stable (i.e. not
453   // dependent on malloc'd pointer values).
454   llvm::sort(GV, [](const Chain *G1, const Chain *G2) {
455     if (G1->size() != G2->size())
456       return G1->size() > G2->size();
457     if (G1->requiresFixup() != G2->requiresFixup())
458       return G1->requiresFixup() > G2->requiresFixup();
459     // Make sure startsBefore() produces a stable final order.
460     assert((G1 == G2 || (G1->startsBefore(G2) ^ G2->startsBefore(G1))) &&
461            "Starts before not total order!");
462     return G1->startsBefore(G2);
463   });
464 
465   Color PreferredColor = Parity < 0 ? Color::Even : Color::Odd;
466   while (Chain *G = getAndEraseNext(PreferredColor, GV)) {
467     // Start off by assuming we'll color to our own preferred color.
468     Color C = PreferredColor;
469     if (Parity == 0)
470       // But if we really don't care, use the chain's preferred color.
471       C = G->getPreferredColor();
472 
473     LLVM_DEBUG(dbgs() << " - Parity=" << Parity
474                       << ", Color=" << ColorNames[(int)C] << "\n");
475 
476     // If we'll need a fixup FMOV, don't bother. Testing has shown that this
477     // happens infrequently and when it does it has at least a 50% chance of
478     // slowing code down instead of speeding it up.
479     if (G->requiresFixup() && C != G->getPreferredColor()) {
480       C = G->getPreferredColor();
481       LLVM_DEBUG(dbgs() << " - " << G->str()
482                         << " - not worthwhile changing; "
483                            "color remains "
484                         << ColorNames[(int)C] << "\n");
485     }
486 
487     Changed |= colorChain(G, C, MBB);
488 
489     Parity += (C == Color::Even) ? G->size() : -G->size();
490     PreferredColor = Parity < 0 ? Color::Even : Color::Odd;
491   }
492 
493   return Changed;
494 }
495 
496 int AArch64A57FPLoadBalancing::scavengeRegister(Chain *G, Color C,
497                                                 MachineBasicBlock &MBB) {
498   // Can we find an appropriate register that is available throughout the life
499   // of the chain? Simulate liveness backwards until the end of the chain.
500   LiveRegUnits Units(*TRI);
501   Units.addLiveOuts(MBB);
502   MachineBasicBlock::iterator I = MBB.end();
503   MachineBasicBlock::iterator ChainEnd = G->end();
504   while (I != ChainEnd) {
505     --I;
506     Units.stepBackward(*I);
507   }
508 
509   // Check which register units are alive throughout the chain.
510   MachineBasicBlock::iterator ChainBegin = G->begin();
511   assert(ChainBegin != ChainEnd && "Chain should contain instructions");
512   do {
513     --I;
514     Units.accumulate(*I);
515   } while (I != ChainBegin);
516 
517   // Make sure we allocate in-order, to get the cheapest registers first.
518   unsigned RegClassID = ChainBegin->getDesc().operands()[0].RegClass;
519   auto Ord = RCI.getOrder(TRI->getRegClass(RegClassID));
520   for (auto Reg : Ord) {
521     if (!Units.available(Reg))
522       continue;
523     if (C == getColor(Reg))
524       return Reg;
525   }
526 
527   return -1;
528 }
529 
530 bool AArch64A57FPLoadBalancing::colorChain(Chain *G, Color C,
531                                            MachineBasicBlock &MBB) {
532   bool Changed = false;
533   LLVM_DEBUG(dbgs() << " - colorChain(" << G->str() << ", "
534                     << ColorNames[(int)C] << ")\n");
535 
536   // Try and obtain a free register of the right class. Without a register
537   // to play with we cannot continue.
538   int Reg = scavengeRegister(G, C, MBB);
539   if (Reg == -1) {
540     LLVM_DEBUG(dbgs() << "Scavenging (thus coloring) failed!\n");
541     return false;
542   }
543   LLVM_DEBUG(dbgs() << " - Scavenged register: " << printReg(Reg, TRI) << "\n");
544 
545   std::map<unsigned, unsigned> Substs;
546   for (MachineInstr &I : *G) {
547     if (!G->contains(I) && (&I != G->getKill() || G->isKillImmutable()))
548       continue;
549 
550     // I is a member of G, or I is a mutable instruction that kills G.
551 
552     std::vector<unsigned> ToErase;
553     for (auto &U : I.operands()) {
554       if (U.isReg() && U.isUse() && Substs.find(U.getReg()) != Substs.end()) {
555         Register OrigReg = U.getReg();
556         U.setReg(Substs[OrigReg]);
557         if (U.isKill())
558           // Don't erase straight away, because there may be other operands
559           // that also reference this substitution!
560           ToErase.push_back(OrigReg);
561       } else if (U.isRegMask()) {
562         for (auto J : Substs) {
563           if (U.clobbersPhysReg(J.first))
564             ToErase.push_back(J.first);
565         }
566       }
567     }
568     // Now it's safe to remove the substs identified earlier.
569     for (auto J : ToErase)
570       Substs.erase(J);
571 
572     // Only change the def if this isn't the last instruction.
573     if (&I != G->getKill()) {
574       MachineOperand &MO = I.getOperand(0);
575 
576       bool Change = TransformAll || getColor(MO.getReg()) != C;
577       if (G->requiresFixup() && &I == G->getLast())
578         Change = false;
579 
580       if (Change) {
581         Substs[MO.getReg()] = Reg;
582         MO.setReg(Reg);
583 
584         Changed = true;
585       }
586     }
587   }
588   assert(Substs.size() == 0 && "No substitutions should be left active!");
589 
590   if (G->getKill()) {
591     LLVM_DEBUG(dbgs() << " - Kill instruction seen.\n");
592   } else {
593     // We didn't have a kill instruction, but we didn't seem to need to change
594     // the destination register anyway.
595     LLVM_DEBUG(dbgs() << " - Destination register not changed.\n");
596   }
597   return Changed;
598 }
599 
600 void AArch64A57FPLoadBalancing::scanInstruction(
601     MachineInstr *MI, unsigned Idx, std::map<unsigned, Chain *> &ActiveChains,
602     std::vector<std::unique_ptr<Chain>> &AllChains) {
603   // Inspect "MI", updating ActiveChains and AllChains.
604 
605   if (isMul(MI)) {
606 
607     for (auto &I : MI->uses())
608       maybeKillChain(I, Idx, ActiveChains);
609     for (auto &I : MI->defs())
610       maybeKillChain(I, Idx, ActiveChains);
611 
612     // Create a new chain. Multiplies don't require forwarding so can go on any
613     // unit.
614     Register DestReg = MI->getOperand(0).getReg();
615 
616     LLVM_DEBUG(dbgs() << "New chain started for register "
617                       << printReg(DestReg, TRI) << " at " << *MI);
618 
619     auto G = std::make_unique<Chain>(MI, Idx, getColor(DestReg));
620     ActiveChains[DestReg] = G.get();
621     AllChains.push_back(std::move(G));
622 
623   } else if (isMla(MI)) {
624 
625     // It is beneficial to keep MLAs on the same functional unit as their
626     // accumulator operand.
627     Register DestReg = MI->getOperand(0).getReg();
628     Register AccumReg = MI->getOperand(3).getReg();
629 
630     maybeKillChain(MI->getOperand(1), Idx, ActiveChains);
631     maybeKillChain(MI->getOperand(2), Idx, ActiveChains);
632     if (DestReg != AccumReg)
633       maybeKillChain(MI->getOperand(0), Idx, ActiveChains);
634 
635     if (ActiveChains.find(AccumReg) != ActiveChains.end()) {
636       LLVM_DEBUG(dbgs() << "Chain found for accumulator register "
637                         << printReg(AccumReg, TRI) << " in MI " << *MI);
638 
639       // For simplicity we only chain together sequences of MULs/MLAs where the
640       // accumulator register is killed on each instruction. This means we don't
641       // need to track other uses of the registers we want to rewrite.
642       //
643       // FIXME: We could extend to handle the non-kill cases for more coverage.
644       if (MI->getOperand(3).isKill()) {
645         // Add to chain.
646         LLVM_DEBUG(dbgs() << "Instruction was successfully added to chain.\n");
647         ActiveChains[AccumReg]->add(MI, Idx, getColor(DestReg));
648         // Handle cases where the destination is not the same as the accumulator.
649         if (DestReg != AccumReg) {
650           ActiveChains[DestReg] = ActiveChains[AccumReg];
651           ActiveChains.erase(AccumReg);
652         }
653         return;
654       }
655 
656       LLVM_DEBUG(
657           dbgs() << "Cannot add to chain because accumulator operand wasn't "
658                  << "marked <kill>!\n");
659       maybeKillChain(MI->getOperand(3), Idx, ActiveChains);
660     }
661 
662     LLVM_DEBUG(dbgs() << "Creating new chain for dest register "
663                       << printReg(DestReg, TRI) << "\n");
664     auto G = std::make_unique<Chain>(MI, Idx, getColor(DestReg));
665     ActiveChains[DestReg] = G.get();
666     AllChains.push_back(std::move(G));
667 
668   } else {
669 
670     // Non-MUL or MLA instruction. Invalidate any chain in the uses or defs
671     // lists.
672     for (auto &I : MI->uses())
673       maybeKillChain(I, Idx, ActiveChains);
674     for (auto &I : MI->defs())
675       maybeKillChain(I, Idx, ActiveChains);
676 
677   }
678 }
679 
680 void AArch64A57FPLoadBalancing::
681 maybeKillChain(MachineOperand &MO, unsigned Idx,
682                std::map<unsigned, Chain*> &ActiveChains) {
683   // Given an operand and the set of active chains (keyed by register),
684   // determine if a chain should be ended and remove from ActiveChains.
685   MachineInstr *MI = MO.getParent();
686 
687   if (MO.isReg()) {
688 
689     // If this is a KILL of a current chain, record it.
690     if (MO.isKill() && ActiveChains.find(MO.getReg()) != ActiveChains.end()) {
691       LLVM_DEBUG(dbgs() << "Kill seen for chain " << printReg(MO.getReg(), TRI)
692                         << "\n");
693       ActiveChains[MO.getReg()]->setKill(MI, Idx, /*Immutable=*/MO.isTied());
694     }
695     ActiveChains.erase(MO.getReg());
696 
697   } else if (MO.isRegMask()) {
698 
699     for (auto I = ActiveChains.begin(), E = ActiveChains.end();
700          I != E;) {
701       if (MO.clobbersPhysReg(I->first)) {
702         LLVM_DEBUG(dbgs() << "Kill (regmask) seen for chain "
703                           << printReg(I->first, TRI) << "\n");
704         I->second->setKill(MI, Idx, /*Immutable=*/true);
705         ActiveChains.erase(I++);
706       } else
707         ++I;
708     }
709 
710   }
711 }
712 
713 Color AArch64A57FPLoadBalancing::getColor(unsigned Reg) {
714   if ((TRI->getEncodingValue(Reg) % 2) == 0)
715     return Color::Even;
716   else
717     return Color::Odd;
718 }
719 
720 // Factory function used by AArch64TargetMachine to add the pass to the passmanager.
721 FunctionPass *llvm::createAArch64A57FPLoadBalancing() {
722   return new AArch64A57FPLoadBalancing();
723 }
724