diff --git a/llvm/lib/Transforms/IPO/PartialInlining.cpp b/llvm/lib/Transforms/IPO/PartialInlining.cpp
index cead7b84c3fc8..f2707afe195c4 100644
--- a/llvm/lib/Transforms/IPO/PartialInlining.cpp
+++ b/llvm/lib/Transforms/IPO/PartialInlining.cpp
@@ -412,9 +412,9 @@ PartialInlinerImpl::computeOutliningColdRegionsInfo(
   bool ColdCandidateFound = false;
   BasicBlock *CurrEntry = EntryBlock;
   std::vector<BasicBlock *> DFS;
-  DenseMap<BasicBlock *, bool> VisitedMap;
+  SmallPtrSet<BasicBlock *, 8> VisitedSet;
   DFS.push_back(CurrEntry);
-  VisitedMap[CurrEntry] = true;
+  VisitedSet.insert(CurrEntry);
 
   // Use Depth First Search on the basic blocks to find CFG edges that are
   // considered cold.
@@ -432,9 +432,8 @@ PartialInlinerImpl::computeOutliningColdRegionsInfo(
         BBProfileCount(ThisBB) < MinBlockCounterExecution)
       continue;
     for (auto SI = succ_begin(ThisBB); SI != succ_end(ThisBB); ++SI) {
-      if (VisitedMap[*SI])
+      if (!VisitedSet.insert(*SI).second)
         continue;
-      VisitedMap[*SI] = true;
       DFS.push_back(*SI);
       // If branch isn't cold, we skip to the next one.
       BranchProbability SuccProb = BPI.getEdgeProbability(ThisBB, *SI);
@@ -492,7 +491,7 @@ PartialInlinerImpl::computeOutliningColdRegionsInfo(
       // at inner regions because the outer region may have live-exit
       // variables.
       for (auto *BB : DominateVector)
-        VisitedMap[BB] = true;
+        VisitedSet.insert(BB);
 
       // ReturnBlock here means the block after the outline call
       BasicBlock *ReturnBlock = ExitBlock->getSingleSuccessor();