apply DivergenceAnalysis for SLU

[drcut]

Summary

Apply non-trivial unswitch to the branches that do not diverge after analysis.

[drcut]

@bcahoon This revision is based on your previous implementation (https://reviews.llvm.org/D109762). In the parent revision(https://reviews.llvm.org/D127547) we use FuncPass to re-implement non-trivial LoopUnswitch, thus, we do not invoke DivergenceAnalysis for each loop. Would you please have are look for whether these modifications (implement non-trivial LoopUnswitch with FuncPass and integrate DivergenceAnalysis in the new FuncPass) work on your test cases? Thanks

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/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"
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,
+    DivergenceInfo *DI) {
   // 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(DI) {
+    // only non-divergent branch can be unswitched
+    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,
+             DivergenceInfo* DI) {
   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, DI))
     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, /*DivergenceInfo*/ nullptr))
     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));
auto DestroyLoopCB = [&LAM](Loop &L, StringRef Name) {
       LAM.clear(L, Name);
     };
     Changed |= unswitchLoop(
         L, AM.getResult<DominatorTreeAnalysis>(F),
         AM.getResult<LoopAnalysis>(F), AM.getResult<AssumptionAnalysis>(F),
         AM.getResult<AAManager>(F), AM.getResult<TargetIRAnalysis>(F),
         /*Trivial=*/false, /*NonTrivial=*/true,
         UnswitchCB, &AM.getResult<ScalarEvolutionAnalysis>(F),
-        MSSAU.hasValue() ? MSSAU.getPointer() : nullptr, DestroyLoopCB);
+        MSSAU.hasValue() ? MSSAU.getPointer() : nullptr, DestroyLoopCB,
+        &AM.getResult<DivergenceAnalysis>(F));
   }
   if (!Changed)
     return PreservedAnalyses::all();
 
   if (MSSA && VerifyMemorySSA)
     MSSA->verifyMemorySSA();
 
   // Historically this pass has had issues with the dominator tree so verify it
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,
+                              /*DivergenceInfo*/ nullptr);
 
   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 }