Index: llvm/include/llvm/Transforms/IPO/Attributor.h
===================================================================
--- llvm/include/llvm/Transforms/IPO/Attributor.h
+++ llvm/include/llvm/Transforms/IPO/Attributor.h
@@ -2023,13 +2023,13 @@
   bool isAssumedToCauseUB() const { return getAssumed(); }
 
   /// Return true if "undefined behavior" is assumed for a specific instruction.
-  virtual bool isAssumedToCauseUB(Instruction *I) const = 0;
+  virtual bool isAssumedToCauseUB(const Instruction *I) const = 0;
 
   /// Return true if "undefined behavior" is known.
   bool isKnownToCauseUB() const { return getKnown(); }
 
   /// Return true if "undefined behavior" is known for a specific instruction.
-  virtual bool isKnownToCauseUB(Instruction *I) const = 0;
+  virtual bool isKnownToCauseUB(Attributor &A, const Instruction *I) const = 0;
 
   /// Return an IR position, see struct IRPosition.
   const IRPosition &getIRPosition() const override { return *this; }
@@ -2748,7 +2748,8 @@
   /// Return an assumed constant for the assocaited value a program point \p
   /// CtxI.
   Optional<ConstantInt *>
-  getAssumedConstantInt(Attributor &A, const Instruction *CtxI = nullptr) const {
+  getAssumedConstantInt(Attributor &A,
+                        const Instruction *CtxI = nullptr) const {
     ConstantRange RangeV = getAssumedConstantRange(A, CtxI);
     if (auto *C = RangeV.getSingleElement())
       return cast<ConstantInt>(
Index: llvm/lib/Transforms/IPO/Attributor.cpp
===================================================================
--- llvm/lib/Transforms/IPO/Attributor.cpp
+++ llvm/lib/Transforms/IPO/Attributor.cpp
@@ -2215,7 +2215,6 @@
   AAUndefinedBehaviorImpl(const IRPosition &IRP) : AAUndefinedBehavior(IRP) {}
 
   /// See AbstractAttribute::updateImpl(...).
-  // through a pointer (i.e. also branches etc.)
   ChangeStatus updateImpl(Attributor &A) override {
     const size_t UBPrevSize = KnownUBInsts.size();
     const size_t NoUBPrevSize = AssumedNoUBInsts.size();
@@ -2292,11 +2291,16 @@
     return ChangeStatus::UNCHANGED;
   }
 
-  bool isKnownToCauseUB(Instruction *I) const override {
-    return KnownUBInsts.count(I);
+  bool isKnownToCauseUB(Attributor &A, const Instruction *I) const override {
+    MustBeExecutedContextExplorer &Explorer =
+      A.getInfoCache().getMustBeExecutedContextExplorer();
+    for (auto It : Explorer.range(I))
+      if (KnownUBInsts.count(It))
+        return true;
+    return false;
   }
 
-  bool isAssumedToCauseUB(Instruction *I) const override {
+  bool isAssumedToCauseUB(const Instruction *I) const override {
     // In simple words, if an instruction is not in the assumed to _not_
     // cause UB, then it is assumed UB (that includes those
     // in the KnownUBInsts set). The rest is boilerplate
@@ -3158,9 +3162,15 @@
   /// See AbstractAttribute::initialize(...).
   void initialize(Attributor &A) override {
     const Function *F = getAssociatedFunction();
-    if (F && !F->isDeclaration()) {
-      ToBeExploredFrom.insert(&F->getEntryBlock().front());
-      assumeLive(A, F->getEntryBlock());
+    if (!F || F->isDeclaration())
+      return;
+    const BasicBlock &EntryBlock = F->getEntryBlock();
+    const Instruction &Front = EntryBlock.front();
+    const auto &AAUB =
+        A.getAAFor<AAUndefinedBehavior>(*this, IRPosition::function(*F));
+    if (!AAUB.isKnownToCauseUB(A, &Front)) {
+      ToBeExploredFrom.insert(&Front);
+      assumeLive(A, EntryBlock);
     }
   }
 
@@ -3445,6 +3455,26 @@
       break;
     }
 
+    // Keep only successors that are not assumed to cause UB.
+
+    // Note: Instead of removing the successors, it's cleaner
+    // to make a copy in which we keep only the valid ones.
+    SmallVector<const Instruction *, 8> AliveSuccessorsCopy;
+    AliveSuccessorsCopy.reserve(AliveSuccessors.size());
+    const Function *F = I->getFunction();
+    const auto &AAUB =
+        A.getAAFor<AAUndefinedBehavior>(*this, IRPosition::function(*F));
+    for (const Instruction *AliveSuccessor : AliveSuccessors) {
+      bool Assumed = false;
+      if (AAUB.isKnownToCauseUB(A, AliveSuccessor) ||
+          (Assumed = AAUB.isAssumedToCauseUB(AliveSuccessor))) {
+        UsedAssumedInformation = Assumed;
+      } else {
+        AliveSuccessorsCopy.push_back(AliveSuccessor);
+      }
+    }
+    AliveSuccessors = std::move(AliveSuccessorsCopy);
+
     if (UsedAssumedInformation) {
       NewToBeExploredFrom.insert(I);
     } else {