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,9 +2164,15 @@
 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 {
+    dbgs() << "AAUndefinedBehavior::updateImpl()\n";
     const size_t UBPrevSize = KnownUBInsts.size();
     const size_t NoUBPrevSize = AssumedNoUBInsts.size();
 
@@ -2242,11 +2248,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
@@ -2312,6 +2321,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 +3107,16 @@
   /// 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);
+      dbgs() << "AAIsDeadFunction::initialize()\n";
+      assumeLive(A, EntryBlock);
     }
   }
 
@@ -3113,6 +3130,7 @@
 
   /// See AbstractAttribute::manifest(...).
   ChangeStatus manifest(Attributor &A) override {
+    dbgs() << "AAIsDeadFunction::manifest(...)\n";
     assert(getState().isValidState() &&
            "Attempted to manifest an invalid state!");
 
@@ -3218,6 +3236,7 @@
   /// Assume \p BB is (partially) live now and indicate to the Attributor \p A
   /// that internal function called from \p BB should now be looked at.
   bool assumeLive(Attributor &A, const BasicBlock &BB) {
+    dbgs() << "Insert: " << BB << "\n";
     if (!AssumedLiveBlocks.insert(&BB).second)
       return false;
 
@@ -3384,6 +3403,33 @@
       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) {
+      dbgs() << "AliveSuccessor: " << *AliveSuccessor << "\n";
+      bool Assumed = false;
+      if (AAUB.isKnownToCauseUB(AliveSuccessor) ||
+          (Assumed = AAUB.isAssumedToCauseUB(AliveSuccessor))) {
+        dbgs() << "Remove\n";
+        if (Assumed)
+          UsedAssumedInformation = true;
+        // Remove the parent BB.
+        const BasicBlock *BB = AliveSuccessor->getParent();
+        dbgs() << *BB << "\n";
+        AssumedLiveBlocks.erase(BB);
+      } else {
+        AliveSuccessorsCopy.push_back(AliveSuccessor);
+      }
+    }
+    AliveSuccessors = std::move(AliveSuccessorsCopy);
+
     if (UsedAssumedInformation) {
       NewToBeExploredFrom.insert(I);
     } else {