[AArch64] Improve load/store optimizer to handle LDUR + LDR.

[ab]

Summary

This patch allows the mixing of scaled and unscaled load/stores to form
load/store pairs.

PR24465

[mcrosier]

I haven't done any testing beyond basic lit tests. Hopefully, I'll have correctness/perf testing in flight by tomorrow. Regardless, I think this is in good enough shape to review.

[kristof.beyls]

test/CodeGen/AArch64/ldp-stp-scaled-unscaled-pairs.ll

1 ↗

(On Diff #32449)

My preference is to use -march=aarch64 instead of -march=arm64 as much as possible, as aarch64 is the official name for the architecture.
Is there a specific reason -mcpu=cortex-a57 needs to be added to the test line? I don't have a strong opinion on whether it should be there or not, just wondering.

81 ↗

(On Diff #32449)

Looking at this test case, I see that before this patch, the following code is produced:

ldur    x8, [x0, #-8]
ldr      x9, [x0]

. If I'm not mistaken, ldur x8, [x0, #-8] has the same functionality as ldr x8, [x0, #-1]? If so, wouldn't it be better to make sure
we produce ldr instead of ldur in the first place?
If we would do that, and there still is a good reason to have special code to convert LDR + LDUR into LDP, I guess none of the above test cases really show that (although I haven't investigated every single test case in detail)?

[kristof.beyls]

test/CodeGen/AArch64/ldp-stp-scaled-unscaled-pairs.ll

81 ↗

(On Diff #32449)

D'oh! The ARMARM clearly states that the scaled immediate offsets in ldr x8, [x0, #imm] can only be positive/unsigned. Please ignore my comment above!

[mcrosier]

Addressed Kristof's comments. Included full diff.

[mcrosier]

Fix test7 which I modified/broke, while doing additional testing.

[mzolotukhin]

Hi Chad,

I checked on a small testcase, and with this patch we do merge STUR and STR. There is one potential issue though: in some cases we intentionally split STP into two STUR/STR, as there is a big performance penalty if STP crosses a cache line (see e.g. performSTORECombine in AArch64ISelLowering.cpp and getMemoryOpCost in AArch64TargetTransformInfo.cpp). So, I think we might want to perform this combining only for non-temporal or known-to-be-well-aligned memory accesses. What do you think, does it make sense?

Thanks,
Michael

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

455

This looks strange. Is it expected that scaled and unscaled offset differ by MemSize^2? (in one case you multiply by MemSize, in the other - divide)

708

Same here.

[mcrosier]

Hi Michael,
So, you're saying that performSTORECombine splits 16B stores that cross line boundaries for performance reasons. AFAICT, this combine is only enabled with Cyclone. Later the AArch64LoadStore optimizer runs and combines these split stores undoing the optimization performed during ISelLowering.

If I understand things correctly, I think this makes sense. However, I'd probably implement this in a separate patch and test it on other subtargets (e.g., A57). Sound reasonable?

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

455

My assumption is that a scaled offset is always the unscaled offset divide by the MemSize and conversely the unscaled offset is always the scaled offset divided by the MemSize; I believe this is valid.

I.e.,
Unscaled Offset = Scaled Offset * MemSize;
Scaled Offset = Unscaled Offset / MemSize;

What did I miss? I'm trying to simply the logic so that both offsets are either scaled or unscaled, but not a mix of the two. I'm guessing you're concerned my conversion is incorrect.

[mzolotukhin]

Hi Chad,

As far as I understand, currently ISelLowering splits the unaligned stores, from which we happen to get STUR and STR, which we can't combine to STP without this patch. With this patch, we'll be able to merge them back, so we'll undo that optimization.

I think that being able to combine STUR/STR to STP is good by itself, but from the patch alone we'll probably get regressions (AFAIR, the biggest one was matmul from LNT testsuite). That's my only concern, but I don't mind landing the patch and addressing this issue later if we actually observe it.

Also, I can't provide a high-quality review of this code, as I'm not very familiar with it, so I'd appreciate other people's feedback on the patch.

Thanks!

Michael

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

455

Yep, I was concerned that it was incorrect, but I just read the code incorrectly. Sorry for that!
(I thought that we do PairedOffset = UnscaledOffset*MemSize in one case and PairedOffset = UnscaleOffset/MemSize in the other).

[ab]

This LGTM (modulo nits), but let's see what the testsuite says first.

In D12116#231366, @mzolotukhin wrote:

As far as I understand, currently ISelLowering splits the unaligned stores, from which we happen to get STUR and STR, which we can't combine to STP without this patch. With this patch, we'll be able to merge them back, so we'll undo that optimization.

But this issue isn't specific to this patch, right? If the unaligned store was split to STR+STR we would have generated an STP even before this change. I agree we'll need to do something about this though, but separately, and for both mixed and non-mixed STR/STUR pairs.

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

450

MI -> Paired, FirstMI -> I ?

607–610

What about removing CanMergeOpc? It's assigned to several times with different values, and all could be folded into ifs/returns.

619–620

Same: fold the condition into the if, remove CanMergeOpc?

627–629

I see this is the only use of the get*From* helpers. What about:

getMatchingPairOpcode(Opc) == getMatchingPairOpcode(PairOpc)

That'll let us get rid of the added functions.

630–633

Same: fold the condition into the return, remove CanMergeOpc?

721

To make sure I understand: the previous code was "incorrect" in checking MI rather than FirstMI, but it didn't make a difference because we didn't allow mismatched opcodes, right?

test/CodeGen/AArch64/ldp-stp-scaled-unscaled-pairs.ll

8 ↗

(On Diff #32821)

Kill the "entry:"s ?

[mcrosier]

Thanks for the feedback, Ahmed.

In D12116#236954, @ab wrote:

This LGTM (modulo nits), but let's see what the testsuite says first.

Unfortunately, I don't have a setup that can test with the testsuite at the moment. That is in the works...

This transformation is very narrow and did not hit anything in Spec2000. Therefore, I'm going to move onto other work with a higher ROI. However, feel free to push this one along.

In D12116#231366, @mzolotukhin wrote:

As far as I understand, currently ISelLowering splits the unaligned stores, from which we happen to get STUR and STR, which we can't combine to STP without this patch. With this patch, we'll be able to merge them back, so we'll undo that optimization.

But this issue isn't specific to this patch, right? If the unaligned store was split to STR+STR we would have generated an STP even before this change. I agree we'll need to do something about this though, but separately, and for both mixed and non-mixed STR/STUR pairs.

That is absolutely correct! This patch only applies to the very narrow case of STR/STUR. The common case is when we're pairing STUR/STUR, which has already been committed.

[mcrosier]

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

721

I'm not sure the checking was "incorrect", but your latter comment is correct; because we didn't allow mismatches opcodes it didn't make a difference. However, because we're converting MI to be either Scaled or Unscaled to match FirstMI, we need to use IsUscaled in the context of this patch.

[mzolotukhin]

Hi Chad, Ahmed,

I ran some testing for this patch and found a bug there. The issue is that when we scale offset, we need to make sure that the original value is divisible by the scale factor. Here is a test to illustrate the issue:

define <2 x i64> @test4(i64* %p) nounwind {
  %a1 = bitcast i64* %p to <2 x i64>*
  %tmp1 = load <2 x i64>, < 2 x i64>* %a1, align 8           # Load <p[0], p[1]>

  %add.ptr2 = getelementptr inbounds i64, i64* %p, i64 3
  %a2 = bitcast i64* %add.ptr2 to <2 x i64>*
  %tmp2 = load <2 x i64>, <2 x i64>* %a2, align 8            # Load <p[3], p[4]>
  %add = add nsw <2 x i64> %tmp1, %tmp2
  ret <2 x i64> %add
}

The current patch will combine these two loads, which is incorrect.

Michael

[ab]

Thanks Michael. Chad, I'll fix that and resubmit, thanks for the patch!

[ab]

And this should catch the unscaled offset not being a multiple of the scale (with abundant asserts).

[mcrosier]

Thanks, Ahmed.

[mzolotukhin]

Hi,

I just finished LNT+externals run, and saw no failures. Thanks for the fix!

Michael

[mcrosier]

Thanks for seeing this one through, Ahmed. LGTM.

[mcrosier]

Committed r246769.

test/CodeGen/AArch64/ldp-stp-scaled-unscaled-pairs.ll

1 ↗

(On Diff #33810)

BTW, I added -aarch64-neon-syntax=apple to the final commit, so that this will pass on all platforms.

[mcrosier]

I speculatively reverted r246769 in r246782 as the compiler looks to be ICEing on Multisource/Benchmarks/tramp3d-v4.

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

static int getMemSize(MachineInstr *MI) {
   case AArch64::LDURXi:
     return 8;
   case AArch64::LDRSWui:
   case AArch64::LDURSWi:
     return 4;
   }
 }
 
+// Given a scaled load/store opcode, return the equivalent unscaled opcode.
+static unsigned getUnscaledFromScaled(unsigned Opc) {
+  switch (Opc) {
+  default:
+    llvm_unreachable("Unexpected scaled opcode!");
+  case AArch64::STRSui:
+    return AArch64::STURSi;
+  case AArch64::STRDui:
+    return AArch64::STURDi;
+  case AArch64::STRQui:
+    return AArch64::STURQi;
+  case AArch64::STRWui:
+    return AArch64::STURWi;
+  case AArch64::STRXui:
+    return AArch64::STURXi;
+  case AArch64::LDRSui:
+    return AArch64::LDURSi;
+  case AArch64::LDRDui:
+    return AArch64::LDURDi;
+  case AArch64::LDRQui:
+    return AArch64::LDURQi;
+  case AArch64::LDRWui:
+    return AArch64::LDURWi;
+  case AArch64::LDRXui:
+    return AArch64::LDURXi;
+  case AArch64::LDRSWui:
+    return AArch64::LDURSWi;
+  }
+}
+
+// Given an unscaled load/store opcode, return the equivalent scaled opcode.
+static unsigned getScaledFromUnscaled(unsigned Opc) {
+  switch (Opc) {
+  default:
+    llvm_unreachable("Unexpected unscaled opcode!");
+  case AArch64::STURSi:
+    return AArch64::STRSui;
+  case AArch64::STURDi:
+    return AArch64::STRDui;
+  case AArch64::STURQi:
+    return AArch64::STRQui;
+  case AArch64::STURWi:
+    return AArch64::STRWui;
+  case AArch64::STURXi:
+    return AArch64::STRXui;
+  case AArch64::LDURSi:
+    return AArch64::LDRSui;
+  case AArch64::LDURDi:
+    return AArch64::LDRDui;
+  case AArch64::LDURQi:
+    return AArch64::LDRQui;
+  case AArch64::LDURWi:
+    return AArch64::LDRWui;
+  case AArch64::LDURXi:
+    return AArch64::LDRXui;
+  case AArch64::LDURSWi:
+    return AArch64::LDRSWui;
+  }
+}
+
 static unsigned getMatchingNonSExtOpcode(unsigned Opc,
                                          bool *IsValidLdStrOpc = nullptr) {
   if (IsValidLdStrOpc)
     *IsValidLdStrOpc = true;
   switch (Opc) {
   default:
     if (IsValidLdStrOpc)
       *IsValidLdStrOpc = false;
AArch64LoadStoreOpt::mergePairedInsns(MachineBasicBlock::iterator I,
   // Insert our new paired instruction after whichever of the paired
   // instructions MergeForward indicates.
   MachineBasicBlock::iterator InsertionPoint = MergeForward ? Paired : 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 ? getLdStBaseOp(Paired) : getLdStBaseOp(I);
 
+  int Offset = getLdStOffsetOp(I).getImm();
+  int PairedOffset = getLdStOffsetOp(Paired).getImm();
+  bool PairedIsUnscaled = isUnscaledLdSt(Paired->getOpcode());
+
+  // We're trying to pair instructions that differ in how they are scaled.
+  // If FirstMI is scaled then scale the offset of MI accordingly.
+  // Otherwise, do the opposite (i.e., make Paired's offset unscaled).
+  if (IsUnscaled != PairedIsUnscaled) {
+    int MemSize = getMemSize(Paired);
+    PairedOffset =
+        PairedIsUnscaled ? PairedOffset / MemSize : PairedOffset * MemSize;
+  }
+
   // Which register is Rt and which is Rt2 depends on the offset order.
   MachineInstr *RtMI, *Rt2MI;
-  if (getLdStOffsetOp(I).getImm() ==
-      getLdStOffsetOp(Paired).getImm() + OffsetStride) {
+  if (Offset == PairedOffset + OffsetStride) {
     RtMI = Paired;
     Rt2MI = I;
     // Here we swapped the assumption made for SExtIdx.
     // I.e., we turn ldp I, Paired into ldp Paired, I.
     // Update the index accordingly.
     if (SExtIdx != -1)
       SExtIdx = (SExtIdx + 1) % 2;
   } else {
     RtMI = I;
     Rt2MI = Paired;
   }
-  // Handle Unscaled
+  // Scale the immediate offset, if necessary.
   int OffsetImm = getLdStOffsetOp(RtMI).getImm();
-  if (IsUnscaled && EnableAArch64UnscaledMemOp)
-    OffsetImm /= OffsetStride;
+  if (isUnscaledLdSt(RtMI->getOpcode()) && EnableAArch64UnscaledMemOp)
+    OffsetImm /= getMemSize(RtMI);
 
   // Construct the new instruction.
   MachineInstrBuilder MIB = BuildMI(*I->getParent(), InsertionPoint,
                                     I->getDebugLoc(), TII->get(NewOpc))
                                 .addOperand(getLdStRegOp(RtMI))
                                 .addOperand(getLdStRegOp(Rt2MI))
                                 .addOperand(BaseRegOp)
                                 .addImm(OffsetImm);
static bool mayAlias(MachineInstr *MIa,
                      const AArch64InstrInfo *TII) {
   for (auto &MIb : MemInsns)
     if (mayAlias(MIa, MIb, TII))
       return true;
 
   return false;
 }
 
+static bool canMergeOpc(unsigned Opc, unsigned PairOpc, LdStPairFlags &Flags) {
+  bool CanMergeOpc = Opc == PairOpc;
+  // Opcodes match nothing more to check.
+  if (CanMergeOpc)
+    return true;
+
+  // Try to match a signed-extended load/store with a zero-extended load/store.
+  Flags.setSExtIdx(-1);
+  bool IsValidLdStrOpc;
+  unsigned NonSExtOpc = getMatchingNonSExtOpcode(Opc, &IsValidLdStrOpc);
+  assert(IsValidLdStrOpc &&
+         "Given Opc should be a Load or Store with an immediate");
+  // Opc will be the first instruction in the pair.
+  CanMergeOpc = NonSExtOpc == getMatchingNonSExtOpcode(PairOpc);
+  if (CanMergeOpc) {
+    Flags.setSExtIdx(NonSExtOpc == (unsigned)Opc ? 1 : 0);
+    return true;
+  }
+
+  // Try to match an unscaled load/store with a scaled load/store.
+  bool IsUnscaled = isUnscaledLdSt(Opc);
+  unsigned NewOpc =
+      IsUnscaled ? getScaledFromUnscaled(Opc) : getUnscaledFromScaled(Opc);
+  CanMergeOpc = NewOpc == PairOpc;
+  if (CanMergeOpc)
+    return true;
+
+  return false;
+}
+
 /// findMatchingInsn - Scan the instructions looking for a load/store that can
 /// be combined with the current instruction into a load/store pair.
 MachineBasicBlock::iterator
 AArch64LoadStoreOpt::findMatchingInsn(MachineBasicBlock::iterator I,
                                       LdStPairFlags &Flags,
                                       unsigned Limit) {
   MachineBasicBlock::iterator E = I->getParent()->end();
   MachineBasicBlock::iterator MBBI = I;
for (unsigned Count = 0; MBBI != E && Count < Limit; ++MBBI) {
     // Skip DBG_VALUE instructions. Otherwise debug info can affect the
     // optimization by changing how far we scan.
     if (MI->isDebugValue())
       continue;
 
     // Now that we know this is a real instruction, count it.
     ++Count;
 
-    bool CanMergeOpc = Opc == MI->getOpcode();
-    Flags.setSExtIdx(-1);
-    if (!CanMergeOpc) {
-      bool IsValidLdStrOpc;
-      unsigned NonSExtOpc = getMatchingNonSExtOpcode(Opc, &IsValidLdStrOpc);
-      assert(IsValidLdStrOpc &&
-             "Given Opc should be a Load or Store with an immediate");
-      // Opc will be the first instruction in the pair.
-      Flags.setSExtIdx(NonSExtOpc == (unsigned)Opc ? 1 : 0);
-      CanMergeOpc = NonSExtOpc == getMatchingNonSExtOpcode(MI->getOpcode());
-    }
-
-    if (CanMergeOpc && getLdStOffsetOp(MI).isImm()) {
+    if (canMergeOpc(Opc, MI->getOpcode(), Flags) &&
+        getLdStOffsetOp(MI).isImm()) {
       assert(MI->mayLoadOrStore() && "Expected memory operation.");
       // If we've found another instruction with the same opcode, check to see
       // if the base and offset are compatible with our starting instruction.
       // These instructions all have scaled immediate operands, so we just
       // check for +1/-1. Make sure to check the new instruction offset is
       // actually an immediate and not a symbolic reference destined for
       // a relocation.
       //
       // Pairwise instructions have a 7-bit signed offset field. Single insns
       // have a 12-bit unsigned offset field. To be a valid combine, the
       // final offset must be in range.
       unsigned MIBaseReg = getLdStBaseOp(MI).getReg();
       int MIOffset = getLdStOffsetOp(MI).getImm();
+
+      // We're trying to pair instructions that differ in how they are scaled.
+      // If FirstMI is scaled then scale the offset of MI accordingly.
+      // Otherwise, do the opposite (i.e., make MI's offset unscaled).
+      bool MIIsUnscaled = isUnscaledLdSt(MI);
+      if (IsUnscaled != MIIsUnscaled) {
+        int MemSize = getMemSize(MI);
+        MIOffset = MIIsUnscaled ? MIOffset / MemSize : MIOffset * MemSize;
+      }
+
       if (BaseReg == MIBaseReg && ((Offset == MIOffset + OffsetStride) ||
                                    (Offset + OffsetStride == MIOffset))) {
         int MinOffset = Offset < MIOffset ? Offset : MIOffset;
         // If this is a volatile load/store that otherwise matched, stop looking
         // as something is going on that we don't have enough information to
         // safely transform. Similarly, stop if we see a hint to avoid pairs.
         if (MI->hasOrderedMemoryRef() || TII->isLdStPairSuppressed(MI))
           return E;
         // If the resultant immediate offset of merging these instructions
         // is out of range for a pairwise instruction, bail and keep looking.
-        bool MIIsUnscaled = isUnscaledLdSt(MI);
-        if (!inBoundsForPair(MIIsUnscaled, MinOffset, OffsetStride)) {
+        if (!inBoundsForPair(IsUnscaled, MinOffset, OffsetStride)) {
           trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
           MemInsns.push_back(MI);
           continue;
         }
         // If the alignment requirements of the paired (scaled) instruction
         // can't express the offset of the unscaled input, bail and keep
         // looking.
         if (IsUnscaled && EnableAArch64UnscaledMemOp &&