AMDGPU/SI: Handle infinite loop for the structurizer to work with CFG with infinite loops.

[cfang]

Summary

The current StructurizeCFG pass only works for CFG with one exit. AMDGPUUnifyDivergentExitNodes combines multiple "return" blocks and/or "unreachable" blocks
to one exit block for the Structurizer to work. However, infinite loop is another kind of special "exit", and if we don't handle it, the case of multiple exits will prevent the structurizer from working.

In this work, for each infinite loop, we add a dummy edge to the "return" block, and thus the AMDGPUUnifyDivergentExitNodes pass will work with infinite loops.
This will make CFG with infinite loops be structurized.

[nhaehnle]

So this does not cover all possible cases of "irreducible" infinite loops, e.g. consider the case with basic blocks A and B where both end with a conditional branch that can go to either A or B. I.e., there are no unconditional branches in the loop.

A more robust approach would be to also transform conditional branches like this:

br i1 %cc, label %A, label %B

becomes

  br i1 true, label %DummyReturnBB, label %local_dummy
local_dummy:
  br i1 %cc, label %A, label %b

lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

194 ↗

(On Diff #144805)

Shouldn't this use eraseFromParent()?

[cfang]

Handle the case that the infinite loop is controlled by a conditional branch. In such case, the two edges of the branch
are both backedges.

We create a "TransitionBB" to hold the conditional branch. And at the place of the original conditional branch,
we create a new conditional branch with the conditional always true, the the true edge connects to the TransitionBB, and the false edge to
the dummy return block.

Thanks to Nicolai for pointing out the case as well as the approach.

[cfang]

In D46340#1086444, @nhaehnle wrote:

So this does not cover all possible cases of "irreducible" infinite loops, e.g. consider the case with basic blocks A and B where both end with a conditional branch that can go to either A or B. I.e., there are no unconditional branches in the loop.

A more robust approach would be to also transform conditional branches like this:

br i1 %cc, label %A, label %B

becomes

  br i1 true, label %DummyReturnBB, label %local_dummy
local_dummy:
  br i1 %cc, label %A, label %b

Thank you so much for pointing out this case and providing the algorithm. Done as suggested.

lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

194 ↗

(On Diff #144805)

Right. There are other places in the same file need to do the same thing using eraseFromParent(). Will do in a separate patch.

[nhaehnle]

Thanks. This change introduces more branches, but since it should only do so in an infinite loop it should be fine.

I did notice some minor style problems though.

lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

186–192 ↗

(On Diff #145114)

The indentation looks wrong here. Please check with clang-format/clang-format-diff (I prefer to use clang-format-diff, because otherwise you're quite likely to get notifications about formatting errors in code that you didn't actually change).

199 ↗

(On Diff #145114)

Typo: *conditional

[cfang]

lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

186–192 ↗

(On Diff #145114)

Nice catch! Corrected. Thanks.

199 ↗

(On Diff #145114)

Thanks.

[cfang]

Correct indentation error and typo.

[cfang]

Any additional comments/suggestions? Thanks!

[cfang]

Ping.

We need this fix to unblock some important tasks. So if there is no additional comments, we need the permission
to integrate the patch. Thanks.

[nhaehnle]

LGTM

llvm/trunk/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

bool AMDGPUUnifyDivergentExitNodes::runOnFunction(Function &F) {
 
   DivergenceAnalysis &DA = getAnalysis<DivergenceAnalysis>();
 
   // Loop over all of the blocks in a function, tracking all of the blocks that
   // return.
   SmallVector<BasicBlock *, 4> ReturningBlocks;
   SmallVector<BasicBlock *, 4> UnreachableBlocks;
 
+  // Dummy return block for infinite loop.
+  BasicBlock *DummyReturnBB = nullptr;
+
   for (BasicBlock *BB : PDT.getRoots()) {
     if (isa<ReturnInst>(BB->getTerminator())) {
       if (!isUniformlyReached(DA, *BB))
         ReturningBlocks.push_back(BB);
     } else if (isa<UnreachableInst>(BB->getTerminator())) {
       if (!isUniformlyReached(DA, *BB))
         UnreachableBlocks.push_back(BB);
+    } else if (BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator())) {
+
+      ConstantInt *BoolTrue = ConstantInt::getTrue(F.getContext());
+      if (DummyReturnBB == nullptr) {
+        DummyReturnBB = BasicBlock::Create(F.getContext(),
+                                           "DummyReturnBlock", &F);
+        Type *RetTy = F.getReturnType();
+        Value *RetVal = RetTy->isVoidTy() ? nullptr : UndefValue::get(RetTy);
+        ReturnInst::Create(F.getContext(), RetVal, DummyReturnBB);
+        ReturningBlocks.push_back(DummyReturnBB);
+      }
+
+      if (BI->isUnconditional()) {
+        BasicBlock *LoopHeaderBB = BI->getSuccessor(0);
+        BI->eraseFromParent(); // Delete the unconditional branch.
+        // Add a new conditional branch with a dummy edge to the return block.
+        BranchInst::Create(LoopHeaderBB, DummyReturnBB, BoolTrue, BB);
+      } else { // Conditional branch.
+        // Create a new transition block to hold the conditional branch.
+        BasicBlock *TransitionBB = BasicBlock::Create(F.getContext(),
+                                                      "TransitionBlock", &F);
+
+        // Move BI from BB to the new transition block.
+        BI->removeFromParent();
+        TransitionBB->getInstList().push_back(BI);
+
+        // Create a branch that will always branch to the transition block.
+        BranchInst::Create(TransitionBB, DummyReturnBB, BoolTrue, BB);
+      }
     }
   }
 
   if (!UnreachableBlocks.empty()) {
     BasicBlock *UnreachableBlock = nullptr;
 
     if (UnreachableBlocks.size() == 1) {
       UnreachableBlock = UnreachableBlocks.front();

llvm/trunk/test/CodeGen/AMDGPU/branch-relaxation.ll

endif:
   ; This call makes the return block long enough that it doesn't get copied.
   call void @llvm.amdgcn.s.sleep(i32 5);
   ret void
 }
 
 ; si_mask_branch
 
 ; GCN-LABEL: {{^}}analyze_mask_branch:
-; GCN: v_cmp_lt_f32_e32 vcc
-; GCN-NEXT: s_and_saveexec_b64 [[MASK:s\[[0-9]+:[0-9]+\]]], vcc
+; GCN: v_cmp_nlt_f32_e32 vcc
+; GCN-NEXT: s_and_saveexec_b64 [[TEMP_MASK:s\[[0-9]+:[0-9]+\]]], vcc
+; GCN-NEXT: s_xor_b64  [[MASK:s\[[0-9]+:[0-9]+\]]], exec, [[TEMP_MASK]]
+; GCN-NEXT: ; mask branch [[FLOW:BB[0-9]+_[0-9]+]]
+
+; GCN: [[FLOW]]: ; %Flow
+; GCN-NEXT: s_or_saveexec_b64 [[TEMP_MASK1:s\[[0-9]+:[0-9]+\]]], [[MASK]]
+; GCN-NEXT: s_xor_b64 exec, exec, [[TEMP_MASK1]]
 ; GCN-NEXT: ; mask branch [[RET:BB[0-9]+_[0-9]+]]
 
-; GCN-NEXT: [[LOOP_BODY:BB[0-9]+_[0-9]+]]: ; %loop_body
+; GCN: [[LOOP_BODY:BB[0-9]+_[0-9]+]]: ; %loop_body
 ; GCN: ;;#ASMSTART
 ; GCN: v_nop_e64
 ; GCN: v_nop_e64
 ; GCN: v_nop_e64
 ; GCN: v_nop_e64
 ; GCN: v_nop_e64
 ; GCN: v_nop_e64
 ; GCN: ;;#ASMEND
+; GCN: s_cbranch_vccz [[RET]]
 
 ; GCN-NEXT: [[LONGBB:BB[0-9]+_[0-9]+]]: ; %loop_body
 ; GCN-NEXT: ; in Loop: Header=[[LOOP_BODY]] Depth=1
 ; GCN-NEXT: s_getpc_b64 vcc
 ; GCN-NEXT: s_sub_u32 vcc_lo, vcc_lo, ([[LONGBB]]+4)-[[LOOP_BODY]]
 ; GCN-NEXT: s_subb_u32 vcc_hi, vcc_hi, 0
 ; GCN-NEXT: s_setpc_b64 vcc
 
-; GCN-NEXT: [[RET]]: ; %ret
-; GCN-NEXT: s_or_b64 exec, exec, [[MASK]]
-; GCN: buffer_store_dword
+; GCN-NEXT: [[RET]]: ; %UnifiedReturnBlock
 ; GCN-NEXT: s_endpgm
 define amdgpu_kernel void @analyze_mask_branch() #0 {
 entry:
   %reg = call float asm sideeffect "v_mov_b32_e64 $0, 0", "=v"()
   %cmp0 = fcmp ogt float %reg, 0.000000e+00
   br i1 %cmp0, label %loop, label %ret
 
 loop:

llvm/trunk/test/CodeGen/AMDGPU/cf-loop-on-constant.ll

 ; RUN: llc -march=amdgcn -verify-machineinstrs < %s | FileCheck -check-prefix=GCN %s
 ; RUN: llc -march=amdgcn -verify-machineinstrs -O0 < %s
 
 ; GCN-LABEL: {{^}}test_loop:
+; GCN: s_and_b64 vcc, exec, -1
 ; GCN: [[LABEL:BB[0-9+]_[0-9]+]]: ; %for.body{{$}}
 ; GCN: ds_read_b32
 ; GCN: ds_write_b32
-; GCN: s_branch [[LABEL]]
+; GCN: s_cbranch_vccnz [[LABEL]]
 ; GCN: s_endpgm
 define amdgpu_kernel void @test_loop(float addrspace(3)* %ptr, i32 %n) nounwind {
 entry:
   %cmp = icmp eq i32 %n, -1
   br i1 %cmp, label %for.exit, label %for.body
 
 for.exit:
   ret void

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

 ; RUN: llc -march=amdgcn -verify-machineinstrs < %s | FileCheck -check-prefix=SI %s
-; RUN: llc -march=amdgcn -mcpu=tonga -mattr=-flat-for-global -verify-machineinstrs < %s | FileCheck -check-prefix=SI %s
+; RUN: opt -mtriple=amdgcn-- -S -amdgpu-unify-divergent-exit-nodes -verify %s | FileCheck -check-prefix=IR %s
 
 ; SI-LABEL: {{^}}infinite_loop:
 ; SI: v_mov_b32_e32 [[REG:v[0-9]+]], 0x3e7
-; SI: BB0_1:
+; SI: [[LOOP:BB[0-9]+_[0-9]+]]:  ; %loop
 ; SI: s_waitcnt lgkmcnt(0)
 ; SI: buffer_store_dword [[REG]]
-; SI: s_branch BB0_1
+; SI: s_branch [[LOOP]]
 define amdgpu_kernel void @infinite_loop(i32 addrspace(1)* %out) {
 entry:
-  br label %for.body
+  br label %loop
 
-for.body:                                         ; preds = %entry, %for.body
+loop:
   store i32 999, i32 addrspace(1)* %out, align 4
-  br label %for.body
+  br label %loop
 }
 
+
+; IR-LABEL: @infinite_loop_ret(
+; IR:  br i1 %cond, label %loop, label %UnifiedReturnBlock
+
+; IR: loop:
+; IR: store i32 999, i32 addrspace(1)* %out, align 4
+; IR: br i1 true, label %loop, label %UnifiedReturnBlock
+
+; IR: UnifiedReturnBlock:
+; IR:  ret void
+
+
+; SI-LABEL: {{^}}infinite_loop_ret:
+; SI: s_cbranch_execz [[RET:BB[0-9]+_[0-9]+]]
+
+; SI: v_mov_b32_e32 [[REG:v[0-9]+]], 0x3e7
+; SI: [[LOOP:BB[0-9]+_[0-9]+]]:  ; %loop
+; SI: s_and_b64 vcc, exec, -1
+; SI: s_waitcnt lgkmcnt(0)
+; SI: buffer_store_dword [[REG]]
+; SI: s_cbranch_vccnz [[LOOP]]
+
+; SI: [[RET]]:  ; %UnifiedReturnBlock
+; SI: s_endpgm
+define amdgpu_kernel void @infinite_loop_ret(i32 addrspace(1)* %out) {
+entry:
+  %tmp = tail call i32 @llvm.amdgcn.workitem.id.x()
+  %cond = icmp eq i32 %tmp, 1
+  br i1 %cond, label %loop, label %return
+
+loop:
+  store i32 999, i32 addrspace(1)* %out, align 4
+  br label %loop
+
+return:
+  ret void
+}
+
+
+; IR-LABEL: @infinite_loops(
+; IR: br i1 undef, label %loop1, label %loop2
+
+; IR: loop1:
+; IR: store i32 999, i32 addrspace(1)* %out, align 4
+; IR: br i1 true, label %loop1, label %DummyReturnBlock
+
+; IR: loop2:
+; IR: store i32 888, i32 addrspace(1)* %out, align 4
+; IR: br i1 true, label %loop2, label %DummyReturnBlock
+
+; IR: DummyReturnBlock:
+; IR: ret void
+
+
+; SI-LABEL: {{^}}infinite_loops:
+
+; SI: v_mov_b32_e32 [[REG1:v[0-9]+]], 0x3e7
+; SI: s_and_b64 vcc, exec, -1
+
+; SI: [[LOOP1:BB[0-9]+_[0-9]+]]:  ; %loop1
+; SI: s_waitcnt lgkmcnt(0)
+; SI: buffer_store_dword [[REG1]]
+; SI: s_cbranch_vccnz [[LOOP1]]
+; SI: s_branch [[RET:BB[0-9]+_[0-9]+]]
+
+; SI: v_mov_b32_e32 [[REG2:v[0-9]+]], 0x378
+; SI: s_and_b64 vcc, exec, -1
+
+; SI: [[LOOP2:BB[0-9]+_[0-9]+]]:  ; %loop2
+; SI: s_waitcnt lgkmcnt(0)
+; SI: buffer_store_dword [[REG2]]
+; SI: s_cbranch_vccnz [[LOOP2]]
+
+; SI: [[RET]]:  ; %DummyReturnBlock
+; SI: s_endpgm
+define amdgpu_kernel void @infinite_loops(i32 addrspace(1)* %out) {
+entry:
+  br i1 undef, label %loop1, label %loop2
+
+loop1:
+  store i32 999, i32 addrspace(1)* %out, align 4
+  br label %loop1
+
+loop2:
+  store i32 888, i32 addrspace(1)* %out, align 4
+  br label %loop2
+}
+
+
+
+; IR-LABEL: @infinite_loop_nest_ret(
+; IR: br i1 %cond1, label %outer_loop, label %UnifiedReturnBlock
+
+; IR: outer_loop:
+; IR: br label %inner_loop
+
+; IR: inner_loop:
+; IR: store i32 999, i32 addrspace(1)* %out, align 4
+; IR: %cond3 = icmp eq i32 %tmp, 3
+; IR: br i1 true, label %TransitionBlock, label %UnifiedReturnBlock
+
+; IR: TransitionBlock:
+; IR: br i1 %cond3, label %inner_loop, label %outer_loop
+
+; IR: UnifiedReturnBlock:
+; IR: ret void
+
+; SI-LABEL: {{^}}infinite_loop_nest_ret:
+; SI: s_cbranch_execz [[RET:BB[0-9]+_[0-9]+]]
+
+; SI: s_mov_b32
+; SI: [[OUTER_LOOP:BB[0-9]+_[0-9]+]]:  ; %outer_loop
+
+; SI: [[INNER_LOOP:BB[0-9]+_[0-9]+]]:  ; %inner_loop
+; SI: s_waitcnt expcnt(0)
+; SI: v_mov_b32_e32 [[REG:v[0-9]+]], 0x3e7
+; SI: v_cmp_ne_u32_e32
+; SI: s_waitcnt lgkmcnt(0)
+; SI: buffer_store_dword [[REG]]
+
+; SI: s_andn2_b64 exec
+; SI: s_cbranch_execnz [[INNER_LOOP]]
+
+; SI: s_andn2_b64 exec
+; SI: s_cbranch_execnz [[OUTER_LOOP]]
+
+; SI: [[RET]]:  ; %UnifiedReturnBlock
+; SI: s_endpgm
+define amdgpu_kernel void @infinite_loop_nest_ret(i32 addrspace(1)* %out) {
+entry:
+  %tmp = tail call i32 @llvm.amdgcn.workitem.id.x()
+  %cond1 = icmp eq i32 %tmp, 1
+  br i1 %cond1, label %outer_loop, label %return
+
+outer_loop:
+ ; %cond2 = icmp eq i32 %tmp, 2
+ ; br i1 %cond2, label %outer_loop, label %inner_loop
+ br label %inner_loop
+
+inner_loop:                                     ; preds = %LeafBlock, %LeafBlock1
+  store i32 999, i32 addrspace(1)* %out, align 4
+  %cond3 = icmp eq i32 %tmp, 3
+  br i1 %cond3, label %inner_loop, label %outer_loop
+
+return:
+  ret void
+}
+
+declare i32 @llvm.amdgcn.workitem.id.x()

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

 ; FIXME: Should fold to unconditional branch?
 ; GCN: ; implicit-def
 ; GCN: s_cbranch_vccz
 
 ; GCN: ds_read_b32
 
 ; GCN: [[BB9:BB[0-9]+_[0-9]+]]: ; %bb9
 ; GCN-NEXT: ; =>This Inner Loop Header: Depth=1
-; GCN-NEXT: s_branch [[BB9]]
+; GCN-NEXT: s_cbranch_vccnz [[BB9]]
 define amdgpu_kernel void @reduced_nested_loop_conditions(i64 addrspace(3)* nocapture %arg) #0 {
 bb:
   %tmp = tail call i32 @llvm.amdgcn.workitem.id.x() #1
   %tmp1 = getelementptr inbounds i64, i64 addrspace(3)* %arg, i32 %tmp
   %tmp2 = load volatile i64, i64 addrspace(3)* %tmp1
   br label %bb5
 
 bb3:                                              ; preds = %bb9

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

 ; FUNC-LABEL: {{^}}loop_land_info_assert:
 ; SI: s_cmp_gt_i32
 ; SI-NEXT: s_cbranch_scc0 [[ENDPGM:BB[0-9]+_[0-9]+]]
 
 ; SI: s_cmpk_gt_i32
 ; SI-NEXT: s_cbranch_scc1 [[ENDPGM]]
 
 ; SI: [[INFLOOP:BB[0-9]+_[0-9]+]]
-; SI: s_branch [[INFLOOP]]
+; SI: s_cbranch_vccnz [[INFLOOP]]
 
 ; SI: [[ENDPGM]]:
 ; SI: s_endpgm
 define amdgpu_kernel void @loop_land_info_assert(i32 %c0, i32 %c1, i32 %c2, i32 %c3, i32 %x, i32 %y, i1 %arg) nounwind {
 entry:
   %cmp = icmp sgt i32 %c0, 0
   br label %while.cond.outer