Index: llvm/include/llvm/Transforms/IPO/Attributor.h
===================================================================
--- llvm/include/llvm/Transforms/IPO/Attributor.h
+++ llvm/include/llvm/Transforms/IPO/Attributor.h
@@ -1737,13 +1737,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
@@ -1976,6 +1976,11 @@
 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 {
@@ -2060,11 +2065,15 @@
     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
@@ -2130,6 +2139,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;
 
@@ -2706,6 +2716,7 @@
       ToBeExploredFrom.insert(&F->getEntryBlock().front());
       assumeLive(A, F->getEntryBlock());
     }
+    MyA = &A;
   }
 
   /// See AbstractAttribute::getAsStr().
@@ -2870,6 +2881,12 @@
     if (!getAssumed())
       return false;
 
+    const Function *F = I->getFunction();
+    const auto &AAUB =
+        MyA->getAAFor<AAUndefinedBehavior>(*this, IRPosition::function(*F));
+    if (AAUB.isAssumedToCauseUB(I))
+      return true;
+
     // If it is not in AssumedLiveBlocks then it for sure dead.
     // Otherwise, it can still be after noreturn call in a live block.
     if (!AssumedLiveBlocks.count(I->getParent()))
@@ -2913,6 +2930,8 @@
     return true;
   }
 
+  Attributor *MyA;
+
   /// Collection of instructions that need to be explored again, e.g., we
   /// did assume they do not transfer control to (one of their) successors.
   SmallSetVector<const Instruction *, 8> ToBeExploredFrom;