[MachineScheduler] Order more stores by ascending address

[Allen]

Summary

According D125377, we order STP Q's by ascending address. While on some
targets, paired 128 bit loads and stores are slow, so the STP will split
into STRQ and STUR, so I hope these stores will also be ordered.
Also add subtarget feature ascend-store-address to control the aggressive order.

[Allen]

This comment was removed by Allen.

[dmgreen]

Do you have more details on where and when you expect this to be beneficial? It may not look like it from the review, but we did a fair amount of testing and benchmarking was done on D125377 on all kinds of CPUs (in-order vs out-of-order, little vs big, etc) to make sure that it was an improvement or benign on the cases we tried. I don't think I would be against this - so long as we had a decent reason to do so. It does constrain the scheduling graph though, so we shouldn't do so unnecessarily.

[Allen]

In D126700#3559712, @dmgreen wrote:

Do you have more details on where and when you expect this to be beneficial? It may not look like it from the review, but we did a fair amount of testing and benchmarking was done on D125377 on all kinds of CPUs (in-order vs out-of-order, little vs big, etc) to make sure that it was an improvement or benign on the cases we tried. I don't think I would be against this - so long as we had a decent reason to do so. It does constrain the scheduling graph though, so we shouldn't do so unnecessarily.

Thanks @dmgreen, base on the following test case , I'll gain benefits base on the tsv110 target (kunpeng 920)
BTW：I have another version, which only do some refactor(keep the logic unchanged), as we don't add new store instructions STRQ and STUR, see detail on https://reviews.llvm.org/D126700?id=433075, which is better ?

double array[0x8000];

void foo (double array[], int size) {
  for (int i=0; i<size; i++) {
     array[i] = 2.0;
  }
  return;
}

int main (){
   
  for (int i=0; i<100000; i++) {
    foo (array, sizeof (array)/sizeof (array[0]));  
  }
  return 0;
}

[dmgreen]

OK I see, thanks. In this version on godbolt I only see ordered STPQ's, perhaps you have some downstream differences that would alter the codegen?
https://godbolt.org/z/1nbr7Eh8r

With this kind of thing we always have two options. Either we can do it unconditionally for all cpus, or we can add a subtarget tuning feature to steer the codegen for specific cpus. There are a number already like FeatureFuseAES that control fusing in the machine scheduler. Adding another to control a more aggressive ordering in AArch64PostRASchedStrategy sounds perfectly sensible to me.

The other option is to do this for all cpus. The general rule there is it should be something that we expect to be beneficial or benign on all/most cpus that would run -mcpu=generic code. I'm not strongly against that for this patch, we need to check that with benchmarks and I did see a case in testing which changed the code from:

STR      q1,[sp,#0x40]     
LDP      q1,q2,[x25,#-0x10]
STR      q0,[sp,#0x60]     
LDUR     q0,[x24,#-0x10]   
FDIV     v0.2D,v0.2D,v1.2D 
STP      q2,q1,[sp,#0x20]  
...
=>
STR      q1,[sp,#0x40]     
LDP      q1,q2,[x25,#-0x10]
STP      q2,q1,[sp,#0x20]  
STR      q0,[sp,#0x60]     
LDUR     q0,[x24,#-0x10]   
FDIV     v0.2D,v0.2D,v1.2D 
...

Those clumped-together stores were apparently a little slower. Perhaps it was just unlucky, and it is only one case. I would guess that in-order memory operations are a little easier for prefetchers than ones that bobble around, but I don't have any evidence of that, and prefetchers are usually pretty sophisticated nowadays. It would take a complex case for something like that to matter, and complex cases can easily defy expectation.

So - do you have any other benchmark results, on any other cpus? Or as an alternative, do you think that adding a subtarget feature is an OK solution?

[Allen]

In D126700#3564241, @dmgreen wrote:

OK I see, thanks. In this version on godbolt I only see ordered STPQ's, perhaps you have some downstream differences that would alter the codegen?
So - do you have any other benchmark results, on any other cpus? Or as an alternative, do you think that adding a subtarget feature is an OK solution?

Many thanks @dmgreen for your detailed reply, and I found the newest upstream already sort the stores by ascending address for tsv110 target. while on my downstream, the assemble is a bit similar with cortex-a57 target. https://godbolt.org/z/bWeKq59ha

So it seems add a subtarget feature is a more reasonable, I'll try that.

[Allen]

Also add subtarget feature ascend-store-address to control the aggressive order.

[fhahn]

A subtarget feature sounds like a good way to gate this, thanks! It makes sense to only do this for CPUs where it's proven to be profitable IMO.

[dmgreen]

llvm/lib/Target/AArch64/AArch64.td

219 ↗

(On Diff #435447)

Capitalize Schedule. Maybe reword as "Schedule vector stores by ascending address".

llvm/lib/Target/AArch64/AArch64InstrInfo.h

106

Can you add a /// doc string

497–498

I think you can probably move these into AArch64InstrInfo, near to isPairedLdSt

llvm/lib/Target/AArch64/AArch64MachineScheduler.cpp

28

Why do we need to check it is an immediate?

[Allen]

Adress review comment, thanks @dmgreen @fhahn

[dmgreen]

Thanks for the update

llvm/lib/Target/AArch64/AArch64MachineScheduler.cpp

28

Why is this needed?

[Allen]

llvm/lib/Target/AArch64/AArch64.td

219 ↗

(On Diff #435447)

Thanks, apply your comment.

llvm/lib/Target/AArch64/AArch64InstrInfo.h

497–498

Done, thanks!

llvm/lib/Target/AArch64/AArch64MachineScheduler.cpp

28

this is because some IR used to match the str q0, [x8, :lo12:seed_lo] for example, and we can't get its offset

28

if we don't check this, it will crash to use getImm() get the immediate with the above str q0, [x8, :lo12:seed_lo]

[dmgreen]

Thanks. LGTM

llvm/lib/Target/AArch64/AArch64.td

219 ↗

(On Diff #435447)

A q store is a store of a vector. This subtarget feature doesn't alter the scalar scores like x/w/s.

llvm/lib/Target/AArch64/AArch64MachineScheduler.cpp

28

OK I see, it is a TargetFlags. That makes sense. It doesn't look like we ever generate them for STQP.

llvm/lib/Target/AArch64/AArch64InstrInfo.h

public:
   static bool isPreLd(const MachineInstr &MI);
 
   /// Returns whether the instruction is a pre-indexed store.
   static bool isPreSt(const MachineInstr &MI);
 
   /// Returns whether the instruction is a pre-indexed load/store.
   static bool isPreLdSt(const MachineInstr &MI);
 
+  static bool isPairedLdSt(const MachineInstr &MI);
+
   /// Returns whether the instruction is FP or NEON.
   static bool isFpOrNEON(const MachineInstr &MI);
 
   /// Returns whether the instruction is in Q form (128 bit operands)
   static bool isQForm(const MachineInstr &MI);
 
   /// Returns the index for the immediate for a given instruction.
   static unsigned getLoadStoreImmIdx(unsigned Opc);
case AArch64::PTRUE_D:
     return true;
   default:
     return false;
   }
 }
 
 /// Return opcode to be used for indirect calls.
 unsigned getBLRCallOpcode(const MachineFunction &MF);
+const MachineOperand &getLdStBaseOp(const MachineInstr &MI);
+const MachineOperand &getLdStOffsetOp(const MachineInstr &MI);
 
 // struct TSFlags {
 #define TSFLAG_ELEMENT_SIZE_TYPE(X)      (X)       // 3-bits
 #define TSFLAG_DESTRUCTIVE_INST_TYPE(X) ((X) << 3) // 4-bits
 #define TSFLAG_FALSE_LANE_TYPE(X)       ((X) << 7) // 2-bits
 #define TSFLAG_INSTR_FLAGS(X)           ((X) << 9) // 2-bits
 // }
 

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

case AArch64::STRQpre:
     return true;
   }
 }
 
 bool AArch64InstrInfo::isPreLdSt(const MachineInstr &MI) {
   return isPreLd(MI) || isPreSt(MI);
 }
 
+bool AArch64InstrInfo::isPairedLdSt(const MachineInstr &MI) {
+  switch (MI.getOpcode()) {
+  default:
+    return false;
+  case AArch64::LDPSi:
+  case AArch64::LDPSWi:
+  case AArch64::LDPDi:
+  case AArch64::LDPQi:
+  case AArch64::LDPWi:
+  case AArch64::LDPXi:
+  case AArch64::STPSi:
+  case AArch64::STPDi:
+  case AArch64::STPQi:
+  case AArch64::STPWi:
+  case AArch64::STPXi:
+  case AArch64::STGPi:
+    return true;
+  }
+}
+
 static const TargetRegisterClass *getRegClass(const MachineInstr &MI,
                                               Register Reg) {
   if (MI.getParent() == nullptr)
     return nullptr;
   const MachineFunction *MF = MI.getParent()->getParent();
   return MF ? MF->getRegInfo().getRegClassOrNull(Reg) : nullptr;
 }
 

 
 unsigned llvm::getBLRCallOpcode(const MachineFunction &MF) {
   if (MF.getSubtarget<AArch64Subtarget>().hardenSlsBlr())
     return AArch64::BLRNoIP;
   else
     return AArch64::BLR;
 }
 
+namespace llvm {
+const MachineOperand &getLdStBaseOp(const MachineInstr &MI) {
+  unsigned Idx =
+      AArch64InstrInfo::isPairedLdSt(MI) || AArch64InstrInfo::isPreLdSt(MI) ? 2
+                                                                            : 1;
+  return MI.getOperand(Idx);
+}
+
+const MachineOperand &getLdStOffsetOp(const MachineInstr &MI) {
+  unsigned Idx =
+      AArch64InstrInfo::isPairedLdSt(MI) || AArch64InstrInfo::isPreLdSt(MI) ? 3
+                                                                            : 2;
+  return MI.getOperand(Idx);
+}
+} // namespace llvm
+
 #define GET_INSTRINFO_HELPERS
 #define GET_INSTRMAP_INFO
 #include "AArch64GenInstrInfo.inc"

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

case AArch64::ST2GOffset:
     return AArch64::ST2GPostIndex;
   case AArch64::STZ2GOffset:
     return AArch64::STZ2GPostIndex;
   case AArch64::STGPi:
     return AArch64::STGPpost;
   }
 }
 
-static bool isPairedLdSt(const MachineInstr &MI) {
-  switch (MI.getOpcode()) {
-  default:
-    return false;
-  case AArch64::LDPSi:
-  case AArch64::LDPSWi:
-  case AArch64::LDPDi:
-  case AArch64::LDPQi:
-  case AArch64::LDPWi:
-  case AArch64::LDPXi:
-  case AArch64::STPSi:
-  case AArch64::STPDi:
-  case AArch64::STPQi:
-  case AArch64::STPWi:
-  case AArch64::STPXi:
-  case AArch64::STGPi:
-    return true;
-  }
-}
-
 static bool isPreLdStPairCandidate(MachineInstr &FirstMI, MachineInstr &MI) {
 
   unsigned OpcA = FirstMI.getOpcode();
   unsigned OpcB = MI.getOpcode();
 
   switch (OpcA) {
   default:
     return false;
static bool isPreLdStPairCandidate(MachineInstr &FirstMI, MachineInstr &MI) {
   case AArch64::LDRXpre:
     return (OpcB == AArch64::LDRXui) || (OpcB == AArch64::LDURXi);
   }
 }
 
 // Returns the scale and offset range of pre/post indexed variants of MI.
 static void getPrePostIndexedMemOpInfo(const MachineInstr &MI, int &Scale,
                                        int &MinOffset, int &MaxOffset) {
-  bool IsPaired = isPairedLdSt(MI);
+  bool IsPaired = AArch64InstrInfo::isPairedLdSt(MI);
   bool IsTagStore = isTagStore(MI);
   // ST*G and all paired ldst have the same scale in pre/post-indexed variants
   // as in the "unsigned offset" variant.
   // All other pre/post indexed ldst instructions are unscaled.
   Scale = (IsTagStore || IsPaired) ? AArch64InstrInfo::getMemScale(MI) : 1;
 
   if (IsPaired) {
     MinOffset = -64;
     MaxOffset = 63;
   } else {
     MinOffset = -256;
     MaxOffset = 255;
   }
 }
 
 static MachineOperand &getLdStRegOp(MachineInstr &MI,
                                     unsigned PairedRegOp = 0) {
   assert(PairedRegOp < 2 && "Unexpected register operand idx.");
   bool IsPreLdSt = AArch64InstrInfo::isPreLdSt(MI);
   if (IsPreLdSt)
     PairedRegOp += 1;
-  unsigned Idx = isPairedLdSt(MI) || IsPreLdSt ? PairedRegOp : 0;
-  return MI.getOperand(Idx);
-}
-
-static const MachineOperand &getLdStBaseOp(const MachineInstr &MI) {
-  unsigned Idx = isPairedLdSt(MI) || AArch64InstrInfo::isPreLdSt(MI) ? 2 : 1;
-  return MI.getOperand(Idx);
-}
-
-static const MachineOperand &getLdStOffsetOp(const MachineInstr &MI) {
-  unsigned Idx = isPairedLdSt(MI) || AArch64InstrInfo::isPreLdSt(MI) ? 3 : 2;
+  unsigned Idx =
+      AArch64InstrInfo::isPairedLdSt(MI) || IsPreLdSt ? PairedRegOp : 0;
   return MI.getOperand(Idx);
 }
 
 static bool isLdOffsetInRangeOfSt(MachineInstr &LoadInst,
                                   MachineInstr &StoreInst,
                                   const AArch64InstrInfo *TII) {
   assert(isMatchingStore(LoadInst, StoreInst) && "Expect only matched ld/st.");
   int LoadSize = TII->getMemScale(LoadInst);
AArch64LoadStoreOpt::mergeUpdateInsn(MachineBasicBlock::iterator I,
   if (Update->getOpcode() == AArch64::SUBXri)
     Value = -Value;
 
   unsigned NewOpc = IsPreIdx ? getPreIndexedOpcode(I->getOpcode())
                              : getPostIndexedOpcode(I->getOpcode());
   MachineInstrBuilder MIB;
   int Scale, MinOffset, MaxOffset;
   getPrePostIndexedMemOpInfo(*I, Scale, MinOffset, MaxOffset);
-  if (!isPairedLdSt(*I)) {
+  if (!AArch64InstrInfo::isPairedLdSt(*I)) {
     // Non-paired instruction.
     MIB = BuildMI(*I->getParent(), I, I->getDebugLoc(), TII->get(NewOpc))
               .add(getLdStRegOp(*Update))
               .add(getLdStRegOp(*I))
               .add(getLdStBaseOp(*I))
               .addImm(Value / Scale)
               .setMemRefs(I->memoperands())
               .setMIFlags(I->mergeFlagsWith(*Update));
MachineBasicBlock::iterator AArch64LoadStoreOpt::findMatchingUpdateInsnForward(
 
   // If the base register overlaps a source/destination register, we can't
   // merge the update. This does not apply to tag store instructions which
   // ignore the address part of the source register.
   // This does not apply to STGPi as well, which does not have unpredictable
   // behavior in this case unlike normal stores, and always performs writeback
   // after reading the source register value.
   if (!isTagStore(MemMI) && MemMI.getOpcode() != AArch64::STGPi) {
-    bool IsPairedInsn = isPairedLdSt(MemMI);
+    bool IsPairedInsn = AArch64InstrInfo::isPairedLdSt(MemMI);
     for (unsigned i = 0, e = IsPairedInsn ? 2 : 1; i != e; ++i) {
       Register DestReg = getLdStRegOp(MemMI, i).getReg();
       if (DestReg == BaseReg || TRI->isSubRegister(BaseReg, DestReg))
         return E;
     }
   }
 
   // Track which register units have been modified and used between the first
MachineBasicBlock::iterator AArch64LoadStoreOpt::findMatchingUpdateInsnBackward(
 
   // If the load/store is the first instruction in the block, there's obviously
   // not any matching update. Ditto if the memory offset isn't zero.
   if (MBBI == B || Offset != 0)
     return E;
   // If the base register overlaps a destination register, we can't
   // merge the update.
   if (!isTagStore(MemMI)) {
-    bool IsPairedInsn = isPairedLdSt(MemMI);
+    bool IsPairedInsn = AArch64InstrInfo::isPairedLdSt(MemMI);
     for (unsigned i = 0, e = IsPairedInsn ? 2 : 1; i != e; ++i) {
       Register DestReg = getLdStRegOp(MemMI, i).getReg();
       if (DestReg == BaseReg || TRI->isSubRegister(BaseReg, DestReg))
         return E;
     }
   }
 
   const bool BaseRegSP = BaseReg == AArch64::SP;

llvm/lib/Target/AArch64/AArch64MachineScheduler.cpp

 //===- AArch64MachineScheduler.cpp - MI Scheduler for AArch64 -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "AArch64MachineScheduler.h"
+#include "AArch64InstrInfo.h"
 #include "MCTargetDesc/AArch64MCTargetDesc.h"
 
 using namespace llvm;
 
+static bool needReorderStoreMI(const MachineInstr *MI) {
+  if (!MI)
+    return false;
+
+  // TODO: Add more type store instructions with immediate offset.
+  switch (MI->getOpcode()) {
+  default:
+    return false;
+  case AArch64::STPQi:
+    return getLdStOffsetOp(*MI).getType() == MachineOperand::MO_Immediate;
+  }
+
+  return false;
+}
+
+// Return true if two stores with same base address may overlap writes
+static bool mayOverlapWrite(const MachineInstr &MI0, const MachineInstr &MI1,
+                            int64_t &Off0, int64_t &Off1) {
+  const MachineOperand &Base0 = getLdStBaseOp(MI0);
+  const MachineOperand &Base1 = getLdStBaseOp(MI1);
+
+  // May overlapping writes if two store instructions without same base
+  if (!Base0.isIdenticalTo(Base1))
+    return true;
+
+  int StoreSize0 = AArch64InstrInfo::getMemScale(MI0);
+  int StoreSize1 = AArch64InstrInfo::getMemScale(MI1);
+  Off0 = AArch64InstrInfo::hasUnscaledLdStOffset(MI0.getOpcode())
+             ? getLdStOffsetOp(MI0).getImm()
+             : getLdStOffsetOp(MI0).getImm() * StoreSize0;
+  Off1 = AArch64InstrInfo::hasUnscaledLdStOffset(MI1.getOpcode())
+             ? getLdStOffsetOp(MI1).getImm()
+             : getLdStOffsetOp(MI1).getImm() * StoreSize1;
+
+  const MachineInstr &MI = (Off0 < Off1) ? MI0 : MI1;
+  int Multiples = AArch64InstrInfo::isPairedLdSt(MI) ? 2 : 1;
+  int StoreSize = AArch64InstrInfo::getMemScale(MI) * Multiples;
+
+  return llabs(Off0 - Off1) < StoreSize;
+}
+
 bool AArch64PostRASchedStrategy::tryCandidate(SchedCandidate &Cand,
                                               SchedCandidate &TryCand) {
   bool OriginalResult = PostGenericScheduler::tryCandidate(Cand, TryCand);
 
   if (Cand.isValid()) {
     MachineInstr *Instr0 = TryCand.SU->getInstr();
     MachineInstr *Instr1 = Cand.SU->getInstr();
-    // When dealing with two STPqi's.
-    if (Instr0 && Instr1 && Instr0->getOpcode() == Instr1->getOpcode () &&
-        Instr0->getOpcode() == AArch64::STPQi)
-    {
-      MachineOperand &Base0 = Instr0->getOperand(2);
-      MachineOperand &Base1 = Instr1->getOperand(2);
-      int64_t Off0 = Instr0->getOperand(3).getImm();
-      int64_t Off1 = Instr1->getOperand(3).getImm();
-      // With the same base address and non-overlapping writes.
-      if (Base0.isIdenticalTo(Base1) && llabs (Off0 - Off1) >= 2) {
-        TryCand.Reason = NodeOrder;
-        // Order them by ascending offsets.
-        return Off0 < Off1;
-      }
-    }
-  }
+
+    if (!needReorderStoreMI(Instr0) || !needReorderStoreMI(Instr1))
+      return OriginalResult;
+
+    int64_t Off0, Off1;
+    if (!mayOverlapWrite(*Instr0, *Instr1, Off0, Off1)) {
 
   return OriginalResult;
 }