[llvm] [LowerTypeTests] Optimize buildBitSet (PR #157386)

[Vitaly Buka via llvm-commits]

https://github.com/vitalybuka created https://github.com/llvm/llvm-project/pull/157386

`buildBitSet` had a loop trough entire GlobalLayout to pickup matching offsets.

The patch maps all offsets to correspondign
`TypeId`, so we pass prepare list of offsets into
`buildBitSet`.

On one large internal binary, `LowerTypeTests`
took 58% of ThinLTO link time before the patch.
After the patch just 7% (absolute saving is 200s).


>From d4e35c3b28e217a868950b0df194660fdb2b1c48 Mon Sep 17 00:00:00 2001
From: Vitaly Buka <vitalybuka at google.com>
Date: Sun, 7 Sep 2025 22:31:58 -0700
Subject: [PATCH] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20initia?=
 =?UTF-8?q?l=20version?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Created using spr 1.3.6
---
 .../llvm/Transforms/IPO/LowerTypeTests.h      | 15 ++---
 llvm/lib/Transforms/IPO/LowerTypeTests.cpp    | 63 +++++++++++--------
 .../Transforms/IPO/LowerTypeTests.cpp         |  5 +-
 3 files changed, 43 insertions(+), 40 deletions(-)

diff --git a/llvm/include/llvm/Transforms/IPO/LowerTypeTests.h b/llvm/include/llvm/Transforms/IPO/LowerTypeTests.h
index 253d13f81857b..a34cbaf72675b 100644
--- a/llvm/include/llvm/Transforms/IPO/LowerTypeTests.h
+++ b/llvm/include/llvm/Transforms/IPO/LowerTypeTests.h
@@ -64,15 +64,12 @@ struct BitSetBuilder {
   uint64_t Min = std::numeric_limits<uint64_t>::max();
   uint64_t Max = 0;
 
-  BitSetBuilder() = default;
-
-  void addOffset(uint64_t Offset) {
-    if (Min > Offset)
-      Min = Offset;
-    if (Max < Offset)
-      Max = Offset;
-
-    Offsets.push_back(Offset);
+  explicit BitSetBuilder(ArrayRef<uint64_t> Offsets) : Offsets(Offsets) {
+    if (!Offsets.empty()) {
+      auto [MinIt, MaxIt] = std::minmax_element(Offsets.begin(), Offsets.end());
+      Min = *MinIt;
+      Max = *MaxIt;
+    }
   }
 
   LLVM_ABI BitSetInfo build();
diff --git a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
index 821a9d82ddb0d..be6cba38d8b3c 100644
--- a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
+++ b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
@@ -573,28 +573,11 @@ class LowerTypeTestsModule {
 };
 } // end anonymous namespace
 
-/// Build a bit set for TypeId using the object layouts in
-/// GlobalLayout.
-static BitSetInfo
-buildBitSet(Metadata *TypeId,
-            const DenseMap<GlobalTypeMember *, uint64_t> &GlobalLayout) {
-  BitSetBuilder BSB;
-
+/// Build a bit set for list of offsets.
+static BitSetInfo buildBitSet(ArrayRef<uint64_t> Offsets) {
   // Compute the byte offset of each address associated with this type
   // identifier.
-  for (const auto &GlobalAndOffset : GlobalLayout) {
-    for (MDNode *Type : GlobalAndOffset.first->types()) {
-      if (Type->getOperand(1) != TypeId)
-        continue;
-      uint64_t Offset =
-          cast<ConstantInt>(
-              cast<ConstantAsMetadata>(Type->getOperand(0))->getValue())
-              ->getZExtValue();
-      BSB.addOffset(GlobalAndOffset.second + Offset);
-    }
-  }
-
-  return BSB.build();
+  return BitSetBuilder(Offsets).build();
 }
 
 /// Build a test that bit BitOffset mod sizeof(Bits)*8 is set in
@@ -1161,21 +1144,47 @@ void LowerTypeTestsModule::importFunction(Function *F,
   F->setVisibility(Visibility);
 }
 
-void LowerTypeTestsModule::lowerTypeTestCalls(
-    ArrayRef<Metadata *> TypeIds, Constant *CombinedGlobalAddr,
-    const DenseMap<GlobalTypeMember *, uint64_t> &GlobalLayout) {
-  // For each type identifier in this disjoint set...
+static auto
+buildBitSets(ArrayRef<Metadata *> TypeIds,
+             const DenseMap<GlobalTypeMember *, uint64_t> &GlobalLayout) {
+  DenseMap<Metadata *, SmallVector<uint64_t, 16>> OffsetsByTypeID;
+  // Pre-populate the map with interesting type identifiers.
+  for (Metadata *TypeId : TypeIds)
+    OffsetsByTypeID[TypeId];
+  for (const auto &[Mem, MemOff] : GlobalLayout) {
+    for (MDNode *Type : Mem->types()) {
+      auto It = OffsetsByTypeID.find(Type->getOperand(1));
+      if (It == OffsetsByTypeID.end())
+        continue;
+      uint64_t Offset =
+          cast<ConstantInt>(
+              cast<ConstantAsMetadata>(Type->getOperand(0))->getValue())
+              ->getZExtValue();
+      It->second.push_back(MemOff + Offset);
+    }
+  }
+
+  SmallVector<std::pair<Metadata *, BitSetInfo>> BitSets;
+  BitSets.reserve(TypeIds.size());
   for (Metadata *TypeId : TypeIds) {
-    // Build the bitset.
-    BitSetInfo BSI = buildBitSet(TypeId, GlobalLayout);
+    BitSets.emplace_back(TypeId, buildBitSet(OffsetsByTypeID[TypeId]));
     LLVM_DEBUG({
       if (auto MDS = dyn_cast<MDString>(TypeId))
         dbgs() << MDS->getString() << ": ";
       else
         dbgs() << "<unnamed>: ";
-      BSI.print(dbgs());
+      BitSets.back().second.print(dbgs());
     });
+  }
+
+  return BitSets;
+}
 
+void LowerTypeTestsModule::lowerTypeTestCalls(
+    ArrayRef<Metadata *> TypeIds, Constant *CombinedGlobalAddr,
+    const DenseMap<GlobalTypeMember *, uint64_t> &GlobalLayout) {
+  // For each type identifier in this disjoint set...
+  for (const auto &[TypeId, BSI] : buildBitSets(TypeIds, GlobalLayout)) {
     ByteArrayInfo *BAI = nullptr;
     TypeIdLowering TIL;
 
diff --git a/llvm/unittests/Transforms/IPO/LowerTypeTests.cpp b/llvm/unittests/Transforms/IPO/LowerTypeTests.cpp
index fc41b036ffb6c..3d6b7efbccd91 100644
--- a/llvm/unittests/Transforms/IPO/LowerTypeTests.cpp
+++ b/llvm/unittests/Transforms/IPO/LowerTypeTests.cpp
@@ -51,10 +51,7 @@ TEST(LowerTypeTests, BitSetBuilder) {
   };
 
   for (auto &&T : BSBTests) {
-    BitSetBuilder BSB;
-    for (auto Offset : T.Offsets)
-      BSB.addOffset(Offset);
-
+    BitSetBuilder BSB(T.Offsets);
     BitSetInfo BSI = BSB.build();
 
     EXPECT_EQ(T.Bits, BSI.Bits);