Ignore/Drop droppable uses for code-sinking in InstCombine

[Tyker]

Summary

This patch allows code-sinking in InstCombine to be performed when instruction have uses in llvm.assume.

Use are considered droppable when it is preferable to modify the User such that the use disappears rather than to prevent a transformation because of the use.
for now uses are considered droppable if they are in an llvm.assume.

[jdoerfert]

Please add a commit message on this one and maybe remove the WIP tag as it is supposed to be reviewed and actually working, right?

[Tyker]

In D73832#1853677, @jdoerfert wrote:

Please add a commit message on this one and maybe remove the WIP tag as it is supposed to be reviewed and actually working, right?

the title kind says it all. yes it is working as expected and should be ready for review.

code-sinking in InstCombine is far from as aggressive as it could be. is it also done somewhere else ?
i think that instead of a single use in the successor block we could just code-sink to the successor block dominate all undroppable uses of the instruction if it exists.
or maybe find the highest common dominator of all undroppable uses and if it is not the current basic block code-sink the instruction to it.

[jdoerfert]

In D73832#1853708, @Tyker wrote:

In D73832#1853677, @jdoerfert wrote:

Please add a commit message on this one and maybe remove the WIP tag as it is supposed to be reviewed and actually working, right?

the title kind says it all. yes it is working as expected and should be ready for review.

Arguably, "droppable uses" is a new concept you need to explain in the commit message.

code-sinking in InstCombine is far from as aggressive as it could be. is it also done somewhere else ?
i think that instead of a single use in the successor block we could just code-sink to the successor block dominate all undroppable uses of the instruction if it exists.
or maybe find the highest common dominator of all undroppable uses and if it is not the current basic block code-sink the instruction to it.

There is for sure potential, I think we have another sinking pass but I'd need to check.

[Tyker]

rebased

[jdoerfert]

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

263

We do we disallow empty BundleOpInfo here?

llvm/test/Transforms/InstCombine/assume.ll

274

Why don't we sink this one anymore?

425

This test opens up the question if we want to sink simple instructions when we need to destroy assumes. I guess that is unrelated to this patch though.

450

Could you add a test case with a FIXME that show we drop too many assumes when we sink, e.g. the ones that are still dominated.

[Tyker]

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

263

i am not sure i understood the question.

we will return true for empty bundles.

this is a simple utility to check if the operand bundle of an llvm.assume still holds any information. it is used in a few places because instcombine was removing instructions like call void @llvm.assume(i1 true) ["nonnull"(%p)] because the argument to the call is known to be true so the assume was considered useless which is not true anymore.

llvm/test/Transforms/InstCombine/assume.ll

274

in this case there is a single use which is droppable (in an llvm.assume). so hasOneUse return true, but getSingleUndroppableUse returns null. so the old passe will sink the instruction whereas the new won't.

this doesn't matter however because this instruction is only used in an @llvm.assume so it is already dead and the backend will it cleanup.

[jdoerfert]

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

263

Oh, empty bundles like call void @llvm.assume(i1 true) ["cold"()] ?
Maybe we should put a comment here because I will forget again.

Nit: no != nullptr

[Tyker]

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

263

call void @llvm.assume(i1 true) ["cold"()] is not considered empty.

this function returns true for truly empty/absent bundles like call void @llvm.assume(i1 true).
or for bundles for which there was WasOn arguments but they have all been dropped. for example if we had call void @llvm.assume(i1 true) ["nonnull"(%p)] and the use of %p get gets dropped. all the only information the assume contains is that a value used to be nonnull, which is basically useless. so the function returns true for those bundles as well.

the general rule is this function returns true if the operand bundle of the assume doesn't hold any usable information. so if this function returns true and the argument of the assume is always true the whole assume should be removed.

i realize that this function probably deserve a more detailed comment.

[lebedev.ri]

llvm/lib/IR/AsmWriter.cpp

2557

Would it be better to ensure that we can't end up in such
seemingly-degenerate state in the first place,
by cleaning up bundles?

[Tyker]

llvm/lib/IR/AsmWriter.cpp

2557

i agree than not having that state would be better but i don't think it is possible without rebuilding instruction after we dropped uses in them.

an other solution could be to mark dropped uses with an undef value instead of a nullptr.

[Tyker]

llvm/lib/IR/AsmWriter.cpp

2557

i tried to find a better solution to represent "dropped" values. but not solution seems to be good.

• nullptr is problematic because many place assume that operands are never null.
• rebuilding instructions seems like too expensive for what we are trying to do.

maybe adding a new kind of value to represent it, that is ignored everywhere. but changing the IR for this seems like overkill.

any thought ?

[jdoerfert]

llvm/lib/IR/AsmWriter.cpp

2557

I like the undef solution if we clarify in the LangRef that an assume on undef is ignored.

If we could change the bundle tag we could set it to "ignore", and then pass undef.

llvm/test/Transforms/InstCombine/assume.ll

274

I see.

[Tyker]

rebased this patch on the changes of previous patches in the stack.

[Tyker]

rebased

[jdoerfert]

llvm/lib/Transforms/Utils/Local.cpp

415

Can you split the isAssumeWithEmptyBundle changes. They look good to me and can be commited, assuming you can add a simple test case showing that instcombine/instsimplify/... will remove certain assume(true) but not all.

For the Instcombine sutff I would probably prefer another opinion (@lebedev.ri @spatel @nikic)

[Tyker]

splitted into this patch and D76147

[jdoerfert]

I like this, let's see if someone else want to comment.

[jdoerfert]

LGTM.

[nikic]

Just a general question from my side: In what direction is this heading?

Is this particular transform (InstCombine code sinking) something where assumes were observed to be problematic in practice? Or is this more of a proof-of-concept that will be extended to other one-use checks in InstCombine? (There's a lot of them, and I'm slightly concerned about the amount of boilerplate that would have to be introduced to handle this everywhere.)

[Tyker]

In D73832#1933035, @nikic wrote:

Just a general question from my side: In what direction is this heading?

Is this particular transform (InstCombine code sinking) something where assumes were observed to be problematic in practice? Or is this more of a proof-of-concept that will be extended to other one-use checks in InstCombine? (There's a lot of them, and I'm slightly concerned about the amount of boilerplate that would have to be introduced to handle this everywhere.)

this is intended as a proof of concept, code-sinking seemed like an easy place to start. I agree with you that extending it to all of instcombine not going to be easy particularly on the testing side.
but we are going to add more assumes with patches like D75825 (disabled by default for now). so we need to minimize the negative impact of assume on other transformations. it is easier to do for assume with operand bundles rather than boolean assumes.

[MaskRay]

This patch caused an assertion failure on UserParent = PN->getIncomingBlock(*I->use_begin()); .

This can be triggered when a bootstrapped clang builds llvm-symbolizer with ThinLTO with minimized bitcode files clang -c -fthinlto-index=AutoUpgrade.o.thinlto.bc AutoUpgrade.bc -O3 (llvm/lib/IR/AutoUpgrade.cpp).

I am trying to figure out a minimal reproduce.

[MaskRay]

In D73832#1942542, @MaskRay wrote:

This patch caused an assertion failure on UserParent = PN->getIncomingBlock(*I->use_begin()); .

This can be triggered when a bootstrapped clang builds llvm-symbolizer with ThinLTO with minimized bitcode files clang -c -fthinlto-index=AutoUpgrade.o.thinlto.bc AutoUpgrade.bc -O3 (llvm/lib/IR/AutoUpgrade.cpp).

I am trying to figure out a minimal reproduce.

Pushed a probable fix: 4c52d51e784ca7603969f2a664e5430d41b9d5dc.

[nikic]

For the record, this change has caused a small but measurable compile-time regression (http://llvm-compile-time-tracker.com/compare.php?from=b539f18c565656cdd49fc0c37efbaa8e584b2d65&to=f1a9efabcb9bc37b663b0e03ed3d5a5aa7cc055e&stat=instructions). It's not a problem for this change taken by itself, but may be a problem if these APIs are going to see significantly wider use.

[Tyker]

In D73832#1942924, @MaskRay wrote:

Pushed a probable fix: 4c52d51e784ca7603969f2a664e5430d41b9d5dc.

thank you

In D73832#1943096, @nikic wrote:

For the record, this change has caused a small but measurable compile-time regression (http://llvm-compile-time-tracker.com/compare.php?from=b539f18c565656cdd49fc0c37efbaa8e584b2d65&to=f1a9efabcb9bc37b663b0e03ed3d5a5aa7cc055e&stat=instructions). It's not a problem for this change taken by itself, but may be a problem if these APIs are going to see significantly wider use.

i added a patch to minimize the overhead by storing inside the Use whether it is droppable. D76923.

llvm/include/llvm/Transforms/Utils/KnowledgeRetention.h

 inline bool hasAttributeInAssume(CallInst &AssumeCI, Value *IsOn,
                                  Attribute::AttrKind Kind,
                                  uint64_t *ArgVal = nullptr,
                                  AssumeQuery AQR = AssumeQuery::Highest) {
   return hasAttributeInAssume(
       AssumeCI, IsOn, Attribute::getNameFromAttrKind(Kind), ArgVal, AQR);
 }
 
+/// Check if the assume still holds any information.
+/// If this returns false this assume can be dropped without losing anyting.
+bool isAssumeWithEmptyBundle(CallInst &Assume);
+
 /// TODO: Add an function to create/fill a map from the bundle when users intend
 /// to make many different queries on the same bundles. to be used for example
 /// in the Attributor.
 
 //===----------------------------------------------------------------------===//
 // Utilities for testing
 //===----------------------------------------------------------------------===//
 

llvm/lib/IR/AsmWriter.cpp

if (!Call->hasOperandBundles())
     return;
 
   Out << " [ ";
 
   bool FirstBundle = true;
   for (unsigned i = 0, e = Call->getNumOperandBundles(); i != e; ++i) {
     OperandBundleUse BU = Call->getOperandBundleAt(i);
 
+    /// Bundles containing dropped values should be ignored.
+    if (any_of(BU.Inputs, [](const Use &U) { return !U.get(); }))
+      continue;
+
     if (!FirstBundle)
       Out << ", ";
     FirstBundle = false;
 
     Out << '"';
     printEscapedString(BU.getTagName(), Out);
     Out << '"';
 

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
-#include "llvm/IR/IntrinsicsX86.h"
-#include "llvm/IR/IntrinsicsARM.h"
 #include "llvm/IR/IntrinsicsAArch64.h"
-#include "llvm/IR/IntrinsicsNVPTX.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
+#include "llvm/IR/IntrinsicsARM.h"
+#include "llvm/IR/IntrinsicsNVPTX.h"
 #include "llvm/IR/IntrinsicsPowerPC.h"
+#include "llvm/IR/IntrinsicsX86.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/Statepoint.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Support/AtomicOrdering.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/InstCombine/InstCombineWorklist.h"
+#include "llvm/Transforms/Utils/KnowledgeRetention.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/SimplifyLibCalls.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <cstring>
 #include <utility>
 #include <vector>
if (match(IIOperand, m_ICmp(Pred, m_Instruction(LHS), m_Zero())) &&
       // TODO: apply nonnull return attributes to calls and invokes
       // TODO: apply range metadata for range check patterns?
     }
 
     // If there is a dominating assume with the same condition as this one,
     // then this one is redundant, and should be removed.
     KnownBits Known(1);
     computeKnownBits(IIOperand, Known, 0, II);
-    if (Known.isAllOnes())
+    if (Known.isAllOnes() && isAssumeWithEmptyBundle(*II))
       return eraseInstFromFunction(*II);
 
     // Update the cache of affected values for this assumption (we might be
     // here because we just simplified the condition).
     AC.updateAffectedValues(II);
     break;
   }
   case Intrinsic::experimental_gc_relocate: {

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

Instruction *InstCombiner::visitFreeze(FreezeInst &I) {
   return nullptr;
 }
 
 /// Try to move the specified instruction from its current block into the
 /// beginning of DestBlock, which can only happen if it's safe to move the
 /// instruction past all of the instructions between it and the end of its
 /// block.
 static bool TryToSinkInstruction(Instruction *I, BasicBlock *DestBlock) {
-  assert(I->hasOneUse() && "Invariants didn't hold!");
+  assert(I->getSingleUndroppableUse() && "Invariants didn't hold!");
   BasicBlock *SrcBlock = I->getParent();
 
   // Cannot move control-flow-involving, volatile loads, vaarg, etc.
   if (isa<PHINode>(I) || I->isEHPad() || I->mayHaveSideEffects() ||
       I->isTerminator())
     return false;
 
   // Do not sink static or dynamic alloca instructions. Static allocas must
static bool TryToSinkInstruction(Instruction *I, BasicBlock *DestBlock) {
   // the end of block that could change the value.
   if (I->mayReadFromMemory()) {
     for (BasicBlock::iterator Scan = I->getIterator(),
                               E = I->getParent()->end();
          Scan != E; ++Scan)
       if (Scan->mayWriteToMemory())
         return false;
   }
+
+  I->dropDroppableUses([DestBlock](const Use *U) {
+    if (auto *I = dyn_cast<Instruction>(U->getUser()))
+      return I->getParent() != DestBlock;
+    return true;
+  });
+  /// FIXME: We could remove droppable uses that are not dominated by
+  /// the new position.
+
   BasicBlock::iterator InsertPos = DestBlock->getFirstInsertionPt();
   I->moveBefore(&*InsertPos);
   ++NumSunkInst;
 
   // Also sink all related debug uses from the source basic block. Otherwise we
   // get debug use before the def. Attempt to salvage debug uses first, to
   // maximise the range variables have location for. If we cannot salvage, then
   // mark the location undef: we know it was supposed to receive a new location
if (ExpensiveCombines && !I->use_empty() && Ty->isIntOrIntVectorTy()) {
         if (isInstructionTriviallyDead(I, &TLI))
           eraseInstFromFunction(*I);
         MadeIRChange = true;
         continue;
       }
     }
 
     // See if we can trivially sink this instruction to a successor basic block.
-    if (EnableCodeSinking && I->hasOneUse()) {
+    if (EnableCodeSinking)
+      if (Use *SingleUse = I->getSingleUndroppableUse()) {
         BasicBlock *BB = I->getParent();
-      Instruction *UserInst = cast<Instruction>(*I->user_begin());
+        Instruction *UserInst = cast<Instruction>(SingleUse->getUser());
         BasicBlock *UserParent;
 
         // Get the block the use occurs in.
         if (PHINode *PN = dyn_cast<PHINode>(UserInst))
           UserParent = PN->getIncomingBlock(*I->use_begin());
         else
           UserParent = UserInst->getParent();
 
         if (UserParent != BB) {
           bool UserIsSuccessor = false;
           // See if the user is one of our successors.
-        for (succ_iterator SI = succ_begin(BB), E = succ_end(BB); SI != E; ++SI)
+          for (succ_iterator SI = succ_begin(BB), E = succ_end(BB); SI != E;
+               ++SI)
             if (*SI == UserParent) {
               UserIsSuccessor = true;
               break;
             }
 
-        // If the user is one of our immediate successors, and if that successor
-        // only has us as a predecessors (we'd have to split the critical edge
-        // otherwise), we can keep going.
+          // If the user is one of our immediate successors, and if that
+          // successor only has us as a predecessors (we'd have to split the
+          // critical edge otherwise), we can keep going.
           if (UserIsSuccessor && UserParent->getUniquePredecessor()) {
             // Okay, the CFG is simple enough, try to sink this instruction.
             if (TryToSinkInstruction(I, UserParent)) {
               LLVM_DEBUG(dbgs() << "IC: Sink: " << *I << '\n');
               MadeIRChange = true;
               // We'll add uses of the sunk instruction below, but since sinking
               // can expose opportunities for it's *operands* add them to the
               // worklist
               for (Use &U : I->operands())
                 if (Instruction *OpI = dyn_cast<Instruction>(U.get()))
                   Worklist.push(OpI);
             }
           }
         }
       }
 
     // Now that we have an instruction, try combining it to simplify it.
     Builder.SetInsertPoint(I);
     Builder.SetCurrentDebugLocation(I->getDebugLoc());
 
 #ifndef NDEBUG
     std::string OrigI;
 #endif

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

 } // namespace
 
 CallInst *llvm::BuildAssumeFromInst(const Instruction *I, Module *M) {
   AssumeBuilderState Builder(M);
   Builder.addInstruction(I);
   return Builder.build();
 }
 
+Value *getValueFromBundleOpInfo(CallInst &Assume,
+                                const CallBase::BundleOpInfo &BOI,
+                                unsigned Idx) {
+  assert(BOI.End - BOI.Begin > Idx && "index out of range");
+  return (Assume.op_begin() + BOI.Begin + Idx)->get();
+};
+
 #ifndef NDEBUG
 
 static bool isExistingAttribute(StringRef Name) {
   return StringSwitch<bool>(Name)
 #define GET_ATTR_NAMES
 #define ATTRIBUTE_ALL(ENUM_NAME, DISPLAY_NAME) .Case(#DISPLAY_NAME, true)
 #include "llvm/IR/Attributes.inc"
       .Default(false);
else
     Lookup = std::prev(
         llvm::upper_bound(Assume.bundle_op_infos(), AttrName,
                           [](StringRef LHS, const CallBase::BundleOpInfo &BOI) {
                             assert(isExistingAttribute(BOI.Tag->getKey()) &&
                                    "this attribute doesn't exist");
                             return LHS < BOI.Tag->getKey();
                           }));
 
-  auto getValueFromBundleOpInfo = [&Assume](const CallBase::BundleOpInfo &BOI,
-                                            unsigned Idx) {
-    assert(BOI.End - BOI.Begin > Idx && "index out of range");
-    return (Assume.op_begin() + BOI.Begin + Idx)->get();
-  };
-
   if (Lookup == Assume.bundle_op_info_end() ||
       Lookup->Tag->getKey() != AttrName)
     return false;
   if (IsOn) {
     if (Lookup->End - Lookup->Begin < BOIE_WasOn)
       return false;
     while (true) {
       if (Lookup == Assume.bundle_op_info_end() ||
           Lookup->Tag->getKey() != AttrName)
         return false;
-      if (getValueFromBundleOpInfo(*Lookup, BOIE_WasOn) == IsOn)
+      if (getValueFromBundleOpInfo(Assume, *Lookup, BOIE_WasOn) == IsOn)
         break;
       if (AQR == AssumeQuery::Highest &&
           Lookup == Assume.bundle_op_info_begin())
         return false;
       Lookup = Lookup + (AQR == AssumeQuery::Lowest ? 1 : -1);
     }
   }
 
   if (Lookup->End - Lookup->Begin < BOIE_Argument)
     return true;
   if (ArgVal)
-    *ArgVal =
-        cast<ConstantInt>(getValueFromBundleOpInfo(*Lookup, BOIE_Argument))
+    *ArgVal = cast<ConstantInt>(
+                  getValueFromBundleOpInfo(Assume, *Lookup, BOIE_Argument))
                   ->getZExtValue();
   return true;
 }
 
+bool llvm::isAssumeWithEmptyBundle(CallInst &CI) {
+  IntrinsicInst &Assume = cast<IntrinsicInst>(CI);
+  assert(Assume.getIntrinsicID() == Intrinsic::assume &&
+         "this function is intended to be used on llvm.assume");
+  return none_of(
+      Assume.bundle_op_infos(), [&Assume](const CallBase::BundleOpInfo &BOI) {
+        return BOI.Begin == BOI.End ||
+               getValueFromBundleOpInfo(Assume, BOI, BOIE_WasOn) != nullptr;
+      });
+}
+
 PreservedAnalyses AssumeBuilderPass::run(Function &F,
                                          FunctionAnalysisManager &AM) {
   for (Instruction &I : instructions(F))
     if (Instruction *Assume = BuildAssumeFromInst(&I))
       Assume->insertBefore(&I);
   return PreservedAnalyses::all();
 }

llvm/lib/Transforms/Utils/Local.cpp

 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Utils/KnowledgeRetention.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 #include <algorithm>
 #include <cassert>
 #include <climits>
 #include <cstdint>
 #include <iterator>
 #include <map>
 #include <utility>
if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
     // Lifetime intrinsics are dead when their right-hand is undef.
     if (II->isLifetimeStartOrEnd())
       return isa<UndefValue>(II->getArgOperand(1));
 
     // Assumptions are dead if their condition is trivially true.  Guards on
     // true are operationally no-ops.  In the future we can consider more
     // sophisticated tradeoffs for guards considering potential for check
     // widening, but for now we keep things simple.
-    if (II->getIntrinsicID() == Intrinsic::assume ||
+    if ((II->getIntrinsicID() == Intrinsic::assume &&
+         isAssumeWithEmptyBundle(*II)) ||
         II->getIntrinsicID() == Intrinsic::experimental_guard) {
       if (ConstantInt *Cond = dyn_cast<ConstantInt>(II->getArgOperand(0)))
         return !Cond->isZero();
 
       return false;
     }
   }
 

llvm/test/Transforms/InstCombine/assume.ll

 
 ; Make sure the above canonicalization does not trigger
 ; if the assume is control dependent on something else
 
 define i1 @nonnull3(i32** %a, i1 %control) {
 ; CHECK-LABEL: @nonnull3(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[LOAD:%.*]] = load i32*, i32** [[A:%.*]], align 8
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32* [[LOAD]], null
 ; CHECK-NEXT:    br i1 [[CONTROL:%.*]], label [[TAKEN:%.*]], label [[NOT_TAKEN:%.*]]
 ; CHECK:       taken:
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32* [[LOAD]], null
 ; CHECK-NEXT:    tail call void @llvm.assume(i1 [[CMP]])
 ; CHECK-NEXT:    ret i1 false
 ; CHECK:       not_taken:
 ; CHECK-NEXT:    [[RVAL_2:%.*]] = icmp sgt i32* [[LOAD]], null
 ; CHECK-NEXT:    ret i1 [[RVAL_2]]
 ;
 entry:
   %load = load i32*, i32** %a

 ;
   %shuf = shufflevector <4 x i8> %x, <4 x i8> undef, <4 x i32> <i32 1, i32 1, i32 2, i32 3>
   %t2 = bitcast <4 x i8> %shuf to i32
   %t3 = icmp ult i32 %t2, 65536
   call void @llvm.assume(i1 %t3)
   ret i32 %t2
 }
 
+define i1 @nonnull3A(i32** %a, i1 %control) {
+; CHECK-LABEL: @nonnull3A(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32*, i32** [[A:%.*]], align 8
+; CHECK-NEXT:    br i1 [[CONTROL:%.*]], label [[TAKEN:%.*]], label [[NOT_TAKEN:%.*]]
+; CHECK:       taken:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32* [[LOAD]], null
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    ret i1 true
+; CHECK:       not_taken:
+; CHECK-NEXT:    [[RVAL_2:%.*]] = icmp sgt i32* [[LOAD]], null
+; CHECK-NEXT:    ret i1 [[RVAL_2]]
+;
+entry:
+  %load = load i32*, i32** %a
+  %cmp = icmp ne i32* %load, null
+  br i1 %control, label %taken, label %not_taken
+taken:
+  tail call void @llvm.assume(i1 %cmp)
+  ret i1 %cmp
+not_taken:
+  tail call void @llvm.assume(i1 %cmp)
+  %rval.2 = icmp sgt i32* %load, null
+  ret i1 %rval.2
+}
+
+define i1 @nonnull3B(i32** %a, i1 %control) {
+; CHECK-LABEL: @nonnull3B(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[CONTROL:%.*]], label [[TAKEN:%.*]], label [[NOT_TAKEN:%.*]]
+; CHECK:       taken:
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32*, i32** [[A:%.*]], align 8
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32* [[LOAD]], null
+; CHECK-NEXT:    tail call void @llvm.assume(i1 true) [ "nonnull"(i32* [[LOAD]]) ]
+; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK:       not_taken:
+; CHECK-NEXT:    ret i1 [[CONTROL]]
+;
+entry:
+  %load = load i32*, i32** %a
+  %cmp = icmp ne i32* %load, null
+  br i1 %control, label %taken, label %not_taken
+taken:
+  tail call void @llvm.assume(i1 true) ["nonnull"(i32* %load)]
+  ret i1 %cmp
+not_taken:
+  tail call void @llvm.assume(i1 %cmp)
+  ret i1 %control
+}
+
 declare void @llvm.dbg.value(metadata, metadata, metadata)
 
 !llvm.dbg.cu = !{!0}
 !llvm.module.flags = !{!5, !6, !7, !8}
 
 !0 = distinct !DICompileUnit(language: DW_LANG_C, file: !3, producer: "Me", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug, enums: null, retainedTypes: null, imports: null)
 !1 = !DILocalVariable(name: "", arg: 1, scope: !2, file: null, line: 1, type: null)
 !2 = distinct !DISubprogram(name: "debug", linkageName: "debug", scope: null, file: null, line: 0, type: null, isLocal: false, isDefinition: true, scopeLine: 1, flags: DIFlagPrototyped, isOptimized: true, unit: !0)

llvm/unittests/Transforms/Utils/KnowledgeRetentionTest.cpp

Tests.push_back(std::make_pair(
         AssertHasTheRightValue(Assume, I->getOperand(2),
                                Attribute::AttrKind::Dereferenceable, 4, true);
         AssertHasTheRightValue(Assume, I->getOperand(3),
                                Attribute::AttrKind::Alignment, 16, true);
         AssertHasTheRightValue(Assume, I->getOperand(3),
                                Attribute::AttrKind::Dereferenceable, 12, true);
       }));
 
-  /// Keep this test last as it modifies the function.
+  /// The tests below modify the code keep them last to not interfere with
+  /// others.
+  Tests.push_back(std::make_pair(
+      "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
+      "8 noalias %P1)\n",
+      [](Instruction *I) {
+        CallInst *Assume = BuildAssumeFromInst(I);
+        Assume->insertBefore(I);
+        Value *Val = I->getOperand(0);
+        AssertMatchesExactlyAttributes(Assume, Val,
+                                       "(nonnull|align|dereferenceable)");
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(1),
+                                       "(noalias|align)");
+
+        Val->dropDroppableUses();
+        AssertMatchesExactlyAttributes(Assume, Val, "");
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(1),
+                                       "(noalias|align)");
+      }));
+  Tests.push_back(
+      std::make_pair("%in_assume = icmp ne i32* %P, null\n", [](Instruction *I) {}));
+  Tests.push_back(
+      std::make_pair("call void @llvm.assume(i1 %in_assume)\n", [](Instruction *I) {
+        I->getOperand(0)->dropDroppableUses();
+        ASSERT_EQ(I->getOperand(0), ConstantInt::getTrue(I->getContext()));
+      }));
   Tests.push_back(std::make_pair(
       "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
       "8 noalias %P1)\n",
       [](Instruction *I) {
         CallInst *Assume = BuildAssumeFromInst(I);
         Assume->insertBefore(I);
         Value *New = I->getFunction()->getArg(3);
         Value *Old = I->getOperand(0);
         AssertMatchesExactlyAttributes(Assume, New, "");
         AssertMatchesExactlyAttributes(Assume, Old,
                                        "(nonnull|align|dereferenceable)");
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(1),
+                                       "(noalias|align)");
         Old->replaceAllUsesWith(New);
         AssertMatchesExactlyAttributes(Assume, New,
                                        "(nonnull|align|dereferenceable)");
         AssertMatchesExactlyAttributes(Assume, Old, "");
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(1),
+                                       "(noalias|align)");
       }));
   RunTest(Head, Tail, Tests);
 }