[AMDGPU] Fix and extend vccz workarounds

[foad]

Summary

We have workarounds for two different cases where vccz can get out of
sync with the value in vcc. This fixes them in two ways:

1. Fix the case where the def of vcc was in a previous basic block, by

pessimistically assuming that vccz might be incorrect at a basic block
boundary.

2. Fix the handling of pre-existing waitcnt instructions by calling

generateWaitcntInstBefore before examining ScoreBrackets to determine
whether there's an outstanding smem read operation.

[dstuttard]

LGTM

[arsenm]

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

1478

Could you have DefVCCLo && DefVCCHi?

[foad]

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

1478

I'm not sure what you're asking. But this particular bit of logic (if defines vcc_lo or vcc_hi ... else if defines vcc ...) hasn't changed, it has just moved up from line 1488, and it was previously discussed here: https://reviews.llvm.org/D69661#inline-627569

[arsenm]

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

1478

I mean call definesRegister 2x instead of 3x

[foad]

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

1478

No, because a def of $vcc does not satisfy Inst.definesRegister(AMDGPU::VCC_LO) or Inst.definesRegister(AMDGPU::VCC_HI). And I'm relying on that here, otherwise the "else if" part would be dead code.

[piotr]

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

1435–1442

Nit: braces around multi-line blocks?

[foad]

Add braces.

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

bool SIInsertWaitcnts::insertWaitcntInBlock(MachineFunction &MF,
                                             WaitcntBrackets &ScoreBrackets) {
   bool Modified = false;
 
   LLVM_DEBUG({
     dbgs() << "*** Block" << Block.getNumber() << " ***";
     ScoreBrackets.dump();
   });
 
-  // Assume VCCZ is correct at basic block boundaries, unless and until we need
-  // to handle cases where that is not true.
+  // Track the correctness of vccz through this basic block. There are two
+  // reasons why it might be incorrect; see ST->hasReadVCCZBug() and
+  // ST->partialVCCWritesUpdateVCCZ().
   bool VCCZCorrect = true;
+  if (ST->hasReadVCCZBug()) {
+    // vccz could be incorrect at a basic block boundary if a predecessor wrote
+    // to vcc and then issued an smem load.
+    VCCZCorrect = false;
+  } else if (!ST->partialVCCWritesUpdateVCCZ()) {
+    // vccz could be incorrect at a basic block boundary if a predecessor wrote
+    // to vcc_lo or vcc_hi.
+    VCCZCorrect = false;
+  }
 
   // Walk over the instructions.
   MachineInstr *OldWaitcntInstr = nullptr;
 
   for (MachineBasicBlock::instr_iterator Iter = Block.instr_begin(),
                                          E = Block.instr_end();
        Iter != E;) {
     MachineInstr &Inst = *Iter;
 
     // Track pre-existing waitcnts from earlier iterations.
     if (Inst.getOpcode() == AMDGPU::S_WAITCNT ||
         (Inst.getOpcode() == AMDGPU::S_WAITCNT_VSCNT &&
          Inst.getOperand(0).isReg() &&
          Inst.getOperand(0).getReg() == AMDGPU::SGPR_NULL)) {
       if (!OldWaitcntInstr)
         OldWaitcntInstr = &Inst;
       ++Iter;
       continue;
     }
 
-    // We might need to restore vccz to its correct value for either of two
-    // different reasons; see ST->hasReadVCCZBug() and
-    // ST->partialVCCWritesUpdateVCCZ().
-    bool RestoreVCCZ = false;
-    if (readsVCCZ(Inst)) {
-      if (!VCCZCorrect) {
-        // Restore vccz if it's not known to be correct already.
-        RestoreVCCZ = true;
-      } else if (ST->hasReadVCCZBug()) {
+    // Generate an s_waitcnt instruction to be placed before Inst, if needed.
+    Modified |= generateWaitcntInstBefore(Inst, ScoreBrackets, OldWaitcntInstr);
+    OldWaitcntInstr = nullptr;
+
+    bool RestoreVCCZ = !VCCZCorrect && readsVCCZ(Inst);
+
+    // Don't examine operands unless we need to track vccz correctness.
+    if (ST->hasReadVCCZBug() || !ST->partialVCCWritesUpdateVCCZ()) {
+      if (Inst.definesRegister(AMDGPU::VCC_LO) ||
+          Inst.definesRegister(AMDGPU::VCC_HI)) {
+        // Up to gfx9, writes to vcc_lo and vcc_hi don't update vccz.
+        if (!ST->partialVCCWritesUpdateVCCZ())
+          VCCZCorrect = false;
+      } else if (Inst.definesRegister(AMDGPU::VCC)) {
         // There is a hardware bug on CI/SI where SMRD instruction may corrupt
         // vccz bit, so when we detect that an instruction may read from a
         // corrupt vccz bit, we need to:
         // 1. Insert s_waitcnt lgkm(0) to wait for all outstanding SMRD
         //    operations to complete.
         // 2. Restore the correct value of vccz by writing the current value
         //    of vcc back to vcc.
-        if (ScoreBrackets.getScoreLB(LGKM_CNT) <
+        if (ST->hasReadVCCZBug() &&
+            ScoreBrackets.getScoreLB(LGKM_CNT) <
                 ScoreBrackets.getScoreUB(LGKM_CNT) &&
             ScoreBrackets.hasPendingEvent(SMEM_ACCESS)) {
-          // Restore vccz if there's an outstanding smem read, which could
-          // complete and clobber vccz at any time.
-          RestoreVCCZ = true;
+          // Writes to vcc while there's an outstanding smem read may get
+          // clobbered as soon as any read completes.
+          VCCZCorrect = false;
+        } else {
+          // Writes to vcc will fix any incorrect value in vccz.
+          VCCZCorrect = true;
         }
       }
     }
 
     if (TII->isSMRD(Inst)) {
       for (const MachineMemOperand *Memop : Inst.memoperands()) {
         const Value *Ptr = Memop->getValue();
         SLoadAddresses.insert(std::make_pair(Ptr, Inst.getParent()));
       }
-    }
-
-    if (!ST->partialVCCWritesUpdateVCCZ()) {
-      if (Inst.definesRegister(AMDGPU::VCC_LO) ||
-          Inst.definesRegister(AMDGPU::VCC_HI)) {
-        // Up to gfx9, writes to vcc_lo and vcc_hi don't update vccz.
+      if (ST->hasReadVCCZBug()) {
+        // This smem read could complete and clobber vccz at any time.
         VCCZCorrect = false;
-      } else if (Inst.definesRegister(AMDGPU::VCC)) {
-        // Writes to vcc will fix any incorrect value in vccz.
-        VCCZCorrect = true;
       }
     }
 
-    // Generate an s_waitcnt instruction to be placed before
-    // cur_Inst, if needed.
-    Modified |= generateWaitcntInstBefore(Inst, ScoreBrackets, OldWaitcntInstr);
-    OldWaitcntInstr = nullptr;
-
     updateEventWaitcntAfter(Inst, &ScoreBrackets);
 
 #if 0 // TODO: implement resource type check controlled by options with ub = LB.
     // If this instruction generates a S_SETVSKIP because it is an
     // indexed resource, and we are on Tahiti, then it will also force
     // an S_WAITCNT vmcnt(0)
     if (RequireCheckResourceType(Inst, context)) {
       // Force the score to as if an S_WAITCNT vmcnt(0) is emitted.

llvm/test/CodeGen/AMDGPU/infinite-loop.ll

 ; SI-NEXT:    s_mov_b32 s3, 0xf000
 ; SI-NEXT:    s_mov_b32 s2, -1
 ; SI-NEXT:    v_mov_b32_e32 v0, 0x3e7
 ; SI-NEXT:    s_and_b64 vcc, exec, -1
 ; SI-NEXT:  BB1_2: ; %loop
 ; SI-NEXT:    ; =>This Inner Loop Header: Depth=1
 ; SI-NEXT:    s_waitcnt lgkmcnt(0)
 ; SI-NEXT:    buffer_store_dword v0, off, s[0:3], 0
+; SI-NEXT:    s_mov_b64 vcc, vcc
 ; SI-NEXT:    s_cbranch_vccnz BB1_2
 ; SI-NEXT:  BB1_3: ; %UnifiedReturnBlock
 ; SI-NEXT:    s_endpgm
 ; IR-LABEL: @infinite_loop_ret(
 ; IR-NEXT:  entry:
 ; IR-NEXT:    [[TMP:%.*]] = tail call i32 @llvm.amdgcn.workitem.id.x()
 ; IR-NEXT:    [[COND:%.*]] = icmp eq i32 [[TMP]], 1
 ; IR-NEXT:    br i1 [[COND]], label [[LOOP:%.*]], label [[UNIFIEDRETURNBLOCK:%.*]]

 ; SI-NEXT:    s_mov_b32 s3, 0xf000
 ; SI-NEXT:    s_mov_b32 s2, -1
 ; SI-NEXT:    v_mov_b32_e32 v0, 0x378
 ; SI-NEXT:    s_and_b64 vcc, exec, -1
 ; SI-NEXT:  BB2_2: ; %loop2
 ; SI-NEXT:    ; =>This Inner Loop Header: Depth=1
 ; SI-NEXT:    s_waitcnt lgkmcnt(0)
 ; SI-NEXT:    buffer_store_dword v0, off, s[0:3], 0
+; SI-NEXT:    s_mov_b64 vcc, vcc
 ; SI-NEXT:    s_cbranch_vccnz BB2_2
 ; SI-NEXT:  ; %bb.3: ; %Flow
 ; SI-NEXT:    s_mov_b64 s[2:3], 0
 ; SI-NEXT:  BB2_4: ; %Flow2
 ; SI-NEXT:    s_and_b64 vcc, exec, s[2:3]
 ; SI-NEXT:    s_waitcnt lgkmcnt(0)
 ; SI-NEXT:    s_mov_b64 vcc, vcc
 ; SI-NEXT:    s_cbranch_vccz BB2_7
 ; SI-NEXT:  ; %bb.5:
 ; SI-NEXT:    s_mov_b32 s3, 0xf000
 ; SI-NEXT:    s_mov_b32 s2, -1
 ; SI-NEXT:    s_waitcnt expcnt(0)
 ; SI-NEXT:    v_mov_b32_e32 v0, 0x3e7
 ; SI-NEXT:    s_and_b64 vcc, exec, 0
 ; SI-NEXT:  BB2_6: ; %loop1
 ; SI-NEXT:    ; =>This Inner Loop Header: Depth=1
 ; SI-NEXT:    buffer_store_dword v0, off, s[0:3], 0
+; SI-NEXT:    s_mov_b64 vcc, vcc
 ; SI-NEXT:    s_cbranch_vccz BB2_6
 ; SI-NEXT:  BB2_7: ; %DummyReturnBlock
 ; SI-NEXT:    s_endpgm
 ; IR-LABEL: @infinite_loops(
 ; IR-NEXT:  entry:
 ; IR-NEXT:    br i1 undef, label [[LOOP1:%.*]], label [[LOOP2:%.*]]
 ; IR:       loop1:
 ; IR-NEXT:    store volatile i32 999, i32 addrspace(1)* [[OUT:%.*]], align 4

llvm/test/CodeGen/AMDGPU/nested-loop-conditions.ll

 ; GCN-NEXT:    s_and_b64 vcc, exec, s[4:5]
 ; GCN-NEXT:    s_cbranch_vccz BB0_7
 ; GCN-NEXT:  ; %bb.5: ; %bb8
 ; GCN-NEXT:    s_waitcnt lgkmcnt(0)
 ; GCN-NEXT:    ds_read_b32 v0, v0
 ; GCN-NEXT:    s_and_b64 vcc, exec, 0
 ; GCN-NEXT:  BB0_6: ; %bb9
 ; GCN-NEXT:    ; =>This Inner Loop Header: Depth=1
+; GCN-NEXT:    s_mov_b64 vcc, vcc
 ; GCN-NEXT:    s_cbranch_vccz BB0_6
 ; GCN-NEXT:  BB0_7: ; %DummyReturnBlock
 ; GCN-NEXT:    s_endpgm
 ; IR-LABEL: @reduced_nested_loop_conditions(
 ; IR-NEXT:  bb:
 ; IR-NEXT:    [[MY_TMP:%.*]] = tail call i32 @llvm.amdgcn.workitem.id.x() #4
 ; IR-NEXT:    [[MY_TMP1:%.*]] = getelementptr inbounds i64, i64 addrspace(3)* [[ARG:%.*]], i32 [[MY_TMP]]
 ; IR-NEXT:    [[MY_TMP2:%.*]] = load volatile i64, i64 addrspace(3)* [[MY_TMP1]]

llvm/test/CodeGen/AMDGPU/si-annotate-cf.ll

 ; SI-NEXT:    s_cbranch_vccz BB3_13
 ; SI-NEXT:  ; %bb.10: ; %for.cond.preheader
 ; SI-NEXT:    s_cmpk_lt_i32 s8, 0x3e8
 ; SI-NEXT:    s_cbranch_scc0 BB3_13
 ; SI-NEXT:  ; %bb.11: ; %for.body
 ; SI-NEXT:    s_and_b64 vcc, exec, 0
 ; SI-NEXT:  BB3_12: ; %self.loop
 ; SI-NEXT:    ; =>This Inner Loop Header: Depth=1
+; SI-NEXT:    s_mov_b64 vcc, vcc
 ; SI-NEXT:    s_cbranch_vccz BB3_12
 ; SI-NEXT:  BB3_13: ; %DummyReturnBlock
 ; SI-NEXT:    s_endpgm
 ;
 ; FLAT-LABEL: loop_land_info_assert:
 ; FLAT:       ; %bb.0: ; %entry
 ; FLAT-NEXT:    s_mov_b32 s7, 0xf000
 ; FLAT-NEXT:    s_mov_b32 s6, -1

 ; FLAT-NEXT:    s_cbranch_vccz BB3_13
 ; FLAT-NEXT:  ; %bb.10: ; %for.cond.preheader
 ; FLAT-NEXT:    s_cmpk_lt_i32 s8, 0x3e8
 ; FLAT-NEXT:    s_cbranch_scc0 BB3_13
 ; FLAT-NEXT:  ; %bb.11: ; %for.body
 ; FLAT-NEXT:    s_and_b64 vcc, exec, 0
 ; FLAT-NEXT:  BB3_12: ; %self.loop
 ; FLAT-NEXT:    ; =>This Inner Loop Header: Depth=1
+; FLAT-NEXT:    s_mov_b64 vcc, vcc
 ; FLAT-NEXT:    s_cbranch_vccz BB3_12
 ; FLAT-NEXT:  BB3_13: ; %DummyReturnBlock
 ; FLAT-NEXT:    s_endpgm
 entry:
   %cmp = icmp sgt i32 %c0, 0
   br label %while.cond.outer
 
 while.cond.outer:

llvm/test/CodeGen/AMDGPU/si-annotate-cfg-loop-assert.ll

 ; CHECK-NEXT:    v_cmp_eq_u32_e64 s[0:1], s1, 0
 ; CHECK-NEXT:    s_or_b64 s[0:1], s[2:3], s[0:1]
 ; CHECK-NEXT:    s_and_b64 vcc, exec, s[0:1]
 ; CHECK-NEXT:    s_cbranch_vccnz BB0_3
 ; CHECK-NEXT:  ; %bb.1: ; %bb9
 ; CHECK-NEXT:    s_and_b64 vcc, exec, 0
 ; CHECK-NEXT:  BB0_2: ; %bb10
 ; CHECK-NEXT:    ; =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    s_mov_b64 vcc, vcc
 ; CHECK-NEXT:    s_cbranch_vccz BB0_2
 ; CHECK-NEXT:  BB0_3: ; %DummyReturnBlock
 ; CHECK-NEXT:    s_endpgm
 bb:
   %tmp = icmp ne i32 %arg, 0
   %tmp7 = icmp ne i32 %arg1, 0
   %tmp8 = and i1 %tmp, %tmp7
   br i1 %tmp8, label %bb9, label %bb11

llvm/test/CodeGen/AMDGPU/vccz-corrupt-bug-workaround.mir

body: |
   bb.1:
 
 ...
 ---
 # Test vcc definition in a previous basic block.
 
 # CHECK-LABEL: name: vcc_def_pred
 # CHECK: bb.1:
+# SI:  $vcc = S_MOV_B64 $vcc
+# GFX9:  $vcc = S_MOV_B64 $vcc
 # CHECK: S_CBRANCH_VCCZ %bb.2, implicit $vcc
 
 name: vcc_def_pred
 body: |
   bb.0:
     $vcc = S_MOV_B64 0
   bb.1:
     S_CBRANCH_VCCZ %bb.2, implicit $vcc

 ...
 
 # CHECK-LABEL: name: load_def_wait_nop_use
 # SI: S_WAITCNT 0
 # SI-NEXT: $sgpr0 = S_LOAD_DWORD_IMM $sgpr0_sgpr1, 0, 0, 0
 # SI-NEXT: $vcc = S_MOV_B64 0
 # SI-NEXT: S_WAITCNT 127
 # SI-NEXT: S_NOP 0
-# FIXME should have $vcc = S_MOV_B64 $vcc
+# SI-NEXT: $vcc = S_MOV_B64 $vcc
 # SI-NEXT: S_CBRANCH_VCCZ %bb.1, implicit $vcc
 name: load_def_wait_nop_use
 body: |
   bb.0:
     $sgpr0 = S_LOAD_DWORD_IMM $sgpr0_sgpr1, 0, 0, 0
     $vcc = S_MOV_B64 0
     S_WAITCNT 127
     S_NOP 0

 
 ---
 # CHECK-LABEL: name: def_load_wait_nop_use
 # SI: S_WAITCNT 0
 # SI-NEXT: $vcc = S_MOV_B64 0
 # SI-NEXT: $sgpr0 = S_LOAD_DWORD_IMM $sgpr0_sgpr1, 0, 0, 0
 # SI-NEXT: S_WAITCNT 127
 # SI-NEXT: S_NOP 0
-# FIXME should have $vcc = S_MOV_B64 $vcc
+# SI-NEXT: $vcc = S_MOV_B64 $vcc
 # SI-NEXT: S_CBRANCH_VCCZ %bb.1, implicit $vcc
 name: def_load_wait_nop_use
 body: |
   bb.0:
     $vcc = S_MOV_B64 0
     $sgpr0 = S_LOAD_DWORD_IMM $sgpr0_sgpr1, 0, 0, 0
     S_WAITCNT 127
     S_NOP 0