[AArch64] Refactor the loads and stores optimizer

[evandro]

Summary

Move remaining inline matching of instructions of some optimizations into separate functions, like in the other optimizations. Otherwise, NFC.

[junbuml]

LGTM

[fhahn]

LGTM

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

585

Nit: should this be at the previous line?

[evandro]

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

585

:-}

[junbuml]

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

1646–1647

Please keep this comment.

[junbuml]

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

1646–1647

Sorry. it doesn't seem so really needed now.

[mcrosier]

Please run clang-format, but otherwise LGTM.

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

struct AArch64LoadStoreOpt : public MachineFunctionPass {
   bool tryToMergeZeroStInst(MachineBasicBlock::iterator &MBBI);
 
   // Find and pair ldr/str instructions.
   bool tryToPairLdStInst(MachineBasicBlock::iterator &MBBI);
 
   // Find and promote load instructions which read directly from store.
   bool tryToPromoteLoadFromStore(MachineBasicBlock::iterator &MBBI);
 
+  // Find and merge a base register updates before or after a ld/st instruction.
+  bool tryToMergeLdStUpdate(MachineBasicBlock::iterator &MBBI);
+
   bool optimizeBlock(MachineBasicBlock &MBB, bool EnableNarrowZeroStOpt);
 
   bool runOnMachineFunction(MachineFunction &Fn) override;
 
   MachineFunctionProperties getRequiredProperties() const override {
     return MachineFunctionProperties().set(
         MachineFunctionProperties::Property::NoVRegs);
   }

 
 static bool isPromotableZeroStoreInst(MachineInstr &MI) {
   unsigned Opc = MI.getOpcode();
   return (Opc == AArch64::STRWui || Opc == AArch64::STURWi ||
           isNarrowStore(Opc)) &&
          getLdStRegOp(MI).getReg() == AArch64::WZR;
 }
 
+static bool isPromotableLoadFromStore(MachineInstr &MI)
+{
+  switch (MI.getOpcode()) {
+  default:
+    return false;
+  // Scaled instructions.
+  case AArch64::LDRBBui:
+  case AArch64::LDRHHui:
+  case AArch64::LDRWui:
+  case AArch64::LDRXui:
+  // Unscaled instructions.
+  case AArch64::LDURBBi:
+  case AArch64::LDURHHi:
+  case AArch64::LDURWi:
+  case AArch64::LDURXi:
+    return true;
+  }
+}
+
+static bool isMergeableLdStUpdate(MachineInstr &MI) {
+  unsigned Opc = MI.getOpcode();
+  switch (Opc) {
+  default:
+    return false;
+  // Scaled instructions.
+  case AArch64::STRSui:
+  case AArch64::STRDui:
+  case AArch64::STRQui:
+  case AArch64::STRXui:
+  case AArch64::STRWui:
+  case AArch64::STRHHui:
+  case AArch64::STRBBui:
+  case AArch64::LDRSui:
+  case AArch64::LDRDui:
+  case AArch64::LDRQui:
+  case AArch64::LDRXui:
+  case AArch64::LDRWui:
+  case AArch64::LDRHHui:
+  case AArch64::LDRBBui:
+  // Unscaled instructions.
+  case AArch64::STURSi:
+  case AArch64::STURDi:
+  case AArch64::STURQi:
+  case AArch64::STURWi:
+  case AArch64::STURXi:
+  case AArch64::LDURSi:
+  case AArch64::LDURDi:
+  case AArch64::LDURQi:
+  case AArch64::LDURWi:
+  case AArch64::LDURXi:
+  // Paired instructions.
+  case AArch64::LDPSi:
+  case AArch64::LDPSWi:
+  case AArch64::LDPDi:
+  case AArch64::LDPQi:
+  case AArch64::LDPWi:
+  case AArch64::LDPXi:
+  case AArch64::STPSi:
+  case AArch64::STPDi:
+  case AArch64::STPQi:
+  case AArch64::STPWi:
+  case AArch64::STPXi:
+    // Make sure this is a reg+imm (as opposed to an address reloc).
+    if (!getLdStOffsetOp(MI).isImm())
+      return false;
+
+    return true;
+  }
+}
+
 MachineBasicBlock::iterator
 AArch64LoadStoreOpt::mergeNarrowZeroStores(MachineBasicBlock::iterator I,
                                            MachineBasicBlock::iterator MergeMI,
                                            const LdStPairFlags &Flags) {
   assert(isPromotableZeroStoreInst(*I) && isPromotableZeroStoreInst(*MergeMI) &&
          "Expected promotable zero stores.");
 
   MachineBasicBlock::iterator NextI = I;
MIB = BuildMI(*I->getParent(), I, I->getDebugLoc(), TII->get(NewOpc))
               .add(getLdStRegOp(*I, 0))
               .add(getLdStRegOp(*I, 1))
               .add(getLdStBaseOp(*I))
               .addImm(Value / Scale)
               .setMemRefs(I->memoperands_begin(), I->memoperands_end());
   }
   (void)MIB;
 
-  if (IsPreIdx)
+  if (IsPreIdx) {
+    ++NumPreFolded;
     DEBUG(dbgs() << "Creating pre-indexed load/store.");
-  else
+  } else {
+    ++NumPostFolded;
     DEBUG(dbgs() << "Creating post-indexed load/store.");
+  }
   DEBUG(dbgs() << "    Replacing instructions:\n    ");
   DEBUG(I->print(dbgs()));
   DEBUG(dbgs() << "    ");
   DEBUG(Update->print(dbgs()));
   DEBUG(dbgs() << "  with instruction:\n    ");
   DEBUG(((MachineInstr *)MIB)->print(dbgs()));
   DEBUG(dbgs() << "\n");
 
if (Paired != E) {
     // Keeping the iterator straight is a pain, so we let the merge routine tell
     // us what the next instruction is after it's done mucking about.
     MBBI = mergePairedInsns(MBBI, Paired, Flags);
     return true;
   }
   return false;
 }
 
+bool
+AArch64LoadStoreOpt::tryToMergeLdStUpdate(MachineBasicBlock::iterator &MBBI) {
+  MachineInstr &MI = *MBBI;
+  MachineBasicBlock::iterator E = MI.getParent()->end();
+  MachineBasicBlock::iterator Update;
+
+  // Look forward to try to form a post-index instruction. For example,
+  // ldr x0, [x20]
+  // add x20, x20, #32
+  //   merged into:
+  // ldr x0, [x20], #32
+  Update = findMatchingUpdateInsnForward(MBBI, 0, UpdateLimit);
+  if (Update != E) {
+    // Merge the update into the ld/st.
+    MBBI = mergeUpdateInsn(MBBI, Update, /*IsPreIdx=*/false);
+    return true;
+  }
+
+  // Don't know how to handle unscaled pre/post-index versions below, so bail.
+  if (TII->isUnscaledLdSt(MI.getOpcode()))
+    return false;
+
+  // Look back to try to find a pre-index instruction. For example,
+  // add x0, x0, #8
+  // ldr x1, [x0]
+  //   merged into:
+  // ldr x1, [x0, #8]!
+  Update = findMatchingUpdateInsnBackward(MBBI, UpdateLimit);
+  if (Update != E) {
+    // Merge the update into the ld/st.
+    MBBI = mergeUpdateInsn(MBBI, Update, /*IsPreIdx=*/true);
+    return true;
+  }
+
+  // The immediate in the load/store is scaled by the size of the memory
+  // operation. The immediate in the add we're looking for,
+  // however, is not, so adjust here.
+  int UnscaledOffset = getLdStOffsetOp(MI).getImm() * getMemScale(MI);
+
+  // Look forward to try to find a post-index instruction. For example,
+  // ldr x1, [x0, #64]
+  // add x0, x0, #64
+  //   merged into:
+  // ldr x1, [x0, #64]!
+  Update = findMatchingUpdateInsnForward(MBBI, UnscaledOffset, UpdateLimit);
+  if (Update != E) {
+    // Merge the update into the ld/st.
+    MBBI = mergeUpdateInsn(MBBI, Update, /*IsPreIdx=*/true);
+    return true;
+  }
+
+  return false;
+}
+
 bool AArch64LoadStoreOpt::optimizeBlock(MachineBasicBlock &MBB,
                                         bool EnableNarrowZeroStOpt) {
   bool Modified = false;
   // Four tranformations to do here:
   // 1) Find loads that directly read from stores and promote them by
   //    replacing with mov instructions. If the store is wider than the load,
   //    the load will be replaced with a bitfield extract.
   //      e.g.,
   //        str w1, [x0, #4]
   //        ldrh w2, [x0, #6]
   //        ; becomes
   //        str w1, [x0, #4]
   //        lsr w2, w1, #16
   for (MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
        MBBI != E;) {
-    MachineInstr &MI = *MBBI;
-    switch (MI.getOpcode()) {
-    default:
-      // Just move on to the next instruction.
-      ++MBBI;
-      break;
-    // Scaled instructions.
-    case AArch64::LDRBBui:
-    case AArch64::LDRHHui:
-    case AArch64::LDRWui:
-    case AArch64::LDRXui:
-    // Unscaled instructions.
-    case AArch64::LDURBBi:
-    case AArch64::LDURHHi:
-    case AArch64::LDURWi:
-    case AArch64::LDURXi:
-      if (tryToPromoteLoadFromStore(MBBI)) {
-        Modified = true;
-        break;
-      }
+    if (isPromotableLoadFromStore(*MBBI) && tryToPromoteLoadFromStore(MBBI))
+    else
       ++MBBI;
-      break;
-    }
   }
   // 2) Merge adjacent zero stores into a wider store.
   //      e.g.,
   //        strh wzr, [x0]
   //        strh wzr, [x0, #2]
   //        ; becomes
   //        str wzr, [x0]
   //      e.g.,
   //        str wzr, [x0]
   //        str wzr, [x0, #4]
   //        ; becomes
   //        str xzr, [x0]
+  if (EnableNarrowZeroStOpt)
     for (MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
-       EnableNarrowZeroStOpt && MBBI != E;) {
-    if (isPromotableZeroStoreInst(*MBBI)) {
-      if (tryToMergeZeroStInst(MBBI)) {
+         MBBI != E;) {
+      if (isPromotableZeroStoreInst(*MBBI) && tryToMergeZeroStInst(MBBI))
         Modified = true;
-      } else
-        ++MBBI;
-    } else
+      else
         ++MBBI;
     }
-
   // 3) Find loads and stores that can be merged into a single load or store
   //    pair instruction.
   //      e.g.,
   //        ldr x0, [x2]
   //        ldr x1, [x2, #8]
   //        ; becomes
   //        ldp x0, x1, [x2]
   for (MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
        MBBI != E;) {
     if (TII->isPairableLdStInst(*MBBI) && tryToPairLdStInst(MBBI))
       Modified = true;
     else
       ++MBBI;
   }
   // 4) Find base register updates that can be merged into the load or store
   //    as a base-reg writeback.
   //      e.g.,
   //        ldr x0, [x2]
   //        add x2, x2, #4
   //        ; becomes
   //        ldr x0, [x2], #4
   for (MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
        MBBI != E;) {
-    MachineInstr &MI = *MBBI;
-    // Do update merging. It's simpler to keep this separate from the above
-    // switchs, though not strictly necessary.
-    unsigned Opc = MI.getOpcode();
-    switch (Opc) {
-    default:
-      // Just move on to the next instruction.
-      ++MBBI;
-      break;
-    // Scaled instructions.
-    case AArch64::STRSui:
-    case AArch64::STRDui:
-    case AArch64::STRQui:
-    case AArch64::STRXui:
-    case AArch64::STRWui:
-    case AArch64::STRHHui:
-    case AArch64::STRBBui:
-    case AArch64::LDRSui:
-    case AArch64::LDRDui:
-    case AArch64::LDRQui:
-    case AArch64::LDRXui:
-    case AArch64::LDRWui:
-    case AArch64::LDRHHui:
-    case AArch64::LDRBBui:
-    // Unscaled instructions.
-    case AArch64::STURSi:
-    case AArch64::STURDi:
-    case AArch64::STURQi:
-    case AArch64::STURWi:
-    case AArch64::STURXi:
-    case AArch64::LDURSi:
-    case AArch64::LDURDi:
-    case AArch64::LDURQi:
-    case AArch64::LDURWi:
-    case AArch64::LDURXi:
-    // Paired instructions.
-    case AArch64::LDPSi:
-    case AArch64::LDPSWi:
-    case AArch64::LDPDi:
-    case AArch64::LDPQi:
-    case AArch64::LDPWi:
-    case AArch64::LDPXi:
-    case AArch64::STPSi:
-    case AArch64::STPDi:
-    case AArch64::STPQi:
-    case AArch64::STPWi:
-    case AArch64::STPXi: {
-      // Make sure this is a reg+imm (as opposed to an address reloc).
-      if (!getLdStOffsetOp(MI).isImm()) {
-        ++MBBI;
-        break;
-      }
-      // Look forward to try to form a post-index instruction. For example,
-      // ldr x0, [x20]
-      // add x20, x20, #32
-      //   merged into:
-      // ldr x0, [x20], #32
-      MachineBasicBlock::iterator Update =
-          findMatchingUpdateInsnForward(MBBI, 0, UpdateLimit);
-      if (Update != E) {
-        // Merge the update into the ld/st.
-        MBBI = mergeUpdateInsn(MBBI, Update, /*IsPreIdx=*/false);
-        Modified = true;
-        ++NumPostFolded;
-        break;
-      }
-
-      // Don't know how to handle unscaled pre/post-index versions below, so
-      // move to the next instruction.
-      if (TII->isUnscaledLdSt(Opc)) {
-        ++MBBI;
-        break;
-      }
-
-      // Look back to try to find a pre-index instruction. For example,
-      // add x0, x0, #8
-      // ldr x1, [x0]
-      //   merged into:
-      // ldr x1, [x0, #8]!
-      Update = findMatchingUpdateInsnBackward(MBBI, UpdateLimit);
-      if (Update != E) {
-        // Merge the update into the ld/st.
-        MBBI = mergeUpdateInsn(MBBI, Update, /*IsPreIdx=*/true);
-        Modified = true;
-        ++NumPreFolded;
-        break;
-      }
-      // The immediate in the load/store is scaled by the size of the memory
-      // operation. The immediate in the add we're looking for,
-      // however, is not, so adjust here.
-      int UnscaledOffset = getLdStOffsetOp(MI).getImm() * getMemScale(MI);
-
-      // Look forward to try to find a post-index instruction. For example,
-      // ldr x1, [x0, #64]
-      // add x0, x0, #64
-      //   merged into:
-      // ldr x1, [x0, #64]!
-      Update = findMatchingUpdateInsnForward(MBBI, UnscaledOffset, UpdateLimit);
-      if (Update != E) {
-        // Merge the update into the ld/st.
-        MBBI = mergeUpdateInsn(MBBI, Update, /*IsPreIdx=*/true);
-        Modified = true;
-        ++NumPreFolded;
-        break;
-      }
-
-      // Nothing found. Just move to the next instruction.
+    if (isMergeableLdStUpdate(*MBBI) && tryToMergeLdStUpdate(MBBI))
+    else
       ++MBBI;
-      break;
-    }
-    }
   }
 
   return Modified;
 }
 
 bool AArch64LoadStoreOpt::runOnMachineFunction(MachineFunction &Fn) {
   if (skipFunction(*Fn.getFunction()))
     return false;