[AArch64] Merge scaled and unscaled narrow zero stores

[guy-david]

No description provided.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Guy David (guy-david)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/136705.diff

2 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp (+10-9)
• (modified) llvm/test/CodeGen/AArch64/str-narrow-zero-merge.mir (+4-8)

diff --git a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp b/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
index 7c47492cf1a8e..57540eeb5ef92 100644
--- a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
+++ b/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
@@ -892,11 +892,10 @@ AArch64LoadStoreOpt::mergeNarrowZeroStores(MachineBasicBlock::iterator I,
     OffsetImm = IOffsetInBytes;
 
   int NewOpcode = getMatchingWideOpcode(Opc);
-  bool FinalIsScaled = !TII->hasUnscaledLdStOffset(NewOpcode);
-
-  // Adjust final offset if the result opcode is a scaled store.
-  if (FinalIsScaled) {
-    int NewOffsetStride = FinalIsScaled ? TII->getMemScale(NewOpcode) : 1;
+  // Adjust final offset on scaled stores because the new instruction
+  // has a different scale.
+  if (!TII->hasUnscaledLdStOffset(NewOpcode)) {
+    int NewOffsetStride = TII->getMemScale(NewOpcode);
     assert(((OffsetImm % NewOffsetStride) == 0) &&
            "Offset should be a multiple of the store memory scale");
     OffsetImm = OffsetImm / NewOffsetStride;
@@ -906,7 +905,7 @@ AArch64LoadStoreOpt::mergeNarrowZeroStores(MachineBasicBlock::iterator I,
   DebugLoc DL = I->getDebugLoc();
   MachineBasicBlock *MBB = I->getParent();
   MachineInstrBuilder MIB;
-  MIB = BuildMI(*MBB, InsertionPoint, DL, TII->get(getMatchingWideOpcode(Opc)))
+  MIB = BuildMI(*MBB, InsertionPoint, DL, TII->get(NewOpcode))
             .addReg(isNarrowStore(Opc) ? AArch64::WZR : AArch64::XZR)
             .add(BaseRegOp)
             .addImm(OffsetImm)
@@ -1539,10 +1538,12 @@ static bool areCandidatesToMergeOrPair(MachineInstr &FirstMI, MachineInstr &MI,
   if (!PairIsValidLdStrOpc)
     return false;
 
-  // FIXME: We don't support merging narrow stores with mixed scaled/unscaled
-  // offsets.
+  // Narrow stores do not have pair-wise opcodes, so constraint their merging
+  // iff they are zero stores.
   if (isNarrowStore(OpcA) || isNarrowStore(OpcB))
-    return false;
+    return getLdStRegOp(FirstMI).getReg() == AArch64::WZR &&
+           getLdStRegOp(MI).getReg() == AArch64::WZR &&
+           TII->getMemScale(FirstMI) == TII->getMemScale(MI);
 
   // The STR<S,D,Q,W,X>pre - STR<S,D,Q,W,X>ui and
   // LDR<S,D,Q,W,X,SW>pre-LDR<S,D,Q,W,X,SW>ui
diff --git a/llvm/test/CodeGen/AArch64/str-narrow-zero-merge.mir b/llvm/test/CodeGen/AArch64/str-narrow-zero-merge.mir
index e995c402c50a8..8ba7c1bcf305f 100644
--- a/llvm/test/CodeGen/AArch64/str-narrow-zero-merge.mir
+++ b/llvm/test/CodeGen/AArch64/str-narrow-zero-merge.mir
@@ -29,8 +29,7 @@ name:  merge_scaled_str_with_unscaled_8
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: merge_scaled_str_with_unscaled_8
-    ; CHECK: STRBBui $wzr, $x0, 4 :: (store (s8))
-    ; CHECK-NEXT: STURBBi $wzr, $x0, 5 :: (store (s8))
+    ; CHECK: STRHHui $wzr, $x0, 2 :: (store (s8))
     ; CHECK-NEXT: RET undef $lr
     STRBBui $wzr, $x0, 4 :: (store (s8))
     STURBBi $wzr, $x0, 5 :: (store (s8))
@@ -41,8 +40,7 @@ name:  merge_unscaled_str_with_scaled_8
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: merge_unscaled_str_with_scaled_8
-    ; CHECK: STURBBi $wzr, $x0, 4 :: (store (s8))
-    ; CHECK-NEXT: STRBBui $wzr, $x0, 5 :: (store (s8))
+    ; CHECK: STURHHi $wzr, $x0, 4 :: (store (s8))
     ; CHECK-NEXT: RET undef $lr
     STURBBi $wzr, $x0, 4 :: (store (s8))
     STRBBui $wzr, $x0, 5 :: (store (s8))
@@ -75,8 +73,7 @@ name:  merge_scaled_str_with_unscaled_16
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: merge_scaled_str_with_unscaled_16
-    ; CHECK: STRHHui $wzr, $x0, 2 :: (store (s16))
-    ; CHECK-NEXT: STURHHi $wzr, $x0, 6 :: (store (s16))
+    ; CHECK: STRWui $wzr, $x0, 1 :: (store (s16))
     ; CHECK-NEXT: RET undef $lr
     STRHHui $wzr, $x0, 2 :: (store (s16))
     STURHHi $wzr, $x0, 6 :: (store (s16))
@@ -87,8 +84,7 @@ name:  merge_unscaled_str_with_scaled_16
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: merge_unscaled_str_with_scaled_16
-    ; CHECK: STURHHi $wzr, $x0, 4 :: (store (s16))
-    ; CHECK-NEXT: STRHHui $wzr, $x0, 3 :: (store (s16))
+    ; CHECK: STURWi $wzr, $x0, 4 :: (store (s16))
     ; CHECK-NEXT: RET undef $lr
     STURHHi $wzr, $x0, 4 :: (store (s16))
     STRHHui $wzr, $x0, 3 :: (store (s16))

[davemgreen]

Does this come up often, or has ISel usually combined them previously? It seems OK so long as the wzr stores are treated OK by the rest of the code.

[guy-david]

With recent LLVM, I no longer see any cases of this happening on a very large bitcode repository, so it's possible ISel combines the stores and nothing in the codegen creates this pattern. Nonetheless, I think it's good to have it.

[davemgreen]

Yeah it comes up but quite rarely. Can you add some extra test cases for edge conditions? Otherwise LGTM.

[davemgreen]

Can you add some test cases for the limits of address ranges? I think it might be 4095 for scaled loads and 255 for unscaled? Otherwise LGTM thanks.

[guy-david]

Can you explain a bit more? Doesn't the new offset always end up being less than or equal to the offset in the first instruction?

[davemgreen]

Can you explain a bit more? Doesn't the new offset always end up being less than or equal to the offset in the first instruction?

Yeah I just wanted to make sure there were not any problems - especially with unscaled and scaled representing the offsets differently with negative offsets and whatnot. We've had them go wrong in the past. This patch doesn't touch them directly but so I think they should be OK, but it's good to make sure they are tested around the extremes.