[wasm] Fix performance bottleneck in DisjointAllocationPool

When compiling modules with many functions, the list of regions in the
{DisjointAllocationPool} can become quite large if the functions die in
a random order (which they typically do, since the order of Liftoff
compilation is different than the order to TurboFan compilation; which
work stealing, both are nondeterministic).
Iterating the list of regions in the {DisjointAllocationPool} was thus
linear in the number of regions, which is linear in the number of
functions of the module. Since we insert new regions one by one, overall
runtime was quadratic.

This CL fixes this by switching from a linked list to a std::set.
Merging a new region is thus logarithmic instead of linear, and overall
we are {n*log(n)} instead of {n^2}.

Note: For {AllocateInRegion} we still need to linearly iterate all
regions that overlap the requested region, but this has not shown to be
a problem so far.

R=ahaas@chromium.org

Bug: v8:10432
Change-Id: I193e56c2abab782e386194fbe64dadfa250916f7
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2154797
Commit-Queue: Clemens Backes <clemensb@chromium.org>
Reviewed-by: Andreas Haas <ahaas@chromium.org>
Cr-Commit-Position: refs/heads/master@{#67303}

src/base/address-region.h

@@ -15,6 +15,14 @@ namespace base {
 // Helper class representing an address region of certain size.
 class AddressRegion {
  public:
+  // Function object that compares the start address of two regions. Usable as
+  // compare function on std data structures and algorithms.
+  struct StartAddressLess {
+    bool operator()(base::AddressRegion a, base::AddressRegion b) const {
+      return a.begin() < b.begin();
+    }
+  };
+
   using Address = uintptr_t;
 
   AddressRegion() = default;

src/wasm/wasm-code-manager.cc

@@ -47,47 +47,59 @@ namespace wasm {
 
 using trap_handler::ProtectedInstructionData;
 
-base::AddressRegion DisjointAllocationPool::Merge(base::AddressRegion region) {
-  auto dest_it = regions_.begin();
-  auto dest_end = regions_.end();
-
-  // Skip over dest regions strictly before {region}.
-  while (dest_it != dest_end && dest_it->end() < region.begin()) ++dest_it;
-
-  // After last dest region: insert and done.
-  if (dest_it == dest_end) {
-    regions_.push_back(region);
-    return region;
-  }
-
-  // Adjacent (from below) to dest: merge and done.
-  if (dest_it->begin() == region.end()) {
-    base::AddressRegion merged_region{region.begin(),
-                                      region.size() + dest_it->size()};
-    DCHECK_EQ(merged_region.end(), dest_it->end());
-    *dest_it = merged_region;
+base::AddressRegion DisjointAllocationPool::Merge(
+    base::AddressRegion new_region) {
+  // Find the possible insertion position by identifying the first region whose
+  // start address is not less than that of {new_region}. Since there cannot be
+  // any overlap between regions, this also means that the start of {above} is
+  // bigger or equal than the *end* of {new_region}.
+  auto above = regions_.lower_bound(new_region);
+  DCHECK(above == regions_.end() || above->begin() >= new_region.end());
+
+  // Check whether to merge with {above}.
+  if (above != regions_.end() && new_region.end() == above->begin()) {
+    base::AddressRegion merged_region{new_region.begin(),
+                                      new_region.size() + above->size()};
+    DCHECK_EQ(merged_region.end(), above->end());
+    // Check whether to also merge with the region below.
+    if (above != regions_.begin()) {
+      auto below = above;
+      --below;
+      if (below->end() == new_region.begin()) {
+        merged_region = {below->begin(), below->size() + merged_region.size()};
+        regions_.erase(below);
+      }
+    }
+    auto insert_pos = regions_.erase(above);
+    regions_.insert(insert_pos, merged_region);
     return merged_region;
   }
 
-  // Before dest: insert and done.
-  if (dest_it->begin() > region.end()) {
-    regions_.insert(dest_it, region);
-    return region;
+  // No element below, and not adjavent to {above}: insert and done.
+  if (above == regions_.begin()) {
+    regions_.insert(above, new_region);
+    return new_region;
   }
 
-  // Src is adjacent from above. Merge and check whether the merged region is
-  // now adjacent to the next region.
-  DCHECK_EQ(dest_it->end(), region.begin());
-  dest_it->set_size(dest_it->size() + region.size());
-  DCHECK_EQ(dest_it->end(), region.end());
-  auto next_dest = dest_it;
-  ++next_dest;
-  if (next_dest != dest_end && dest_it->end() == next_dest->begin()) {
-    dest_it->set_size(dest_it->size() + next_dest->size());
-    DCHECK_EQ(dest_it->end(), next_dest->end());
-    regions_.erase(next_dest);
+  auto below = above;
+  --below;
+  // Sanity check:
+  DCHECK(above == regions_.end() || below->end() < above->begin());
+
+  // Adjacent to {below}: merge and done.
+  if (below->end() == new_region.begin()) {
+    base::AddressRegion merged_region{below->begin(),
+                                      below->size() + new_region.size()};
+    DCHECK_EQ(merged_region.end(), new_region.end());
+    regions_.erase(below);
+    regions_.insert(above, merged_region);
+    return merged_region;
   }
-  return *dest_it;
+
+  // Not adjacent to any existing region: insert between {below} and {above}.
+  DCHECK_LT(below->end(), new_region.begin());
+  regions_.insert(above, new_region);
+  return new_region;
 }
 
 base::AddressRegion DisjointAllocationPool::Allocate(size_t size) {
@@ -97,24 +109,31 @@ base::AddressRegion DisjointAllocationPool::Allocate(size_t size) {
 
 base::AddressRegion DisjointAllocationPool::AllocateInRegion(
     size_t size, base::AddressRegion region) {
-  for (auto it = regions_.begin(), end = regions_.end(); it != end; ++it) {
+  // Get an iterator to the first contained region whose start address is not
+  // smaller than the start address of {region}. Start the search from the
+  // region one before that (the last one whose start address is smaller).
+  auto it = regions_.lower_bound(region);
+  if (it != regions_.begin()) --it;
+
+  for (auto end = regions_.end(); it != end; ++it) {
     base::AddressRegion overlap = it->GetOverlap(region);
     if (size > overlap.size()) continue;
     base::AddressRegion ret{overlap.begin(), size};
-    if (size == it->size()) {
-      // We use the full region --> erase the region from {regions_}.
-      regions_.erase(it);
-    } else if (ret.begin() == it->begin()) {
-      // We return a region at the start --> shrink remaining region from front.
-      *it = base::AddressRegion{it->begin() + size, it->size() - size};
-    } else if (ret.end() == it->end()) {
+    base::AddressRegion old = *it;
+    auto insert_pos = regions_.erase(it);
+    if (size == old.size()) {
+      // We use the full region --> nothing to add back.
+    } else if (ret.begin() == old.begin()) {
+      // We return a region at the start --> shrink old region from front.
+      regions_.insert(insert_pos, {old.begin() + size, old.size() - size});
+    } else if (ret.end() == old.end()) {
       // We return a region at the end --> shrink remaining region.
-      *it = base::AddressRegion{it->begin(), it->size() - size};
+      regions_.insert(insert_pos, {old.begin(), old.size() - size});
     } else {
-      // We return something in the middle --> split the remaining region.
-      regions_.insert(
-          it, base::AddressRegion{it->begin(), ret.begin() - it->begin()});
-      *it = base::AddressRegion{ret.end(), it->end() - ret.end()};
+      // We return something in the middle --> split the remaining region
+      // (insert the region with smaller address first).
+      regions_.insert(insert_pos, {old.begin(), ret.begin() - old.begin()});
+      regions_.insert(insert_pos, {ret.end(), old.end() - ret.end()});
     }
     return ret;
   }

src/wasm/wasm-code-manager.h

@@ -6,9 +6,9 @@
 #define V8_WASM_WASM_CODE_MANAGER_H_
 
 #include <atomic>
-#include <list>
 #include <map>
 #include <memory>
+#include <set>
 #include <unordered_set>
 #include <utility>
 #include <vector>
@@ -88,10 +88,9 @@ class V8_EXPORT_PRIVATE DisjointAllocationPool final {
   explicit DisjointAllocationPool(base::AddressRegion region)
       : regions_({region}) {}
 
-  // Merge the parameter region into this object while preserving ordering of
-  // the regions. The assumption is that the passed parameter is not
-  // intersecting this object - for example, it was obtained from a previous
-  // Allocate. Returns the merged region.
+  // Merge the parameter region into this object. The assumption is that the
+  // passed parameter is not intersecting this object - for example, it was
+  // obtained from a previous Allocate. Returns the merged region.
   base::AddressRegion Merge(base::AddressRegion);
 
   // Allocate a contiguous region of size {size}. Return an empty pool on
@@ -103,10 +102,11 @@ class V8_EXPORT_PRIVATE DisjointAllocationPool final {
   base::AddressRegion AllocateInRegion(size_t size, base::AddressRegion);
 
   bool IsEmpty() const { return regions_.empty(); }
-  const std::list<base::AddressRegion>& regions() const { return regions_; }
+
+  const auto& regions() const { return regions_; }
 
  private:
-  std::list<base::AddressRegion> regions_;
+  std::set<base::AddressRegion, base::AddressRegion::StartAddressLess> regions_;
 };
 
 class V8_EXPORT_PRIVATE WasmCode final {

test/unittests/wasm/wasm-code-manager-unittest.cc

@@ -71,8 +71,8 @@ TEST_F(DisjointAllocationPoolTest, SimpleExtract) {
   a.Merge(b);
   CheckPool(a, {{1, 4}});
   CHECK_EQ(a.regions().size(), 1);
-  CHECK_EQ(a.regions().front().begin(), 1);
-  CHECK_EQ(a.regions().front().end(), 5);
+  CHECK_EQ(a.regions().begin()->begin(), 1);
+  CHECK_EQ(a.regions().begin()->end(), 5);
 }
 
 TEST_F(DisjointAllocationPoolTest, ExtractAll) {
@@ -111,6 +111,18 @@ TEST_F(DisjointAllocationPoolTest, Merging) {
   CheckPool(a, {{10, 15}});
 }
 
+TEST_F(DisjointAllocationPoolTest, MergingFirst) {
+  DisjointAllocationPool a = Make({{10, 5}, {20, 5}});
+  a.Merge({5, 5});
+  CheckPool(a, {{5, 10}, {20, 5}});
+}
+
+TEST_F(DisjointAllocationPoolTest, MergingAbove) {
+  DisjointAllocationPool a = Make({{10, 5}, {25, 5}});
+  a.Merge({20, 5});
+  CheckPool(a, {{10, 5}, {20, 10}});
+}
+
 TEST_F(DisjointAllocationPoolTest, MergingMore) {
   DisjointAllocationPool a = Make({{10, 5}, {20, 5}, {30, 5}});
   a.Merge({15, 5});