Index: llvm/include/llvm/Transforms/IPO/Attributor.h
===================================================================
--- llvm/include/llvm/Transforms/IPO/Attributor.h
+++ llvm/include/llvm/Transforms/IPO/Attributor.h
@@ -2025,13 +2025,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(const Instruction *I) const = 0;
 
   /// Return an IR position, see struct IRPosition.
   const IRPosition &getIRPosition() const override { return *this; }
Index: llvm/lib/Transforms/IPO/Attributor.cpp
===================================================================
--- llvm/lib/Transforms/IPO/Attributor.cpp
+++ llvm/lib/Transforms/IPO/Attributor.cpp
@@ -2164,8 +2164,12 @@
 struct AAUndefinedBehaviorImpl : public AAUndefinedBehavior {
   AAUndefinedBehaviorImpl(const IRPosition &IRP) : AAUndefinedBehavior(IRP) {}
 
+  void initialize(Attributor &A) override {
+    AAUndefinedBehavior::initialize(A);
+    Explorer = &(A.getInfoCache().getMustBeExecutedContextExplorer());
+  }
+
   /// 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();
@@ -2242,11 +2246,14 @@
     return ChangeStatus::UNCHANGED;
   }
 
-  bool isKnownToCauseUB(Instruction *I) const override {
-    return KnownUBInsts.count(I);
+  bool isKnownToCauseUB(const Instruction *I) const override {
+    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
@@ -2275,7 +2282,7 @@
     if (KnownUBInsts.empty())
       return ChangeStatus::UNCHANGED;
     for (Instruction *I : KnownUBInsts)
-      A.changeToUnreachableAfterManifest(I);
+      //A.changeToUnreachableAfterManifest(I);
     return ChangeStatus::CHANGED;
   }
 
@@ -2312,6 +2319,7 @@
   SmallPtrSet<Instruction *, 8> KnownUBInsts;
 
 private:
+  MustBeExecutedContextExplorer *Explorer;
   /// A set of all the (live) instructions that are assumed to _not_ cause UB.
   SmallPtrSet<Instruction *, 8> AssumedNoUBInsts;
 
@@ -3097,9 +3105,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(&Front)) {
+      ToBeExploredFrom.insert(&Front);
+      assumeLive(A, EntryBlock);
     }
   }
 
@@ -3384,6 +3398,30 @@
       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(AliveSuccessor) ||
+          (Assumed = AAUB.isAssumedToCauseUB(AliveSuccessor))) {
+        if (Assumed)
+          UsedAssumedInformation = true;
+        // Remove the parent BB.
+        const BasicBlock *BB = AliveSuccessor->getParent();
+        AssumedLiveBlocks.erase(BB);
+      } else {
+        AliveSuccessorsCopy.push_back(AliveSuccessor);
+      }
+    }
+    AliveSuccessors = std::move(AliveSuccessorsCopy);
+
     if (UsedAssumedInformation) {
       NewToBeExploredFrom.insert(I);
     } else {