[AMDGPU] Only raise wave priority if there is a long enough sequence of VALU instructions.

[tsymalla]

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

135

You could early opt-out when VALUInstsThreshold == 0 (at the beginning), is that correct?

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

30–31

Should we be counting cycles instead of instructions?

140

You're counting VALU instructions here and above?

198–199

Why not construct directly at the insert point?

[kosarev]

Updated as suggested.

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

30–31

We don't know it yet, I'm afraid. For the couple real use cases that we have for the issue that this pass tries to address counting instructions looks sufficient. This being coupled with that proper cycle counting might be not very trivial, we may be at the risk of over-engineering here.

135

We could, but I'm not sure I know how that might be useful in practice.

140

Yes, that's not good. Combined the two new loops into one and simplified related code. Thanks.

198–199

Well, the instruction we create here is not just an auxiliary value that we coincidentally happen to use in both the cases, if that's what you mean. That instruction must be spent in all cases and we do want it be the same instruction.

[kosarev]

Ping.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

198–199

But you can construct it in the right place rather than constructing the instruction and then inserting after. You have the insert point you want, you can just construct it there? It's unusual to need insert/insertAfter

[kosarev]

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

198–199

I'm not sure I see how this answers the point. Granted that constructing the instruction at the points of insertion is possible, but as I said in this case we have reasons to create it in a single place.

[foad]

I'm not sure how to review this. Can you explain why this heuristic makes sense intuitively? Or do you have any benchmarks to back it up? Or preferably both? :)

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

115

If you're ignoring loops, can you described succinctly what you *are* counting? Is it something like the minimum number of VALU instructions along any path from the start of the function to the VMEM load in question?

[foad]

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

198–199

I agree with Matt that it would be more normal to construct-and-insert at the same time, though I don't feel very strongly about it. I think I suggested before that you could pass the insertion point into BuildSetprioMI. In this case the insertion point would be something like MBBInfos[MBB].LastVMEMLoad ? std::next(MBBInfos[MBB].LastVMEMLoad) : MBB->begin().

[kosarev]

Reworked as suggested.

[kosarev]

In D124671#3521606, @foad wrote:

I'm not sure how to review this. Can you explain why this heuristic makes sense intuitively? Or do you have any benchmarks to back it up? Or preferably both? :)

Right, that's the hardest part of it, isn't it. Because we don't have much use cases provided, the idea was to start with replicating more or less the same counting logic we have in the other compiler, and then adjust things using feedback from people who can give it some proper testing.

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

115

Good point. Amended the comment.

198–199

OK, changed.

[nhaehnle]

I don't quite follow the cross-basic-block logic. Given what I understand of the goal of the heuristic, I expect it to be: "insert a s_setprio 0 before the first long section of dense VALU that can happen after a VMEM load (if any)".

Instead of MBBInfo::LastVMEMLoad I can see a MBBInfo::LastVALUSequenceBreak and an MBBInfo::PastVMEMLoad boolean to indicate whether a VMEM load can have been issued at the end of each basic block.

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

30–31

Could this also be configured via a function attribute? Those are less problematic than command-line options in a driver context.

155–156

I believe NumVALUInsts should be reset to 0 here. It should probably also be reset at a number of other events, in particular DS instructions.

The reasoning is that we want to lower priority just before running a long dense block of VALU, so that other waves have a better chance of running address calculation VALU.

[kosarev]

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

155–156

So would something like this work for the purpose?

if (SIInstrInfo::isVALU(MI))
  ++NumVALUInsts;
else
  NumVALUInsts = 0;

[kosarev]

Updated to count VALU instructions that follow VMEM loads and to add
support for threshold function attributes.

[kosarev]

In D124671#3698865, @nhaehnle wrote:

I don't quite follow the cross-basic-block logic. Given what I understand of the goal of the heuristic, I expect it to be: "insert a s_setprio 0 before the first long section of dense VALU that can happen after a VMEM load (if any)".

Instead of MBBInfo::LastVMEMLoad I can see a MBBInfo::LastVALUSequenceBreak and an MBBInfo::PastVMEMLoad boolean to indicate whether a VMEM load can have been issued at the end of each basic block.

I think knowing if we are at a point past a VMEM load wouldn't help as there may be another VMEM load down the control flow. And a similar problem with the last sequence break as at where the sequence finally becomes 'long enough', we can have several points at which it begins.

If I'm not wrong that there is no harm in lowering the priority immediately after VMEM loads, it seems it's still easier implementation-wise to track them and not other things. The updated version just makes sure we only consider VMEM loads that are followed by VALU sequences of the required length. Then the rest of the logic remains the same.

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

30–31

Done.

155–156

Updated to break the sequence on DS instructions.

[kosarev]

Update commit description.

[kosarev]

Update commit title as well.

[kosarev]

Ping.

[kosarev]

Ping.

[nhaehnle]

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

170–171

Why are you only counting if AtStart? The idea was: find places where a VMEM load is followed by a long sequence of dense VALU. Lower priority between that VMEM load and the dense VALU.

[nhaehnle]

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

168–172

If the winner is MaxNumVALUInstsInMiddle, I believe the relevant VMEMLoad should be the VMEMLoad that appeared just before the corresponding sequence of VALU instead of the last one.

I think that instead of Info.LastVMEMLoad, we'd want to have Info.LastVMEMLoadBeforeLongVALU (feel free to think of a better name), which is only set here at the end if the threshold is exceeded. The loop above would keep track of the most recently seen VMEMLoad as well as the one corresponding to the longest VALU sequence in the middle so far.

170–171

Thank you for the offline clarification, I understand this better now.

[kosarev]

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

168–172

The ticket says the priority is supposed to be lowered after the last VMEM load, and this updated implementation has been prepared in assumption that counting VALUs doesn't affect that principle. So if this point stands, then I understand we should not be interested in any VALU sequences preceding VMEM loads, even if long enough.

In terms of the code, as we reset MaxNumVALUInstsInMiddle every time we ran into a VMEM load, I'm not sure I see how an in-the-middle sequence followed by a VMEM load can possibly be the winner.

Info.LastVMEMLoad storing the actual last VMEM load in the block should not be a problem because it is the Info.MayReachVMEMLoad flag that the following analysis takes into account. And speaking of naming, I'm not perfectly happy with the the name of that flag as it actually means 'may reach any of the last VMEM loads that precede a long-enough sequence of VALU instructions'. Would appreciate if anyone can suggest a better alternative of a practical length.

[nhaehnle]

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

168–172

I did think the policy was supposed to be "lower priority before the longest sequence of VALU instructions, if the sequence length crosses a certain threshold". But we can start with "lower after last VMEM load, if subsequent VALU sequence is long enough". So then, you're right and the code is okay as-is.

[kosarev]

@nhaehnle Are we fine with accepting and submitting this? Or, do we want to give this some testing beforehand? Also not quite sure if the pass is supposed to be enabled to by default.

[nhaehnle]

Yes, I think it's fine.

llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp

 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/Allocator.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "amdgpu-set-wave-priority"
 
+static cl::opt<unsigned> VALUInstsThreshold(
+    "amdgpu-set-wave-priority-valu-insts-threshold",
+    cl::desc("VALU instruction count threshold for adjusting wave priority"),
+    cl::init(100), cl::Hidden);
+
 namespace {
 
 struct MBBInfo {
   MBBInfo() = default;
   bool MayReachVMEMLoad = false;
+  unsigned NumVALUInstsAtBeginning = 0;
+  unsigned NumVALUInstsAtEnd = 0;
+  MachineInstr *LastVMEMLoad = nullptr;
 };
 
 using MBBInfoSet = DenseMap<const MachineBasicBlock *, MBBInfo>;
 
 class AMDGPUSetWavePriority : public MachineFunctionPass {
 public:
   static char ID;
 
bool AMDGPUSetWavePriority::runOnMachineFunction(MachineFunction &MF) {
 
   Function &F = MF.getFunction();
   if (skipFunction(F) || !AMDGPU::isEntryFunctionCC(F.getCallingConv()))
     return false;
 
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   TII = ST.getInstrInfo();
 
+  // Count VALU instructions that may be executed by the time each of the
+  // function's blocks receive control. We currently assume that
+  // backedges/loops, branch probabilities and other details can be
+  // ignored.
   MBBInfoSet MBBInfos;
+  ReversePostOrderTraversal<MachineFunction *> RPOT(&MF);
+  for (const MachineBasicBlock *MBB : RPOT) {
+    unsigned NumVALUInsts = 0;
+    for (const MachineBasicBlock *Pred : MBB->predecessors())
+      NumVALUInsts = std::max(NumVALUInsts, MBBInfos[Pred].NumVALUInstsAtEnd);
+
+    MBBInfo &Info = MBBInfos[MBB];
+    Info.NumVALUInstsAtBeginning = NumVALUInsts;
+    Info.NumVALUInstsAtEnd =
+        NumVALUInsts + count_if(*MBB, [](const MachineInstr &MI) {
+          return SIInstrInfo::isVALU(MI);
+        });
+  }
+
+  // Find VMEM loads that may be executed before the number of executed
+  // VALU instructions hits the specified threshold.
   SmallVector<const MachineBasicBlock *, 16> Worklist;
   for (MachineBasicBlock &MBB : MF) {
-    if (any_of(MBB, isVMEMLoad))
+    MBBInfo &Info = MBBInfos[&MBB];
+    unsigned NumVALUInsts = Info.NumVALUInstsAtBeginning;
+    MachineInstr *LastVMEMLoad = nullptr;
+    for (MachineInstr &MI : MBB) {
+      if (NumVALUInsts >= VALUInstsThreshold)
+        break;
+      if (isVMEMLoad(MI))
+        LastVMEMLoad = &MI;
+      if (SIInstrInfo::isVALU(MI))
+        ++NumVALUInsts;
+    }
+    if (LastVMEMLoad) {
+      Info.LastVMEMLoad = LastVMEMLoad;
       Worklist.push_back(&MBB);
     }
+  }
 
-  // Mark blocks from which control may reach VMEM loads.
+  // Mark blocks from which control may reach the VMEM loads.
   while (!Worklist.empty()) {
     const MachineBasicBlock *MBB = Worklist.pop_back_val();
     MBBInfo &Info = MBBInfos[MBB];
     if (!Info.MayReachVMEMLoad) {
       Info.MayReachVMEMLoad = true;
       Worklist.append(MBB->pred_begin(), MBB->pred_end());
     }
   }
 
   MachineBasicBlock &Entry = MF.front();
   if (!MBBInfos[&Entry].MayReachVMEMLoad)
     return false;
 
   // Raise the priority at the beginning of the shader.
   MachineBasicBlock::iterator I = Entry.begin(), E = Entry.end();
   while (I != E && !SIInstrInfo::isVALU(*I) && !I->isTerminator())
     ++I;
   Entry.insert(I, BuildSetprioMI(MF, HighPriority));
 
   // Lower the priority on edges where control leaves blocks from which
-  // VMEM loads are reachable.
+  // the VMEM loads are reachable.
   SmallSet<MachineBasicBlock *, 16> PriorityLoweringBlocks;
   for (MachineBasicBlock &MBB : MF) {
     if (MBBInfos[&MBB].MayReachVMEMLoad) {
       if (MBB.succ_empty())
         PriorityLoweringBlocks.insert(&MBB);
       continue;
     }
 
     if (CanLowerPriorityDirectlyInPredecessors(MBB, MBBInfos)) {
       for (MachineBasicBlock *Pred : MBB.predecessors()) {
         if (MBBInfos[Pred].MayReachVMEMLoad)
           PriorityLoweringBlocks.insert(Pred);
       }
       continue;
     }
 
     // Where lowering the priority in predecessors is not possible, the
     // block receiving control either was not part of a loop in the first
     // place or the loop simplification/canonicalization pass should have
     // already tried to split the edge and insert a preheader, and if for
     // whatever reason it failed to do so, then this leaves us with the
     // only option of lowering the priority within the loop.
     PriorityLoweringBlocks.insert(&MBB);
   }
 
   for (MachineBasicBlock *MBB : PriorityLoweringBlocks) {
-    MachineBasicBlock::iterator I = MBB->end(), B = MBB->begin();
-    while (I != B) {
-      if (isVMEMLoad(*--I)) {
-        ++I;
-        break;
-      }
+    MachineInstr *Setprio = BuildSetprioMI(MF, LowPriority);
+    if (MachineInstr *LastVMEMLoad = MBBInfos[MBB].LastVMEMLoad) {
+      MBB->insertAfter(MachineBasicBlock::instr_iterator(LastVMEMLoad),
+                       Setprio);
+      continue;
     }
-    MBB->insert(I, BuildSetprioMI(MF, LowPriority));
+    MBB->insert(MBB->begin(), Setprio);
   }
 
   return true;
 }

llvm/test/CodeGen/AMDGPU/set-wave-priority.ll

-; RUN: llc -mtriple=amdgcn -amdgpu-set-wave-priority=true -o - %s | \
-; RUN:   FileCheck %s
+; RUN: llc -mtriple=amdgcn -amdgpu-set-wave-priority=true \
+; RUN:   -amdgpu-set-wave-priority-valu-insts-threshold=4 -o - %s | FileCheck %s
 
 ; CHECK-LABEL: no_setprio:
 ; CHECK-NOT:       s_setprio
 ; CHECK:           ; return to shader part epilog
 define amdgpu_ps <2 x float> @no_setprio() {
   ret <2 x float> <float 0.0, float 0.0>
 }
 
exit:
   ret <2 x float> %mul2
 
 another_load:
   %v2 = call <2 x float> @llvm.amdgcn.struct.buffer.load.v2f32(<4 x i32> %p, i32 0, i32 0, i32 1, i32 0)
   %sum = fadd <2 x float> %v, %v2
   ret <2 x float> %sum
 }
 
+; CHECK-LABEL: valu_insts_threshold:
+; CHECK:           s_setprio 3
+; CHECK:           buffer_load_dwordx2
+; CHECK-NEXT:      s_setprio 0
+; CHECK-COUNT-4:   v_add_f32_e32
+; CHECK:           s_cbranch_scc0 [[A:.*]]
+; CHECK:       {{.*}}:  ; %b
+; CHECK-NEXT:      buffer_load_dwordx2
+; CHECK:           s_branch [[END:.*]]
+; CHECK:       [[A]]:  ; %a
+; CHECK:           s_branch [[END]]
+; CHECK:       [[END]]:
+define amdgpu_ps <2 x float> @valu_insts_threshold(<4 x i32> inreg %p, i32 inreg %i) {
+  %v = call <2 x float> @llvm.amdgcn.struct.buffer.load.v2f32(<4 x i32> %p, i32 0, i32 0, i32 0, i32 0)
+  %add = fadd <2 x float> %v, %v
+  %add2 = fadd <2 x float> %add, %add
+
+  %cond = icmp eq i32 %i, 0
+  br i1 %cond, label %a, label %b
+
+a:
+  ret <2 x float> %add2
+
+b:
+  %v2 = call <2 x float> @llvm.amdgcn.struct.buffer.load.v2f32(<4 x i32> %p, i32 0, i32 1, i32 0, i32 0)
+  %sub = fsub <2 x float> %add2, %v2
+  ret <2 x float> %sub
+}
+
 declare <2 x float> @llvm.amdgcn.struct.buffer.load.v2f32(<4 x i32>, i32, i32, i32, i32) nounwind