[AArch64] Merge two adjacent str WZR into str XZR

[junbuml]

Summary

This change merges adjacent 32 bit zero stores into a 64 bit zero store.
e.g.,

str wzr, [x0]
str wzr, [x0, #4]

becomes

str xzr, [x0]

Therefore, four adjacent 32 bit zero stores will be a single stp.
e.g.,

str wzr, [x0]
str wzr, [x0, #4]
str wzr, [x0, #8]
str wzr, [x0, #12]

becomes

stp xzr, xzr, [x0]

[t.p.northover]

Hi Jun,

I think this looks fine as code. Just a couple of minor grammatical nits before committing.

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

636

I think you mean "isPromotableZeroStore" here.

1106

Again, "IsMergeableZeroStore" conveys the meaning better.

[junbuml]

Addressed Tim's comments.
Thanks Tim for the very quick review!

[t.p.northover]

Thanks Jun. Looks fine to me now.

Tim.

[junbuml]

Applied it to current truck.

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

static unsigned getMatchingWideOpcode(unsigned Opc) {
   case AArch64::STRBBui:
     return AArch64::STRHHui;
   case AArch64::STRHHui:
     return AArch64::STRWui;
   case AArch64::STURBBi:
     return AArch64::STURHHi;
   case AArch64::STURHHi:
     return AArch64::STURWi;
+  case AArch64::STURWi:
+    return AArch64::STURXi;
+  case AArch64::STRWui:
+    return AArch64::STRXui;
   case AArch64::LDRHHui:
   case AArch64::LDRSHWui:
     return AArch64::LDRWui;
   case AArch64::LDURHHi:
   case AArch64::LDURSHWi:
     return AArch64::LDURWi;
   case AArch64::LDRBBui:
   case AArch64::LDRSBWui:
int UnscaledStOffset = isUnscaledLdSt(StoreInst)
                              : getLdStOffsetOp(StoreInst).getImm() * StoreSize;
   int UnscaledLdOffset = isUnscaledLdSt(LoadInst)
                              ? getLdStOffsetOp(LoadInst).getImm()
                              : getLdStOffsetOp(LoadInst).getImm() * LoadSize;
   return (UnscaledStOffset <= UnscaledLdOffset) &&
          (UnscaledLdOffset + LoadSize <= (UnscaledStOffset + StoreSize));
 }
 
+static bool isPromotableZeroStore(MachineInstr *MI) {
+  MachineOperand MO = getLdStRegOp(MI);
+  unsigned Opc = MI->getOpcode();
+  return ((isNarrowStore(Opc) || Opc == AArch64::STURWi ||
+           Opc == AArch64::STRWui) &&
+          MO.isReg() && MO.getReg() == AArch64::WZR);
+}
+
 MachineBasicBlock::iterator
 AArch64LoadStoreOpt::mergePairedInsns(MachineBasicBlock::iterator I,
                                       MachineBasicBlock::iterator Paired,
                                       const LdStPairFlags &Flags) {
   MachineBasicBlock::iterator NextI = I;
   ++NextI;
   // If NextI is the second of the two instructions to be merged, we need
   // to skip one further. Either way we merge will invalidate the iterator,
if (isNarrowLoad(Opc)) {
     // Erase the old instructions.
     I->eraseFromParent();
     Paired->eraseFromParent();
     return NextI;
   }
 
   // Construct the new instruction.
   MachineInstrBuilder MIB;
-  if (isNarrowStore(Opc)) {
+  if (isPromotableZeroStore(I)) {
     // Change the scaled offset from small to large type.
     if (!IsUnscaled) {
       assert(((OffsetImm & 1) == 0) && "Unexpected offset to merge");
       OffsetImm /= 2;
     }
     MIB = BuildMI(*I->getParent(), InsertionPoint, I->getDebugLoc(),
                   TII->get(getMatchingWideOpcode(Opc)))
-              .addOperand(getLdStRegOp(I))
+              .addReg(isNarrowStore(I) ? AArch64::WZR : AArch64::XZR)
               .addOperand(BaseRegOp)
               .addImm(OffsetImm)
               .setMemRefs(I->mergeMemRefsWith(*Paired));
   } else {
     // Handle Unscaled
     if (IsUnscaled)
       OffsetImm /= OffsetStride;
     MIB = BuildMI(*I->getParent(), InsertionPoint, I->getDebugLoc(),
AArch64LoadStoreOpt::findMatchingInsn(MachineBasicBlock::iterator I,
   ++MBBI;
 
   unsigned Opc = FirstMI->getOpcode();
   bool MayLoad = FirstMI->mayLoad();
   bool IsUnscaled = isUnscaledLdSt(FirstMI);
   unsigned Reg = getLdStRegOp(FirstMI).getReg();
   unsigned BaseReg = getLdStBaseOp(FirstMI).getReg();
   int Offset = getLdStOffsetOp(FirstMI).getImm();
-  bool IsNarrowStore = isNarrowStore(Opc);
+  bool IsMergeableZeroStore = isPromotableZeroStore(FirstMI);
 
-  // For narrow stores, find only the case where the stored value is WZR.
-  if (IsNarrowStore && Reg != AArch64::WZR)
+  // For narrow stores, merge only when storing zero.
+  if (isNarrowStore(Opc) && !IsMergeableZeroStore)
     return E;
 
   // Early exit if the first instruction modifies the base register.
   // e.g., ldr x0, [x0]
   if (FirstMI->modifiesRegister(BaseReg, TRI))
     return E;
 
   // Early exit if the offset is not possible to match. (6 bits of positive
   // range, plus allow an extra one in case we find a later insn that matches
   // with Offset-1)
   int OffsetStride = IsUnscaled ? getMemScale(FirstMI) : 1;
-  if (!(isNarrowLoad(Opc) || IsNarrowStore) &&
+  if (!(isNarrowLoad(Opc) || IsMergeableZeroStore) &&
       !inBoundsForPair(IsUnscaled, Offset, OffsetStride))
     return E;
 
   // Track which registers have been modified and used between the first insn
   // (inclusive) and the second insn.
   ModifiedRegs.reset();
   UsedRegs.reset();
 
if (CanMergeOpc && getLdStOffsetOp(MI).isImm()) {
         bool IsNarrowLoad = isNarrowLoad(MI->getOpcode());
         if (!IsNarrowLoad &&
             !inBoundsForPair(MIIsUnscaled, MinOffset, OffsetStride)) {
           trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
           MemInsns.push_back(MI);
           continue;
         }
 
-        if (IsNarrowLoad || IsNarrowStore) {
+        if (IsNarrowLoad || IsMergeableZeroStore) {
           // 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) {
             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()) ||
-            (IsNarrowStore && Reg != getLdStRegOp(MI).getReg())) {
+            (IsMergeableZeroStore && 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
if (TII->isLdStPairSuppressed(MI))
     return false;
 
   // Look ahead up to LdStLimit instructions for a pairable instruction.
   LdStPairFlags Flags;
   MachineBasicBlock::iterator Paired = findMatchingInsn(MBBI, Flags, LdStLimit);
   if (Paired != E) {
     if (isNarrowLoad(MI)) {
       ++NumNarrowLoadsPromoted;
-    } else if (isNarrowStore(MI)) {
+    } else if (isPromotableZeroStore(MI)) {
       ++NumZeroStoresPromoted;
     } else {
       ++NumPairCreated;
       if (isUnscaledLdSt(MI))
         ++NumUnscaledPairCreated;
     }
 
     // Merge the loads into a pair. Keeping the iterator straight is a
default:
       break;
     // Scaled instructions.
     case AArch64::LDRBBui:
     case AArch64::LDRHHui:
     case AArch64::LDRSBWui:
     case AArch64::LDRSHWui:
     case AArch64::STRBBui:
     case AArch64::STRHHui:
+    case AArch64::STRWui:
     // Unscaled instructions.
     case AArch64::LDURBBi:
     case AArch64::LDURHHi:
     case AArch64::LDURSBWi:
     case AArch64::LDURSHWi:
     case AArch64::STURBBi:
-    case AArch64::STURHHi: {
+    case AArch64::STURHHi:
+    case AArch64::STURWi: {
       if (tryToMergeLdStInst(MBBI)) {
         Modified = true;
         break;
       }
       ++MBBI;
       break;
     }
     }

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

store i16 0, i16* %arrayidx
   %add = add nsw i32 %n, 1
   %idxprom1 = sext i32 %add to i64
   %arrayidx2 = getelementptr inbounds i16, i16* %P, i64 %idxprom1
   store i16 0, i16* %arrayidx2
   ret void
 }
 
 ; CHECK-LABEL: Strh_zero_4
-; CHECK: stp wzr, wzr
+; CHECK: str xzr
 define void @Strh_zero_4(i16* nocapture %P, i32 %n) {
 entry:
   %idxprom = sext i32 %n to i64
   %arrayidx = getelementptr inbounds i16, i16* %P, i64 %idxprom
   store i16 0, i16* %arrayidx
   %add = add nsw i32 %n, 1
   %idxprom1 = sext i32 %add to i64
   %arrayidx2 = getelementptr inbounds i16, i16* %P, i64 %idxprom1
   store i16 0, i16* %arrayidx2
   %add3 = add nsw i32 %n, 2
   %idxprom4 = sext i32 %add3 to i64
   %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
+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
+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
+  store i32 0, i32* %arrayidx2
+  %add3 = add nsw i32 %n, 2
+  %idxprom4 = sext i32 %add3 to i64
+  %arrayidx5 = getelementptr inbounds i32, i32* %P, i64 %idxprom4
+  store i32 0, i32* %arrayidx5
+  %add6 = add nsw i32 %n, 3
+  %idxprom7 = sext i32 %add6 to i64
+  %arrayidx8 = getelementptr inbounds i32, i32* %P, i64 %idxprom7
+  store i32 0, i32* %arrayidx8
+  ret void
+}
+
 ; CHECK-LABEL: Sturb_zero
 ; CHECK: sturh wzr
 define void @Sturb_zero(i8* nocapture %P, i32 %n) #0 {
 entry:
   %sub = add nsw i32 %n, -2
   %idxprom = sext i32 %sub to i64
   %arrayidx = getelementptr inbounds i8, i8* %P, i64 %idxprom
   store i8 0, i8* %arrayidx
entry:
   %sub1 = add nsw i32 %n, -3
   %idxprom2 = sext i32 %sub1 to i64
   %arrayidx3 = getelementptr inbounds i16, i16* %P, i64 %idxprom2
   store i16 0, i16* %arrayidx3
   ret void
 }
 
 ; CHECK-LABEL: Sturh_zero_4
-; CHECK: stp wzr, wzr
+; CHECK: stur xzr
 define void @Sturh_zero_4(i16* nocapture %P, i32 %n) {
 entry:
   %sub = add nsw i32 %n, -3
   %idxprom = sext i32 %sub to i64
   %arrayidx = getelementptr inbounds i16, i16* %P, i64 %idxprom
   store i16 0, i16* %arrayidx
   %sub1 = add nsw i32 %n, -4
   %idxprom2 = sext i32 %sub1 to i64
   %arrayidx3 = getelementptr inbounds i16, i16* %P, i64 %idxprom2
   store i16 0, i16* %arrayidx3
   %sub4 = add nsw i32 %n, -2
   %idxprom5 = sext i32 %sub4 to i64
   %arrayidx6 = getelementptr inbounds i16, i16* %P, i64 %idxprom5
   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
+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
+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
+  %arrayidx3 = getelementptr inbounds i32, i32* %P, i64 %idxprom2
+  store i32 0, i32* %arrayidx3
+  %sub4 = add nsw i32 %n, -2
+  %idxprom5 = sext i32 %sub4 to i64
+  %arrayidx6 = getelementptr inbounds i32, i32* %P, i64 %idxprom5
+  store i32 0, i32* %arrayidx6
+  %sub7 = add nsw i32 %n, -1
+  %idxprom8 = sext i32 %sub7 to i64
+  %arrayidx9 = getelementptr inbounds i32, i32* %P, i64 %idxprom8
+  store i32 0, i32* %arrayidx9
+  ret void
+}
+