[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

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

771

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

774

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

767

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

772

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)
+  int64_t ScaledIAddress =
+      AArch64InstrInfo::getLdStOffsetOp(*I).getImm() * OffsetStride;
+  int64_t ScaledMergeMIAddress =
+      AArch64InstrInfo::getLdStOffsetOp(*MergeMI).getImm() *
+      MergeMIOffsetStride;
+  if (ScaledIAddress == ScaledMergeMIAddress + OffsetStride)
     RtMI = &*MergeMI;
   else
     RtMI = &*I;
 
   int OffsetImm = AArch64InstrInfo::getLdStOffsetOp(*RtMI).getImm();
+  bool IsScaledRtMI = !TII->hasUnscaledLdStOffset(RtMI->getOpcode());
   // Change the scaled offset from small to large type.
-  if (IsScaled) {
+  if (IsScaledRtMI) {
     assert(((OffsetImm & 1) == 0) && "Unexpected offset to merge");
     OffsetImm /= 2;
   }
 
   // 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)
             .add(BaseRegOp)
             .addImm(OffsetImm)
             .cloneMergedMemRefs({&*I, &*MergeMI})
             .setMIFlags(I->mergeFlagsWith(*MergeMI));

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

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=aarch64-linux-gnu | FileCheck %s
+
+%struct.d = type { i16, i16, i16 }
+
+define void @merge_unscaled_str_with_scaled() {
+; CHECK-LABEL: merge_unscaled_str_with_scaled:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    sub sp, sp, #16
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    str xzr, [sp, #64]
+; CHECK-NEXT:    add sp, sp, #16
+; CHECK-NEXT:    ret
+entry:
+  %ptr.1 = alloca i32, align 4
+  %ptr.2 = alloca %struct.d, align 2
+  store i32 0, ptr %ptr.1, align 4
+  %0 = getelementptr inbounds i8, ptr %ptr.2, i64 2
+  store i32 0, ptr %0, align 2
+  ret void
+}
+
+define void @merge_scaled_str_with_scaled_str(){
+; CHECK-LABEL: merge_scaled_str_with_scaled_str:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    sub sp, sp, #16
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    str xzr, [sp, #8]
+; CHECK-NEXT:    add sp, sp, #16
+; CHECK-NEXT:    ret
+  %ptr.1 = alloca i32, align 4
+  %ptr.2 = alloca i32, align 4
+  store i32 0, ptr %ptr.1, align 4
+  store i32 0, ptr %ptr.2, align 4
+  ret void
+}
+
+define void @merge_scaled_str_with_scaled_str_2() {
+; CHECK-LABEL: merge_scaled_str_with_scaled_str_2:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    sub sp, sp, #16
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    str wzr, [sp, #12]
+; CHECK-NEXT:    add sp, sp, #16
+; CHECK-NEXT:    ret
+entry:
+  %g = alloca i16, align 2
+  %h = alloca i16, align 2
+  store i16 0, ptr %g, align 2
+  store i16 0, ptr %h, align 2
+  ret void
+}