[NewPM] Disable non-trivial loop-unswitch on targets with divergence

[sameerds]

Summary

Unswitching a loop on a non-trivial divergent branch is expensive
since it serializes the execution of both version of the
loop. But identifying a divergent branch needs divergence analysis,
which is a function level analysis.

The legacy pass manager handles this dependency by isolating such a
loop transform and rerunning the required function analyses. This
functionality is currently missing in the new pass manager, and there
is no safe way for the SimpleLoopUnswitch pass to depend on
DivergenceAnalysis. So we conservatively assume that all non-trivial
branches are divergent if the target has divergence.

[rampitec]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2907

Is this a correct condition? Double negative "!NonTrivial" = Trivial, right? Are we not doing trivial loop unswitching?

[tra]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2907

This should be a bit easier to read:

if (!(EnableNonTrivialUnswitch || 
      (NonTrivial && !TTI.hasBranchDivergence())))

Or, maybe extract part of the condition into a bool ShouldUnswitch = NonTrivial && !TTI.hasBranchDivergence().

[sameerds]

Improved the readability of a potentially triple-negative boolean condition

[sameerds]

In D98958#2640544, @sameerds wrote:

Improved the readability of a potentially triple-negative boolean condition

The linter pointed out two overlong lines in the new comment. I've fixed that locally in my working directory. But I didn't want to trigger yet another build by uploading that NFC change.

[tra]

llvm/test/Transforms/LoopUnswitch/AMDGPU/divergent-unswitch.ll

16

Perhaps the test should be replaced with a test that verifies that unswitch does *not* happen.

[rampitec]

LGTM modulo Artem's comment about test for unswitch to not happen.

[sameerds]

llvm/test/Transforms/LoopUnswitch/AMDGPU/divergent-unswitch.ll

16

I am not sure what you mean. I am guessing that the way phab presents this change might be to blame. There are two files: divergent-unswitch.ll (this file) and a new uniform-unswitch.ll. This file checks that the divergent unswitch does not happen, and passes.

The part that carries this comment is moved to uniform-unswitch.ll, where it is checked to ensure that the uniform unswitch does happen. But the new behaviour conservatily treats all non-trivial branches as divergent, and hence the new test is marked as XFAIL.

[tra]

llvm/test/Transforms/LoopUnswitch/AMDGPU/divergent-unswitch.ll

16

I've missed that you've split the files.

the new behaviour conservatily treats all non-trivial branches as divergent, and hence the new test is marked as XFAIL.

You may want to add a comment explaining that next to that XFAIL. Tests that test something that does not happen are hard to understand without additional details.

I guess I can rephrase my suggestion as "could we replace XFAIL with a positive test that unswitching does not happen on AMDGPU unless it's force-enabled by the flag"?

If this is a temporary arrangement and we do plan to eventually allow unswitching for uniform loops, XFAIL+a comment may do.

[sameerds]

llvm/test/Transforms/LoopUnswitch/AMDGPU/divergent-unswitch.ll

16

If this is a temporary arrangement and we do plan to eventually allow unswitching for uniform loops, XFAIL+a comment may do.

Correct. The test uniform-unswitch.ll does have a comment point to a bug report after the XFAIL line. It is not highlighted in this diff because Phabricator treats it as existing code that got moved to a new file.

But one can ask how long is "eventually". The deeper issue here is that this can be classified as a functional regression in the new pass manager. The old manager correctly handled the dependence of a loop transform on a function analysis, but the new pass manager has no such infrastructure. This test is supposed to pass (uniform non-trivial branches should get unswitched) when we bring that functionality to the new pass manager.

[tra]

llvm/test/Transforms/LoopUnswitch/AMDGPU/divergent-unswitch.ll

16

SGTM.

llvm/test/Transforms/LoopUnswitch/AMDGPU/uniform-unswitch.ll

2

Please add a comment/TODO explaining that the unswitch is expected to happen in principle, but we still need <something> to make it work.

[RKSimon]

llvm/test/Transforms/LoopUnswitch/AMDGPU/uniform-unswitch.ll

2

@sameerds This is breaking a windows buildbot as its actually passing.....

http://lab.llvm.org:8011/#/builders/123/builds/3602

[aeubanks]

llvm/test/Transforms/LoopUnswitch/AMDGPU/uniform-unswitch.ll

2

If I had to guess, some bots are still using the legacy PM.
This test should either use -passes=default<O3> (new PM) or -O3 -enable-new-pm=0 (legacy PM).

[rampitec]

Actually test unexpectedly passes with my local Linux build too:

********************
Unexpectedly Passed Tests (1):
  LLVM :: Transforms/LoopUnswitch/AMDGPU/uniform-unswitch.ll

[sameerds]

llvm/test/Transforms/LoopUnswitch/AMDGPU/uniform-unswitch.ll

2

Thanks to Richard Smith (@rsmith ?) for fixing this. The test was updated to force -enable-new-pm=1.

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

static bool unswitchLoop(Loop &L, DominatorTree &DT, LoopInfo &LI,
   // Try trivial unswitch first before loop over other basic blocks in the loop.
   if (unswitchAllTrivialConditions(L, DT, LI, SE, MSSAU)) {
     // If we unswitched successfully we will want to clean up the loop before
     // processing it further so just mark it as unswitched and return.
     UnswitchCB(/*CurrentLoopValid*/ true, {});
     return true;
   }
 
-  // If we're not doing non-trivial unswitching, we're done. We both accept
-  // a parameter but also check a local flag that can be used for testing
-  // a debugging.
-  if (!NonTrivial && !EnableNonTrivialUnswitch)
+  // Check whether we should continue with non-trivial conditions.
+  // EnableNonTrivialUnswitch: Global variable that forces non-trivial
+  //                           unswitching for testing and debugging.
+  // NonTrivial: Parameter that enables non-trivial unswitching for this
+  //             invocation of the transform. But this should be allowed only
+  //             for targets without branch divergence.
+  //
+  // FIXME: If divergence analysis becomes available to a loop
+  // transform, we should allow unswitching for non-trivial uniform
+  // branches even on targets that have divergence.
+  // https://bugs.llvm.org/show_bug.cgi?id=48819
+  bool ContinueWithNonTrivial =
+      EnableNonTrivialUnswitch || (NonTrivial && !TTI.hasBranchDivergence());
+  if (!ContinueWithNonTrivial)
     return false;
 
   // Skip non-trivial unswitching for optsize functions.
   if (L.getHeader()->getParent()->hasOptSize())
     return false;
 
   // Skip non-trivial unswitching for loops that cannot be cloned.
   if (!L.isSafeToClone())

llvm/test/Transforms/LoopUnswitch/AMDGPU/divergent-unswitch.ll

-; RUN: opt -mtriple=amdgcn-- -O3 -S -enable-new-pm=0 %s | FileCheck %s
-
-; This fails with the new pass manager:
-; https://bugs.llvm.org/show_bug.cgi?id=48819
-
-; Check that loop unswitch happened and condition hoisted out of the loop.
-; Condition is uniform so all targets should perform unswitching.
-
-; CHECK-LABEL: {{^}}define amdgpu_kernel void @uniform_unswitch
-; CHECK: entry:
-; CHECK-NEXT: [[LOOP_COND:%[a-z0-9]+]] = icmp
-; CHECK-NEXT: [[IF_COND:%[a-z0-9]+]] = icmp eq i32 %x, 123456
-; CHECK-NEXT: and i1 [[LOOP_COND]], [[IF_COND]]
-; CHECK-NEXT: br i1
-
-define amdgpu_kernel void @uniform_unswitch(i32 * nocapture %out, i32 %n, i32 %x) {
-entry:
-  %cmp6 = icmp sgt i32 %n, 0
-  br i1 %cmp6, label %for.body.lr.ph, label %for.cond.cleanup
-
-for.body.lr.ph:                                   ; preds = %entry
-  %cmp1 = icmp eq i32 %x, 123456
-  br label %for.body
-
-for.cond.cleanup.loopexit:                        ; preds = %for.inc
-  br label %for.cond.cleanup
-
-for.cond.cleanup:                                 ; preds = %for.cond.cleanup.loopexit, %entry
-  ret void
-
-for.body:                                         ; preds = %for.inc, %for.body.lr.ph
-  %i.07 = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %for.inc ]
-  br i1 %cmp1, label %if.then, label %for.inc
-
-if.then:                                          ; preds = %for.body
-  %arrayidx = getelementptr inbounds i32, i32 * %out, i32 %i.07
-  store i32 %i.07, i32 * %arrayidx, align 4
-  br label %for.inc
-
-for.inc:                                          ; preds = %for.body, %if.then
-  %inc = add nuw nsw i32 %i.07, 1
-  %exitcond = icmp eq i32 %inc, %n
-  br i1 %exitcond, label %for.cond.cleanup.loopexit, label %for.body
-}
+; RUN: opt -mtriple=amdgcn-- -O3 -S %s | FileCheck %s
 
 ; Check that loop unswitch does not happen if condition is divergent.
 
 ; CHECK-LABEL: {{^}}define amdgpu_kernel void @divergent_unswitch
 ; CHECK: entry:
 ; CHECK: icmp
 ; CHECK: [[IF_COND:%[a-z0-9]+]] = icmp {{.*}} 567890
 ; CHECK: br label

llvm/test/Transforms/LoopUnswitch/AMDGPU/uniform-unswitch.ll

-; RUN: opt -mtriple=amdgcn-- -O3 -S -enable-new-pm=0 %s | FileCheck %s
+; RUN: opt -mtriple=amdgcn-- -O3 -S %s | FileCheck %s
+; XFAIL: *
+
+; Check that loop unswitch happened and condition hoisted out of the loop.
+; Condition is uniform so even targets with divergence should perform unswitching.
 
 ; This fails with the new pass manager:
 ; https://bugs.llvm.org/show_bug.cgi?id=48819
-
-; Check that loop unswitch happened and condition hoisted out of the loop.
-; Condition is uniform so all targets should perform unswitching.
+; The correct behaviour (allow uniform non-trivial branches to be
+; unswitched on all targets) requires access to the function-level
+; divergence analysis from a loop transform, which is currently not
+; supported in the new pass manager.
 
 ; CHECK-LABEL: {{^}}define amdgpu_kernel void @uniform_unswitch
 ; CHECK: entry:
 ; CHECK-NEXT: [[LOOP_COND:%[a-z0-9]+]] = icmp
 ; CHECK-NEXT: [[IF_COND:%[a-z0-9]+]] = icmp eq i32 %x, 123456
 ; CHECK-NEXT: and i1 [[LOOP_COND]], [[IF_COND]]
 ; CHECK-NEXT: br i1
 
if.then:                                          ; preds = %for.body
   br label %for.inc
 
 for.inc:                                          ; preds = %for.body, %if.then
   %inc = add nuw nsw i32 %i.07, 1
   %exitcond = icmp eq i32 %inc, %n
   br i1 %exitcond, label %for.cond.cleanup.loopexit, label %for.body
 }
 
-; Check that loop unswitch does not happen if condition is divergent.
-
-; CHECK-LABEL: {{^}}define amdgpu_kernel void @divergent_unswitch
-; CHECK: entry:
-; CHECK: icmp
-; CHECK: [[IF_COND:%[a-z0-9]+]] = icmp {{.*}} 567890
-; CHECK: br label
-; CHECK: br i1 [[IF_COND]]
-
-define amdgpu_kernel void @divergent_unswitch(i32 * nocapture %out, i32 %n) {
-entry:
-  %cmp9 = icmp sgt i32 %n, 0
-  br i1 %cmp9, label %for.body.lr.ph, label %for.cond.cleanup
-
-for.body.lr.ph:                                   ; preds = %entry
-  %call = tail call i32 @llvm.amdgcn.workitem.id.x() #0
-  %cmp2 = icmp eq i32 %call, 567890
-  br label %for.body
-
-for.cond.cleanup.loopexit:                        ; preds = %for.inc
-  br label %for.cond.cleanup
-
-for.cond.cleanup:                                 ; preds = %for.cond.cleanup.loopexit, %entry
-  ret void
-
-for.body:                                         ; preds = %for.inc, %for.body.lr.ph
-  %i.010 = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %for.inc ]
-  br i1 %cmp2, label %if.then, label %for.inc
-
-if.then:                                          ; preds = %for.body
-  %arrayidx = getelementptr inbounds i32, i32 * %out, i32 %i.010
-  store i32 %i.010, i32 * %arrayidx, align 4
-  br label %for.inc
-
-for.inc:                                          ; preds = %for.body, %if.then
-  %inc = add nuw nsw i32 %i.010, 1
-  %exitcond = icmp eq i32 %inc, %n
-  br i1 %exitcond, label %for.cond.cleanup.loopexit, label %for.body
-}
-
 declare i32 @llvm.amdgcn.workitem.id.x() #0
 
 attributes #0 = { nounwind readnone }