AMDGPU: Fix divergence analysis of control flow intrinsics

[arsenm]

Summary

The mask results of these should be uniform. The trickier part is the
dummy booleans used as IR glue need to be treated as divergent. This
should make the divergence analysis results correct for the IR the DAG
is constructed from.

This should allow us to eliminate requiresUniformRegister, which has
an expensive, recursive scan over all users looking for control flow
intrinsics. This should avoid recent compile time regressions.

[rampitec]

Let's Alex review, but could you run a PSDB in the meanwhile?

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

712–739

I think this works for the def cases, but the uses don't currently have the analog in DivergenceAnalysis for values that have to be uniform uses

[nhaehnle]

This looks plausible enough, but feels like the kind of change that should ideally be run through a thorough set of real execution checks.

[arsenm]

Now with more passing

[arsenm]

Split removal into separate patch

[arsenm]

I think the only obstacle to eliminating requiresUniformRegister is due to the treatment of phis with always uniform inputs. This case for example, the LCSSA phi in the return block was incorrectly concluded to be divergent, despite there only being one always uniform input

Printing analysis 'Legacy Divergence Analysis' for function 'atomic_nand_i32_lds':
DIVERGENT: i32 addrspace(3)* %ptr

           :
DIVERGENT:       %1 = load i32, i32 addrspace(3)* %ptr, align 4
                 br label %atomicrmw.start

           atomicrmw.start:
                 %phi.broken = phi i64 [ %4, %atomicrmw.start ], [ 0, %0 ]
DIVERGENT:       %loaded = phi i32 [ %1, %0 ], [ %newloaded, %atomicrmw.start ]
DIVERGENT:       %2 = and i32 %loaded, 4
DIVERGENT:       %new = xor i32 %2, -1
DIVERGENT:       %3 = cmpxchg i32 addrspace(3)* %ptr, i32 %loaded, i32 %new seq_cst seq_cst
DIVERGENT:       %success = extractvalue { i32, i1 } %3, 1
DIVERGENT:       %newloaded = extractvalue { i32, i1 } %3, 0
                 %4 = call i64 @llvm.amdgcn.if.break.i64(i1 %success, i64 %phi.broken)
DIVERGENT:       %5 = call i1 @llvm.amdgcn.loop.i64(i64 %4)
DIVERGENT:       br i1 %5, label %atomicrmw.end, label %atomicrmw.start

           atomicrmw.end:
DIVERGENT:       %newloaded.lcssa = phi i32 [ %newloaded, %atomicrmw.start ]
DIVERGENT:       %.lcssa = phi i64 [ %4, %atomicrmw.start ]
DIVERGENT:       call void @llvm.amdgcn.end.cf.i64(i64 %.lcssa)
DIVERGENT:       ret i32 %newloaded.lcssa

[alex-t]

I agree that current expensive walk on IR should be removed.
I initially consider it as a temporary solution that allows to switch to the new concept
of assignment correct register classes keeping in mind to eventually find another way to workaround CF results exceptions.

What I conceptually don't like here is mixing the "divergence" and register class notions.
The mask produced by the CF intrinsics is always scalar - not same as always uniform.
The isAlwaysUniform is queried by the DA in the divergence propagation.
I'm afraid we can get some subtle bugs claiming CF intrinsic "result 1" always uniform.
Although I cannot yet provide any example to illustrate my concerns.

I would suggest to run extended testing: MESA tests, BLAS tests.... something else that is tremendous enough.

[arsenm]

In D73815#1857088, @alex-t wrote:

I agree that current expensive walk on IR should be removed.
I initially consider it as a temporary solution that allows to switch to the new concept
of assignment correct register classes keeping in mind to eventually find another way to workaround CF results exceptions.

What I conceptually don't like here is mixing the "divergence" and register class notions.
The mask produced by the CF intrinsics is always scalar - not same as always uniform.

Scalar is the only modeled concept of uniform we have. We don't currently try to model workgroup uniform, so these should be the same

[alex-t]

In D73815#1857100, @arsenm wrote:

In D73815#1857088, @alex-t wrote:

I agree that current expensive walk on IR should be removed.
I initially consider it as a temporary solution that allows to switch to the new concept
of assignment correct register classes keeping in mind to eventually find another way to workaround CF results exceptions.

What I conceptually don't like here is mixing the "divergence" and register class notions.
The mask produced by the CF intrinsics is always scalar - not same as always uniform.

Scalar is the only modeled concept of uniform we have. We don't currently try to model workgroup uniform, so these should be the same

It should work if we can guarantee that the values produced by the CF intrinsics are only used by the another CF intrinsics.
I mean something like AND/OR the uniform mask with some divergent value and then feeding result to another CF...
but normally this should not happen.
Once again: I like getting rid of requiresUniformRegister but insist on the extended testing.
Also, the trouble with LCSSA use of the uniform mask out of the divergent loop need to be solved anyway.

[arsenm]

In D73815#1857176, @alex-t wrote:

In D73815#1857100, @arsenm wrote:

In D73815#1857088, @alex-t wrote:

I agree that current expensive walk on IR should be removed.
I initially consider it as a temporary solution that allows to switch to the new concept
of assignment correct register classes keeping in mind to eventually find another way to workaround CF results exceptions.

What I conceptually don't like here is mixing the "divergence" and register class notions.
The mask produced by the CF intrinsics is always scalar - not same as always uniform.

Scalar is the only modeled concept of uniform we have. We don't currently try to model workgroup uniform, so these should be the same

It should work if we can guarantee that the values produced by the CF intrinsics are only used by the another CF intrinsics.
I mean something like AND/OR the uniform mask with some divergent value and then feeding result to another CF...
but normally this should not happen.

In no context should this ever happen. The control flow intrinsics are not general use, and are just a hack to get some control flow representation through the DAG. We should not be trying to worry about handling some bug case where somehow a divergent operation was introduced between a control flow intrinsic and the user. That is a bug and it should break. We shouldn't need to bend over backwards to support this case.

Once again: I like getting rid of requiresUniformRegister but insist on the extended testing.
Also, the trouble with LCSSA use of the uniform mask out of the divergent loop need to be solved anyway.

I've split that into an independent problem. The control flow intrinsics can still have correct divergence information without yet ripping out requiresUniformRegister

[arsenm]

096cd991ee90875603a9cacf3b460ac677258539

llvm/lib/Target/AMDGPU/AMDGPUSearchableTables.td

 def : SourceOfDivergence<int_amdgcn_mfma_f32_32x32x2f32>;
 def : SourceOfDivergence<int_amdgcn_mfma_f32_32x32x4f16>;
 def : SourceOfDivergence<int_amdgcn_mfma_f32_32x32x8f16>;
 def : SourceOfDivergence<int_amdgcn_mfma_i32_32x32x4i8>;
 def : SourceOfDivergence<int_amdgcn_mfma_i32_32x32x8i8>;
 def : SourceOfDivergence<int_amdgcn_mfma_f32_32x32x2bf16>;
 def : SourceOfDivergence<int_amdgcn_mfma_f32_32x32x4bf16>;
 
+// The dummy boolean output is divergent from the IR's perspective,
+// but the mask results are uniform. These produce a divergent and
+// uniform result, so the returned struct is collectively divergent.
+// isAlwaysUniform can override the extract of the uniform component.
+def : SourceOfDivergence<int_amdgcn_if>;
+def : SourceOfDivergence<int_amdgcn_else>;
+def : SourceOfDivergence<int_amdgcn_loop>;
+
 foreach intr = AMDGPUImageDimAtomicIntrinsics in
 def : SourceOfDivergence<intr>;

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

bool GCNTTIImpl::isAlwaysUniform(const Value *V) const {
   if (const IntrinsicInst *Intrinsic = dyn_cast<IntrinsicInst>(V)) {
     switch (Intrinsic->getIntrinsicID()) {
     default:
       return false;
     case Intrinsic::amdgcn_readfirstlane:
     case Intrinsic::amdgcn_readlane:
     case Intrinsic::amdgcn_icmp:
     case Intrinsic::amdgcn_fcmp:
+    case Intrinsic::amdgcn_if_break:
       return true;
     }
   }
 
   const ExtractValueInst *ExtValue = dyn_cast<ExtractValueInst>(V);
   if (!ExtValue)
     return false;
 
-  if (const CallInst *CI = dyn_cast<CallInst>(ExtValue->getOperand(0))) {
-    // If we have inline asm returning mixed SGPR and VGPR results, we inferred
-    // divergent for the overall struct return. We need to override it in the
-    // case we're extracting an SGPR component here.
-    if (isa<InlineAsm>(CI->getCalledValue()))
-      return !isInlineAsmSourceOfDivergence(CI, ExtValue->getIndices());
-  }
+  const CallInst *CI = dyn_cast<CallInst>(ExtValue->getOperand(0));
+  if (!CI)
+    return false;
+
+  if (const IntrinsicInst *Intrinsic = dyn_cast<IntrinsicInst>(CI)) {
+    switch (Intrinsic->getIntrinsicID()) {
+    default:
+      return false;
+    case Intrinsic::amdgcn_if:
+    case Intrinsic::amdgcn_else: {
+      ArrayRef<unsigned> Indices = ExtValue->getIndices();
+      return Indices.size() == 1 && Indices[0] == 1;
+    }
+    }
+  }
+
 
   return false;
 }
 
 bool GCNTTIImpl::collectFlatAddressOperands(SmallVectorImpl<int> &OpIndexes,
                                             Intrinsic::ID IID) const {
   switch (IID) {
   case Intrinsic::amdgcn_atomic_inc:
   case Intrinsic::amdgcn_atomic_dec:

llvm/test/Analysis/DivergenceAnalysis/AMDGPU/control-flow-intrinsics.ll

+; RUN: opt -mtriple=amdgcn-mesa-mesa3d -analyze -divergence -use-gpu-divergence-analysis %s | FileCheck %s
+
+; Tests control flow intrinsics that should be treated as uniform
+
+; CHECK: Printing analysis 'Legacy Divergence Analysis' for function 'test_if_break':
+; CHECK: DIVERGENT: %cond = icmp eq i32 %arg0, 0
+; CHECK-NOT: DIVERGENT
+; CHECK: ret void
+define amdgpu_ps void @test_if_break(i32 %arg0, i64 inreg %saved) {
+entry:
+  %cond = icmp eq i32 %arg0, 0
+  %break = call i64 @llvm.amdgcn.if.break.i64.i64(i1 %cond, i64 %saved)
+  store volatile i64 %break, i64 addrspace(1)* undef
+  ret void
+}
+
+; CHECK: Printing analysis 'Legacy Divergence Analysis' for function 'test_if':
+; CHECK: DIVERGENT: %cond = icmp eq i32 %arg0, 0
+; CHECK-NEXT: DIVERGENT: %if = call { i1, i64 } @llvm.amdgcn.if.i64(i1 %cond)
+; CHECK-NEXT: DIVERGENT: %if.bool = extractvalue { i1, i64 } %if, 0
+; CHECK-NOT: DIVERGENT
+; CHECK: DIVERGENT: %if.bool.ext = zext i1 %if.bool to i32
+define void @test_if(i32 %arg0) {
+entry:
+  %cond = icmp eq i32 %arg0, 0
+  %if = call { i1, i64 } @llvm.amdgcn.if.i64(i1 %cond)
+  %if.bool = extractvalue { i1, i64 } %if, 0
+  %if.mask = extractvalue { i1, i64 } %if, 1
+  %if.bool.ext = zext i1 %if.bool to i32
+  store volatile i32 %if.bool.ext, i32 addrspace(1)* undef
+  store volatile i64 %if.mask, i64 addrspace(1)* undef
+  ret void
+}
+
+; The result should still be treated as divergent, even with a uniform source.
+; CHECK: Printing analysis 'Legacy Divergence Analysis' for function 'test_if_uniform':
+; CHECK-NOT: DIVERGENT
+; CHECK: DIVERGENT: %if = call { i1, i64 } @llvm.amdgcn.if.i64(i1 %cond)
+; CHECK-NEXT: DIVERGENT: %if.bool = extractvalue { i1, i64 } %if, 0
+; CHECK-NOT: DIVERGENT
+; CHECK: DIVERGENT: %if.bool.ext = zext i1 %if.bool to i32
+define amdgpu_ps void @test_if_uniform(i32 inreg %arg0) {
+entry:
+  %cond = icmp eq i32 %arg0, 0
+  %if = call { i1, i64 } @llvm.amdgcn.if.i64(i1 %cond)
+  %if.bool = extractvalue { i1, i64 } %if, 0
+  %if.mask = extractvalue { i1, i64 } %if, 1
+  %if.bool.ext = zext i1 %if.bool to i32
+  store volatile i32 %if.bool.ext, i32 addrspace(1)* undef
+  store volatile i64 %if.mask, i64 addrspace(1)* undef
+  ret void
+}
+
+; CHECK: Printing analysis 'Legacy Divergence Analysis' for function 'test_loop_uniform':
+; CHECK: DIVERGENT: %loop = call i1 @llvm.amdgcn.loop.i64(i64 %mask)
+define amdgpu_ps void @test_loop_uniform(i64 inreg %mask) {
+entry:
+  %loop = call i1 @llvm.amdgcn.loop.i64(i64 %mask)
+  %loop.ext = zext i1 %loop to i32
+  store volatile i32 %loop.ext, i32 addrspace(1)* undef
+  ret void
+}
+
+; CHECK: Printing analysis 'Legacy Divergence Analysis' for function 'test_else':
+; CHECK: DIVERGENT: %else = call { i1, i64 } @llvm.amdgcn.else.i64.i64(i64 %mask)
+; CHECK: DIVERGENT:       %else.bool = extractvalue { i1, i64 } %else, 0
+; CHECK: {{^[ \t]+}}%else.mask = extractvalue { i1, i64 } %else, 1
+define amdgpu_ps void @test_else(i64 inreg %mask) {
+entry:
+  %else = call { i1, i64 } @llvm.amdgcn.else.i64.i64(i64 %mask)
+  %else.bool = extractvalue { i1, i64 } %else, 0
+  %else.mask = extractvalue { i1, i64 } %else, 1
+  %else.bool.ext = zext i1 %else.bool to i32
+  store volatile i32 %else.bool.ext, i32 addrspace(1)* undef
+  store volatile i64 %else.mask, i64 addrspace(1)* undef
+  ret void
+}
+
+; This case is probably always broken
+; CHECK: Printing analysis 'Legacy Divergence Analysis' for function 'test_else_divergent_mask':
+; CHECK: DIVERGENT: %if = call { i1, i64 } @llvm.amdgcn.else.i64.i64(i64 %mask)
+; CHECK-NEXT: DIVERGENT: %if.bool = extractvalue { i1, i64 } %if, 0
+; CHECK-NOT: DIVERGENT
+; CHECK: DIVERGENT: %if.bool.ext = zext i1 %if.bool to i32
+define void @test_else_divergent_mask(i64 %mask) {
+entry:
+  %if = call { i1, i64 } @llvm.amdgcn.else.i64.i64(i64 %mask)
+  %if.bool = extractvalue { i1, i64 } %if, 0
+  %if.mask = extractvalue { i1, i64 } %if, 1
+  %if.bool.ext = zext i1 %if.bool to i32
+  store volatile i32 %if.bool.ext, i32 addrspace(1)* undef
+  store volatile i64 %if.mask, i64 addrspace(1)* undef
+  ret void
+}
+
+declare { i1, i64 } @llvm.amdgcn.if.i64(i1) #0
+declare { i1, i64 } @llvm.amdgcn.else.i64.i64(i64) #0
+declare i64 @llvm.amdgcn.if.break.i64.i64(i1, i64) #1
+declare i1 @llvm.amdgcn.loop.i64(i64) #1
+
+attributes #0 = { convergent nounwind }
+attributes #1 = { convergent nounwind readnone }