[llvm] [BOLT][Passes] use ADT containers for instrumentation spanning tree. (PR #192289)

[David CARLIER via llvm-commits]

https://github.com/devnexen created https://github.com/llvm/llvm-project/pull/192289

Swap `std::unordered_map<…, std::set<…>>` for
`DenseMap<…, SmallVector<…>>` in `Instrumentation::instrumentFunction` and switch read paths from `STOutSet[&BB]` to `find()`. This removes per-set heap allocations, stops inserting empty buckets on every probe, and replaces linear `is_contained()` scans over a red-black tree with linear scans over inline `SmallVector` storage (most basic blocks have at most a couple of spanning-tree out-edges). NFC.

>From 03a8830f164f6025c38611fb10d1e22546234fc1 Mon Sep 17 00:00:00 2001
From: David Carlier <devnexen at gmail.com>
Date: Wed, 15 Apr 2026 17:56:45 +0100
Subject: [PATCH] [BOLT][Passes] use ADT containers for instrumentation
 spanning tree.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Swap `std::unordered_map<…, std::set<…>>` for
`DenseMap<…, SmallVector<…>>` in `Instrumentation::instrumentFunction`
and switch read paths from `STOutSet[&BB]` to `find()`. This removes
per-set heap allocations, stops inserting empty buckets on every probe,
and replaces linear `is_contained()` scans over a red-black tree with
linear scans over inline `SmallVector` storage (most basic blocks have
at most a couple of spanning-tree out-edges). NFC.
---
 bolt/lib/Passes/Instrumentation.cpp | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/bolt/lib/Passes/Instrumentation.cpp b/bolt/lib/Passes/Instrumentation.cpp
index 418643713c54e..4a64cd49090cf 100644
--- a/bolt/lib/Passes/Instrumentation.cpp
+++ b/bolt/lib/Passes/Instrumentation.cpp
@@ -411,8 +411,8 @@ void Instrumentation::instrumentFunction(BinaryFunction &Function,
   // Exit basic blocks are always instrumented so we start the traversal with
   // a minimum number of defined variables to make the equation solvable.
   std::stack<std::pair<const BinaryBasicBlock *, BinaryBasicBlock *>> Stack;
-  std::unordered_map<const BinaryBasicBlock *,
-                     std::set<const BinaryBasicBlock *>>
+  DenseMap<const BinaryBasicBlock *,
+                     SmallVector<const BinaryBasicBlock *>>
       STOutSet;
   for (auto BBI = Function.getLayout().block_rbegin();
        BBI != Function.getLayout().block_rend(); ++BBI) {
@@ -440,7 +440,7 @@ void Instrumentation::instrumentFunction(BinaryFunction &Function,
 
       VisitedSet.insert(BB);
       if (Pred)
-        STOutSet[Pred].insert(BB);
+        STOutSet[Pred].push_back(BB);
 
       for (BinaryBasicBlock *SuccBB : BB->successors())
         Stack.push(std::make_pair(BB, SuccBB));
@@ -507,7 +507,9 @@ void Instrumentation::instrumentFunction(BinaryFunction &Function,
       }
       if (TargetFunc) {
         // Do not instrument edges in the spanning tree
-        if (llvm::is_contained(STOutSet[&BB], TargetBB)) {
+        auto STIt = STOutSet.find(&BB);
+        if (STIt != STOutSet.end() &&
+            llvm::is_contained(STIt->second, TargetBB)) {
           auto L = BC.scopeLock();
           createEdgeDescription(*FuncDesc, Function, FromOffset, BBToID[&BB],
                                 Function, ToOffset, BBToID[TargetBB],
@@ -524,7 +526,9 @@ void Instrumentation::instrumentFunction(BinaryFunction &Function,
       if (IsJumpTable) {
         for (BinaryBasicBlock *&Succ : BB.successors()) {
           // Do not instrument edges in the spanning tree
-          if (llvm::is_contained(STOutSet[&BB], &*Succ)) {
+          auto STIt = STOutSet.find(&BB);
+          if (STIt != STOutSet.end() &&
+              llvm::is_contained(STIt->second, &*Succ)) {
             auto L = BC.scopeLock();
             createEdgeDescription(*FuncDesc, Function, FromOffset, BBToID[&BB],
                                   Function, Succ->getInputOffset(),
@@ -567,7 +571,8 @@ void Instrumentation::instrumentFunction(BinaryFunction &Function,
       FromOffset = *BC.MIB->getOffset(*LastInstr);
 
       // Do not instrument edges in the spanning tree
-      if (llvm::is_contained(STOutSet[&BB], FTBB)) {
+      auto STIt = STOutSet.find(&BB);
+      if (STIt != STOutSet.end() && llvm::is_contained(STIt->second, FTBB)) {
         auto L = BC.scopeLock();
         createEdgeDescription(*FuncDesc, Function, FromOffset, BBToID[&BB],
                               Function, FTBB->getInputOffset(), BBToID[FTBB],
@@ -585,7 +590,8 @@ void Instrumentation::instrumentFunction(BinaryFunction &Function,
   if (!opts::ConservativeInstrumentation) {
     for (auto BBI = Function.begin(), BBE = Function.end(); BBI != BBE; ++BBI) {
       BinaryBasicBlock &BB = *BBI;
-      if (STOutSet[&BB].size() == 0)
+      auto STIt = STOutSet.find(&BB);
+      if (STIt == STOutSet.end() || STIt->second.empty())
         instrumentLeafNode(BB, BB.begin(), IsLeafFunction, *FuncDesc,
                            BBToID[&BB]);
     }