[AArch64] Handle missing store pair opportunity

[junbuml]

Summary

This change will handle missing store pair opportunity where the first store
instruction stores zero followed by the non-zero store. For example, this change
will convert :

  
str wzr, [x8]
str w1, [x8, #4]

into:

stp wzr, w1, [x8]

[mcrosier]

LGTM, with a few minor nits.

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

1608

I tend to add a comment next to boolean arguments (e.g., "/*FindNarrowMerge=*/ true").

1644

/*FindNarrowMerge=*/ false

test/CodeGen/AArch64/arm64-narrow-ldst-merge.ll

355–356

Please put a space between the semicolon and the CHECK directive. Also, please remove the space between the CHECK directive and the colon below. I'm not sure if FileCheck considers "CHECK :" as a directive.

369–370

Please put a space between the semicolon and the CHECK directive.

370

Please put a space between the semicolon and the CHECK directive.

[junbuml]

Address Chad's comment.

[flyingforyou]

LGTM, too. with a few minor nits.

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

1608

Please put a space.

"/*FindNarrowMerge = */ true"

1644

Please put a space /*FindNarrowMerge = */ false.

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

struct AArch64LoadStoreOpt : public MachineFunctionPass {
   // Track which registers have been modified and used.
   BitVector ModifiedRegs, UsedRegs;
 
   // Scan the instructions looking for a load/store that can be combined
   // with the current instruction into a load/store pair.
   // Return the matching instruction if one is found, else MBB->end().
   MachineBasicBlock::iterator findMatchingInsn(MachineBasicBlock::iterator I,
                                                LdStPairFlags &Flags,
-                                               unsigned Limit);
+                                               unsigned Limit,
+                                               bool FindNarrowMerge);
 
   // Scan the instructions looking for a store that writes to the address from
   // which the current load instruction reads. Return true if one is found.
   bool findMatchingStore(MachineBasicBlock::iterator I, unsigned Limit,
                          MachineBasicBlock::iterator &StoreI);
 
   // Merge the two instructions indicated into a wider instruction.
   MachineBasicBlock::iterator
if (isNarrowLoad(Opc)) {
     DEBUG((BitExtMI2)->print(dbgs()));
     DEBUG(dbgs() << "\n");
 
     // Erase the old instructions.
     I->eraseFromParent();
     MergeMI->eraseFromParent();
     return NextI;
   }
-  assert(isPromotableZeroStoreInst(I) && "Expected promotable zero store");
+  assert(isPromotableZeroStoreInst(I) && isPromotableZeroStoreInst(MergeMI) &&
+         "Expected promotable zero store");
 
   // Construct the new instruction.
   MachineInstrBuilder MIB;
   MIB = BuildMI(*MBB, InsertionPoint, DL, TII->get(getMatchingWideOpcode(Opc)))
             .addReg(isNarrowStore(Opc) ? AArch64::WZR : AArch64::XZR)
             .addOperand(BaseRegOp)
             .addImm(OffsetImm)
             .setMemRefs(I->mergeMemRefsWith(*MergeMI));
static bool canMergeOpc(unsigned OpcA, unsigned OpcB, LdStPairFlags &Flags,
 
   // FIXME: Can we also match a mixed sext/zext unscaled/scaled pair?
 }
 
 /// Scan the instructions looking for a load/store that can be combined with the
 /// current instruction into a wider equivalent or a load/store pair.
 MachineBasicBlock::iterator
 AArch64LoadStoreOpt::findMatchingInsn(MachineBasicBlock::iterator I,
-                                      LdStPairFlags &Flags, unsigned Limit) {
+                                      LdStPairFlags &Flags, unsigned Limit,
+                                      bool FindNarrowMerge) {
   MachineBasicBlock::iterator E = I->getParent()->end();
   MachineBasicBlock::iterator MBBI = I;
   MachineInstr *FirstMI = I;
   ++MBBI;
 
   unsigned Opc = FirstMI->getOpcode();
   bool MayLoad = FirstMI->mayLoad();
   bool IsUnscaled = TII->isUnscaledLdSt(FirstMI);
if (canMergeOpc(Opc, MI->getOpcode(), Flags, TII) &&
       if (BaseReg == MIBaseReg && ((Offset == MIOffset + OffsetStride) ||
                                    (Offset + OffsetStride == MIOffset))) {
         int MinOffset = Offset < MIOffset ? Offset : MIOffset;
         // If this is a volatile load/store that otherwise matched, stop looking
         // as something is going on that we don't have enough information to
         // safely transform. Similarly, stop if we see a hint to avoid pairs.
         if (MI->hasOrderedMemoryRef() || TII->isLdStPairSuppressed(MI))
           return E;
-        // If the resultant immediate offset of merging these instructions
-        // is out of range for a pairwise instruction, bail and keep looking.
-        bool IsNarrowLoad = isNarrowLoad(MI->getOpcode());
-        if (!IsNarrowLoad &&
-            !inBoundsForPair(IsUnscaled, MinOffset, OffsetStride)) {
+
+        if (FindNarrowMerge) {
+          // If the alignment requirements of the scaled wide load/store
+          // instruction can't express the offset of the scaled narrow input,
+          // bail and keep looking. For promotable zero stores, allow only when
+          // the stored value is the same (i.e., WZR).
+          if ((!IsUnscaled && alignTo(MinOffset, 2) != MinOffset) ||
+              (IsPromotableZeroStore && Reg != getLdStRegOp(MI).getReg())) {
             trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
             MemInsns.push_back(MI);
             continue;
           }
-
-        if (IsNarrowLoad || IsPromotableZeroStore) {
-          // If the alignment requirements of the scaled wide load/store
-          // instruction can't express the offset of the scaled narrow
-          // input, bail and keep looking.
-          if (!IsUnscaled && alignTo(MinOffset, 2) != MinOffset) {
+        } else {
+          // If the resultant immediate offset of merging these instructions
+          // is out of range for a pairwise instruction, bail and keep looking.
+          if (!inBoundsForPair(IsUnscaled, MinOffset, OffsetStride)) {
             trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
             MemInsns.push_back(MI);
             continue;
           }
-        } else {
           // If the alignment requirements of the paired (scaled) instruction
           // can't express the offset of the unscaled input, bail and keep
           // looking.
           if (IsUnscaled && (alignTo(MinOffset, OffsetStride) != MinOffset)) {
             trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
             MemInsns.push_back(MI);
             continue;
           }
         }
         // If the destination register of the loads is the same register, bail
         // and keep looking. A load-pair instruction with both destination
         // registers the same is UNPREDICTABLE and will result in an exception.
-        // For narrow stores, allow only when the stored value is the same
-        // (i.e., WZR).
-        if ((MayLoad && Reg == getLdStRegOp(MI).getReg()) ||
-            (IsPromotableZeroStore && Reg != getLdStRegOp(MI).getReg())) {
+        if (MayLoad && Reg == getLdStRegOp(MI).getReg()) {
           trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
           MemInsns.push_back(MI);
           continue;
         }
 
         // If the Rt of the second instruction was not modified or used between
         // the two instructions and none of the instructions between the second
         // and first alias with the second, we can combine the second into the
bool AArch64LoadStoreOpt::tryToMergeLdStInst(
   // For promotable zero stores, the stored value should be WZR.
   if (isPromotableZeroStoreOpcode(MI) &&
       getLdStRegOp(MI).getReg() != AArch64::WZR)
     return false;
 
   // Look ahead up to LdStLimit instructions for a mergable instruction.
   LdStPairFlags Flags;
   MachineBasicBlock::iterator MergeMI =
-      findMatchingInsn(MBBI, Flags, LdStLimit);
+      findMatchingInsn(MBBI, Flags, LdStLimit, /*FindNarrowMerge=*/true);
   if (MergeMI != E) {
     if (isNarrowLoad(MI)) {
       ++NumNarrowLoadsPromoted;
     } else if (isPromotableZeroStoreInst(MI)) {
       ++NumZeroStoresPromoted;
     }
     // 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.
bool AArch64LoadStoreOpt::tryToPairLdStInst(MachineBasicBlock::iterator &MBBI) {
   bool IsUnscaled = TII->isUnscaledLdSt(MI);
   int Offset = getLdStOffsetOp(MI).getImm();
   int OffsetStride = IsUnscaled ? getMemScale(MI) : 1;
   if (!inBoundsForPair(IsUnscaled, Offset, OffsetStride))
     return false;
 
   // Look ahead up to LdStLimit instructions for a pairable instruction.
   LdStPairFlags Flags;
-  MachineBasicBlock::iterator Paired = findMatchingInsn(MBBI, Flags, LdStLimit);
+  MachineBasicBlock::iterator Paired =
+      findMatchingInsn(MBBI, Flags, LdStLimit, /*FindNarrowMerge=*/false);
   if (Paired != E) {
     ++NumPairCreated;
     if (TII->isUnscaledLdSt(MI))
       ++NumUnscaledPairCreated;
     // 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;

test/CodeGen/AArch64/arm64-narrow-ldst-merge.ll

entry:
   %arrayidx5 = getelementptr inbounds i16, i16* %P, i64 %idxprom4
   store i16 0, i16* %arrayidx5
   %add6 = add nsw i32 %n, 3
   %idxprom7 = sext i32 %add6 to i64
   %arrayidx8 = getelementptr inbounds i16, i16* %P, i64 %idxprom7
   store i16 0, i16* %arrayidx8
   ret void
 }
 
-;CHECK-LABEL: Strw_zero
-;CHECK : str xzr
+; CHECK-LABEL: Strw_zero
+; CHECK: str xzr
 define void @Strw_zero(i32* nocapture %P, i32 %n) {
 entry:
   %idxprom = sext i32 %n to i64
   %arrayidx = getelementptr inbounds i32, i32* %P, i64 %idxprom
   store i32 0, i32* %arrayidx
   %add = add nsw i32 %n, 1
   %idxprom1 = sext i32 %add to i64
   %arrayidx2 = getelementptr inbounds i32, i32* %P, i64 %idxprom1
   store i32 0, i32* %arrayidx2
   ret void
 }
 
-;CHECK-LABEL: Strw_zero_4
-;CHECK : stp xzr
+; CHECK-LABEL: Strw_zero_nonzero
+; CHECK: stp wzr, w1
+define void @Strw_zero_nonzero(i32* nocapture %P, i32 %n)  {
+entry:
+  %idxprom = sext i32 %n to i64
+  %arrayidx = getelementptr inbounds i32, i32* %P, i64 %idxprom
+  store i32 0, i32* %arrayidx
+  %add = add nsw i32 %n, 1
+  %idxprom1 = sext i32 %add to i64
+  %arrayidx2 = getelementptr inbounds i32, i32* %P, i64 %idxprom1
+  store i32 %n, i32* %arrayidx2
+  ret void
+}
+
+; CHECK-LABEL: Strw_zero_4
+; CHECK: stp xzr
 define void @Strw_zero_4(i32* nocapture %P, i32 %n) {
 entry:
   %idxprom = sext i32 %n to i64
   %arrayidx = getelementptr inbounds i32, i32* %P, i64 %idxprom
   store i32 0, i32* %arrayidx
   %add = add nsw i32 %n, 1
   %idxprom1 = sext i32 %add to i64
   %arrayidx2 = getelementptr inbounds i32, i32* %P, i64 %idxprom1
entry:
   store i16 0, i16* %arrayidx6
   %sub7 = add nsw i32 %n, -1
   %idxprom8 = sext i32 %sub7 to i64
   %arrayidx9 = getelementptr inbounds i16, i16* %P, i64 %idxprom8
   store i16 0, i16* %arrayidx9
   ret void
 }
 
-;CHECK-LABEL: Sturw_zero
-;CHECK : stur xzr
+; CHECK-LABEL: Sturw_zero
+; CHECK: stur xzr
 define void @Sturw_zero(i32* nocapture %P, i32 %n) {
 entry:
   %sub = add nsw i32 %n, -3
   %idxprom = sext i32 %sub to i64
   %arrayidx = getelementptr inbounds i32, i32* %P, i64 %idxprom
   store i32 0, i32* %arrayidx
   %sub1 = add nsw i32 %n, -4
   %idxprom2 = sext i32 %sub1 to i64
   %arrayidx3 = getelementptr inbounds i32, i32* %P, i64 %idxprom2
   store i32 0, i32* %arrayidx3
   ret void
 }
 
-;CHECK-LABEL: Sturw_zero_4
-;CHECK : str xzr
+; CHECK-LABEL: Sturw_zero_4
+; CHECK: stp xzr, xzr
 define void @Sturw_zero_4(i32* nocapture %P, i32 %n) {
 entry:
   %sub = add nsw i32 %n, -3
   %idxprom = sext i32 %sub to i64
   %arrayidx = getelementptr inbounds i32, i32* %P, i64 %idxprom
   store i32 0, i32* %arrayidx
   %sub1 = add nsw i32 %n, -4
   %idxprom2 = sext i32 %sub1 to i64