lib/Target/ARM64/ARM64LoadStoreOptimizer.cpp

   1 //===-- ARM64LoadStoreOptimizer.cpp - ARM64 load/store opt. pass --*- C++ -*-=//
   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 contains a pass that performs load / store related peephole
  11 // optimizations. This pass should be run after register allocation.
  12 //
  13 //===----------------------------------------------------------------------===//
  14
  15 #include "ARM64InstrInfo.h"
  16 #include "MCTargetDesc/ARM64AddressingModes.h"
  17 #include "llvm/ADT/BitVector.h"
  18 #include "llvm/CodeGen/MachineBasicBlock.h"
  19 #include "llvm/CodeGen/MachineFunctionPass.h"
  20 #include "llvm/CodeGen/MachineInstr.h"
  21 #include "llvm/CodeGen/MachineInstrBuilder.h"
  22 #include "llvm/Target/TargetInstrInfo.h"
  23 #include "llvm/Target/TargetMachine.h"
  24 #include "llvm/Target/TargetRegisterInfo.h"
  25 #include "llvm/Support/CommandLine.h"
  26 #include "llvm/Support/Debug.h"
  27 #include "llvm/Support/ErrorHandling.h"
  28 #include "llvm/Support/raw_ostream.h"
  29 #include "llvm/ADT/Statistic.h"
  30 using namespace llvm;
  31
  32 #define DEBUG_TYPE "arm64-ldst-opt"
  33
  34 /// ARM64AllocLoadStoreOpt - Post-register allocation pass to combine
  35 /// load / store instructions to form ldp / stp instructions.
  36
  37 STATISTIC(NumPairCreated, "Number of load/store pair instructions generated");
  38 STATISTIC(NumPostFolded, "Number of post-index updates folded");
  39 STATISTIC(NumPreFolded, "Number of pre-index updates folded");
  40 STATISTIC(NumUnscaledPairCreated,
  41           "Number of load/store from unscaled generated");
  42
  43 static cl::opt<unsigned> ScanLimit("arm64-load-store-scan-limit", cl::init(20),
  44                                    cl::Hidden);
  45
  46 // Place holder while testing unscaled load/store combining
  47 static cl::opt<bool>
  48 EnableARM64UnscaledMemOp("arm64-unscaled-mem-op", cl::Hidden,
  49                          cl::desc("Allow ARM64 unscaled load/store combining"),
  50                          cl::init(true));
  51
  52 namespace {
  53 struct ARM64LoadStoreOpt : public MachineFunctionPass {
  54   static char ID;
  55   ARM64LoadStoreOpt() : MachineFunctionPass(ID) {}
  56
  57   const ARM64InstrInfo *TII;
  58   const TargetRegisterInfo *TRI;
  59
  60   // Scan the instructions looking for a load/store that can be combined
  61   // with the current instruction into a load/store pair.
  62   // Return the matching instruction if one is found, else MBB->end().
  63   // If a matching instruction is found, mergeForward is set to true if the
  64   // merge is to remove the first instruction and replace the second with
  65   // a pair-wise insn, and false if the reverse is true.
  66   MachineBasicBlock::iterator findMatchingInsn(MachineBasicBlock::iterator I,
  67                                                bool &mergeForward,
  68                                                unsigned Limit);
  69   // Merge the two instructions indicated into a single pair-wise instruction.
  70   // If mergeForward is true, erase the first instruction and fold its
  71   // operation into the second. If false, the reverse. Return the instruction
  72   // following the first instruction (which may change during processing).
  73   MachineBasicBlock::iterator
  74   mergePairedInsns(MachineBasicBlock::iterator I,
  75                    MachineBasicBlock::iterator Paired, bool mergeForward);
  76
  77   // Scan the instruction list to find a base register update that can
  78   // be combined with the current instruction (a load or store) using
  79   // pre or post indexed addressing with writeback. Scan forwards.
  80   MachineBasicBlock::iterator
  81   findMatchingUpdateInsnForward(MachineBasicBlock::iterator I, unsigned Limit,
  82                                 int Value);
  83
  84   // Scan the instruction list to find a base register update that can
  85   // be combined with the current instruction (a load or store) using
  86   // pre or post indexed addressing with writeback. Scan backwards.
  87   MachineBasicBlock::iterator
  88   findMatchingUpdateInsnBackward(MachineBasicBlock::iterator I, unsigned Limit);
  89
  90   // Merge a pre-index base register update into a ld/st instruction.
  91   MachineBasicBlock::iterator
  92   mergePreIdxUpdateInsn(MachineBasicBlock::iterator I,
  93                         MachineBasicBlock::iterator Update);
  94
  95   // Merge a post-index base register update into a ld/st instruction.
  96   MachineBasicBlock::iterator
  97   mergePostIdxUpdateInsn(MachineBasicBlock::iterator I,
  98                          MachineBasicBlock::iterator Update);
  99
 100   bool optimizeBlock(MachineBasicBlock &MBB);
 101
 102   bool runOnMachineFunction(MachineFunction &Fn) override;
 103
 104   const char *getPassName() const override {
 105     return "ARM64 load / store optimization pass";
 106   }
 107
 108 private:
 109   int getMemSize(MachineInstr *MemMI);
 110 };
 111 char ARM64LoadStoreOpt::ID = 0;
 112 }
 113
 114 static bool isUnscaledLdst(unsigned Opc) {
 115   switch (Opc) {
 116   default:
 117     return false;
 118   case ARM64::STURSi:
 119     return true;
 120   case ARM64::STURDi:
 121     return true;
 122   case ARM64::STURQi:
 123     return true;
 124   case ARM64::STURWi:
 125     return true;
 126   case ARM64::STURXi:
 127     return true;
 128   case ARM64::LDURSi:
 129     return true;
 130   case ARM64::LDURDi:
 131     return true;
 132   case ARM64::LDURQi:
 133     return true;
 134   case ARM64::LDURWi:
 135     return true;
 136   case ARM64::LDURXi:
 137     return true;
 138   }
 139 }
 140
 141 // Size in bytes of the data moved by an unscaled load or store
 142 int ARM64LoadStoreOpt::getMemSize(MachineInstr *MemMI) {
 143   switch (MemMI->getOpcode()) {
 144   default:
 145     llvm_unreachable("Opcode has has unknown size!");
 146   case ARM64::STRSui:
 147   case ARM64::STURSi:
 148     return 4;
 149   case ARM64::STRDui:
 150   case ARM64::STURDi:
 151     return 8;
 152   case ARM64::STRQui:
 153   case ARM64::STURQi:
 154     return 16;
 155   case ARM64::STRWui:
 156   case ARM64::STURWi:
 157     return 4;
 158   case ARM64::STRXui:
 159   case ARM64::STURXi:
 160     return 8;
 161   case ARM64::LDRSui:
 162   case ARM64::LDURSi:
 163     return 4;
 164   case ARM64::LDRDui:
 165   case ARM64::LDURDi:
 166     return 8;
 167   case ARM64::LDRQui:
 168   case ARM64::LDURQi:
 169     return 16;
 170   case ARM64::LDRWui:
 171   case ARM64::LDURWi:
 172     return 4;
 173   case ARM64::LDRXui:
 174   case ARM64::LDURXi:
 175     return 8;
 176   }
 177 }
 178
 179 static unsigned getMatchingPairOpcode(unsigned Opc) {
 180   switch (Opc) {
 181   default:
 182     llvm_unreachable("Opcode has no pairwise equivalent!");
 183   case ARM64::STRSui:
 184   case ARM64::STURSi:
 185     return ARM64::STPSi;
 186   case ARM64::STRDui:
 187   case ARM64::STURDi:
 188     return ARM64::STPDi;
 189   case ARM64::STRQui:
 190   case ARM64::STURQi:
 191     return ARM64::STPQi;
 192   case ARM64::STRWui:
 193   case ARM64::STURWi:
 194     return ARM64::STPWi;
 195   case ARM64::STRXui:
 196   case ARM64::STURXi:
 197     return ARM64::STPXi;
 198   case ARM64::LDRSui:
 199   case ARM64::LDURSi:
 200     return ARM64::LDPSi;
 201   case ARM64::LDRDui:
 202   case ARM64::LDURDi:
 203     return ARM64::LDPDi;
 204   case ARM64::LDRQui:
 205   case ARM64::LDURQi:
 206     return ARM64::LDPQi;
 207   case ARM64::LDRWui:
 208   case ARM64::LDURWi:
 209     return ARM64::LDPWi;
 210   case ARM64::LDRXui:
 211   case ARM64::LDURXi:
 212     return ARM64::LDPXi;
 213   }
 214 }
 215
 216 static unsigned getPreIndexedOpcode(unsigned Opc) {
 217   switch (Opc) {
 218   default:
 219     llvm_unreachable("Opcode has no pre-indexed equivalent!");
 220   case ARM64::STRSui:    return ARM64::STRSpre;
 221   case ARM64::STRDui:    return ARM64::STRDpre;
 222   case ARM64::STRQui:    return ARM64::STRQpre;
 223   case ARM64::STRWui:    return ARM64::STRWpre;
 224   case ARM64::STRXui:    return ARM64::STRXpre;
 225   case ARM64::LDRSui:    return ARM64::LDRSpre;
 226   case ARM64::LDRDui:    return ARM64::LDRDpre;
 227   case ARM64::LDRQui:    return ARM64::LDRQpre;
 228   case ARM64::LDRWui:    return ARM64::LDRWpre;
 229   case ARM64::LDRXui:    return ARM64::LDRXpre;
 230   }
 231 }
 232
 233 static unsigned getPostIndexedOpcode(unsigned Opc) {
 234   switch (Opc) {
 235   default:
 236     llvm_unreachable("Opcode has no post-indexed wise equivalent!");
 237   case ARM64::STRSui:
 238     return ARM64::STRSpost;
 239   case ARM64::STRDui:
 240     return ARM64::STRDpost;
 241   case ARM64::STRQui:
 242     return ARM64::STRQpost;
 243   case ARM64::STRWui:
 244     return ARM64::STRWpost;
 245   case ARM64::STRXui:
 246     return ARM64::STRXpost;
 247   case ARM64::LDRSui:
 248     return ARM64::LDRSpost;
 249   case ARM64::LDRDui:
 250     return ARM64::LDRDpost;
 251   case ARM64::LDRQui:
 252     return ARM64::LDRQpost;
 253   case ARM64::LDRWui:
 254     return ARM64::LDRWpost;
 255   case ARM64::LDRXui:
 256     return ARM64::LDRXpost;
 257   }
 258 }
 259
 260 MachineBasicBlock::iterator
 261 ARM64LoadStoreOpt::mergePairedInsns(MachineBasicBlock::iterator I,
 262                                     MachineBasicBlock::iterator Paired,
 263                                     bool mergeForward) {
 264   MachineBasicBlock::iterator NextI = I;
 265   ++NextI;
 266   // If NextI is the second of the two instructions to be merged, we need
 267   // to skip one further. Either way we merge will invalidate the iterator,
 268   // and we don't need to scan the new instruction, as it's a pairwise
 269   // instruction, which we're not considering for further action anyway.
 270   if (NextI == Paired)
 271     ++NextI;
 272
 273   bool IsUnscaled = isUnscaledLdst(I->getOpcode());
 274   int OffsetStride = IsUnscaled && EnableARM64UnscaledMemOp ? getMemSize(I) : 1;
 275
 276   unsigned NewOpc = getMatchingPairOpcode(I->getOpcode());
 277   // Insert our new paired instruction after whichever of the paired
 278   // instructions mergeForward indicates.
 279   MachineBasicBlock::iterator InsertionPoint = mergeForward ? Paired : I;
 280   // Also based on mergeForward is from where we copy the base register operand
 281   // so we get the flags compatible with the input code.
 282   MachineOperand &BaseRegOp =
 283       mergeForward ? Paired->getOperand(1) : I->getOperand(1);
 284
 285   // Which register is Rt and which is Rt2 depends on the offset order.
 286   MachineInstr *RtMI, *Rt2MI;
 287   if (I->getOperand(2).getImm() ==
 288       Paired->getOperand(2).getImm() + OffsetStride) {
 289     RtMI = Paired;
 290     Rt2MI = I;
 291   } else {
 292     RtMI = I;
 293     Rt2MI = Paired;
 294   }
 295   // Handle Unscaled
 296   int OffsetImm = RtMI->getOperand(2).getImm();
 297   if (IsUnscaled && EnableARM64UnscaledMemOp)
 298     OffsetImm /= OffsetStride;
 299
 300   // Construct the new instruction.
 301   MachineInstrBuilder MIB = BuildMI(*I->getParent(), InsertionPoint,
 302                                     I->getDebugLoc(), TII->get(NewOpc))
 303                                 .addOperand(RtMI->getOperand(0))
 304                                 .addOperand(Rt2MI->getOperand(0))
 305                                 .addOperand(BaseRegOp)
 306                                 .addImm(OffsetImm);
 307   (void)MIB;
 308
 309   // FIXME: Do we need/want to copy the mem operands from the source
 310   //        instructions? Probably. What uses them after this?
 311
 312   DEBUG(dbgs() << "Creating pair load/store. Replacing instructions:\n    ");
 313   DEBUG(I->print(dbgs()));
 314   DEBUG(dbgs() << "    ");
 315   DEBUG(Paired->print(dbgs()));
 316   DEBUG(dbgs() << "  with instruction:\n    ");
 317   DEBUG(((MachineInstr *)MIB)->print(dbgs()));
 318   DEBUG(dbgs() << "\n");
 319
 320   // Erase the old instructions.
 321   I->eraseFromParent();
 322   Paired->eraseFromParent();
 323
 324   return NextI;
 325 }
 326
 327 /// trackRegDefsUses - Remember what registers the specified instruction uses
 328 /// and modifies.
 329 static void trackRegDefsUses(MachineInstr *MI, BitVector &ModifiedRegs,
 330                              BitVector &UsedRegs,
 331                              const TargetRegisterInfo *TRI) {
 332   for (unsigned i = 0, e = MI->getNumOperands(); i != e; ++i) {
 333     MachineOperand &MO = MI->getOperand(i);
 334     if (MO.isRegMask())
 335       ModifiedRegs.setBitsNotInMask(MO.getRegMask());
 336
 337     if (!MO.isReg())
 338       continue;
 339     unsigned Reg = MO.getReg();
 340     if (MO.isDef()) {
 341       for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
 342         ModifiedRegs.set(*AI);
 343     } else {
 344       assert(MO.isUse() && "Reg operand not a def and not a use?!?");
 345       for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
 346         UsedRegs.set(*AI);
 347     }
 348   }
 349 }
 350
 351 static bool inBoundsForPair(bool IsUnscaled, int Offset, int OffsetStride) {
 352   if (!IsUnscaled && (Offset > 63 || Offset < -64))
 353     return false;
 354   if (IsUnscaled) {
 355     // Convert the byte-offset used by unscaled into an "element" offset used
 356     // by the scaled pair load/store instructions.
 357     int elemOffset = Offset / OffsetStride;
 358     if (elemOffset > 63 || elemOffset < -64)
 359       return false;
 360   }
 361   return true;
 362 }
 363
 364 // Do alignment, specialized to power of 2 and for signed ints,
 365 // avoiding having to do a C-style cast from uint_64t to int when
 366 // using RoundUpToAlignment from include/llvm/Support/MathExtras.h.
 367 // FIXME: Move this function to include/MathExtras.h?
 368 static int alignTo(int Num, int PowOf2) {
 369   return (Num + PowOf2 - 1) & ~(PowOf2 - 1);
 370 }
 371
 372 /// findMatchingInsn - Scan the instructions looking for a load/store that can
 373 /// be combined with the current instruction into a load/store pair.
 374 MachineBasicBlock::iterator
 375 ARM64LoadStoreOpt::findMatchingInsn(MachineBasicBlock::iterator I,
 376                                     bool &mergeForward, unsigned Limit) {
 377   MachineBasicBlock::iterator E = I->getParent()->end();
 378   MachineBasicBlock::iterator MBBI = I;
 379   MachineInstr *FirstMI = I;
 380   ++MBBI;
 381
 382   int Opc = FirstMI->getOpcode();
 383   bool mayLoad = FirstMI->mayLoad();
 384   bool IsUnscaled = isUnscaledLdst(Opc);
 385   unsigned Reg = FirstMI->getOperand(0).getReg();
 386   unsigned BaseReg = FirstMI->getOperand(1).getReg();
 387   int Offset = FirstMI->getOperand(2).getImm();
 388
 389   // Early exit if the first instruction modifies the base register.
 390   // e.g., ldr x0, [x0]
 391   // Early exit if the offset if not possible to match. (6 bits of positive
 392   // range, plus allow an extra one in case we find a later insn that matches
 393   // with Offset-1
 394   if (FirstMI->modifiesRegister(BaseReg, TRI))
 395     return E;
 396   int OffsetStride =
 397       IsUnscaled && EnableARM64UnscaledMemOp ? getMemSize(FirstMI) : 1;
 398   if (!inBoundsForPair(IsUnscaled, Offset, OffsetStride))
 399     return E;
 400
 401   // Track which registers have been modified and used between the first insn
 402   // (inclusive) and the second insn.
 403   BitVector ModifiedRegs, UsedRegs;
 404   ModifiedRegs.resize(TRI->getNumRegs());
 405   UsedRegs.resize(TRI->getNumRegs());
 406   for (unsigned Count = 0; MBBI != E && Count < Limit; ++MBBI) {
 407     MachineInstr *MI = MBBI;
 408     // Skip DBG_VALUE instructions. Otherwise debug info can affect the
 409     // optimization by changing how far we scan.
 410     if (MI->isDebugValue())
 411       continue;
 412
 413     // Now that we know this is a real instruction, count it.
 414     ++Count;
 415
 416     if (Opc == MI->getOpcode() && MI->getOperand(2).isImm()) {
 417       // If we've found another instruction with the same opcode, check to see
 418       // if the base and offset are compatible with our starting instruction.
 419       // These instructions all have scaled immediate operands, so we just
 420       // check for +1/-1. Make sure to check the new instruction offset is
 421       // actually an immediate and not a symbolic reference destined for
 422       // a relocation.
 423       //
 424       // Pairwise instructions have a 7-bit signed offset field. Single insns
 425       // have a 12-bit unsigned offset field. To be a valid combine, the
 426       // final offset must be in range.
 427       unsigned MIBaseReg = MI->getOperand(1).getReg();
 428       int MIOffset = MI->getOperand(2).getImm();
 429       if (BaseReg == MIBaseReg && ((Offset == MIOffset + OffsetStride) ||
 430                                    (Offset + OffsetStride == MIOffset))) {
 431         int MinOffset = Offset < MIOffset ? Offset : MIOffset;
 432         // If this is a volatile load/store that otherwise matched, stop looking
 433         // as something is going on that we don't have enough information to
 434         // safely transform. Similarly, stop if we see a hint to avoid pairs.
 435         if (MI->hasOrderedMemoryRef() || TII->isLdStPairSuppressed(MI))
 436           return E;
 437         // If the resultant immediate offset of merging these instructions
 438         // is out of range for a pairwise instruction, bail and keep looking.
 439         bool MIIsUnscaled = isUnscaledLdst(MI->getOpcode());
 440         if (!inBoundsForPair(MIIsUnscaled, MinOffset, OffsetStride)) {
 441           trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
 442           continue;
 443         }
 444         // If the alignment requirements of the paired (scaled) instruction
 445         // can't express the offset of the unscaled input, bail and keep
 446         // looking.
 447         if (IsUnscaled && EnableARM64UnscaledMemOp &&
 448             (alignTo(MinOffset, OffsetStride) != MinOffset)) {
 449           trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
 450           continue;
 451         }
 452         // If the destination register of the loads is the same register, bail
 453         // and keep looking. A load-pair instruction with both destination
 454         // registers the same is UNPREDICTABLE and will result in an exception.
 455         if (mayLoad && Reg == MI->getOperand(0).getReg()) {
 456           trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
 457           continue;
 458         }
 459
 460         // If the Rt of the second instruction was not modified or used between
 461         // the two instructions, we can combine the second into the first.
 462         if (!ModifiedRegs[MI->getOperand(0).getReg()] &&
 463             !UsedRegs[MI->getOperand(0).getReg()]) {
 464           mergeForward = false;
 465           return MBBI;
 466         }
 467
 468         // Likewise, if the Rt of the first instruction is not modified or used
 469         // between the two instructions, we can combine the first into the
 470         // second.
 471         if (!ModifiedRegs[FirstMI->getOperand(0).getReg()] &&
 472             !UsedRegs[FirstMI->getOperand(0).getReg()]) {
 473           mergeForward = true;
 474           return MBBI;
 475         }
 476         // Unable to combine these instructions due to interference in between.
 477         // Keep looking.
 478       }
 479     }
 480
 481     // If the instruction wasn't a matching load or store, but does (or can)
 482     // modify memory, stop searching, as we don't have alias analysis or
 483     // anything like that to tell us whether the access is tromping on the
 484     // locations we care about. The big one we want to catch is calls.
 485     //
 486     // FIXME: Theoretically, we can do better than that for SP and FP based
 487     // references since we can effectively know where those are touching. It's
 488     // unclear if it's worth the extra code, though. Most paired instructions
 489     // will be sequential, perhaps with a few intervening non-memory related
 490     // instructions.
 491     if (MI->mayStore() || MI->isCall())
 492       return E;
 493     // Likewise, if we're matching a store instruction, we don't want to
 494     // move across a load, as it may be reading the same location.
 495     if (FirstMI->mayStore() && MI->mayLoad())
 496       return E;
 497
 498     // Update modified / uses register lists.
 499     trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
 500
 501     // Otherwise, if the base register is modified, we have no match, so
 502     // return early.
 503     if (ModifiedRegs[BaseReg])
 504       return E;
 505   }
 506   return E;
 507 }
 508
 509 MachineBasicBlock::iterator
 510 ARM64LoadStoreOpt::mergePreIdxUpdateInsn(MachineBasicBlock::iterator I,
 511                                          MachineBasicBlock::iterator Update) {
 512   assert((Update->getOpcode() == ARM64::ADDXri ||
 513           Update->getOpcode() == ARM64::SUBXri) &&
 514          "Unexpected base register update instruction to merge!");
 515   MachineBasicBlock::iterator NextI = I;
 516   // Return the instruction following the merged instruction, which is
 517   // the instruction following our unmerged load. Unless that's the add/sub
 518   // instruction we're merging, in which case it's the one after that.
 519   if (++NextI == Update)
 520     ++NextI;
 521
 522   int Value = Update->getOperand(2).getImm();
 523   assert(ARM64_AM::getShiftValue(Update->getOperand(3).getImm()) == 0 &&
 524          "Can't merge 1 << 12 offset into pre-indexed load / store");
 525   if (Update->getOpcode() == ARM64::SUBXri)
 526     Value = -Value;
 527
 528   unsigned NewOpc = getPreIndexedOpcode(I->getOpcode());
 529   MachineInstrBuilder MIB =
 530       BuildMI(*I->getParent(), I, I->getDebugLoc(), TII->get(NewOpc))
 531           .addOperand(I->getOperand(0))
 532           .addOperand(I->getOperand(1))
 533           .addImm(Value);
 534   (void)MIB;
 535
 536   DEBUG(dbgs() << "Creating pre-indexed load/store.");
 537   DEBUG(dbgs() << "    Replacing instructions:\n    ");
 538   DEBUG(I->print(dbgs()));
 539   DEBUG(dbgs() << "    ");
 540   DEBUG(Update->print(dbgs()));
 541   DEBUG(dbgs() << "  with instruction:\n    ");
 542   DEBUG(((MachineInstr *)MIB)->print(dbgs()));
 543   DEBUG(dbgs() << "\n");
 544
 545   // Erase the old instructions for the block.
 546   I->eraseFromParent();
 547   Update->eraseFromParent();
 548
 549   return NextI;
 550 }
 551
 552 MachineBasicBlock::iterator
 553 ARM64LoadStoreOpt::mergePostIdxUpdateInsn(MachineBasicBlock::iterator I,
 554                                           MachineBasicBlock::iterator Update) {
 555   assert((Update->getOpcode() == ARM64::ADDXri ||
 556           Update->getOpcode() == ARM64::SUBXri) &&
 557          "Unexpected base register update instruction to merge!");
 558   MachineBasicBlock::iterator NextI = I;
 559   // Return the instruction following the merged instruction, which is
 560   // the instruction following our unmerged load. Unless that's the add/sub
 561   // instruction we're merging, in which case it's the one after that.
 562   if (++NextI == Update)
 563     ++NextI;
 564
 565   int Value = Update->getOperand(2).getImm();
 566   assert(ARM64_AM::getShiftValue(Update->getOperand(3).getImm()) == 0 &&
 567          "Can't merge 1 << 12 offset into post-indexed load / store");
 568   if (Update->getOpcode() == ARM64::SUBXri)
 569     Value = -Value;
 570
 571   unsigned NewOpc = getPostIndexedOpcode(I->getOpcode());
 572   MachineInstrBuilder MIB =
 573       BuildMI(*I->getParent(), I, I->getDebugLoc(), TII->get(NewOpc))
 574           .addOperand(I->getOperand(0))
 575           .addOperand(I->getOperand(1))
 576           .addImm(Value);
 577   (void)MIB;
 578
 579   DEBUG(dbgs() << "Creating post-indexed load/store.");
 580   DEBUG(dbgs() << "    Replacing instructions:\n    ");
 581   DEBUG(I->print(dbgs()));
 582   DEBUG(dbgs() << "    ");
 583   DEBUG(Update->print(dbgs()));
 584   DEBUG(dbgs() << "  with instruction:\n    ");
 585   DEBUG(((MachineInstr *)MIB)->print(dbgs()));
 586   DEBUG(dbgs() << "\n");
 587
 588   // Erase the old instructions for the block.
 589   I->eraseFromParent();
 590   Update->eraseFromParent();
 591
 592   return NextI;
 593 }
 594
 595 static bool isMatchingUpdateInsn(MachineInstr *MI, unsigned BaseReg,
 596                                  int Offset) {
 597   switch (MI->getOpcode()) {
 598   default:
 599     break;
 600   case ARM64::SUBXri:
 601     // Negate the offset for a SUB instruction.
 602     Offset *= -1;
 603   // FALLTHROUGH
 604   case ARM64::ADDXri:
 605     // Make sure it's a vanilla immediate operand, not a relocation or
 606     // anything else we can't handle.
 607     if (!MI->getOperand(2).isImm())
 608       break;
 609     // Watch out for 1 << 12 shifted value.
 610     if (ARM64_AM::getShiftValue(MI->getOperand(3).getImm()))
 611       break;
 612     // If the instruction has the base register as source and dest and the
 613     // immediate will fit in a signed 9-bit integer, then we have a match.
 614     if (MI->getOperand(0).getReg() == BaseReg &&
 615         MI->getOperand(1).getReg() == BaseReg &&
 616         MI->getOperand(2).getImm() <= 255 &&
 617         MI->getOperand(2).getImm() >= -256) {
 618       // If we have a non-zero Offset, we check that it matches the amount
 619       // we're adding to the register.
 620       if (!Offset || Offset == MI->getOperand(2).getImm())
 621         return true;
 622     }
 623     break;
 624   }
 625   return false;
 626 }
 627
 628 MachineBasicBlock::iterator
 629 ARM64LoadStoreOpt::findMatchingUpdateInsnForward(MachineBasicBlock::iterator I,
 630                                                  unsigned Limit, int Value) {
 631   MachineBasicBlock::iterator E = I->getParent()->end();
 632   MachineInstr *MemMI = I;
 633   MachineBasicBlock::iterator MBBI = I;
 634   const MachineFunction &MF = *MemMI->getParent()->getParent();
 635
 636   unsigned DestReg = MemMI->getOperand(0).getReg();
 637   unsigned BaseReg = MemMI->getOperand(1).getReg();
 638   int Offset = MemMI->getOperand(2).getImm() *
 639                TII->getRegClass(MemMI->getDesc(), 0, TRI, MF)->getSize();
 640
 641   // If the base register overlaps the destination register, we can't
 642   // merge the update.
 643   if (DestReg == BaseReg || TRI->isSubRegister(BaseReg, DestReg))
 644     return E;
 645
 646   // Scan forward looking for post-index opportunities.
 647   // Updating instructions can't be formed if the memory insn already
 648   // has an offset other than the value we're looking for.
 649   if (Offset != Value)
 650     return E;
 651
 652   // Track which registers have been modified and used between the first insn
 653   // (inclusive) and the second insn.
 654   BitVector ModifiedRegs, UsedRegs;
 655   ModifiedRegs.resize(TRI->getNumRegs());
 656   UsedRegs.resize(TRI->getNumRegs());
 657   ++MBBI;
 658   for (unsigned Count = 0; MBBI != E; ++MBBI) {
 659     MachineInstr *MI = MBBI;
 660     // Skip DBG_VALUE instructions. Otherwise debug info can affect the
 661     // optimization by changing how far we scan.
 662     if (MI->isDebugValue())
 663       continue;
 664
 665     // Now that we know this is a real instruction, count it.
 666     ++Count;
 667
 668     // If we found a match, return it.
 669     if (isMatchingUpdateInsn(MI, BaseReg, Value))
 670       return MBBI;
 671
 672     // Update the status of what the instruction clobbered and used.
 673     trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
 674
 675     // Otherwise, if the base register is used or modified, we have no match, so
 676     // return early.
 677     if (ModifiedRegs[BaseReg] || UsedRegs[BaseReg])
 678       return E;
 679   }
 680   return E;
 681 }
 682
 683 MachineBasicBlock::iterator
 684 ARM64LoadStoreOpt::findMatchingUpdateInsnBackward(MachineBasicBlock::iterator I,
 685                                                   unsigned Limit) {
 686   MachineBasicBlock::iterator B = I->getParent()->begin();
 687   MachineBasicBlock::iterator E = I->getParent()->end();
 688   MachineInstr *MemMI = I;
 689   MachineBasicBlock::iterator MBBI = I;
 690   const MachineFunction &MF = *MemMI->getParent()->getParent();
 691
 692   unsigned DestReg = MemMI->getOperand(0).getReg();
 693   unsigned BaseReg = MemMI->getOperand(1).getReg();
 694   int Offset = MemMI->getOperand(2).getImm();
 695   unsigned RegSize = TII->getRegClass(MemMI->getDesc(), 0, TRI, MF)->getSize();
 696
 697   // If the load/store is the first instruction in the block, there's obviously
 698   // not any matching update. Ditto if the memory offset isn't zero.
 699   if (MBBI == B || Offset != 0)
 700     return E;
 701   // If the base register overlaps the destination register, we can't
 702   // merge the update.
 703   if (DestReg == BaseReg || TRI->isSubRegister(BaseReg, DestReg))
 704     return E;
 705
 706   // Track which registers have been modified and used between the first insn
 707   // (inclusive) and the second insn.
 708   BitVector ModifiedRegs, UsedRegs;
 709   ModifiedRegs.resize(TRI->getNumRegs());
 710   UsedRegs.resize(TRI->getNumRegs());
 711   --MBBI;
 712   for (unsigned Count = 0; MBBI != B; --MBBI) {
 713     MachineInstr *MI = MBBI;
 714     // Skip DBG_VALUE instructions. Otherwise debug info can affect the
 715     // optimization by changing how far we scan.
 716     if (MI->isDebugValue())
 717       continue;
 718
 719     // Now that we know this is a real instruction, count it.
 720     ++Count;
 721
 722     // If we found a match, return it.
 723     if (isMatchingUpdateInsn(MI, BaseReg, RegSize))
 724       return MBBI;
 725
 726     // Update the status of what the instruction clobbered and used.
 727     trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
 728
 729     // Otherwise, if the base register is used or modified, we have no match, so
 730     // return early.
 731     if (ModifiedRegs[BaseReg] || UsedRegs[BaseReg])
 732       return E;
 733   }
 734   return E;
 735 }
 736
 737 bool ARM64LoadStoreOpt::optimizeBlock(MachineBasicBlock &MBB) {
 738   bool Modified = false;
 739   // Two tranformations to do here:
 740   // 1) Find loads and stores that can be merged into a single load or store
 741   //    pair instruction.
 742   //      e.g.,
 743   //        ldr x0, [x2]
 744   //        ldr x1, [x2, #8]
 745   //        ; becomes
 746   //        ldp x0, x1, [x2]
 747   // 2) Find base register updates that can be merged into the load or store
 748   //    as a base-reg writeback.
 749   //      e.g.,
 750   //        ldr x0, [x2]
 751   //        add x2, x2, #4
 752   //        ; becomes
 753   //        ldr x0, [x2], #4
 754
 755   for (MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
 756        MBBI != E;) {
 757     MachineInstr *MI = MBBI;
 758     switch (MI->getOpcode()) {
 759     default:
 760       // Just move on to the next instruction.
 761       ++MBBI;
 762       break;
 763     case ARM64::STRSui:
 764     case ARM64::STRDui:
 765     case ARM64::STRQui:
 766     case ARM64::STRXui:
 767     case ARM64::STRWui:
 768     case ARM64::LDRSui:
 769     case ARM64::LDRDui:
 770     case ARM64::LDRQui:
 771     case ARM64::LDRXui:
 772     case ARM64::LDRWui:
 773     // do the unscaled versions as well
 774     case ARM64::STURSi:
 775     case ARM64::STURDi:
 776     case ARM64::STURQi:
 777     case ARM64::STURWi:
 778     case ARM64::STURXi:
 779     case ARM64::LDURSi:
 780     case ARM64::LDURDi:
 781     case ARM64::LDURQi:
 782     case ARM64::LDURWi:
 783     case ARM64::LDURXi: {
 784       // If this is a volatile load/store, don't mess with it.
 785       if (MI->hasOrderedMemoryRef()) {
 786         ++MBBI;
 787         break;
 788       }
 789       // Make sure this is a reg+imm (as opposed to an address reloc).
 790       if (!MI->getOperand(2).isImm()) {
 791         ++MBBI;
 792         break;
 793       }
 794       // Check if this load/store has a hint to avoid pair formation.
 795       // MachineMemOperands hints are set by the ARM64StorePairSuppress pass.
 796       if (TII->isLdStPairSuppressed(MI)) {
 797         ++MBBI;
 798         break;
 799       }
 800       // Look ahead up to ScanLimit instructions for a pairable instruction.
 801       bool mergeForward = false;
 802       MachineBasicBlock::iterator Paired =
 803           findMatchingInsn(MBBI, mergeForward, ScanLimit);
 804       if (Paired != E) {
 805         // Merge the loads into a pair. Keeping the iterator straight is a
 806         // pain, so we let the merge routine tell us what the next instruction
 807         // is after it's done mucking about.
 808         MBBI = mergePairedInsns(MBBI, Paired, mergeForward);
 809
 810         Modified = true;
 811         ++NumPairCreated;
 812         if (isUnscaledLdst(MI->getOpcode()))
 813           ++NumUnscaledPairCreated;
 814         break;
 815       }
 816       ++MBBI;
 817       break;
 818     }
 819       // FIXME: Do the other instructions.
 820     }
 821   }
 822
 823   for (MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
 824        MBBI != E;) {
 825     MachineInstr *MI = MBBI;
 826     // Do update merging. It's simpler to keep this separate from the above
 827     // switch, though not strictly necessary.
 828     int Opc = MI->getOpcode();
 829     switch (Opc) {
 830     default:
 831       // Just move on to the next instruction.
 832       ++MBBI;
 833       break;
 834     case ARM64::STRSui:
 835     case ARM64::STRDui:
 836     case ARM64::STRQui:
 837     case ARM64::STRXui:
 838     case ARM64::STRWui:
 839     case ARM64::LDRSui:
 840     case ARM64::LDRDui:
 841     case ARM64::LDRQui:
 842     case ARM64::LDRXui:
 843     case ARM64::LDRWui:
 844     // do the unscaled versions as well
 845     case ARM64::STURSi:
 846     case ARM64::STURDi:
 847     case ARM64::STURQi:
 848     case ARM64::STURWi:
 849     case ARM64::STURXi:
 850     case ARM64::LDURSi:
 851     case ARM64::LDURDi:
 852     case ARM64::LDURQi:
 853     case ARM64::LDURWi:
 854     case ARM64::LDURXi: {
 855       // Make sure this is a reg+imm (as opposed to an address reloc).
 856       if (!MI->getOperand(2).isImm()) {
 857         ++MBBI;
 858         break;
 859       }
 860       // Look ahead up to ScanLimit instructions for a mergable instruction.
 861       MachineBasicBlock::iterator Update =
 862           findMatchingUpdateInsnForward(MBBI, ScanLimit, 0);
 863       if (Update != E) {
 864         // Merge the update into the ld/st.
 865         MBBI = mergePostIdxUpdateInsn(MBBI, Update);
 866         Modified = true;
 867         ++NumPostFolded;
 868         break;
 869       }
 870       // Don't know how to handle pre/post-index versions, so move to the next
 871       // instruction.
 872       if (isUnscaledLdst(Opc)) {
 873         ++MBBI;
 874         break;
 875       }
 876
 877       // Look back to try to find a pre-index instruction. For example,
 878       // add x0, x0, #8
 879       // ldr x1, [x0]
 880       //   merged into:
 881       // ldr x1, [x0, #8]!
 882       Update = findMatchingUpdateInsnBackward(MBBI, ScanLimit);
 883       if (Update != E) {
 884         // Merge the update into the ld/st.
 885         MBBI = mergePreIdxUpdateInsn(MBBI, Update);
 886         Modified = true;
 887         ++NumPreFolded;
 888         break;
 889       }
 890
 891       // Look forward to try to find a post-index instruction. For example,
 892       // ldr x1, [x0, #64]
 893       // add x0, x0, #64
 894       //   merged into:
 895       // ldr x1, [x0, #64]!
 896
 897       // The immediate in the load/store is scaled by the size of the register
 898       // being loaded. The immediate in the add we're looking for,
 899       // however, is not, so adjust here.
 900       int Value = MI->getOperand(2).getImm() *
 901                   TII->getRegClass(MI->getDesc(), 0, TRI, *(MBB.getParent()))
 902                       ->getSize();
 903       Update = findMatchingUpdateInsnForward(MBBI, ScanLimit, Value);
 904       if (Update != E) {
 905         // Merge the update into the ld/st.
 906         MBBI = mergePreIdxUpdateInsn(MBBI, Update);
 907         Modified = true;
 908         ++NumPreFolded;
 909         break;
 910       }
 911
 912       // Nothing found. Just move to the next instruction.
 913       ++MBBI;
 914       break;
 915     }
 916       // FIXME: Do the other instructions.
 917     }
 918   }
 919
 920   return Modified;
 921 }
 922
 923 bool ARM64LoadStoreOpt::runOnMachineFunction(MachineFunction &Fn) {
 924   const TargetMachine &TM = Fn.getTarget();
 925   TII = static_cast<const ARM64InstrInfo *>(TM.getInstrInfo());
 926   TRI = TM.getRegisterInfo();
 927
 928   bool Modified = false;
 929   for (auto &MBB : Fn)
 930     Modified |= optimizeBlock(MBB);
 931
 932   return Modified;
 933 }
 934
 935 // FIXME: Do we need/want a pre-alloc pass like ARM has to try to keep
 936 // loads and stores near one another?
 937
 938 /// createARMLoadStoreOptimizationPass - returns an instance of the load / store
 939 /// optimization pass.
 940 FunctionPass *llvm::createARM64LoadStoreOptimizationPass() {
 941   return new ARM64LoadStoreOpt();
 942 }