[AMDGPU] Handle uniform control flow inside divergent

llvm/lib/Target/AMDGPU/SILowerI1Copies.cpp

@@ -147,12 +147,10 @@ class PhiIncomingAnalysis {
         continue;
       }
 
-      ReachableMap.try_emplace(MBB, false);
-
       // If this block has a divergent terminator and the def block is its
       // post-dominator, the wave may first visit the other successors.
       if (TII->hasDivergentBranch(MBB) && PDT.dominates(&DefBlock, MBB))
-        append_range(Stack, MBB->successors());
+        Stack.push_back(MBB);
     }
 
     while (!Stack.empty()) {
@@ -161,6 +159,12 @@ class PhiIncomingAnalysis {
         append_range(Stack, MBB->successors());
     }
 
+    // Insert remaining incoming blocks.
+    for (auto Incoming : Incomings) {
+      MachineBasicBlock *MBB = Incoming.Block;
+      ReachableMap.try_emplace(MBB, false);
+    }
+
     for (auto &[MBB, IsSource] : ReachableMap) {
       bool HaveReachablePred = false;
       for (MachineBasicBlock *Pred : MBB->predecessors()) {

llvm/test/CodeGen/AMDGPU/uniform-inside-divergent-cfg.ll

@@ -0,0 +1,163 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc -mtriple=amdgpu11.00 -structurizecfg-skip-uniform-regions < %s | FileCheck %s
+
+define amdgpu_kernel void @div_unif_div(ptr addrspace(1) %out, float %ubeta, i32 %m) {
+; CHECK-LABEL: div_unif_div:
+; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:    s_load_b64 s[0:1], s[4:5], 0x2c
+; CHECK-NEXT:    v_dual_mov_b32 v1, 0 :: v_dual_and_b32 v0, 0x3ff, v0
+; CHECK-NEXT:    s_waitcnt lgkmcnt(0)
+; CHECK-NEXT:    s_delay_alu instid0(VALU_DEP_1)
+; CHECK-NEXT:    v_cmp_gt_u32_e32 vcc_lo, s1, v0
+; CHECK-NEXT:    s_mov_b32 s1, 0
+; CHECK-NEXT:    s_and_saveexec_b32 s2, vcc_lo
+; CHECK-NEXT:    s_cbranch_execz .LBB0_8
+; CHECK-NEXT:  ; %bb.1: ; %A
+; CHECK-NEXT:    v_cmp_eq_f32_e64 s0, s0, 0
+; CHECK-NEXT:    s_mov_b32 s1, exec_lo
+; CHECK-NEXT:    s_and_b32 vcc_lo, exec_lo, s0
+; CHECK-NEXT:    s_cbranch_vccnz .LBB0_7
+; CHECK-NEXT:  ; %bb.2: ; %B
+; CHECK-NEXT:    s_mov_b32 s0, exec_lo
+; CHECK-NEXT:    ; implicit-def: $vgpr1
+; CHECK-NEXT:    v_cmpx_lt_u32_e32 2, v0
+; CHECK-NEXT:    s_xor_b32 s0, exec_lo, s0
+; CHECK-NEXT:  ; %bb.3: ; %B2
+; CHECK-NEXT:    v_mul_u32_u24_e32 v1, 3, v0
+; CHECK-NEXT:  ; %bb.4: ; %Flow
+; CHECK-NEXT:    s_and_not1_saveexec_b32 s0, s0
+; CHECK-NEXT:  ; %bb.5: ; %B1
+; CHECK-NEXT:    v_add_nc_u32_e32 v1, 7, v0
+; CHECK-NEXT:  ; %bb.6: ; %Flow3
+; CHECK-NEXT:    s_or_b32 exec_lo, exec_lo, s0
+; CHECK-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
+; CHECK-NEXT:    s_or_b32 s1, s1, exec_lo
+; CHECK-NEXT:    s_branch .LBB0_8
+; CHECK-NEXT:  .LBB0_7:
+; CHECK-NEXT:    v_mov_b32_e32 v1, 1
+; CHECK-NEXT:  .LBB0_8: ; %join
+; CHECK-NEXT:    s_or_b32 exec_lo, exec_lo, s2
+; CHECK-NEXT:    s_and_saveexec_b32 s0, s1
+; CHECK-NEXT:    s_cbranch_execz .LBB0_10
+; CHECK-NEXT:  ; %bb.9: ; %do.store
+; CHECK-NEXT:    s_load_b64 s[0:1], s[4:5], 0x24
+; CHECK-NEXT:    v_lshlrev_b32_e32 v0, 2, v0
+; CHECK-NEXT:    s_waitcnt lgkmcnt(0)
+; CHECK-NEXT:    global_store_b32 v0, v1, s[0:1]
+; CHECK-NEXT:  .LBB0_10: ; %exit
+; CHECK-NEXT:    s_endpgm
+entry:
+  %tid = call i32 @llvm.amdgcn.workitem.id.x()
+  %div = icmp ult i32 %tid, %m                 ; DIVERGENT -> becomes SI_IF
+  br i1 %div, label %A, label %join
+
+A:                                             ; only lanes with tid < m
+  %unif = fcmp une float %ubeta, 0.0           ; UNIFORM -> region gets skipped
+  br i1 %unif, label %B, label %join
+
+B:
+  %div2 = icmp ult i32 %tid, 3                 ; divergent inner if/else
+  br i1 %div2, label %B1, label %B2
+B1:
+  %v1 = add i32 %tid, 7
+  br label %C
+B2:
+  %v2 = mul i32 %tid, 3
+  br label %C
+C:
+  %v = phi i32 [ %v1, %B1 ], [ %v2, %B2 ]
+  br label %join
+
+join:
+  %val  = phi i32 [ 0, %entry ], [ 1, %A ], [ %v, %C ]
+  %pred = phi i1  [ false, %entry ], [ true, %A ], [ true, %C ]
+  br i1 %pred, label %do.store, label %exit
+
+do.store:
+  %gep = getelementptr i32, ptr addrspace(1) %out, i32 %tid
+  store i32 %val, ptr addrspace(1) %gep
+  br label %exit
+
+exit:
+  ret void
+}
+
+define amdgpu_kernel void @unif_div(ptr addrspace(1) %out, float %u1, float %u2, i32 %m) {
+; CHECK-LABEL: unif_div:
+; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:    s_load_b128 s[0:3], s[4:5], 0x2c
+; CHECK-NEXT:    v_dual_mov_b32 v1, 0 :: v_dual_and_b32 v0, 0x3ff, v0
+; CHECK-NEXT:    s_waitcnt lgkmcnt(0)
+; CHECK-NEXT:    s_mov_b32 s3, 0
+; CHECK-NEXT:    s_delay_alu instid0(VALU_DEP_1)
+; CHECK-NEXT:    v_cmp_gt_u32_e32 vcc_lo, s2, v0
+; CHECK-NEXT:    s_and_saveexec_b32 s2, vcc_lo
+; CHECK-NEXT:    s_cbranch_execz .LBB1_7
+; CHECK-NEXT:  ; %bb.1: ; %A
+; CHECK-NEXT:    v_cmp_eq_f32_e64 s0, s0, 0
+; CHECK-NEXT:    s_and_b32 vcc_lo, exec_lo, s0
+; CHECK-NEXT:    s_cbranch_vccnz .LBB1_4
+; CHECK-NEXT:  ; %bb.2: ; %B
+; CHECK-NEXT:    v_cmp_eq_f32_e64 s0, s1, 0
+; CHECK-NEXT:    s_and_b32 vcc_lo, exec_lo, s0
+; CHECK-NEXT:    s_cbranch_vccnz .LBB1_5
+; CHECK-NEXT:  ; %bb.3: ; %C
+; CHECK-NEXT:    v_add_nc_u32_e32 v1, 7, v0
+; CHECK-NEXT:    s_branch .LBB1_6
+; CHECK-NEXT:  .LBB1_4:
+; CHECK-NEXT:    s_mov_b32 s3, exec_lo
+; CHECK-NEXT:    v_mov_b32_e32 v1, 1
+; CHECK-NEXT:    s_branch .LBB1_7
+; CHECK-NEXT:  .LBB1_5: ; %D
+; CHECK-NEXT:    v_mul_u32_u24_e32 v1, 3, v0
+; CHECK-NEXT:  .LBB1_6: ; %join
+; CHECK-NEXT:    s_or_b32 s3, exec_lo, exec_lo
+; CHECK-NEXT:  .LBB1_7: ; %join
+; CHECK-NEXT:    s_or_b32 exec_lo, exec_lo, s2
+; CHECK-NEXT:    s_and_saveexec_b32 s0, s3
+; CHECK-NEXT:    s_cbranch_execz .LBB1_9
+; CHECK-NEXT:  ; %bb.8: ; %do.store
+; CHECK-NEXT:    s_load_b64 s[0:1], s[4:5], 0x24
+; CHECK-NEXT:    v_lshlrev_b32_e32 v0, 2, v0
+; CHECK-NEXT:    s_waitcnt lgkmcnt(0)
+; CHECK-NEXT:    global_store_b32 v0, v1, s[0:1]
+; CHECK-NEXT:  .LBB1_9: ; %exit
+; CHECK-NEXT:    s_endpgm
+entry:
+  %tid = call i32 @llvm.amdgcn.workitem.id.x()
+  %div = icmp ult i32 %tid, %m                 ; DIVERGENT -> SI_IF
+  br i1 %div, label %A, label %join
+
+A:                                             ; only lanes with tid < m
+  %c1 = fcmp une float %u1, 0.0                ; uniform
+  br i1 %c1, label %B, label %join
+
+B:
+  %c2 = fcmp une float %u2, 0.0                ; uniform
+  br i1 %c2, label %C, label %D
+
+C:
+  %vc = add i32 %tid, 7
+  br label %E
+
+D:
+  %vd = mul i32 %tid, 3
+  br label %E
+
+E:                                             ; not a direct successor of A
+  %v = phi i32 [ %vc, %C ], [ %vd, %D ]
+  br label %join
+
+join:
+  %val  = phi i32 [ 0, %entry ], [ 1, %A ], [ %v, %E ]
+  %pred = phi i1  [ false, %entry ], [ true, %A ], [ true, %E ]
+  br i1 %pred, label %do.store, label %exit
+
+do.store:
+  %gep = getelementptr i32, ptr addrspace(1) %out, i32 %tid
+  store i32 %val, ptr addrspace(1) %gep
+  br label %exit
+
+exit:
+  ret void
+}