Bug 1750930 - Make the indirect stub table two-level. r=yury, a=RyanVM

For large applications with multiple threads and large tables of
indirect functions and tiered compilation, the shared (per-module)
indirect stubs set can easily reach one million entries.  However the
insert-one-at-a-time strategy for this set uses linear time, so
merging a new set of stubs into the set will tend to take quadratic
time.  Furthermore, a lock is held by each thread during this slow
process, serializing all the threads.  The result can be very slow
application startup.

As a stopgap (because we may remove the indirect stubs and don't need
to commit to a complex solution right now), we break the indirect
stubs set it into per-tls sets to avoid the very long insertion times
resulting from having one shared set per module.  Set merging remains
quadratic but the since the sets are much smaller it matters much
less.

A followup bug will be filed for a better, permanent solution, should
we need it.

Differential Revision: https://phabricator.services.mozilla.com/D136470

js/src/wasm/WasmCode.cpp

@@ -884,8 +884,10 @@ bool LazyStubTier::createOneEntryStub(uint32_t funcExportIndex,
 
 // This uses the funcIndex as the major key and the tls pointer value as the
 // minor key, the same as the < and == predicates used in RemoveDuplicates.
+// However, since we only ever use this to search tables where every entry has
+// the same tls, there is no actual code for tls comparison here.
 
-auto IndirectStubComparator = [](uint32_t funcIndex, void* tlsData,
+auto IndirectStubComparator = [](uint32_t funcIndex,
                                  const IndirectStub& stub) -> int {
   if (funcIndex < stub.funcIndex) {
     return -1;
@@ -894,12 +896,6 @@ auto IndirectStubComparator = [](uint32_t funcIndex, void* tlsData,
     return 1;
   }
   // Function indices are equal.
-  if (uintptr_t(tlsData) < uintptr_t(stub.tls)) {
-    return -1;
-  }
-  if (uintptr_t(tlsData) > uintptr_t(stub.tls)) {
-    return 1;
-  }
   return 0;
 };
 
@@ -982,24 +978,63 @@ bool LazyStubTier::createManyIndirectStubs(
   }
 
   // Record the runtime info about generated indirect stubs.
-  if (!indirectStubVector_.reserve(indirectStubVector_.length() +
-                                   targets.length())) {
-    return false;
+
+  // Count the number of new slots needed for the different tls values in the
+  // table.  While there may be multiple tls values in the target set, the
+  // typical number is one or two.
+  struct Counter {
+    explicit Counter(void* tls) : tls(tls), counter(0) {}
+    void* tls;
+    size_t counter;
+  };
+  Vector<Counter, 8, SystemAllocPolicy> counters{};
+  for (const auto& target : targets) {
+    size_t i = 0;
+    while (i < counters.length() && target.tls != counters[i].tls) {
+      i++;
+    }
+    if (i == counters.length() && !counters.emplaceBack(target.tls)) {
+      return false;
+    }
+    counters[i].counter++;
+  }
+
+  // Reserve space in the tables, creating new tables as necessary.  Do this
+  // first to avoid OOM while we're midway through installing stubs in the
+  // tables.
+  for (const auto& counter : counters) {
+    auto probe = indirectStubTable_.lookupForAdd(counter.tls);
+    if (!probe) {
+      IndirectStubVector v{};
+      if (!indirectStubTable_.add(probe, counter.tls, std::move(v))) {
+        return false;
+      }
+    }
+    IndirectStubVector& indirectStubVector = probe->value();
+    if (!indirectStubVector.reserve(indirectStubVector.length() +
+                                    counter.counter)) {
+      return false;
+    }
   }
 
+  // We have storage, so now we can commit.
   for (const auto& target : targets) {
     auto stub = IndirectStub{target.functionIdx, lastStubSegmentIndex_,
-                             indirectStubRangeIndex, target.tls};
+                             indirectStubRangeIndex};
+
+    auto probe = indirectStubTable_.lookup(target.tls);
+    MOZ_RELEASE_ASSERT(probe);
+    IndirectStubVector& indirectStubVector = probe->value();
 
     size_t indirectStubIndex;
     MOZ_ALWAYS_FALSE(BinarySearchIf(
-        indirectStubVector_, 0, indirectStubVector_.length(),
+        indirectStubVector, 0, indirectStubVector.length(),
         [&stub](const IndirectStub& otherStub) {
-          return IndirectStubComparator(stub.funcIndex, stub.tls, otherStub);
+          return IndirectStubComparator(stub.funcIndex, otherStub);
         },
         &indirectStubIndex));
-    MOZ_ALWAYS_TRUE(indirectStubVector_.insert(
-        indirectStubVector_.begin() + indirectStubIndex, std::move(stub)));
+    MOZ_ALWAYS_TRUE(indirectStubVector.insert(
+        indirectStubVector.begin() + indirectStubIndex, std::move(stub)));
 
     ++indirectStubRangeIndex;
   }
@@ -1078,16 +1113,20 @@ void* LazyStubTier::lookupInterpEntry(uint32_t funcIndex) const {
 
 void* LazyStubTier::lookupIndirectStub(uint32_t funcIndex, void* tls) const {
   size_t match;
+  auto probe = indirectStubTable_.lookup(tls);
+  if (!probe) {
+    return nullptr;
+  }
+  const IndirectStubVector& indirectStubVector = probe->value();
   if (!BinarySearchIf(
-          indirectStubVector_, 0, indirectStubVector_.length(),
-          [funcIndex, tls](const IndirectStub& stub) {
-            return IndirectStubComparator(funcIndex, tls, stub);
+          indirectStubVector, 0, indirectStubVector.length(),
+          [funcIndex](const IndirectStub& stub) {
+            return IndirectStubComparator(funcIndex, stub);
           },
           &match)) {
     return nullptr;
   }
-
-  const IndirectStub& indirectStub = indirectStubVector_[match];
+  const IndirectStub& indirectStub = indirectStubVector[match];
 
   const LazyStubSegment& segment = *stubSegments_[indirectStub.segmentIndex];
   return segment.base() +

js/src/wasm/WasmCode.h

@@ -573,28 +573,48 @@ struct LazyFuncExport {
 
 using LazyFuncExportVector = Vector<LazyFuncExport, 0, SystemAllocPolicy>;
 
-// IndirectStub provides a mapping between function indices and
-// indirect stubs code ranges.
-
-// The function index is the index of the function *within a specific module*,
+// IndirectStub provides a mapping between a function index and an indirect stub
+// code range.
+//
+// The function index is the index of the function *within its defining module*,
 // not necessarily in the module that owns the stub.  That module's and
-// function's instance is provided by the tls field.
+// function's instance is provided by the tls field of the IndirectStubTable
+// entry within which this IndirectStub is found.
 
 struct IndirectStub {
   size_t funcIndex;
   size_t segmentIndex;
   size_t codeRangeIndex;
-  void* tls;
-  IndirectStub(size_t funcIndex, size_t segmentIndex, size_t codeRangeIndex,
-               TlsData* tls)
+  IndirectStub(size_t funcIndex, size_t segmentIndex, size_t codeRangeIndex)
       : funcIndex(funcIndex),
         segmentIndex(segmentIndex),
-        codeRangeIndex(codeRangeIndex),
-        tls(tls) {}
+        codeRangeIndex(codeRangeIndex) {}
 };
 
+// IndirectStubVector represents a set of IndirectStubs.  These stubs all belong
+// to the same IndirectStubTable entry, and so all have the same tls value.
+//
+// The IndirectStubVector is ordered by IndirectStubComparator (WasmCode.cpp):
+// the sort key is the funcIndex.  The vector is binary-searched by that
+// predicate when an entry is needed.
+//
+// Creating an indirect stub is not an idempotent operation!  There must be NO
+// duplicate entries in the table, or equivalently, an entry that is in the
+// table must always be found by a binary search.
+
 using IndirectStubVector = Vector<IndirectStub, 0, SystemAllocPolicy>;
 
+// An IndirectStubTable represents a set of indirect stubs belonging to a
+// module.  There table is keyed uniquely by tls and there is one
+// IndirectStubVector per tls value represented in the set.
+//
+// While the set is usually very small, its can grow with the product of the
+// number of instances and the number of threads in a system, and we therefore
+// use a hash table.
+
+using IndirectStubTable =
+    HashMap<void*, IndirectStubVector, DefaultHasher<void*>, SystemAllocPolicy>;
+
 // LazyStubTier contains all the necessary information for lazy function entry
 // stubs and indirect stubs that are generated at runtime.
 // None of its data are ever serialized.
@@ -606,15 +626,7 @@ using IndirectStubVector = Vector<IndirectStub, 0, SystemAllocPolicy>;
 class LazyStubTier {
   LazyStubSegmentVector stubSegments_;
   LazyFuncExportVector exports_;
-  // The indirectStubVector_ is totally ordered by IndirectStubComparator (in
-  // WasmCode.cpp): the primary index is the funcIndex, the secondary index the
-  // pointer value of the tls.  The vector is binary-searched by that predicate
-  // when an entry is needed.
-  //
-  // Creating an indirect stub is not an idempotent operation!  There must be NO
-  // duplicate entries in the table, which is another way of saying that an
-  // entry that is in the table must always be found by a lookup.
-  IndirectStubVector indirectStubVector_;
+  IndirectStubTable indirectStubTable_;
   size_t lastStubSegmentIndex_;
 
   [[nodiscard]] bool createManyEntryStubs(const Uint32Vector& funcExportIndices,