diff --git a/llvm/lib/Transforms/Instrumentation/LowerAllowCheckPass.cpp b/llvm/lib/Transforms/Instrumentation/LowerAllowCheckPass.cpp
index 10e908ef73ce5..f1a0411521e01 100644
--- a/llvm/lib/Transforms/Instrumentation/LowerAllowCheckPass.cpp
+++ b/llvm/lib/Transforms/Instrumentation/LowerAllowCheckPass.cpp
@@ -15,6 +15,7 @@
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
@@ -111,31 +112,29 @@ static bool removeUbsanTraps(Function &F, const BlockFrequencyInfo &BFI,
     return ShouldRemoveRandom() || ShouldRemoveHot(*(II->getParent()), cutoff);
   };
 
-  for (BasicBlock &BB : F) {
-    for (Instruction &I : BB) {
-      IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I);
-      if (!II)
-        continue;
-      auto ID = II->getIntrinsicID();
-      switch (ID) {
-      case Intrinsic::allow_ubsan_check:
-      case Intrinsic::allow_runtime_check: {
-        ++NumChecksTotal;
-
-        bool ToRemove = ShouldRemove(II);
-
-        ReplaceWithValue.push_back({
-            II,
-            ToRemove,
-        });
-        if (ToRemove)
-          ++NumChecksRemoved;
-        emitRemark(II, ORE, ToRemove);
-        break;
-      }
-      default:
-        break;
-      }
+  for (Instruction &I : instructions(F)) {
+    IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I);
+    if (!II)
+      continue;
+    auto ID = II->getIntrinsicID();
+    switch (ID) {
+    case Intrinsic::allow_ubsan_check:
+    case Intrinsic::allow_runtime_check: {
+      ++NumChecksTotal;
+
+      bool ToRemove = ShouldRemove(II);
+
+      ReplaceWithValue.push_back({
+          II,
+          ToRemove,
+      });
+      if (ToRemove)
+        ++NumChecksRemoved;
+      emitRemark(II, ORE, ToRemove);
+      break;
+    }
+    default:
+      break;
     }
   }