[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 ↗

(On Diff #47044)

I think you mean "isPromotableZeroStore" here.

1106 ↗

(On Diff #47044)

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.

llvm/trunk/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

static bool isNarrowStore(unsigned Opc) {
   case AArch64::STRBBui:
   case AArch64::STURBBi:
   case AArch64::STRHHui:
   case AArch64::STURHHi:
     return true;
   }
 }
 
-static bool isNarrowStore(MachineInstr *MI) {
-  return isNarrowStore(MI->getOpcode());
-}
-
 static bool isNarrowLoad(unsigned Opc) {
   switch (Opc) {
   default:
     return false;
   case AArch64::LDRHHui:
   case AArch64::LDURHHi:
   case AArch64::LDRBBui:
   case AArch64::LDURBBi:
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 isPromotableZeroStoreOpcode(MachineInstr *MI) {
+  unsigned Opc = MI->getOpcode();
+  return isNarrowStore(Opc) || Opc == AArch64::STRWui || Opc == AArch64::STURWi;
+}
+
+static bool isPromotableZeroStoreInst(MachineInstr *MI) {
+  return (isPromotableZeroStoreOpcode(MI)) &&
+         getLdStRegOp(MI).getReg() == AArch64::WZR;
+}
+
 MachineBasicBlock::iterator
 AArch64LoadStoreOpt::mergeNarrowInsns(MachineBasicBlock::iterator I,
                                       MachineBasicBlock::iterator MergeMI,
                                       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)) {
     DEBUG((BitExtMI2)->print(dbgs()));
     DEBUG(dbgs() << "\n");
 
     // Erase the old instructions.
     I->eraseFromParent();
     MergeMI->eraseFromParent();
     return NextI;
   }
-  assert(isNarrowStore(Opc) && "Expected narrow store");
+  assert(isPromotableZeroStoreInst(I) && "Expected promotable zero store");
 
   // Construct the new instruction.
   MachineInstrBuilder MIB;
   MIB = BuildMI(*MBB, InsertionPoint, DL, TII->get(getMatchingWideOpcode(Opc)))
-            .addOperand(getLdStRegOp(I))
+            .addReg(isNarrowStore(Opc) ? AArch64::WZR : AArch64::XZR)
             .addOperand(BaseRegOp)
             .addImm(OffsetImm)
             .setMemRefs(I->mergeMemRefsWith(*MergeMI));
 
   (void)MIB;
 
   DEBUG(dbgs() << "Creating wider load/store. Replacing instructions:\n    ");
   DEBUG(I->print(dbgs()));
AArch64LoadStoreOpt::findMatchingInsn(MachineBasicBlock::iterator I,
 
   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();
   int OffsetStride = IsUnscaled ? getMemScale(FirstMI) : 1;
-  bool IsNarrowStore = isNarrowStore(Opc);
+  bool IsPromotableZeroStore = isPromotableZeroStoreInst(FirstMI);
 
   // Track which registers have been modified and used between the first insn
   // (inclusive) and the second insn.
   ModifiedRegs.reset();
   UsedRegs.reset();
 
   // Remember any instructions that read/write memory between FirstMI and MI.
   SmallVector<MachineInstr *, 4> MemInsns;
if (canMergeOpc(Opc, MI->getOpcode(), Flags) &&
         bool IsNarrowLoad = isNarrowLoad(MI->getOpcode());
         if (!IsNarrowLoad &&
             !inBoundsForPair(IsUnscaled, MinOffset, OffsetStride)) {
           trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
           MemInsns.push_back(MI);
           continue;
         }
 
-        if (IsNarrowLoad || IsNarrowStore) {
+        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) {
             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())) {
+            (IsPromotableZeroStore && 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::isCandidateToMergeOrPair(MachineInstr *MI) {
   return true;
 }
 
 // Find narrow loads that can be converted into a single wider load with
 // bitfield extract instructions.  Also merge adjacent zero stores into a wider
 // store.
 bool AArch64LoadStoreOpt::tryToMergeLdStInst(
     MachineBasicBlock::iterator &MBBI) {
-  assert((isNarrowLoad(MBBI) || isNarrowStore(MBBI)) && "Expected narrow op.");
+  assert((isNarrowLoad(MBBI) || isPromotableZeroStoreOpcode(MBBI)) &&
+         "Expected narrow op.");
   MachineInstr *MI = MBBI;
   MachineBasicBlock::iterator E = MI->getParent()->end();
 
   if (!isCandidateToMergeOrPair(MI))
     return false;
 
-  // For narrow stores, find only the case where the stored value is WZR.
-  if (isNarrowStore(MI) && getLdStRegOp(MI).getReg() != AArch64::WZR)
+  // 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);
+  MachineBasicBlock::iterator MergeMI =
+      findMatchingInsn(MBBI, Flags, LdStLimit);
   if (MergeMI != E) {
     if (isNarrowLoad(MI)) {
       ++NumNarrowLoadsPromoted;
-    } else if (isNarrowStore(MI)) {
+    } 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.
     MBBI = mergeNarrowInsns(MBBI, MergeMI, Flags);
     return true;
   }
   return false;
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;
     }
     }

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

entry:
   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:
   %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
+}
+