[AArch64] merge scaled and unscaled zero narrow stores.

[zjaffal]

Summary

This patch fixes a crash when a sclaed and unscaled zero stores are merged.

[t.p.northover]

llvm/test/CodeGen/AArch64/str-narrow-zero-merge.ll

11 ↗

(On Diff #523763)

I think the patch needs more work. This offset is clearly not right (it's in the caller's frame).

[zjaffal]

1. Change the way offsets are calculated to handle mixing of scaled and unscaled stores.
2. Change the test to use .mir instead of .ll

[fhahn]

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

765

Instead we could just assign OffsetStrideRtMI in the if above (like RtMI) using the already computed strides?

768

There's no need to check for IsScaled here, as we will just multiply with 1 in the unscaled case

[zjaffal]

1. Assign OffsetImm to the appropriate value without doing recalculations.
2. Remove unecessary checks.

[fhahn]

Could you make sure that we have tests where the first offset is > the second offset and with negative offsets?

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

760

nit: This doesn't calculate the final offset, only selects the offset

766

the code below doesn't adjust the final offset any longer?

[zjaffal]

Fix comments

[fhahn]

LGTM, thanks! Would be good to wait with committing a day or so in case there are additional comments.

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

AArch64LoadStoreOpt::mergeNarrowZeroStores(MachineBasicBlock::iterator I,
   // 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,
   // and we don't need to scan the new instruction, as it's a pairwise
   // instruction, which we're not considering for further action anyway.
   if (NextI == MergeMI)
     NextI = next_nodbg(NextI, E);
 
   unsigned Opc = I->getOpcode();
+  unsigned MergeMIOpc = MergeMI->getOpcode();
   bool IsScaled = !TII->hasUnscaledLdStOffset(Opc);
-  int OffsetStride = IsScaled ? 1 : TII->getMemScale(*I);
+  bool IsMergedMIScaled = !TII->hasUnscaledLdStOffset(MergeMIOpc);
+  int OffsetStride = IsScaled ? TII->getMemScale(*I) : 1;
+  int MergeMIOffsetStride = IsMergedMIScaled ? TII->getMemScale(*MergeMI) : 1;
 
   bool MergeForward = Flags.getMergeForward();
   // Insert our new paired instruction after whichever of the paired
   // instructions MergeForward indicates.
   MachineBasicBlock::iterator InsertionPoint = MergeForward ? MergeMI : I;
   // Also based on MergeForward is from where we copy the base register operand
   // so we get the flags compatible with the input code.
   const MachineOperand &BaseRegOp =
       MergeForward ? AArch64InstrInfo::getLdStBaseOp(*MergeMI)
                    : AArch64InstrInfo::getLdStBaseOp(*I);
 
   // Which register is Rt and which is Rt2 depends on the offset order.
-  MachineInstr *RtMI;
-  if (AArch64InstrInfo::getLdStOffsetOp(*I).getImm() ==
-      AArch64InstrInfo::getLdStOffsetOp(*MergeMI).getImm() + OffsetStride)
-    RtMI = &*MergeMI;
+  int64_t IOffsetInBytes =
+      AArch64InstrInfo::getLdStOffsetOp(*I).getImm() * OffsetStride;
+  int64_t MIOffsetInBytes =
+      AArch64InstrInfo::getLdStOffsetOp(*MergeMI).getImm() *
+      MergeMIOffsetStride;
+  int64_t OffsetImm;
+  // Calculate final offset in bytes
+  if (IOffsetInBytes > MIOffsetInBytes)
+    OffsetImm = MIOffsetInBytes;
   else
-    RtMI = &*I;
+    OffsetImm = IOffsetInBytes;
 
-  int OffsetImm = AArch64InstrInfo::getLdStOffsetOp(*RtMI).getImm();
-  // Change the scaled offset from small to large type.
-  if (IsScaled) {
-    assert(((OffsetImm & 1) == 0) && "Unexpected offset to merge");
-    OffsetImm /= 2;
+  // Adjust final offset if the result opcode is a scaled store.
+  int NewOpcode = getMatchingWideOpcode(Opc);
+  bool FinalIsScaled = !TII->hasUnscaledLdStOffset(NewOpcode);
+  if (FinalIsScaled) {
+    int NewOffsetStride = FinalIsScaled ? TII->getMemScale(NewOpcode) : 1;
+    assert(((OffsetImm % NewOffsetStride) == 0) &&
+           "Offset should be a multiple of the store memory scale");
+    OffsetImm = OffsetImm / NewOffsetStride;
   }
 
   // Construct the new instruction.
   DebugLoc DL = I->getDebugLoc();
   MachineBasicBlock *MBB = I->getParent();
   MachineInstrBuilder MIB;
   MIB = BuildMI(*MBB, InsertionPoint, DL, TII->get(getMatchingWideOpcode(Opc)))
             .addReg(isNarrowStore(Opc) ? AArch64::WZR : AArch64::XZR)

llvm/test/CodeGen/AArch64/str-narrow-zero-merge.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 2
+# RUN: llc -mtriple=aarch64-linux-gnu -verify-machineinstrs -run-pass=aarch64-ldst-opt %s -o - | FileCheck %s
+
+
+---
+name:  merge_unscaled_str_with_unscaled_str_8
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: merge_unscaled_str_with_unscaled_str_8
+    ; CHECK: STURHHi $wzr, $x0, 4 :: (store (s8))
+    ; CHECK-NEXT: RET undef $lr
+    STURBBi $wzr, $x0, 4 :: (store (s8))
+    STURBBi $wzr, $x0, 5 :: (store (s8))
+    RET undef $lr
+...
+---
+name:  merge_scaled_str_with_scaled_str_8
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: merge_scaled_str_with_scaled_str_8
+    ; CHECK: STRHHui $wzr, $x0, 2 :: (store (s8))
+    ; CHECK-NEXT: RET undef $lr
+    STRBBui $wzr, $x0, 4 :: (store (s8))
+    STRBBui $wzr, $x0, 5 :: (store (s8))
+    RET undef $lr
+...
+---
+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-NEXT: RET undef $lr
+    STRBBui $wzr, $x0, 4 :: (store (s8))
+    STURBBi $wzr, $x0, 5 :: (store (s8))
+    RET undef $lr
+...
+---
+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-NEXT: RET undef $lr
+    STURBBi $wzr, $x0, 4 :: (store (s8))
+    STRBBui $wzr, $x0, 5 :: (store (s8))
+    RET undef $lr
+...
+---
+name:  merge_unscaled_str_with_unscaled_str_16
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: merge_unscaled_str_with_unscaled_str_16
+    ; CHECK: STURWi $wzr, $x0, 4 :: (store (s16))
+    ; CHECK-NEXT: RET undef $lr
+    STURHHi $wzr, $x0, 4 :: (store (s16))
+    STURHHi $wzr, $x0, 6 :: (store (s16))
+    RET undef $lr
+...
+---
+name:  merge_scaled_str_with_scaled_str_16
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: merge_scaled_str_with_scaled_str_16
+    ; CHECK: STRWui $wzr, $x0, 2 :: (store (s16))
+    ; CHECK-NEXT: RET undef $lr
+    STRHHui $wzr, $x0, 4 :: (store (s16))
+    STRHHui $wzr, $x0, 5 :: (store (s16))
+    RET undef $lr
+...
+---
+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-NEXT: RET undef $lr
+    STRHHui $wzr, $x0, 2 :: (store (s16))
+    STURHHi $wzr, $x0, 6 :: (store (s16))
+    RET undef $lr
+...
+---
+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-NEXT: RET undef $lr
+    STURHHi $wzr, $x0, 4 :: (store (s16))
+    STRHHui $wzr, $x0, 3 :: (store (s16))
+    RET undef $lr
+...
+---
+name:  merge_unscaled_str_with_unscaled_32
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: merge_unscaled_str_with_unscaled_32
+    ; CHECK: STURXi $xzr, $x0, 4 :: (store (s32))
+    ; CHECK-NEXT: RET undef $lr
+    STURWi $wzr, $x0, 4 :: (store (s32))
+    STURWi $wzr, $x0, 8 :: (store (s32))
+    RET undef $lr
+...
+---
+name:  merge_scaled_str_with_scaled_32
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: merge_scaled_str_with_scaled_32
+    ; CHECK: STRXui $xzr, $x0, 1 :: (store (s32))
+    ; CHECK-NEXT: RET undef $lr
+    STRWui $wzr, $x0, 2 :: (store (s32))
+    STRWui $wzr, $x0, 3 :: (store (s32))
+    RET undef $lr
+...
+---
+name:  merge_scaled_str_with_unscaled_32
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: merge_scaled_str_with_unscaled_32
+    ; CHECK: STRXui $xzr, $x0, 1 :: (store (s32))
+    ; CHECK-NEXT: RET undef $lr
+    STRWui $wzr, $x0, 2 :: (store (s32))
+    STURWi $wzr, $x0, 12 :: (store (s32))
+    RET undef $lr
+...
+---
+name:  merge_unscaled_str_with_scaled_32
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: merge_unscaled_str_with_scaled_32
+    ; CHECK: STURXi $xzr, $x0, 8 :: (store (s32))
+    ; CHECK-NEXT: RET undef $lr
+    STURWi $wzr, $x0, 8 :: (store (s32))
+    STRWui $wzr, $x0, 3 :: (store (s32))
+    RET undef $lr
+...
+---
+name:  merge_scaled_str_with_unscaled_32_lower_address_second
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: merge_scaled_str_with_unscaled_32_lower_address_second
+    ; CHECK: STRXui $xzr, $x0, 1 :: (store (s32))
+    ; CHECK-NEXT: RET undef $lr
+    STRWui $wzr, $x0, 3 :: (store (s32))
+    STURWi $wzr, $x0, 8 :: (store (s32))
+    RET undef $lr
+...
+---
+name:  merge_unscaled_str_with_scaled_32_lower_address_second
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: merge_unscaled_str_with_scaled_32_lower_address_second
+    ; CHECK: STURXi $xzr, $x0, 4 :: (store (s32))
+    ; CHECK-NEXT: RET undef $lr
+    STURWi $wzr, $x0, 8 :: (store (s32))
+    STRWui $wzr, $x0, 1 :: (store (s32))
+    RET undef $lr
+...
+---
+name:  merge_scaled_str_with_unscaled_32_negative_address
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: merge_scaled_str_with_unscaled_32_negative_address
+    ; CHECK: STPWi $wzr, $wzr, $x0, -1 :: (store (s32))
+    ; CHECK-NEXT: RET undef $lr
+    STRWui $wzr, $x0, 0 :: (store (s32))
+    STURWi $wzr, $x0, -4 :: (store (s32))
+    RET undef $lr