[LowerTypeTests] Optimize buildBitSet (llvm#157386)

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

The patch maps all offsets to correspondign
`TypeId`, so we pass prepared 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).

Author: vitalybuka

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();

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;
 

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);