StructurizeCFG: Relax uniformity checks.

[sheredom]

Summary

This change relaxes the checks for hasOnlyUniformBranches such that our region is uniform if:

1. All conditional branches that are direct children are uniform.
2. And either: a. All sub-regions are uniform. b. There is one or less conditional branches among the direct children.

[nhaehnle]

I'm personally not confident enough in this change to enable it unconditionally immediately. I think this should be guarded with an option.

test/CodeGen/AMDGPU/control-flow-optnone.ll

18–19

I think this doesn't fit with what the test is actually trying to test. There's a comment above that says it wants to be changed to optnone. Perhaps better to change the run line accordingly?

[sheredom]

Guarded the change by an option to allow users of the structurizer to turn this functionality on in their own time.

[sheredom]

In D62198#1510400, @nhaehnle wrote:

I'm personally not confident enough in this change to enable it unconditionally immediately. I think this should be guarded with an option.

I've introduced an option and enabled it only for the uniform-regions.ll test.

[nhaehnle]

LGTM

lib/Transforms/Scalar/StructurizeCFG.cpp

for (Instruction &I : *BB) {
         }
         Updater.RewriteUseAfterInsertions(U);
       }
     }
 }
 
 static bool hasOnlyUniformBranches(Region *R, unsigned UniformMDKindID,
                                    const LegacyDivergenceAnalysis &DA) {
+  // Bool for if all sub-regions are uniform.
+  bool SubRegionsAreUniform = true;
+  // Count of how many direct children are conditional.
+  unsigned ConditionalDirectChildren = 0;
+
   for (auto E : R->elements()) {
     if (!E->isSubRegion()) {
       auto Br = dyn_cast<BranchInst>(E->getEntry()->getTerminator());
       if (!Br || !Br->isConditional())
         continue;
 
       if (!DA.isUniform(Br))
         return false;
+
+      // One of our direct children is conditional.
+      ConditionalDirectChildren++;
+
       LLVM_DEBUG(dbgs() << "BB: " << Br->getParent()->getName()
                         << " has uniform terminator\n");
     } else {
       // Explicitly refuse to treat regions as uniform if they have non-uniform
       // subregions. We cannot rely on DivergenceAnalysis for branches in
       // subregions because those branches may have been removed and re-created,
       // so we look for our metadata instead.
       //
       // Warning: It would be nice to treat regions as uniform based only on
       // their direct child basic blocks' terminators, regardless of whether
       // subregions are uniform or not. However, this requires a very careful
       // look at SIAnnotateControlFlow to make sure nothing breaks there.
       for (auto BB : E->getNodeAs<Region>()->blocks()) {
         auto Br = dyn_cast<BranchInst>(BB->getTerminator());
         if (!Br || !Br->isConditional())
           continue;
 
-        if (!Br->getMetadata(UniformMDKindID))
-          return false;
+        if (!Br->getMetadata(UniformMDKindID)) {
+          SubRegionsAreUniform = false;
+          break;
         }
       }
     }
-  return true;
+  }
+
+  // Our region is uniform if:
+  // 1. All conditional branches that are direct children are uniform (checked
+  // above).
+  // 2. And either:
+  //   a. All sub-regions are uniform.
+  //   b. There is one or less conditional branches among the direct children.
+  return SubRegionsAreUniform || (ConditionalDirectChildren <= 1);
 }
 
 /// Run the transformation for each region found
 bool StructurizeCFG::runOnRegion(Region *R, RGPassManager &RGM) {
   if (R->isTopLevelRegion())
     return false;
 
   DA = nullptr;

test/CodeGen/AMDGPU/control-flow-optnone.ll

 ; GCN: s_branch
 
 ; GCN-DAG: v_cmp_lt_i32
 ; GCN-DAG: v_cmp_gt_i32
 ; GCN: s_and_b64
 ; GCN: s_mov_b64 exec
 
 ; GCN: s_or_b64 exec, exec
-; GCN: v_cmp_eq_u32
-; GCN: s_cbranch_vccnz
+; GCN: s_cmp_eq_u32
+; GCN: s_cbranch_scc1
 ; GCN-NEXT: s_branch
 define amdgpu_kernel void @copytoreg_divergent_brcond(i32 %arg, i32 %arg1, i32 %arg2) #0 {
 bb:
   %tmp = tail call i32 @llvm.amdgcn.workitem.id.x()
   %tmp3 = zext i32 %tmp to i64
   %tmp5 = add i64 %tmp3, undef
   %tmp6 = trunc i64 %tmp5 to i32
   %tmp7 = mul nsw i32 %tmp6, %arg2

test/Transforms/StructurizeCFG/AMDGPU/uniform-regions.ll

 
 if:
   br label %end
 
 end:
   ret void
 }
 
+define amdgpu_cs void @uniform_branch_to_nonuniform_subregions(i32 addrspace(4)* %ptr, i32 inreg %data) {
+; CHECK-LABEL: @uniform_branch_to_nonuniform_subregions(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[DATA:%.*]], 42
+; CHECK-NEXT:    br i1 [[C]], label [[UNIFORM_FOR_BODY:%.*]], label [[FOR_BODY:%.*]], !structurizecfg.uniform !0
+; CHECK:       for.body:
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[TMP0:%.*]], [[FLOW1:%.*]] ]
+; CHECK-NEXT:    [[CC:%.*]] = icmp ult i32 [[I]], 4
+; CHECK-NEXT:    br i1 [[CC]], label [[MID_LOOP:%.*]], label [[FLOW1]]
+; CHECK:       mid.loop:
+; CHECK-NEXT:    [[V:%.*]] = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[CC2:%.*]] = icmp eq i32 [[V]], 0
+; CHECK-NEXT:    br i1 [[CC2]], label [[END_LOOP:%.*]], label [[FLOW2:%.*]]
+; CHECK:       Flow1:
+; CHECK-NEXT:    [[TMP0]] = phi i32 [ [[TMP2:%.*]], [[FLOW2]] ], [ undef, [[FOR_BODY]] ]
+; CHECK-NEXT:    [[TMP1:%.*]] = phi i1 [ [[TMP3:%.*]], [[FLOW2]] ], [ true, [[FOR_BODY]] ]
+; CHECK-NEXT:    br i1 [[TMP1]], label [[FOR_END:%.*]], label [[FOR_BODY]]
+; CHECK:       end.loop:
+; CHECK-NEXT:    [[I_INC:%.*]] = add i32 [[I]], 1
+; CHECK-NEXT:    br label [[FLOW2]]
+; CHECK:       Flow2:
+; CHECK-NEXT:    [[TMP2]] = phi i32 [ [[I_INC]], [[END_LOOP]] ], [ undef, [[MID_LOOP]] ]
+; CHECK-NEXT:    [[TMP3]] = phi i1 [ false, [[END_LOOP]] ], [ true, [[MID_LOOP]] ]
+; CHECK-NEXT:    br label [[FLOW1]]
+; CHECK:       for.end:
+; CHECK-NEXT:    br i1 [[CC]], label [[IF:%.*]], label [[FLOW:%.*]]
+; CHECK:       if:
+; CHECK-NEXT:    br label [[FLOW]]
+; CHECK:       uniform.for.body:
+; CHECK-NEXT:    [[UNIFORM_I:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[TMP4:%.*]], [[FLOW4:%.*]] ]
+; CHECK-NEXT:    [[UNIFORM_CC:%.*]] = icmp ult i32 [[UNIFORM_I]], 4
+; CHECK-NEXT:    br i1 [[UNIFORM_CC]], label [[UNIFORM_MID_LOOP:%.*]], label [[FLOW4]]
+; CHECK:       uniform.mid.loop:
+; CHECK-NEXT:    [[UNIFORM_V:%.*]] = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[UNIFORM_CC2:%.*]] = icmp eq i32 [[UNIFORM_V]], 0
+; CHECK-NEXT:    br i1 [[UNIFORM_CC2]], label [[UNIFORM_END_LOOP:%.*]], label [[FLOW5:%.*]]
+; CHECK:       Flow4:
+; CHECK-NEXT:    [[TMP4]] = phi i32 [ [[TMP6:%.*]], [[FLOW5]] ], [ undef, [[UNIFORM_FOR_BODY]] ]
+; CHECK-NEXT:    [[TMP5:%.*]] = phi i1 [ [[TMP7:%.*]], [[FLOW5]] ], [ true, [[UNIFORM_FOR_BODY]] ]
+; CHECK-NEXT:    br i1 [[TMP5]], label [[UNIFORM_FOR_END:%.*]], label [[UNIFORM_FOR_BODY]]
+; CHECK:       uniform.end.loop:
+; CHECK-NEXT:    [[UNIFORM_I_INC:%.*]] = add i32 [[UNIFORM_I]], 1
+; CHECK-NEXT:    br label [[FLOW5]]
+; CHECK:       Flow5:
+; CHECK-NEXT:    [[TMP6]] = phi i32 [ [[UNIFORM_I_INC]], [[UNIFORM_END_LOOP]] ], [ undef, [[UNIFORM_MID_LOOP]] ]
+; CHECK-NEXT:    [[TMP7]] = phi i1 [ false, [[UNIFORM_END_LOOP]] ], [ true, [[UNIFORM_MID_LOOP]] ]
+; CHECK-NEXT:    br label [[FLOW4]]
+; CHECK:       uniform.for.end:
+; CHECK-NEXT:    br i1 [[UNIFORM_CC]], label [[UNIFORM_IF:%.*]], label [[FLOW3:%.*]]
+; CHECK:       uniform.if:
+; CHECK-NEXT:    br label [[FLOW3]]
+; CHECK:       Flow:
+; CHECK-NEXT:    br label [[END:%.*]]
+; CHECK:       Flow3:
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %c = icmp eq i32 %data, 42
+  br i1 %c, label %uniform.for.body, label %for.body
+
+for.body:
+  %i = phi i32 [0, %entry], [%i.inc, %end.loop]
+  %cc = icmp ult i32 %i, 4
+  br i1 %cc, label %mid.loop, label %for.end
+
+mid.loop:
+  %v = call i32 @llvm.amdgcn.workitem.id.x()
+  %cc2 = icmp eq i32 %v, 0
+  br i1 %cc2, label %end.loop, label %for.end
+
+end.loop:
+  %i.inc = add i32 %i, 1
+  br label %for.body
+
+for.end:
+  br i1 %cc, label %if, label %end
+
+if:
+  br label %end
+
+uniform.for.body:
+  %uniform.i = phi i32 [0, %entry], [%uniform.i.inc, %uniform.end.loop]
+  %uniform.cc = icmp ult i32 %uniform.i, 4
+  br i1 %uniform.cc, label %uniform.mid.loop, label %uniform.for.end
+
+uniform.mid.loop:
+  %uniform.v = call i32 @llvm.amdgcn.workitem.id.x()
+  %uniform.cc2 = icmp eq i32 %uniform.v, 0
+  br i1 %uniform.cc2, label %uniform.end.loop, label %uniform.for.end
+
+uniform.end.loop:
+  %uniform.i.inc = add i32 %uniform.i, 1
+  br label %uniform.for.body
+
+uniform.for.end:
+  br i1 %uniform.cc, label %uniform.if, label %end
+
+uniform.if:
+  br label %end
+
+end:
+  ret void
+}
+
 declare i32 @llvm.amdgcn.workitem.id.x()