[NewPM][SimpleLoopUnswitch] Add DivergenceInfo

[bcahoon]

Summary

Enable DivergenceAnalysis with non-trivial Simple Loop Unswitch.
Because DivergenceAnalysis is a function pass, it's not possible
to add the analysis to the pass manager pipeline when it's used
by a loop pass. This patch creates a new DivergenceInfo instance
only when needed by Simple Loop Unswitch. The PostDominatorTree
is needed as well. This adds DivergenceAnalysis when using the new
pass manager only and does not change the legacy pass manager.

[aeubanks]

lgtm but wait for other lgtms

llvm/test/Transforms/SimpleLoopUnswitch/AMDGPU/nontrivial-unswitch.ll

1 ↗

(On Diff #372481)

not a big deal, but we could just do REQUIRES: amdgpu-registered-target right? unless you're planning on adding more of these sorts of tests

[asbirlea]

lgtm.

[aeubanks]

you should probably check that this doesn't affect compile times too much since we're creating some new (potentially expensive) analyses without any sort of caching

[bcahoon]

Added amdgpu-registered-target to test case.

[arsenm]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2803

Constructing your own PostDominatorTree seems not good. Should this be a pass dependency for divergent targets?

[bcahoon]

In D109762#2999963, @aeubanks wrote:

you should probably check that this doesn't affect compile times too much since we're creating some new (potentially expensive) analyses without any sort of caching

I've run a few test cases, but I don't know of a straightforward way to evaluate a large set of tests for compile time. It would be nice to be be able to cache the analysis.

[bcahoon]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2803

I agree that it's not good, but I believe the only way to do this as a dependency/required is to have all the loop passes update the post-dominator tree info when needed (which is what's done for dominators and several other loop analyses).

[aeubanks]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2803

See the loop analyses part of https://llvm.org/docs/NewPassManager.html#using-analyses for limitations on accessing function analyses in a loop pass.

[bcahoon]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2803

Thanks for that reference. If I understand correctly, another approach to adding divergence analysis would be to add an entry to LoopStandardAnalysisResults for DivergenceAnalysisInfo (but, at least initially, without the guarantee that loop passes update it). That way it can be accessed by SimpleLoopUnswtich.

Then, in PassBuilder.cpp add an instance of the DivergeranceAnalysis prior to SimpleLoopUnswitch. This also requires a call to createFunctionToLoopPassAdaptor between DivergenceAnalysis and SimpleLoopUnswitch. We could do this only for targets that have branch divergence so that the pass pipeline remains the same for other targets.

This approach seems most similar to how the legacy pass manager works. My main hesitation is adding DivergenceAnalysis to LoopStandardAnalysisResults.

[aeubanks]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2803

I don't think splitting up the loop pipeline to add in a function pass (to request a function analysis) is ok
also, adding DivergenceAnalysis to LoopStandardAnalysisResults doesn't seem right since expecting all loop passes to update DivergenceAnalysis doesn't seem right

I think the proper way is to relax the constraint that a loop pass can't request a function analysis on demand. Part of the reason we don't allow this for other IR units is potential future concurrency and determinism, but I don't think we will ever do concurrency for loop passes (as in run loop passes on multiple loops in a function concurrently). I believe @asbirlea has mentioned something like this in the past, any thoughts?

Just curious, are you actually seeing benchmarks where this nontrivial loop unswitching matters?

[bcahoon]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2803

Just curious, are you actually seeing benchmarks where this nontrivial loop unswitching matters?

Yes, I've been tracking down the source of some performance regressions that we're seeing after the switch from the legacy pass manager to the new pass manager. I've seen a couple instances where non-trivial loop unswitching does improve performance (though, in one instance, I've noticed that the threshold needs to increase in simple loop unswitch to recover performance).

[sameerds]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2803

Just to add some detail, here's one idea of "splitting

2803

I don't think splitting up the loop pipeline to add in a function pass (to request a function analysis) is ok
also, adding DivergenceAnalysis to LoopStandardAnalysisResults doesn't seem right since expecting all loop passes to update DivergenceAnalysis doesn't seem right

I think the proper way is to relax the constraint that a loop pass can't request a function analysis on demand.

Just FYI, here's one idea of how "splitting the loop pipeline" might look like if we attempted that:

https://lists.llvm.org/pipermail/llvm-dev/2021-February/148619.html

[aeubanks]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2803

I'm currently talking to some people to see if we can do some redesigning and allow loop passes to request function analyses on demand.

[bcahoon]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2803

Great to hear. Let me know if there is anything I can do to help. Do you think it's worthwhile to merge this patch now or wait until the outcome of the redesign?

[aeubanks]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2803

Merging this now is fine. (and still LGTM)

[aeubanks]

I talked to some people and we've decided that the best thing to do would be to refactor out the nontrivial unswitching part into a function pass. Nontrivial unswitching is fairly special in the kinds of transforms it does.

[bcahoon]

In D109762#3059988, @aeubanks wrote:

I talked to some people and we've decided that the best thing to do would be to refactor out the nontrivial unswitching part into a function pass. Nontrivial unswitching is fairly special in the kinds of transforms it does.

Thanks for the update. I've waited to commit this patch because the motivating code that is improved by non-trivial unswitching requires additional changes (that are in the regular LoopUnswitch pass but not the SimpleLoopUnswitch pass).

[sameerds]

In D109762#3059988, @aeubanks wrote:

I talked to some people and we've decided that the best thing to do would be to refactor out the nontrivial unswitching part into a function pass. Nontrivial unswitching is fairly special in the kinds of transforms it does.

Will that always work as expected? The real dependency is that the divergence analysis is not incrementally updated. So even if this is a function pass, we may want to rerun it on the whole function every time it manages to unswitch a loop.

[aeubanks]

In D109762#3069261, @sameerds wrote:

In D109762#3059988, @aeubanks wrote:

I talked to some people and we've decided that the best thing to do would be to refactor out the nontrivial unswitching part into a function pass. Nontrivial unswitching is fairly special in the kinds of transforms it does.

Will that always work as expected? The real dependency is that the divergence analysis is not incrementally updated. So even if this is a function pass, we may want to rerun it on the whole function every time it manages to unswitch a loop.

With this patch, we have to rerun DivergenceAnalysis every time we run the pass on every loop. If we change nontrivial unswitching into a function pass, we can upgrade that to only having to be rerun every time we actually unswitch something.
To be more specific, with the function pass, we'd have to invalidate most analyses anyway after every successful unswitch.

for every loop
  DA = FAM.getAnalysis<DivergenceAnalysis>(); // this will not rerun the analysis if we haven't invalidated below because we didn't successfully unswitch
  if successful unswitch
    auto PA = PreservedAnalyases::none();
    PA.preserve<SCEV>();
    // and all the other existing analyses we preserved
    FAM.invalidate(PA);

[asbirlea]

In D109762#3070676, @aeubanks wrote:

In D109762#3069261, @sameerds wrote:

In D109762#3059988, @aeubanks wrote:

I talked to some people and we've decided that the best thing to do would be to refactor out the nontrivial unswitching part into a function pass. Nontrivial unswitching is fairly special in the kinds of transforms it does.

Will that always work as expected? The real dependency is that the divergence analysis is not incrementally updated. So even if this is a function pass, we may want to rerun it on the whole function every time it manages to unswitch a loop.

With this patch, we have to rerun DivergenceAnalysis every time we run the pass on every loop. If we change nontrivial unswitching into a function pass, we can upgrade that to only having to be rerun every time we actually unswitch something.
To be more specific, with the function pass, we'd have to invalidate most analyses anyway after every successful unswitch.

for every loop
  DA = FAM.getAnalysis<DivergenceAnalysis>(); // this will not rerun the analysis if we haven't invalidated below because we didn't successfully unswitch
  if successful unswitch
    auto PA = PreservedAnalyases::none();
    PA.preserve<SCEV>();
    // and all the other existing analyses we preserved
    FAM.invalidate(PA);

There's also the option to update all analysis inside the new function pass that's doing non-trivial unswitching. This is currently the case with all other analyses (they are updated), so the only one that needs to be added is DA.
This limits the need to update DA to just the new Function pass, vs all Loop passes that are part of the same LPM as SimpleLoopUnswitch.
The alternative is, as you both already said, for DA to be recomputed when it is invalidated.

[Peakulorain]

Has this patch been merged into the main branch yet?

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/CodeMetrics.h"
+#include "llvm/Analysis/DivergenceAnalysis.h"
 #include "llvm/Analysis/GuardUtils.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/LoopAnalysisManager.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopIterator.h"
 #include "llvm/Analysis/LoopPass.h"
 #include "llvm/Analysis/MemorySSA.h"
 #include "llvm/Analysis/MemorySSAUpdater.h"
 #include "llvm/Analysis/MustExecute.h"
+#include "llvm/Analysis/PostDominators.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstrTypes.h"
static int CalculateUnswitchCostMultiplier(
   return CostMultiplier;
 }
 
 static bool unswitchBestCondition(
     Loop &L, DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
     AAResults &AA, TargetTransformInfo &TTI,
     function_ref<void(bool, bool, ArrayRef<Loop *>)> UnswitchCB,
     ScalarEvolution *SE, MemorySSAUpdater *MSSAU,
-    function_ref<void(Loop &, StringRef)> DestroyLoopCB) {
+    function_ref<void(Loop &, StringRef)> DestroyLoopCB,
+    bool UseDivergenceInfo) {
   // Collect all invariant conditions within this loop (as opposed to an inner
   // loop which would be handled when visiting that inner loop).
   SmallVector<std::pair<Instruction *, TinyPtrVector<Value *>>, 4>
       UnswitchCandidates;
 
   // Whether or not we should also collect guards in the loop.
   bool CollectGuards = false;
   if (UnswitchGuards) {
for (auto *ExitBB : ExitBlocks) {
     auto *I = ExitBB->getFirstNonPHI();
     if (isa<CleanupPadInst>(I) || isa<CatchSwitchInst>(I)) {
       LLVM_DEBUG(dbgs() << "Cannot unswitch because of cleanuppad/catchswitch "
                            "in exit block\n");
       return false;
     }
   }
 
+  if (UseDivergenceInfo) {
+    Function *F = L.getHeader()->getParent();
+    PostDominatorTree PDT(*F);
+    DivergenceInfo DI(*F, DT, PDT, LI, TTI, /*KnownReducible*/ true);
+    llvm::erase_if(UnswitchCandidates,
+                   [&](std::pair<Instruction *, TinyPtrVector<Value *>> Cand) {
+                     return DI.isDivergent(*Cand.second[0]);
+                   });
+    if (UnswitchCandidates.empty())
+      return false;
+  }
+
   LLVM_DEBUG(
       dbgs() << "Considering " << UnswitchCandidates.size()
              << " non-trivial loop invariant conditions for unswitching.\n");
 
   // Given that unswitching these terminators will require duplicating parts of
   // the loop, so we need to be able to model that cost. Compute the ephemeral
   // values and set up a data structure to hold per-BB costs. We cache each
   // block's cost so that we don't recompute this when considering different

 /// If `SE` is non-null, we will update that analysis based on the unswitching
 /// done.
 static bool
 unswitchLoop(Loop &L, DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
              AAResults &AA, TargetTransformInfo &TTI, bool Trivial,
              bool NonTrivial,
              function_ref<void(bool, bool, ArrayRef<Loop *>)> UnswitchCB,
              ScalarEvolution *SE, MemorySSAUpdater *MSSAU,
-             function_ref<void(Loop &, StringRef)> DestroyLoopCB) {
+             function_ref<void(Loop &, StringRef)> DestroyLoopCB,
+             bool UseDivergenceInfo) {
   assert(L.isRecursivelyLCSSAForm(DT, LI) &&
          "Loops must be in LCSSA form before unswitching.");
 
   // Must be in loop simplified form: we need a preheader and dedicated exits.
   if (!L.isLoopSimplifyForm())
     return false;
 
   // Try trivial unswitch first before loop over other basic blocks in the loop.
   if (Trivial && 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, false, {});
     return true;
   }
 
   // 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)
+  if (!EnableNonTrivialUnswitch && !NonTrivial)
     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())
     return false;
 
   // For non-trivial unswitching, because it often creates new loops, we rely on
   // the pass manager to iterate on the loops rather than trying to immediately
   // reach a fixed point. There is no substantial advantage to iterating
   // internally, and if any of the new loops are simplified enough to contain
   // trivial unswitching we want to prefer those.
 
   // Try to unswitch the best invariant condition. We prefer this full unswitch to
   // a partial unswitch when possible below the threshold.
   if (unswitchBestCondition(L, DT, LI, AC, AA, TTI, UnswitchCB, SE, MSSAU,
-                            DestroyLoopCB))
+                            DestroyLoopCB, UseDivergenceInfo))
     return true;
 
   // No other opportunities to unswitch.
   return false;
 }
 
 PreservedAnalyses SimpleLoopUnswitchPass::run(Loop &L, LoopAnalysisManager &AM,
                                               LoopStandardAnalysisResults &AR,
PreservedAnalyses SimpleLoopUnswitchPass::run(Loop &L, LoopAnalysisManager &AM,
   if (AR.MSSA) {
     MSSAU = MemorySSAUpdater(AR.MSSA);
     if (VerifyMemorySSA)
       AR.MSSA->verifyMemorySSA();
   }
   if (!unswitchLoop(L, AR.DT, AR.LI, AR.AC, AR.AA, AR.TTI, Trivial, NonTrivial,
                     UnswitchCB, &AR.SE,
                     MSSAU.hasValue() ? MSSAU.getPointer() : nullptr,
-                    DestroyLoopCB))
+                    DestroyLoopCB, AR.TTI.useGPUDivergenceAnalysis()))
     return PreservedAnalyses::all();
 
   if (AR.MSSA && VerifyMemorySSA)
     AR.MSSA->verifyMemorySSA();
 
   // Historically this pass has had issues with the dominator tree so verify it
   // in asserts builds.
   assert(AR.DT.verify(DominatorTree::VerificationLevel::Fast));
bool SimpleLoopUnswitchLegacyPass::runOnLoop(Loop *L, LPPassManager &LPM) {
   auto DestroyLoopCB = [&LPM](Loop &L, StringRef /* Name */) {
     LPM.markLoopAsDeleted(L);
   };
 
   if (VerifyMemorySSA)
     MSSA->verifyMemorySSA();
 
   bool Changed = unswitchLoop(*L, DT, LI, AC, AA, TTI, true, NonTrivial,
-                              UnswitchCB, SE, &MSSAU, DestroyLoopCB);
+                              UnswitchCB, SE, &MSSAU, DestroyLoopCB,
+                              /*UseDivergenceInfo*/ false);
 
   if (VerifyMemorySSA)
     MSSA->verifyMemorySSA();
 
   // Historically this pass has had issues with the dominator tree so verify it
   // in asserts builds.
   assert(DT.verify(DominatorTree::VerificationLevel::Fast));
 

llvm/test/Transforms/SimpleLoopUnswitch/divergent-nontrivial-unswitch.ll

+; RUN: opt -mtriple=amdgcn-unknown-amdhsa -passes='loop(simple-loop-unswitch<nontrivial>),verify<loops>' -S < %s | FileCheck %s
+; REQUIRES: amdgpu-registered-target
+
+; Check that non-trivial loop unswitch occurs on a target with divergence.
+
+; CHECK-LABEL: @nontrivial_unswitch(
+; CHECK: for.body.us:
+; CHECK: if.then.us:
+; CHECK: for.inc.us:
+; CHECK: for.body:
+; CHECK: for.inc:
+define amdgpu_kernel void @nontrivial_unswitch(i32 * nocapture %out, i32 %n, i1 %cond) {
+entry:
+  br label %for.body
+
+for.body:
+  %i = phi i32 [ 0, %entry ], [ %inc, %for.inc ]
+  br i1 %cond, label %if.then, label %for.inc
+
+if.then:
+  %arrayidx = getelementptr inbounds i32, i32 * %out, i32 %i
+  store i32 %i, i32 * %arrayidx, align 4
+  br label %for.inc
+
+for.inc:
+  %inc = add nuw nsw i32 %i, 1
+  %exitcond = icmp eq i32 %inc, %n
+  br i1 %exitcond, label %for.cond.cleanup.loopexit, label %for.body
+
+for.cond.cleanup.loopexit:
+  ret void
+}
+
+
+; Check that loop unswitch does not happen if the condition is divergent.
+
+; CHECK-LABEL: @divergent_unswitch(
+; CHECK: entry:
+; CHECK: [[IF_COND:%[a-z0-9]+]] = icmp {{.*}}, 1
+; CHECK: br label
+; CHECK: for.body:
+; CHECK: br i1 [[IF_COND]]
+
+define amdgpu_kernel void @divergent_unswitch(i32 * nocapture %out, i32 %n) {
+entry:
+  br label %for.body.lr.ph
+
+for.body.lr.ph:
+  %call = tail call i32 @llvm.amdgcn.workitem.id.x() #0
+  %cmp2 = icmp eq i32 %call, 1
+  br label %for.body
+
+for.body:
+  %i = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %for.inc ]
+  br i1 %cmp2, label %if.then, label %for.inc
+
+if.then:
+  %arrayidx = getelementptr inbounds i32, i32 * %out, i32 %i
+  store i32 %i, i32 * %arrayidx, align 4
+  br label %for.inc
+
+for.inc:
+  %inc = add nuw nsw i32 %i, 1
+  %exitcond = icmp eq i32 %inc, %n
+  br i1 %exitcond, label %for.cond.cleanup.loopexit, label %for.body
+
+for.cond.cleanup.loopexit:
+  ret void
+}
+
+declare i32 @llvm.amdgcn.workitem.id.x() #0
+
+attributes #0 = { nounwind readnone }