[AMDGPU] Combine adjacent waitcounts in a single strongest wait

[msearles]

LGTM

lib/Target/AMDGPU/SIInsertWaitcnts.cpp

public:
   void generateSWaitCntInstBefore(MachineInstr &MI,
                                   BlockWaitcntBrackets *ScoreBrackets);
   void updateEventWaitCntAfter(MachineInstr &Inst,
                                BlockWaitcntBrackets *ScoreBrackets);
   void mergeInputScoreBrackets(MachineBasicBlock &Block);
   MachineBasicBlock *loopBottom(const MachineLoop *Loop);
   void insertWaitcntInBlock(MachineFunction &MF, MachineBasicBlock &Block);
   void insertWaitcntBeforeCF(MachineBasicBlock &Block, MachineInstr *Inst);
+  bool isWaitcntStronger(unsigned LHS, unsigned RHS);
+  unsigned combineWaitcnt(unsigned LHS, unsigned RHS);
 };
 
 } // end anonymous namespace
 
 RegInterval BlockWaitcntBrackets::getRegInterval(const MachineInstr *MI,
                                                  const SIInstrInfo *TII,
                                                  const MachineRegisterInfo *MRI,
                                                  const SIRegisterInfo *TRI,

 }
 
 static bool readsVCCZ(const MachineInstr &MI) {
   unsigned Opc = MI.getOpcode();
   return (Opc == AMDGPU::S_CBRANCH_VCCNZ || Opc == AMDGPU::S_CBRANCH_VCCZ) &&
          !MI.getOperand(1).isUndef();
 }
 
+/// \brief Given wait count encodings checks if LHS is stronger than RHS.
+bool SIInsertWaitcnts::isWaitcntStronger(unsigned LHS, unsigned RHS) {
+  if (AMDGPU::decodeVmcnt(IV, LHS) > AMDGPU::decodeVmcnt(IV, RHS))
+    return false;
+  if (AMDGPU::decodeLgkmcnt(IV, LHS) > AMDGPU::decodeLgkmcnt(IV, RHS))
+    return false;
+  if (AMDGPU::decodeExpcnt(IV, LHS) > AMDGPU::decodeExpcnt(IV, RHS))
+    return false;
+  return true;
+}
+
+/// \brief Given wait count encodings create a new encoding which is stronger
+/// or equal to both.
+unsigned SIInsertWaitcnts::combineWaitcnt(unsigned LHS, unsigned RHS) {
+  unsigned VmCnt = std::min(AMDGPU::decodeVmcnt(IV, LHS),
+                            AMDGPU::decodeVmcnt(IV, RHS));
+  unsigned LgkmCnt = std::min(AMDGPU::decodeLgkmcnt(IV, LHS),
+                              AMDGPU::decodeLgkmcnt(IV, RHS));
+  unsigned ExpCnt = std::min(AMDGPU::decodeExpcnt(IV, LHS),
+                             AMDGPU::decodeExpcnt(IV, RHS));
+  return AMDGPU::encodeWaitcnt(IV, VmCnt, ExpCnt, LgkmCnt);
+}
+
 ///  \brief Generate s_waitcnt instruction to be placed before cur_Inst.
 ///  Instructions of a given type are returned in order,
 ///  but instructions of different types can complete out of order.
 ///  We rely on this in-order completion
 ///  and simply assign a score to the memory access instructions.
 ///  We keep track of the active "score bracket" to determine
 ///  if an access of a memory read requires an s_waitcnt
 ///  and if so what the value of each counter is.
if (EmitSwaitcnt != 0) {
       // We don't (yet) track waitcnts that existed prior to the waitcnt
       // pass (we just skip over them); because the waitcnt pass is ignorant
       // of them, it may insert a redundant waitcnt. To avoid this, check
       // the prev instr. If it and the to-be-inserted waitcnt are the
       // same, keep the prev waitcnt and skip the insertion. We assume that
       // whomever. e.g., for memory model, inserted the prev waitcnt really
       // wants it there.
       bool insertSWaitInst = true;
-      if (MI.getIterator() != MI.getParent()->begin()) {
-        MachineInstr *MIPrevInst = &*std::prev(MI.getIterator());
-        if (MIPrevInst &&
-            MIPrevInst->getOpcode() == AMDGPU::S_WAITCNT &&
-            MIPrevInst->getOperand(0).getImm() == Enc) {
+      for (MachineBasicBlock::iterator I = MI.getIterator(),
+                                       B = MI.getParent()->begin();
+           insertSWaitInst && I != B; --I) {
+        if (I ==  MI.getIterator())
+          continue;
+
+        switch (I->getOpcode()) {
+        case AMDGPU::S_WAITCNT:
+          if (isWaitcntStronger(I->getOperand(0).getImm(), Enc))
             insertSWaitInst = false;
+          else if (!OldWaitcnt) {
+            OldWaitcnt = &*I;
+            Enc = combineWaitcnt(I->getOperand(0).getImm(), Enc);
           }
+          break;
+        // TODO: skip over instructions which never require wait.
+        }
+        break;
       }
       if (insertSWaitInst) {
         if (OldWaitcnt && OldWaitcnt->getOpcode() == AMDGPU::S_WAITCNT) {
           OldWaitcnt->getOperand(0).setImm(Enc);
+          if (!OldWaitcnt->getParent())
             MI.getParent()->insert(MI, OldWaitcnt);
 
           DEBUG(dbgs() << "updateWaitcntInBlock\n"
                        << "Old Instr: " << MI << '\n'
                        << "New Instr: " << *OldWaitcnt << '\n');
         } else {
             auto SWaitInst = BuildMI(*MI.getParent(), MI.getIterator(),
                                MI.getDebugLoc(), TII->get(AMDGPU::S_WAITCNT))
                              .addImm(Enc);

test/CodeGen/AMDGPU/llvm.amdgcn.s.waitcnt.ll

 
 ; Test that the intrinsic is merged with automatically generated waits and
 ; emitted as late as possible.
 ;
 ; CHECK-LABEL: {{^}}test2:
 ; CHECK-NOT: s_waitcnt
 ; CHECK: image_load
 ; CHECK-NEXT: v_lshlrev_b32
-; CHECK-NEXT: s_waitcnt
-; CHECK: s_waitcnt vmcnt(0){{$}}
+; CHECK-NEXT: s_waitcnt vmcnt(0) expcnt(0){{$}}
 ; CHECK-NEXT: image_store
 define amdgpu_ps void @test2(<8 x i32> inreg %rsrc, i32 %c) {
   %t = call <4 x float> @llvm.amdgcn.image.load.v4f32.i32.v8i32(i32 %c, <8 x i32> %rsrc, i32 15, i1 0, i1 0, i1 0, i1 0)
   call void @llvm.amdgcn.s.waitcnt(i32 3840) ; 0xf00
   %c.1 = mul i32 %c, 2
   call void @llvm.amdgcn.image.store.v4f32.i32.v8i32(<4 x float> %t, i32 %c.1, <8 x i32> %rsrc, i32 15, i1 0, i1 0, i1 0, i1 0)
   ret void
 }
 
 declare void @llvm.amdgcn.s.waitcnt(i32) #0
 
 declare <4 x float> @llvm.amdgcn.image.load.v4f32.i32.v8i32(i32, <8 x i32>, i32, i1, i1, i1, i1) #1
 declare void @llvm.amdgcn.image.store.v4f32.i32.v8i32(<4 x float>, i32, <8 x i32>, i32, i1, i1, i1, i1) #0
 
 attributes #0 = { nounwind }
 attributes #1 = { nounwind readonly }