diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index 90bbd102868c5..18fa10b09a66a 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3604,23 +3604,6 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
       return eraseInstFromFunction(*II);
     }
 
-    // assume( (load addr) != null ) -> add 'nonnull' metadata to load
-    // (if assume is valid at the load)
-    Instruction *LHS;
-    if (match(IIOperand, m_SpecificICmp(ICmpInst::ICMP_NE, m_Instruction(LHS),
-                                        m_Zero())) &&
-        LHS->getOpcode() == Instruction::Load &&
-        LHS->getType()->isPointerTy() &&
-        isValidAssumeForContext(II, LHS, &DT)) {
-      MDNode *MD = MDNode::get(II->getContext(), {});
-      LHS->setMetadata(LLVMContext::MD_nonnull, MD);
-      LHS->setMetadata(LLVMContext::MD_noundef, MD);
-      return RemoveConditionFromAssume(II);
-
-      // TODO: apply nonnull return attributes to calls and invokes
-      // TODO: apply range metadata for range check patterns?
-    }
-
     for (unsigned Idx = 0; Idx < II->getNumOperandBundles(); Idx++) {
       OperandBundleUse OBU = II->getOperandBundleAt(Idx);
 
@@ -3675,11 +3658,24 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
       if (OBU.getTagName() == "nonnull" && OBU.Inputs.size() == 1) {
         RetainedKnowledge RK = getKnowledgeFromOperandInAssume(
             *cast<AssumeInst>(II), II->arg_size() + Idx);
-        if (!RK || RK.AttrKind != Attribute::NonNull ||
-            !isKnownNonZero(RK.WasOn,
-                            getSimplifyQuery().getWithInstruction(II)))
+        if (!RK || RK.AttrKind != Attribute::NonNull)
           continue;
-        return CallBase::removeOperandBundle(II, OBU.getTagID());
+
+        // Drop assume if we can prove nonnull without it
+        if (isKnownNonZero(RK.WasOn, getSimplifyQuery().getWithInstruction(II)))
+          return CallBase::removeOperandBundle(II, OBU.getTagID());
+
+        // Fold the assume into metadata if it's valid at the load
+        if (auto *LI = dyn_cast<LoadInst>(RK.WasOn);
+            LI &&
+            isValidAssumeForContext(II, LI, &DT, /*AllowEphemerals=*/true)) {
+          MDNode *MD = MDNode::get(II->getContext(), {});
+          LI->setMetadata(LLVMContext::MD_nonnull, MD);
+          LI->setMetadata(LLVMContext::MD_noundef, MD);
+          return CallBase::removeOperandBundle(II, OBU.getTagID());
+        }
+
+        // TODO: apply nonnull return attributes to calls and invokes
       }
     }
 
@@ -3688,14 +3684,13 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
     // call void @llvm.assume(i1 %A)
     // into
     // call void @llvm.assume(i1 true) [ "nonnull"(i32* %PTR) ]
-    if (EnableKnowledgeRetention &&
-        match(IIOperand,
+    if (match(IIOperand,
               m_SpecificICmp(ICmpInst::ICMP_NE, m_Value(A), m_Zero())) &&
         A->getType()->isPointerTy()) {
       if (auto *Replacement = buildAssumeFromKnowledge(
               {RetainedKnowledge{Attribute::NonNull, 0, A}}, Next, &AC, &DT)) {
 
-        Replacement->insertBefore(Next->getIterator());
+        InsertNewInstBefore(Replacement, Next->getIterator());
         AC.registerAssumption(Replacement);
         return RemoveConditionFromAssume(II);
       }
diff --git a/llvm/test/Transforms/InstCombine/assume-icmp-null-select.ll b/llvm/test/Transforms/InstCombine/assume-icmp-null-select.ll
index 108ee8cf80755..38475816f17f7 100644
--- a/llvm/test/Transforms/InstCombine/assume-icmp-null-select.ll
+++ b/llvm/test/Transforms/InstCombine/assume-icmp-null-select.ll
@@ -24,8 +24,7 @@ define ptr @example2(ptr %x) {
 ; CHECK-NEXT:    [[Y:%.*]] = load ptr, ptr [[X:%.*]], align 8
 ; CHECK-NEXT:    [[Y_IS_NULL:%.*]] = icmp eq ptr [[Y]], null
 ; CHECK-NEXT:    [[RES:%.*]] = select i1 [[Y_IS_NULL]], ptr null, ptr [[X]]
-; CHECK-NEXT:    [[NONNULL:%.*]] = icmp ne ptr [[RES]], null
-; CHECK-NEXT:    call void @llvm.assume(i1 [[NONNULL]])
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(ptr [[RES]]) ]
 ; CHECK-NEXT:    ret ptr [[RES]]
 ;
   %y = load ptr, ptr %x, align 8
diff --git a/llvm/test/Transforms/InstCombine/assume-loop-align.ll b/llvm/test/Transforms/InstCombine/assume-loop-align.ll
index 7669d5bae5b08..f6805ba4894f0 100644
--- a/llvm/test/Transforms/InstCombine/assume-loop-align.ll
+++ b/llvm/test/Transforms/InstCombine/assume-loop-align.ll
@@ -60,4 +60,3 @@ declare void @llvm.assume(i1) #1
 
 attributes #0 = { nounwind uwtable }
 attributes #1 = { nounwind }
-
diff --git a/llvm/test/Transforms/InstCombine/assume.ll b/llvm/test/Transforms/InstCombine/assume.ll
index c3a193b9fc6c4..197d680babe11 100644
--- a/llvm/test/Transforms/InstCombine/assume.ll
+++ b/llvm/test/Transforms/InstCombine/assume.ll
@@ -339,28 +339,16 @@ define i1 @nonnull2(ptr %a) {
 ; if the assume is control dependent on something else
 
 define i1 @nonnull3(ptr %a, i1 %control) {
-; FIXME: in the BUNDLES version we could duplicate the load and keep the assume nonnull.
-; DEFAULT-LABEL: @nonnull3(
-; DEFAULT-NEXT:  entry:
-; DEFAULT-NEXT:    [[LOAD:%.*]] = load ptr, ptr [[A:%.*]], align 8
-; DEFAULT-NEXT:    [[CMP:%.*]] = icmp ne ptr [[LOAD]], null
-; DEFAULT-NEXT:    br i1 [[CONTROL:%.*]], label [[TAKEN:%.*]], label [[NOT_TAKEN:%.*]]
-; DEFAULT:       taken:
-; DEFAULT-NEXT:    tail call void @llvm.assume(i1 [[CMP]])
-; DEFAULT-NEXT:    ret i1 false
-; DEFAULT:       not_taken:
-; DEFAULT-NEXT:    [[RVAL_2:%.*]] = icmp sgt ptr [[LOAD]], null
-; DEFAULT-NEXT:    ret i1 [[RVAL_2]]
-;
-; BUNDLES-LABEL: @nonnull3(
-; BUNDLES-NEXT:  entry:
-; BUNDLES-NEXT:    br i1 [[CONTROL:%.*]], label [[TAKEN:%.*]], label [[NOT_TAKEN:%.*]]
-; BUNDLES:       taken:
-; BUNDLES-NEXT:    ret i1 false
-; BUNDLES:       not_taken:
-; BUNDLES-NEXT:    [[LOAD:%.*]] = load ptr, ptr [[A:%.*]], align 8
-; BUNDLES-NEXT:    [[RVAL_2:%.*]] = icmp sgt ptr [[LOAD]], null
-; BUNDLES-NEXT:    ret i1 [[RVAL_2]]
+; FIXME: we could duplicate the load and keep the assume nonnull.
+; CHECK-LABEL: @nonnull3(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[CONTROL:%.*]], label [[TAKEN:%.*]], label [[NOT_TAKEN:%.*]]
+; CHECK:       taken:
+; CHECK-NEXT:    ret i1 false
+; CHECK:       not_taken:
+; CHECK-NEXT:    [[LOAD:%.*]] = load ptr, ptr [[A:%.*]], align 8
+; CHECK-NEXT:    [[RVAL_2:%.*]] = icmp sgt ptr [[LOAD]], null
+; CHECK-NEXT:    ret i1 [[RVAL_2]]
 ;
 entry:
   %load = load ptr, ptr %a
@@ -380,18 +368,11 @@ not_taken:
 ; interrupted by an exception being thrown
 
 define i1 @nonnull4(ptr %a) {
-; DEFAULT-LABEL: @nonnull4(
-; DEFAULT-NEXT:    [[LOAD:%.*]] = load ptr, ptr [[A:%.*]], align 8
-; DEFAULT-NEXT:    tail call void @escape(ptr [[LOAD]])
-; DEFAULT-NEXT:    [[CMP:%.*]] = icmp ne ptr [[LOAD]], null
-; DEFAULT-NEXT:    tail call void @llvm.assume(i1 [[CMP]])
-; DEFAULT-NEXT:    ret i1 false
-;
-; BUNDLES-LABEL: @nonnull4(
-; BUNDLES-NEXT:    [[LOAD:%.*]] = load ptr, ptr [[A:%.*]], align 8
-; BUNDLES-NEXT:    tail call void @escape(ptr [[LOAD]])
-; BUNDLES-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(ptr [[LOAD]]) ]
-; BUNDLES-NEXT:    ret i1 false
+; CHECK-LABEL: @nonnull4(
+; CHECK-NEXT:    [[LOAD:%.*]] = load ptr, ptr [[A:%.*]], align 8
+; CHECK-NEXT:    tail call void @escape(ptr [[LOAD]])
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(ptr [[LOAD]]) ]
+; CHECK-NEXT:    ret i1 false
 ;
   %load = load ptr, ptr %a
   ;; This call may throw!
@@ -515,27 +496,15 @@ define i32 @PR40940(<4 x i8> %x) {
 }
 
 define i1 @nonnull3A(ptr %a, i1 %control) {
-; DEFAULT-LABEL: @nonnull3A(
-; DEFAULT-NEXT:  entry:
-; DEFAULT-NEXT:    [[LOAD:%.*]] = load ptr, ptr [[A:%.*]], align 8
-; DEFAULT-NEXT:    br i1 [[CONTROL:%.*]], label [[TAKEN:%.*]], label [[NOT_TAKEN:%.*]]
-; DEFAULT:       taken:
-; DEFAULT-NEXT:    [[CMP:%.*]] = icmp ne ptr [[LOAD]], null
-; DEFAULT-NEXT:    call void @llvm.assume(i1 [[CMP]])
-; DEFAULT-NEXT:    ret i1 [[CMP]]
-; DEFAULT:       not_taken:
-; DEFAULT-NEXT:    [[RVAL_2:%.*]] = icmp sgt ptr [[LOAD]], null
-; DEFAULT-NEXT:    ret i1 [[RVAL_2]]
-;
-; BUNDLES-LABEL: @nonnull3A(
-; BUNDLES-NEXT:  entry:
-; BUNDLES-NEXT:    br i1 [[CONTROL:%.*]], label [[TAKEN:%.*]], label [[NOT_TAKEN:%.*]]
-; BUNDLES:       taken:
-; BUNDLES-NEXT:    ret i1 true
-; BUNDLES:       not_taken:
-; BUNDLES-NEXT:    [[LOAD:%.*]] = load ptr, ptr [[A:%.*]], align 8
-; BUNDLES-NEXT:    [[RVAL_2:%.*]] = icmp sgt ptr [[LOAD]], null
-; BUNDLES-NEXT:    ret i1 [[RVAL_2]]
+; CHECK-LABEL: @nonnull3A(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[CONTROL:%.*]], label [[TAKEN:%.*]], label [[NOT_TAKEN:%.*]]
+; CHECK:       taken:
+; CHECK-NEXT:    ret i1 true
+; CHECK:       not_taken:
+; CHECK-NEXT:    [[LOAD:%.*]] = load ptr, ptr [[A:%.*]], align 8
+; CHECK-NEXT:    [[RVAL_2:%.*]] = icmp sgt ptr [[LOAD]], null
+; CHECK-NEXT:    ret i1 [[RVAL_2]]
 ;
 entry:
   %load = load ptr, ptr %a
@@ -635,16 +604,8 @@ not_taken:
 }
 
 define void @nonnull_only_ephemeral_use(ptr %p) {
-; DEFAULT-LABEL: @nonnull_only_ephemeral_use(
-; DEFAULT-NEXT:    [[A:%.*]] = load ptr, ptr [[P:%.*]], align 8
-; DEFAULT-NEXT:    [[CMP:%.*]] = icmp ne ptr [[A]], null
-; DEFAULT-NEXT:    tail call void @llvm.assume(i1 [[CMP]])
-; DEFAULT-NEXT:    ret void
-;
-; BUNDLES-LABEL: @nonnull_only_ephemeral_use(
-; BUNDLES-NEXT:    [[A:%.*]] = load ptr, ptr [[P:%.*]], align 8
-; BUNDLES-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(ptr [[A]]) ]
-; BUNDLES-NEXT:    ret void
+; CHECK-LABEL: @nonnull_only_ephemeral_use(
+; CHECK-NEXT:    ret void
 ;
   %a = load ptr, ptr %p
   %cmp = icmp ne ptr %a, null