[AMDGPU/MemOpsCluster] Let mem ops clustering logic also consider number of clustered bytes

[hsmhsm]

Summary

While clustering mem ops, AMDGPU target needs to consider number of clustered bytes
to decide on max number of mem ops that can be clustered. This patch adds support to pass
number of clustered bytes to target mem ops clustering logic.

[hsmhsm]

Taken care of clang formatting.

[arsenm]

Missing tests

llvm/lib/Target/AMDGPU/SIInsertHardClauses.cpp

159–166

It would be better to not depend on the memory operands here, but this belongs in a helper function some kind of not (and this can also sink down to the use)

[hsmhsm]

llvm/lib/Target/AMDGPU/SIInsertHardClauses.cpp

159–166

Hi @arsenm

Did you mean here the helper function as a kind of below?

unsigned getDstMemOperandSize(const MachineInstr *MI) const {
  if (!MI || MI->memoperands_empty())
    return 0;

  return MI->memoperands().front()->getSize();
}

And, use above helper function as below?

unsigned WidthA = getDstMemOperandSize(CI.Last);
unsigned WidthB = getDstMemOperandSize(&MI);

[hsmhsm]

In D80545#2054804, @arsenm wrote:

Missing tests

There is NO functionality related changes in this patch. What kind of tests that you suggest here?

[hsmhsm]

Have taken care of Matt's review comments

[arsenm]

llvm/lib/CodeGen/MachineScheduler.cpp

1576–1577

This won't correctly handle multiple mem operands

llvm/lib/Target/AMDGPU/SIInsertHardClauses.cpp

159–166

Yes, but it's worse to rely on the memory operands here than getting this from the instruction opcode / operand

[hsmhsm]

Have taken care of second set of review comments (by Matt)

[hsmhsm]

Replace few if stmts with ternary expression

[foad]

llvm/lib/Target/AMDGPU/SIInsertHardClauses.cpp

159–166

Right, getMemOperandsWithOffset could be extended to return the width. Some targets already have an internal function getMemOperandsWithOffsetWidth which does that.

176–178

The comment explains that we don't really want to limit the size of the cluster here, so it's probably best to pass in a small dummy value like 2 instead of WidthA + WidthB.

[hsmhsm]

Modified getMemOperandsWithOffset() to pass Width as an argument. Note that
SIInstrInfo::getMemOperandsWithOffset() needs to be extended to compute and
assign Width. It will be implemented in next patch.

[hsmhsm]

getMemOperandsWithOffset() implemented by other targets needs compute Width too.
If it is not computing it, then, default computation is based on memoperands() api.

[hsmhsm]

Compute Width within SIInstrInfo::getMemOperandsWithOffset()

[hsmhsm]

llvm/lib/CodeGen/MachineScheduler.cpp

1576–1577

Taken care

llvm/lib/Target/AMDGPU/SIInsertHardClauses.cpp

159–166

Taken care

159–166

Taken care

176–178

Taken care

[hsmhsm]

PING (since it is bit of high priority)

[foad]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

2045

It's a bit silly that this function has "Width" in its name, and getMemOperandsWithOffset doesn't, but both of them take same Width argument.

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

272–273

Move this below the mayLoadOrStore check. Then, does getDstMemOperand ever fail?

llvm/lib/Target/Hexagon/HexagonInstrInfo.cpp

2972 ↗

(On Diff #266538)

Doesn't AccessSize give you the width already?

llvm/lib/Target/X86/X86InstrInfo.cpp

3700–3701 ↗

(On Diff #266538)

Perhaps MemOp.hasOneMemOperand() would be more appropriate? (Or what should the behaviour be if there is more than one?)

Please at least add a FIXME to get the width without relying on memoperands, or get an opinion from the X86 maintainers if this is OK.

[hsmhsm]

Have taken care of review comments made by Jay.

[hsmhsm]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

2045

You are right, I somehow had missed it. Fixed it.

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

272–273

Fixed it. However, I do not know much about the functionality guarantee of mayLoadOrStore(). So for the purpose of safety, I have added an if stmt to check if DstOp is null or not, and return false if it is null.

llvm/lib/Target/Hexagon/HexagonInstrInfo.cpp

2972 ↗

(On Diff #266538)

You are right, I somehow had missed it. Fixed it.

llvm/lib/Target/X86/X86InstrInfo.cpp

3700–3701 ↗

(On Diff #266538)

Yes, you are right. I have added a FIXME comment for now. I will get in touch with X86 maintainers and will get it fixed as soon as possible.

[hsmhsm]

One of the fixes for review comments resulted in one unit test failure, hence reverted back to original state.

[hsmhsm]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

272–273

Apparently, above if stmt resulted in one unit test failure. Hence reverted it back. And again, assuming here that DstOp will not be null is also NOT guaranteed, We seem to be getting null in many unit test cases, and hence assuming it non-null will again result in many test failures. Now there two possibilities
(1) getDstMemOperand() might not be handling all the cases OR
(2) getDstMemOperand() is handling all the cases, but, it is expected that DstOp will be null in few cases even though, mayLoadOrStore() is true.

[foad]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

275

Hopefully if you follow the advice above, you will find DstOp is never null (because it is only ever null here in cases where getMemOperandsWithOffsetWidth is going to return false anyway).

If you do still find cases where DstOp is null then please debug them to see what has gone wrong.

277–292

I would prefer this logic to be folded into getMemOperandsWithOffsetWidth itself, where we already test for all the different instruction types (isMUBUF and so on). It might make sense for each of the cases in getMemOperandsWithOffsetWidth to set DstOp and then fall through to some common code at the end of the function that sets Width based on DstOp.

2736–2741

This code can be removed now. You should be able to use the Width values returned by getMemOperandsWithOffsetWidth instead. But you can leave that for a later patch if you prefer.

[hsmhsm]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

275

I will hope for the best :)

277–292

In fact, I initially thought about it. However, first look at the code within getMemOperandsWithOffsetWidth() made me to feel on the other way around. I felt that this function requires clean-up, and hence did not want to touch it. Anyway, let me follow your suggestion. I will experiment and prepare a patch.

2736–2741

I would prefer to make this change in a separate patch, after the current patch is successfully merged.

[hsmhsm]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

275

DstOp is null in many unit test-cases. I am not sure, what is going-on here. I need to debug it. It may take some time. I am not sure, if it is even worth considering the current situation and other P1 tickets at my hand. Anyway, let me see what best I can do here.

[hsmhsm]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

275

Hi @foad

Only above review comment to be taken care in the patch, and all other review comments are taken care. However, since we will not be going to use the width value until we will be able to come-up with a safe heuristic for mem ops cluster threshold, we can safely merge this patch, and in the next patch, we can address this review comment. That way, we no need to block this entire big patch to be in hold condition, and we do not need to deal with merge conflicts often and often. What do you say?

[hsmhsm]

Make sure that Width is properly computed within SIInstrInfo::getMemOperandsWithOffsetWidth()

[foad]

Looks OK to me with some minor cleanups (see inline comments).

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

267

Do you need this function at all? Can't you use RI.getOpSize instead?

272–274

Could use RI.getRegClassForReg here.

275

You don't need the outer parentheses here.

346

You can remove some duplicated code by making this an "else", so there is only one "return true" for all BUF instructions.

[hsmhsm]

Thanks Jay, I will take care of these minor comments in a next immediate patch.

llvm/include/llvm/CodeGen/TargetInstrInfo.h

public:
   ///   DAG->addMutation(createLoadClusterDAGMutation(DAG->TII, DAG->TRI));
   /// or
   ///   DAG->addMutation(createStoreClusterDAGMutation(DAG->TII, DAG->TRI));
   /// to TargetPassConfig::createMachineScheduler() to have an effect.
   ///
   /// \p BaseOps1 and \p BaseOps2 are memory operands of two memory operations.
   /// \p NumLoads is the number of loads that will be in the cluster if this
   /// hook returns true.
+  /// \p NumBytes is the number of bytes that will be loaded from all the
+  /// clustered loads if this hook returns true.
   virtual bool shouldClusterMemOps(ArrayRef<const MachineOperand *> BaseOps1,
                                    ArrayRef<const MachineOperand *> BaseOps2,
-                                   unsigned NumLoads) const {
+                                   unsigned NumLoads, unsigned NumBytes) const {
     llvm_unreachable("target did not implement shouldClusterMemOps()");
   }
 
   /// Reverses the branch condition of the specified condition list,
   /// returning false on success and true if it cannot be reversed.
   virtual bool
   reverseBranchCondition(SmallVectorImpl<MachineOperand> &Cond) const {
     return true;

llvm/lib/CodeGen/MachineScheduler.cpp

 
 /// Post-process the DAG to create cluster edges between neighboring
 /// loads or between neighboring stores.
 class BaseMemOpClusterMutation : public ScheduleDAGMutation {
   struct MemOpInfo {
     SUnit *SU;
     SmallVector<const MachineOperand *, 4> BaseOps;
     int64_t Offset;
+    unsigned Width;
 
     MemOpInfo(SUnit *SU, ArrayRef<const MachineOperand *> BaseOps,
-              int64_t Offset)
-        : SU(SU), BaseOps(BaseOps.begin(), BaseOps.end()), Offset(Offset) {}
+              int64_t Offset, unsigned Width)
+        : SU(SU), BaseOps(BaseOps.begin(), BaseOps.end()), Offset(Offset),
+          Width(Width) {}
 
     static bool Compare(const MachineOperand *const &A,
                         const MachineOperand *const &B) {
       if (A->getType() != B->getType())
         return A->getType() < B->getType();
       if (A->isReg())
         return A->getReg() < B->getReg();
       if (A->isFI()) {

 }
 
 } // end namespace llvm
 
 void BaseMemOpClusterMutation::clusterNeighboringMemOps(
     ArrayRef<SUnit *> MemOps, ScheduleDAGInstrs *DAG) {
   SmallVector<MemOpInfo, 32> MemOpRecords;
   for (SUnit *SU : MemOps) {
+    const MachineInstr &MI = *SU->getInstr();
     SmallVector<const MachineOperand *, 4> BaseOps;
     int64_t Offset;
     bool OffsetIsScalable;
-    if (TII->getMemOperandsWithOffset(*SU->getInstr(), BaseOps, Offset,
-                                      OffsetIsScalable, TRI))
-      MemOpRecords.push_back(MemOpInfo(SU, BaseOps, Offset));
+    if (TII->getMemOperandsWithOffset(MI, BaseOps, Offset, OffsetIsScalable,
+                                      TRI)) {
+      unsigned Width =
+          !MI.memoperands_empty() ? MI.memoperands().front()->getSize() : 0;
+      MemOpRecords.push_back(MemOpInfo(SU, BaseOps, Offset, Width));
+    }
 #ifndef NDEBUG
     for (auto *Op : BaseOps)
       assert(Op);
 #endif
   }
   if (MemOpRecords.size() < 2)
     return;
 
   llvm::sort(MemOpRecords);
 
   // At this point, `MemOpRecords` array must hold atleast two mem ops. Try to
   // cluster mem ops collected within `MemOpRecords` array.
   unsigned ClusterLength = 1;
+  unsigned CurrentClusterBytes = MemOpRecords[0].Width;
   for (unsigned Idx = 0, End = MemOpRecords.size(); Idx < (End - 1); ++Idx) {
     // Decision to cluster mem ops is taken based on target dependent logic
     auto MemOpa = MemOpRecords[Idx];
     auto MemOpb = MemOpRecords[Idx + 1];
     ++ClusterLength;
-    if (!TII->shouldClusterMemOps(MemOpa.BaseOps, MemOpb.BaseOps,
-                                  ClusterLength)) {
+    CurrentClusterBytes += MemOpb.Width;
+    if (!TII->shouldClusterMemOps(MemOpa.BaseOps, MemOpb.BaseOps, ClusterLength,
+                                  CurrentClusterBytes)) {
       // Current mem ops pair could not be clustered, reset cluster length, and
       // go to next pair
       ClusterLength = 1;
+      CurrentClusterBytes = MemOpb.Width;
       continue;
     }
 
     SUnit *SUa = MemOpa.SU;
     SUnit *SUb = MemOpb.SU;
     if (SUa->NodeNum > SUb->NodeNum)
       std::swap(SUa, SUb);
 
     // FIXME: Is this check really required?
     if (!DAG->addEdge(SUb, SDep(SUa, SDep::Cluster))) {
       ClusterLength = 1;
+      CurrentClusterBytes = MemOpb.Width;
       continue;
     }
 
     LLVM_DEBUG(dbgs() << "Cluster ld/st SU(" << SUa->NodeNum << ") - SU("
                       << SUb->NodeNum << ")\n");
 
     // Copy successor edges from SUa to SUb. Interleaving computation
     // dependent on SUa can prevent load combining due to register reuse.

llvm/lib/Target/AArch64/AArch64InstrInfo.h

public:
   /// \p Scale, \p Width, \p MinOffset, and \p MaxOffset accordingly.
   ///
   /// For unscaled instructions, \p Scale is set to 1.
   static bool getMemOpInfo(unsigned Opcode, TypeSize &Scale, unsigned &Width,
                            int64_t &MinOffset, int64_t &MaxOffset);
 
   bool shouldClusterMemOps(ArrayRef<const MachineOperand *> BaseOps1,
                            ArrayRef<const MachineOperand *> BaseOps2,
-                           unsigned NumLoads) const override;
+                           unsigned NumLoads, unsigned NumBytes) const override;
 
   void copyPhysRegTuple(MachineBasicBlock &MBB, MachineBasicBlock::iterator I,
                         const DebugLoc &DL, MCRegister DestReg,
                         MCRegister SrcReg, bool KillSrc, unsigned Opcode,
                         llvm::ArrayRef<unsigned> Indices) const;
   void copyGPRRegTuple(MachineBasicBlock &MBB, MachineBasicBlock::iterator I,
                        DebugLoc DL, unsigned DestReg, unsigned SrcReg,
                        bool KillSrc, unsigned Opcode, unsigned ZeroReg,

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

bool AArch64InstrInfo::getMemOperandsWithOffset(
     const MachineInstr &LdSt, SmallVectorImpl<const MachineOperand *> &BaseOps,
     int64_t &Offset, bool &OffsetIsScalable, const TargetRegisterInfo *TRI)
     const {
   if (!LdSt.mayLoadOrStore())
     return false;
 
   const MachineOperand *BaseOp;
   unsigned Width;
   if (!getMemOperandWithOffsetWidth(LdSt, BaseOp, Offset, OffsetIsScalable,
                                     Width, TRI))
     return false;
   BaseOps.push_back(BaseOp);
   return true;
 }
 
 bool AArch64InstrInfo::getMemOperandWithOffsetWidth(
     const MachineInstr &LdSt, const MachineOperand *&BaseOp, int64_t &Offset,
static bool shouldClusterFI(const MachineFrameInfo &MFI, int FI1,
   return FI1 == FI2;
 }
 
 /// Detect opportunities for ldp/stp formation.
 ///
 /// Only called for LdSt for which getMemOperandWithOffset returns true.
 bool AArch64InstrInfo::shouldClusterMemOps(
     ArrayRef<const MachineOperand *> BaseOps1,
-    ArrayRef<const MachineOperand *> BaseOps2, unsigned NumLoads) const {
+    ArrayRef<const MachineOperand *> BaseOps2, unsigned NumLoads,
+    unsigned NumBytes) const {
   assert(BaseOps1.size() == 1 && BaseOps2.size() == 1);
   const MachineOperand &BaseOp1 = *BaseOps1.front();
   const MachineOperand &BaseOp2 = *BaseOps2.front();
   const MachineInstr &FirstLdSt = *BaseOp1.getParent();
   const MachineInstr &SecondLdSt = *BaseOp2.getParent();
   if (BaseOp1.getType() != BaseOp2.getType())
     return false;
 

llvm/lib/Target/AMDGPU/SIInsertHardClauses.cpp

for (auto &MBB : MF) {
           if (!SII->getMemOperandsWithOffset(MI, BaseOps, Dummy1, Dummy2,
                                              TRI)) {
             // We failed to get the base operands, so we'll never clause this
             // instruction with any other, so pretend it's illegal.
             Type = HARDCLAUSE_ILLEGAL;
           }
         }
 
+        unsigned WidthA = CI.Last
+                              ? !CI.Last->memoperands_empty()
+                                    ? CI.Last->memoperands().front()->getSize()
+                                    : 0
+                              : 0;
+        unsigned WidthB = !MI.memoperands_empty()
+                              ? MI.memoperands().front()->getSize()
+                              : 0;
+
         if (CI.Length == 64 ||
             (CI.Length && Type != HARDCLAUSE_INTERNAL &&
              (Type != CI.Type ||
               // Note that we lie to shouldClusterMemOps about the size of the
               // cluster. When shouldClusterMemOps is called from the machine
               // scheduler it limits the size of the cluster to avoid increasing
               // register pressure too much, but this pass runs after register
               // allocation so there is no need for that kind of limit.
-              !SII->shouldClusterMemOps(CI.BaseOps, BaseOps, 2)))) {
+              !SII->shouldClusterMemOps(CI.BaseOps, BaseOps, 2,
+                                        WidthA + WidthB)))) {
           // Finish the current clause.
           Changed |= emitClause(CI, SII);
           CI = ClauseInfo();
         }
 
         if (CI.Length) {
           // Extend the current clause.
           ++CI.Length;
           if (Type != HARDCLAUSE_INTERNAL) {

llvm/lib/Target/AMDGPU/SIInstrInfo.h

public:
   bool
   getMemOperandsWithOffset(const MachineInstr &LdSt,
                            SmallVectorImpl<const MachineOperand *> &BaseOps,
                            int64_t &Offset, bool &OffsetIsScalable,
                            const TargetRegisterInfo *TRI) const final;
 
   bool shouldClusterMemOps(ArrayRef<const MachineOperand *> BaseOps1,
                            ArrayRef<const MachineOperand *> BaseOps2,
-                           unsigned NumLoads) const override;
+                           unsigned NumLoads, unsigned NumBytes) const override;
 
   bool shouldScheduleLoadsNear(SDNode *Load0, SDNode *Load1, int64_t Offset0,
                                int64_t Offset1, unsigned NumLoads) const override;
 
   void copyPhysReg(MachineBasicBlock &MBB, MachineBasicBlock::iterator MI,
                    const DebugLoc &DL, MCRegister DestReg, MCRegister SrcReg,
                    bool KillSrc) const override;
 

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

static bool isStride64(unsigned Opc) {
   case AMDGPU::DS_WRITE2ST64_B32:
   case AMDGPU::DS_WRITE2ST64_B64:
     return true;
   default:
     return false;
   }
 }
 
 bool SIInstrInfo::getMemOperandsWithOffset(
     const MachineInstr &LdSt, SmallVectorImpl<const MachineOperand *> &BaseOps,
     int64_t &Offset, bool &OffsetIsScalable, const TargetRegisterInfo *TRI)
     const {
   if (!LdSt.mayLoadOrStore())
     return false;
 
   unsigned Opc = LdSt.getOpcode();
   OffsetIsScalable = false;
   const MachineOperand *BaseOp, *OffsetOp;
 
   if (isDS(LdSt)) {
     BaseOp = getNamedOperand(LdSt, AMDGPU::OpName::addr);
     OffsetOp = getNamedOperand(LdSt, AMDGPU::OpName::offset);
     if (OffsetOp) {
       // Normal, single offset LDS instruction.
       if (!BaseOp) {
         // DS_CONSUME/DS_APPEND use M0 for the base address.
         // TODO: find the implicit use operand for M0 and use that as BaseOp?
         return false;
       }
       BaseOps.push_back(BaseOp);
       Offset = OffsetOp->getImm();
     } else {
       // The 2 offset instructions use offset0 and offset1 instead. We can treat
       // these as a load with a single offset if the 2 offsets are consecutive.
       // We will use this for some partially aligned loads.
       const MachineOperand *Offset0Op =
           getNamedOperand(LdSt, AMDGPU::OpName::offset0);
       const MachineOperand *Offset1Op =
           getNamedOperand(LdSt, AMDGPU::OpName::offset1);
 
       unsigned Offset0 = Offset0Op->getImm();
       unsigned Offset1 = Offset1Op->getImm();
if (SOffset && SOffset->isReg()) {
         return false;
 
       const MachineOperand *OffsetImm =
         getNamedOperand(LdSt, AMDGPU::OpName::offset);
       BaseOps.push_back(RSrc);
       BaseOps.push_back(SOffset);
       Offset = OffsetImm->getImm();
       return true;
     }
 
     BaseOp = getNamedOperand(LdSt, AMDGPU::OpName::srsrc);
     if (!BaseOp) // e.g. BUFFER_WBINVL1_VOL
       return false;
     BaseOps.push_back(BaseOp);
 
     BaseOp = getNamedOperand(LdSt, AMDGPU::OpName::vaddr);
     if (BaseOp)
static bool memOpsHaveSameBasePtr(const MachineInstr &MI1,
   if (isa<UndefValue>(Base1) || isa<UndefValue>(Base2))
     return false;
 
   return Base1 == Base2;
 }
 
 bool SIInstrInfo::shouldClusterMemOps(ArrayRef<const MachineOperand *> BaseOps1,
                                       ArrayRef<const MachineOperand *> BaseOps2,
-                                      unsigned NumLoads) const {
+                                      unsigned NumLoads,
+                                      unsigned NumBytes) const {
   assert(!BaseOps1.empty() && !BaseOps2.empty());
   const MachineInstr &FirstLdSt = *BaseOps1.front()->getParent();
   const MachineInstr &SecondLdSt = *BaseOps2.front()->getParent();
 
   if (!memOpsHaveSameBasePtr(FirstLdSt, BaseOps1, SecondLdSt, BaseOps2))
     return false;
 
   const MachineOperand *FirstDst = nullptr;
bool SIInstrInfo::checkInstOffsetsDoNotOverlap(const MachineInstr &MIa,
   bool Offset0IsScalable, Offset1IsScalable;
   if (!getMemOperandsWithOffset(MIa, BaseOps0, Offset0, Offset0IsScalable, &RI) ||
       !getMemOperandsWithOffset(MIb, BaseOps1, Offset1, Offset1IsScalable, &RI))
     return false;
 
   if (!memOpsHaveSameBaseOperands(BaseOps0, BaseOps1))
     return false;
 
   if (!MIa.hasOneMemOperand() || !MIb.hasOneMemOperand()) {
     // FIXME: Handle ds_read2 / ds_write2.
     return false;
   }
   unsigned Width0 = MIa.memoperands().front()->getSize();
   unsigned Width1 = MIb.memoperands().front()->getSize();
   return offsetsDoNotOverlap(Width0, Offset0, Width1, Offset1);
 }
 
 bool SIInstrInfo::areMemAccessesTriviallyDisjoint(const MachineInstr &MIa,
                                                   const MachineInstr &MIb) const {
   assert(MIa.mayLoadOrStore() &&
          "MIa must load from or modify a memory location");
   assert(MIb.mayLoadOrStore() &&