Hoist llvm.assume into single predecessor if block otherwise empty

[markus]

Summary

Here is a first go at what was discussed in the thread of https://lists.llvm.org/pipermail/llvm-dev/2021-May/150739.html

Hoist llvm.assume instrinsics (and rewrite condition) from blocks where they would inhibit transformation by SimplifyCFGOpt::simplifyUncondBranch.

[nikic]

This kind of transformation should probably eventually reside in lib/Transforms/Utils/SimplifyCFG.cpp but for now it seemed easier and a smaller change to just leave it in CodeGenPrepare.

I don't think this transform makes a lot of sense in CGP, I'd prefer it to go directly to SimplifyCFG (after which we can drop the CGP code).

However, I'm not sure the whole concept of "hoisting assumes" even makes sense. I'm not sure if we have any passes that can reason about assume(c1 || c2) conditions. We have a lot that can reason about assume(c1 && c2) once it has been canonicalized to assume(c1) assume(c2), but probably nothing that handles ||. The only case in which it could end up being useful is if at a later point, we can show that c1 is actually false, which means we can fold it to just assume(c2), at which point it becomes usable.

Another thing to consider is that assumes with operand bundles would require a more complex transform, to convert the operand bundle form back into a condition form.

I think we would be better off to just drop these assumes. The framing I'm thinking of here is that SimplifyCFG should generally be ignoring ephemeral values in transforms. So if normally SimplifyCFG would drop a block that contains no instructions, then it should also drop a block that only contains ephemeral values.

[jdoerfert]

In D103316#2786946, @nikic wrote:

This kind of transformation should probably eventually reside in lib/Transforms/Utils/SimplifyCFG.cpp but for now it seemed easier and a smaller change to just leave it in CodeGenPrepare.

I don't think this transform makes a lot of sense in CGP, I'd prefer it to go directly to SimplifyCFG (after which we can drop the CGP code).

I agree.

I think we would be better off to just drop these assumes. The framing I'm thinking of here is that SimplifyCFG should generally be ignoring ephemeral values in transforms. So if normally SimplifyCFG would drop a block that contains no instructions, then it should also drop a block that only contains ephemeral values.

I disagree.

if (x)
  assume(y)

is a natural way to spell something for the user.
We should not just drop it without having a good reason.
While we might not use assume(x && y) directly, we will use it if we specialize x to true in some places.

[markus]

Since there seem to be agreement on the placement in SimplifyCFG the code has been moved there.

Wrt assume operand bundles I am thinking that those can always simply be dropped.

Old CodeGenPrepare tests have been temporarily disabled but will be revamped in the SimplifyCFG setting once it is decided if this transformation should hoist or drop assumes (from or in otherwise empty blocks).

[reames]

In D103316#2787008, @jdoerfert wrote:

In D103316#2786946, @nikic wrote:

This kind of transformation should probably eventually reside in lib/Transforms/Utils/SimplifyCFG.cpp but for now it seemed easier and a smaller change to just leave it in CodeGenPrepare.

I don't think this transform makes a lot of sense in CGP, I'd prefer it to go directly to SimplifyCFG (after which we can drop the CGP code).

I agree.

I think we would be better off to just drop these assumes. The framing I'm thinking of here is that SimplifyCFG should generally be ignoring ephemeral values in transforms. So if normally SimplifyCFG would drop a block that contains no instructions, then it should also drop a block that only contains ephemeral values.

I disagree.

if (x)
  assume(y)

is a natural way to spell something for the user.
We should not just drop it without having a good reason.
While we might not use assume(x && y) directly, we will use it if we specialize x to true in some places.

Er, I think you misread the comment you were replying to. The comment was discussing assume(A || B), your response is discussing assume(A && B). Given the context of the discussion, that difference seems important.

My own opinion is closer to that of @nikic. I think we should not be letting assumes block transforms. We should attempt to salvage assume information where we can, but only on a best effort basis. I do think we should generate the assume(A || B) form if that's the logical salvage even if nothing currently can infer from that.

[reames]

In D103316#2789027, @markus wrote:

Wrt assume operand bundles I am thinking that those can always simply be dropped.

I think you can simply ignore operand bundle forms. This is not on by default, and the implementation isn't yet mature. The folks working on that can pay the cost of figuring out what to do for them. As long as your code doesn't crash on operand bundle assumes, I don't really care what it does with them.

[reames]

llvm/lib/CodeGen/CodeGenPrepare.cpp

380 ↗

(On Diff #348737)

Please split the new transform in SimplifyCFG and the deletion of the old CGP code into two patches. The former should be independently worthwhile if structured properly.

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6779

I don't think this is the right framing. We don't want to unconditionally hoist assumes, we only want to hoist them if they'd otherwise inhibit a transformation. (e.g. If the successor block was otherwise empty, we hoist as part of threading the edge to the successors successor.)

llvm/test/Transforms/InstCombine/assume-align.ll

16 ↗

(On Diff #348737)

If we don't predicate the store, the shown transform here does not appear profitable.

[jdoerfert]

[EDIT respond to the assume bundle comment]

In D103316#2800110, @reames wrote:

In D103316#2787008, @jdoerfert wrote:

In D103316#2786946, @nikic wrote:

This kind of transformation should probably eventually reside in lib/Transforms/Utils/SimplifyCFG.cpp but for now it seemed easier and a smaller change to just leave it in CodeGenPrepare.

I don't think this transform makes a lot of sense in CGP, I'd prefer it to go directly to SimplifyCFG (after which we can drop the CGP code).

I agree.

I think we would be better off to just drop these assumes. The framing I'm thinking of here is that SimplifyCFG should generally be ignoring ephemeral values in transforms. So if normally SimplifyCFG would drop a block that contains no instructions, then it should also drop a block that only contains ephemeral values.

I disagree.

if (x)
  assume(y)

is a natural way to spell something for the user.
We should not just drop it without having a good reason.
While we might not use assume(x && y) directly, we will use it if we specialize x to true in some places.

Er, I think you misread the comment you were replying to. The comment was discussing assume(A || B), your response is discussing assume(A && B). Given the context of the discussion, that difference seems important.

My own opinion is closer to that of @nikic. I think we should not be letting assumes block transforms. We should attempt to salvage assume information where we can, but only on a best effort basis. I do think we should generate the assume(A || B) form if that's the logical salvage even if nothing currently can infer from that.

Unsure if I got it right or wrong before but what you said last is what I like too. Salvage assumes in blocks we want to remove if possible. At the same time, "ignore" assumes when the decision for transformations are made.

In D103316#2800112, @reames wrote:

In D103316#2789027, @markus wrote:

Wrt assume operand bundles I am thinking that those can always simply be dropped.

I think you can simply ignore operand bundle forms. This is not on by default, and the implementation isn't yet mature. The folks working on that can pay the cost of figuring out what to do for them. As long as your code doesn't crash on operand bundle assumes, I don't really care what it does with them.

Given that operand bundles for align are the default, people might care: https://clang.godbolt.org/z/Kq8c81T15
That said, I think dropping them is the only reasonable choice we have right now. We don't have a concept of a "pre-condition" that needs to hold in order for the rest to be usable.

[markus]

In D103316#2800153, @jdoerfert wrote:

In D103316#2800110, @reames wrote:

My own opinion is closer to that of @nikic. I think we should not be letting assumes block transforms. We should attempt to salvage assume information where we can, but only on a best effort basis. I do think we should generate the assume(A || B) form if that's the logical salvage even if nothing currently can infer from that.

Unsure if I got it right or wrong before but what you said last is what I like too. Salvage assumes in blocks we want to remove if possible. At the same time, "ignore" assumes when the decision for transformations are made.

So what does this mean in practice?
I interpret it as follows:
We provide a function say bool isBlockEmptyExceptAssumes(BasicBlock *BB) and we update the decision making of all the existing CFG transformations to use that. We provide another function say void hoistAssumesFromBlockKnownToBeOtherwiseEmpty(BasicBlock *BB) that all those CFG transformations need to call during their transformation if they used that fact that any of the blocks checked out as true with the former function.

Now maybe there aren't that many CFG transformations where this is relevant and maybe it integrates easily, I don't know, but it does seem like a somewhat larger change.

llvm/lib/CodeGen/CodeGenPrepare.cpp

380 ↗

(On Diff #348737)

That makes sense. Will do.

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6779

I don't fully understand this comment. The current patch does not unconditionally hoist assumes. It hoists assumes if the block in question is otherwise empty. Thus preventing those assumes from inhibiting other transformations later on. Of course it will require another iteration of CFG considerations but isn't that the way these CFG optimizations are generally structured (one transform enables another transform)?

I guess one way I could interpret this comment is that we should provide a utility function say isBlockEmptyExceptAssumes() that can then be queried in the decision making of all the other CFG transformations. So if a block only contains assumes then treat it as if it was empty when performing those optimizations. Maybe that is doable but it does not seem obvious that it would integrate easily with the existing transformations.

llvm/test/Transforms/InstCombine/assume-align.ll

16 ↗

(On Diff #348737)

I didn't even know that we had predicated scalar stores on IR. Either way is it a generic comment that it would be good if the stores was eventually predicated (in the final emission) or does it imply that predication should happen as part of simplify-cfg now with the added transformation?

[jdoerfert]

In D103316#2802200, @markus wrote:

We provide a function say bool isBlockEmptyExceptAssumes(BasicBlock *BB) and we update the decision making of all the existing CFG transformations to use that. We provide another function say void hoistAssumesFromBlockKnownToBeOtherwiseEmpty(BasicBlock *BB) that all those CFG transformations need to call during their transformation if they used that fact that any of the blocks checked out as true with the former function.

Now maybe there aren't that many CFG transformations where this is relevant and maybe it integrates easily, I don't know, but it does seem like a somewhat larger change.

The above sounds sensible. To keep it small you can use it only in the transformation relevant to you right now. We can update others as we go.

[markus]

Is this the way we want to integrate into SimplifyCFG? Currently only tested with the supplied lit-test.

[bjope]

Not sure I've actually captured all the logic in how we ended up here, considering the original problem was that CodeGenPrepare is too pessimistic when it removes all llvm.assume intrinsics when doing CFG transformations.

The first patch from @markus aimed at improving CodeGenPrepare by making it a bit less aggressive and keeping some llvm.assumes. And then there was an idea to make it even less aggressive by also salvaging some llvm.assume intrinsics.

Now, we have ended up with a patch modifying SimplifyCFG instead. Which obviously doesn't solve the problems we have in CodeGenPrepare. So what is the ultimate goal with this approach? Is the idea that we should remove the CFG transformations in CodeGenPrepare (at least eliminateMostlyEmptyBlocks?) in a follow up patch? Or is the idea to reuse these Utils helpers in CodeGenPrepare when doing such CFG transforms?

I mean, currently the CFG tranformations in CodeGenPrepare might depend on the existence of llvm.assume in the code. That problem still exist unless we also do something in CodeGenPrepare as well, given that the ultimate goal still is that we want to keep/salvage as many llvm.assume intrinsics as possible until ISel.
According to code comments in CodeGenPrepare those transforms are done (at least partially) as a cleanup after early IR passes in codegen such as LSR. And not all targets are running SimplifyCFG between LSR and CodeGenPrepare. Maybe the idea is to also add SimplifyCFG to the codegen pipeline just before CodeGenPrepare as a preparation to remove CodeGenPrepare::eliminateMostlyEmptyBlocks?

[jdoerfert]

I thought the problem was that a block with an assume caused us to miss out on some backed optimization.
The proposed solution was to eliminate blocks that only contain an assume as side effect early as middle-end
optimizations would also benefit from a simplified CFG.

[bjope]

In D103316#2817597, @jdoerfert wrote:

I thought the problem was that a block with an assume caused us to miss out on some backed optimization.
The proposed solution was to eliminate blocks that only contain an assume as side effect early as middle-end
optimizations would also benefit from a simplified CFG.

The problem @markus is working on is that the assumes are removed already in CodeGenPrepare. So there is no assumes left in the IR after CodeGenPrepare, and any AliasAnalysis queries after CodeGenPrepare won't benefit from the assumes.
Thus, any guidance here that isn't moving towards the goal of handling assumes better in CodeGenPrepare (preferrably keeping/salvaging as many assumes as possible) is a bit confusing.

[nikic]

@bjope This SimplifyCFG patch allows us to remove the assume dropping code in CGP subsequently.

llvm/lib/Transforms/Utils/Local.cpp

1112

This needs to be CreateLogicalOr.

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

307

You need isSafeToSpeculativelyExecute() here, not just !mayHaveSideEffects().

[markus]

Thanks @bjope for the input. I do share your concern. While I think we all agree that SimplifyCFG is the right place for this kind of transformation the steps that would follow acceptance and landing of the current patch do not seem to directly lead to being able to remove the unconditional assume elimination from CodeGenPrepare. At least not without possibly introducing other degradation as there may be other transformation that have run in between and now put assumes in inconvenient places. As pointed out by @bjope.

[markus]

Updated to address (some) review comments, clang-tidy warnings and taking the Succ->getSinglePredecessor() condition into account to avoid duplication of instructions.

llvm/include/llvm/Transforms/Utils/Local.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_UTILS_LOCAL_H
 #define LLVM_TRANSFORMS_UTILS_LOCAL_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/Utils/Local.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/Type.h"

 /// successor (BB!). Eliminate the edge between them, moving the instructions in
 /// the predecessor into BB. This deletes the predecessor block.
 void MergeBasicBlockIntoOnlyPred(BasicBlock *BB, DomTreeUpdater *DTU = nullptr);
 
 /// BB is known to contain an unconditional branch, and contains no instructions
 /// other than PHI nodes, potential debug intrinsics and the branch. If
 /// possible, eliminate BB by rewriting all the predecessors to branch to the
 /// successor block and return true. If we can't transform, return false.
-bool TryToSimplifyUncondBranchFromEmptyBlock(BasicBlock *BB,
-                                             DomTreeUpdater *DTU = nullptr);
+bool TryToSimplifyUncondBranchFromEmptyBlock(
+    BasicBlock *BB, DomTreeUpdater *DTU = nullptr,
+    SetVector<Instruction *> *AssumesAndDeps = nullptr);
 
 /// Check for and eliminate duplicate PHI nodes in this block. This doesn't try
 /// to be clever about PHI nodes which differ only in the order of the incoming
 /// values, but instcombine orders them so it usually won't matter.
 bool EliminateDuplicatePHINodes(BasicBlock *BB);
 
 /// This function is used to do simplification of a CFG.  For example, it
 /// adjusts branches to branches to eliminate the extra hop, it eliminates

llvm/lib/Transforms/Utils/Local.cpp

for (unsigned i = 0, e = BBPreds.size(); i != e; ++i) {
       // newly retargeted branch.
       PN->addIncoming(Selected, PredBB);
     }
   }
 
   replaceUndefValuesInPhi(PN, IncomingValues);
 }
 
-bool llvm::TryToSimplifyUncondBranchFromEmptyBlock(BasicBlock *BB,
-                                                   DomTreeUpdater *DTU) {
+bool llvm::TryToSimplifyUncondBranchFromEmptyBlock(
+    BasicBlock *BB, DomTreeUpdater *DTU,
+    SetVector<Instruction *> *AssumesAndDeps) {
   assert(BB != &BB->getParent()->getEntryBlock() &&
          "TryToSimplifyUncondBranchFromEmptyBlock called on entry block!");
 
   // We can't eliminate infinite loops.
   BasicBlock *Succ = cast<BranchInst>(BB->getTerminator())->getSuccessor(0);
   if (BB == Succ) return false;
 
   // Check to see if merging these blocks would cause conflicts for any of the
if (auto *CBI = dyn_cast<CallBrInst>(PredBB->getTerminator())) {
       for (unsigned i = 0, e = CBI->getNumIndirectDests(); i != e; ++i)
         if (Succ == CBI->getIndirectDest(i))
           return false;
     }
   }
 
   LLVM_DEBUG(dbgs() << "Killing Trivial BB: \n" << *BB);
 
+  if (AssumesAndDeps) {
+    // Hoist @llvm.assume into predecessors.
+    for (BasicBlock *PredBB : predecessors(BB)) {
+      auto *PredBranch = dyn_cast<BranchInst>(PredBB->getTerminator());
+      if (!PredBranch || PredBranch->isUnconditional())
+        continue;
+
+      // Now clone AssumesAndDeps into Pred. When we get to a @llvm.assume
+      // update the assume condition to depend on PredBranch->getCondition().
+      SmallDenseMap<Value *, Value *, 8> VMap;
+      auto newIfAvailable = [&VMap](Value *OldV) {
+        if (VMap.count(OldV))
+          return VMap[OldV];
+        else
+          return OldV;
+      };
+      Module *M = BB->getParent()->getParent();
+      IRBuilder<> IRB(M->getContext());
+      IRB.SetInsertPoint(PredBranch);
+
+      Value *Cond = PredBranch->getCondition();
+      auto *PredTrueSucc = PredBranch->getSuccessor(0);
+      if (BB == PredTrueSucc)
+        Cond = IRB.CreateNot(Cond);
+
+      for (auto It = AssumesAndDeps->rbegin(), E = AssumesAndDeps->rend();
+           It != E; ++It) {
+        if (auto *Assume = dyn_cast<AssumeInst>(*It)) {
+          auto *OrCond =
+              IRB.CreateOr(Cond, newIfAvailable(Assume->getOperand(0)));
+          // Note that any Operand Bundle is dropped at this point.
+          IRB.CreateAssumption(OrCond);
+        } else {
+          auto *OldI = *It;
+          auto *NewI = OldI->clone();
+          for (unsigned I = 0, E = OldI->getNumOperands(); I < E; ++I)
+            NewI->setOperand(I, newIfAvailable(OldI->getOperand(I)));
+          VMap[OldI] = NewI;
+          NewI->insertBefore(PredBranch);
+        }
+      }
+    }
+  }
+
   SmallVector<DominatorTree::UpdateType, 32> Updates;
   if (DTU) {
     // All predecessors of BB will be moved to Succ.
     SmallPtrSet<BasicBlock *, 8> PredsOfBB(pred_begin(BB), pred_end(BB));
     SmallPtrSet<BasicBlock *, 8> PredsOfSucc(pred_begin(Succ), pred_end(Succ));
     Updates.reserve(Updates.size() + 2 * PredsOfBB.size() + 1);
     for (auto *PredOfBB : PredsOfBB)
       // This predecessor of BB may already have Succ as a successor.

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

public:
   bool requestResimplify() {
     Resimplify = true;
     return true;
   }
 };
 
 } // end anonymous namespace
 
+static bool
+isBlockEmptyExceptAssumes(BasicBlock *BB,
+                          SetVector<Instruction *> &AssumesAndDeps) {
+  auto *Pred = BB->getSinglePredecessor();
+  if (!Pred)
+    return false;
+
+  for (auto It = BB->rbegin(), E = BB->rend(); It != E; ++It) {
+    auto I = &(*It);
+    // Skip the terminator.
+    if (I->isTerminator())
+      continue;
+
+    // Bail if side-effects.
+    if (I->mayHaveSideEffects() && !isa<AssumeInst>(I))
+      return false;
+
+    // Insert into AssumesAndDeps either if instruction is a @llvm.assume or if
+    // all uses of the instruction are covered by the AssumesAndDeps set.
+    // Otherwise bail.
+    if (isa<AssumeInst>(I))
+      AssumesAndDeps.insert(I);
+    else if (!empty(I->users()) && all_of(I->users(), [&](User *U) {
+               return isa<Instruction>(U)
+                          ? AssumesAndDeps.count(cast<Instruction>(U))
+                          : false;
+             }))
+      AssumesAndDeps.insert(I);
+    else
+      return false;
+  }
+
+  return true;
+}
+
 /// Return true if it is safe to merge these two
 /// terminator instructions together.
 static bool
 SafeToMergeTerminators(Instruction *SI1, Instruction *SI2,
                        SmallSetVector<BasicBlock *, 4> *FailBlocks = nullptr) {
   if (SI1 == SI2)
     return false; // Can't merge with self!
 
bool SimplifyCFGOpt::simplifyUncondBranch(BranchInst *BI,
   // can be eliminated when the pass is invoked later in the back-end.)
   // Note that if BB has only one predecessor then we do not introduce new
   // backedge, so we can eliminate BB.
   bool NeedCanonicalLoop =
       Options.NeedCanonicalLoop &&
       (!LoopHeaders.empty() && BB->hasNPredecessorsOrMore(2) &&
        (is_contained(LoopHeaders, BB) || is_contained(LoopHeaders, Succ)));
   BasicBlock::iterator I = BB->getFirstNonPHIOrDbg(true)->getIterator();
-  if (I->isTerminator() && BB != &BB->getParent()->getEntryBlock() &&
-      !NeedCanonicalLoop && TryToSimplifyUncondBranchFromEmptyBlock(BB, DTU))
+  SetVector<Instruction *> AssumesAndDeps;
+  if (isBlockEmptyExceptAssumes(BB, AssumesAndDeps) &&
+      BB != &BB->getParent()->getEntryBlock() && !NeedCanonicalLoop &&
+      TryToSimplifyUncondBranchFromEmptyBlock(BB, DTU, &AssumesAndDeps))
     return true;
 
   // If the only instruction in the block is a seteq/setne comparison against a
   // constant, try to simplify the block.
   if (ICmpInst *ICI = dyn_cast<ICmpInst>(I))
     if (ICI->isEquality() && isa<ConstantInt>(ICI->getOperand(1))) {
       for (++I; isa<DbgInfoIntrinsic>(I); ++I)
         ;
bool SimplifyCFGOpt::simplifyOnceImpl(BasicBlock *BB) {
   Changed |= removeUndefIntroducingPredecessor(BB, DTU);
 
   // Merge basic blocks into their predecessor if there is only one distinct
   // pred, and if there is only one distinct successor of the predecessor, and
   // if there are no PHI nodes.
   if (MergeBlockIntoPredecessor(BB, DTU))
     return true;
 
   if (SinkCommon && Options.SinkCommonInsts)
     Changed |= SinkCommonCodeFromPredecessors(BB, DTU);
 
   IRBuilder<> Builder(BB);
 
   if (Options.FoldTwoEntryPHINode) {
     // If there is a trivial two-entry PHI node in this basic block, and we can
     // eliminate it, do so now.
     if (auto *PN = dyn_cast<PHINode>(BB->begin()))

llvm/test/Transforms/SimplifyCFG/hoist-assume.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -passes=simplify-cfg < %s 2>&1 | FileCheck %s
+
+declare void @llvm.assume(i1 noundef)
+
+define dso_local void @f(i32 %a, i32 %b) {
+; CHECK-LABEL: @f(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[A:%.*]], 0
+; CHECK-NEXT:    [[TMP0:%.*]] = xor i1 [[COND]], true
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp sgt i32 [[B:%.*]], 5
+; CHECK-NEXT:    [[TMP2:%.*]] = or i1 [[TMP0]], [[TMP1]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[TMP2]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %cond = icmp eq i32 %a, 0
+  br i1 %cond, label %if.then, label %end
+
+if.then:
+  %cond2 = icmp sgt i32 %b, 5
+  call void @llvm.assume(i1 %cond2)
+  br label %end
+
+end:
+  ret void
+}
+
+define dso_local void @g(i32 %a, i32 %b) {
+; CHECK-LABEL: @g(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[A:%.*]], 0
+; CHECK-NEXT:    [[TMP0:%.*]] = xor i1 [[COND]], true
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp sgt i32 [[B:%.*]], 5
+; CHECK-NEXT:    [[TMP2:%.*]] = or i1 [[TMP0]], [[TMP1]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[TMP2]])
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp sle i32 [[B]], 5
+; CHECK-NEXT:    [[TMP4:%.*]] = or i1 [[COND]], [[TMP3]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[TMP4]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %cond = icmp eq i32 %a, 0
+  br i1 %cond, label %if.then, label %if.else
+
+if.then:
+  %cond2 = icmp sgt i32 %b, 5
+  call void @llvm.assume(i1 %cond2)
+  br label %end
+
+if.else:
+  %cond3 = icmp sle i32 %b, 5
+  call void @llvm.assume(i1 %cond3)
+  br label %end
+
+end:
+  ret void
+}