Balance program partitions by size (#2461)

2026-11-06 08:23:19 +08:00 · 2026-10-28 06:10:55 -05:00
parent 62b34e7e01 007a144cb6
commit 0bf0f69dc3
2 changed files with 240 additions and 11 deletions
@@ -36,6 +36,10 @@ struct Symbol {
  // Local definitions are rewritten with hidden visibility.
  std::string PromotedHeader;
  std::vector<unsigned> Refs;
+  // The functions its body calls directly (sorted).
+  std::vector<unsigned> Calls;
+  // A variable defined with "constant", whose loads the optimizer can fold.
+  bool Constant = false;
  std::vector<unsigned> Metadata;
  // The source-module group the emitter wrote before this definition, or
  // empty.
@@ -74,6 +78,8 @@ uint32_t stableHash(StringRef Name) {
  return Hash;
 }

+constexpr unsigned MinGroupedPartitions = 32;
+
 // Places the functions whose root carries a source-module group by group: a
 // module's functions call each other most, so keeping them together leaves
 // fewer calls for the thin link to import across partitions. A group larger
@@ -83,11 +89,11 @@ uint32_t stableHash(StringRef Name) {
 // lighter of two partitions chosen by hashes of its key, after the functions
 // placed by name. The placement depends only on the module text and the
 // partition count. Modules with fewer than MinGroupedPartitions partitions
-// keep the placement by name: with few partitions per core, the balance of
-// the largest partitions decides the build time more than the imports.
+// keep the placement by name, which balanceRoots evens out: with few
+// partitions per core, the balance of the largest partitions decides the
+// build time more than the imports.
 void placeGroups(std::vector<Symbol> &Symbols, ArrayRef<unsigned> Roots,
                 unsigned Count) {
-  constexpr unsigned MinGroupedPartitions = 32;
  if (Count < MinGroupedPartitions)
    return;
  auto Bytes = [](const Symbol &S) -> uint64_t {
@@ -153,6 +159,67 @@ void placeGroups(std::vector<Symbol> &Symbols, ArrayRef<unsigned> Roots,
      Symbols[I].Partition = Where[Keys[I]];
 }

+// In modules of fewer than MinGroupedPartitions partitions, which
+// placeGroups leaves placed by name, the functions of each root (a
+// function and the chain that follows it) move together to another partition
+// when their name's partition would hold more than 1.25 times the mean
+// function bytes: each root probes the next partitions in turn, keeping the
+// first with room, else its own. An edit then recompiles about a mean
+// partition, and the largest partition, which bounds a cold build, stays
+// near the mean unless one root is larger. Roots and their bytes, and so
+// placement, depend only on the module text.
+void balanceRoots(std::vector<Symbol> &Symbols, ArrayRef<unsigned> Roots,
+                  ArrayRef<char> Live, unsigned Count) {
+  if (Count >= MinGroupedPartitions)
+    return;
+  DenseMap<unsigned, uint64_t> RootBytes;
+  uint64_t Total = 0;
+  for (unsigned I = 0; I < Symbols.size(); ++I)
+    if (Symbols[I].Kind == SymbolKind::Function && Live[I])
+      RootBytes[Roots[I]] += Symbols[I].Header.size() + Symbols[I].Body.size();
+  // Each root counts in whole 4 KiB units, so that an edit that grows a root
+  // by less than that leaves every comparison, and every other root's
+  // partition, unchanged.
+  constexpr uint64_t Unit = 4096;
+  for (auto &Entry : RootBytes) {
+    Entry.second = (Entry.second + Unit - 1) / Unit * Unit;
+    Total += Entry.second;
+  }
+  std::vector<unsigned> Order;
+  for (auto &Entry : RootBytes)
+    Order.push_back(Entry.first);
+  llvm::sort(Order, [&](unsigned A, unsigned B) {
+    return Symbols[A].Name < Symbols[B].Name;
+  });
+  uint64_t Cap = Total * 5 / 4 / Count;
+  std::vector<uint64_t> Load(Count, 0);
+  DenseMap<unsigned, unsigned> Placed;
+  // Roots of more than half the cap are placed first, in name order, so
+  // that small roots fill around them; the others follow in name order.
+  auto Place = [&](unsigned Root) {
+    uint64_t Bytes = RootBytes[Root];
+    unsigned Home = stableHash(Symbols[Root].Name) % Count, Chosen = Home;
+    for (unsigned K = 0; K < Count; ++K) {
+      unsigned P = (Home + K) % Count;
+      if (Load[P] + Bytes <= Cap) {
+        Chosen = P;
+        break;
+      }
+    }
+    Load[Chosen] += Bytes;
+    Placed[Root] = Chosen;
+  };
+  for (unsigned Root : Order)
+    if (RootBytes[Root] * 2 > Cap)
+      Place(Root);
+  for (unsigned Root : Order)
+    if (RootBytes[Root] * 2 <= Cap)
+      Place(Root);
+  for (unsigned I = 0; I < Symbols.size(); ++I)
+    if (Symbols[I].Kind == SymbolKind::Function && Live[I])
+      Symbols[I].Partition = Placed[Roots[I]];
+}
+
 bool startsWithLinkage(StringRef Rest) {
  for (StringRef Linkage :
       {"linkonce", "weak", "common", "appending", "extern_weak",
@@ -278,6 +345,7 @@ bool parseGlobal(StringRef Line, Symbol &S) {
    }
    if (Cursor.consume_front("constant ")) {
      Keyword = "constant";
+      S.Constant = true;
      break;
    }
    size_t Space = Cursor.find(' ');
@@ -325,10 +393,13 @@ bool parseGlobal(StringRef Line, Symbol &S) {
 }

 // Collects the symbols (sorted) and metadata nodes (in order of first
-// appearance) that Text references outside quoted strings.
+// appearance) that Text references outside quoted strings, and in Calls the
+// symbols it names as direct callees, which the emitter writes followed by
+// "(".
 void scanReferences(StringRef Text, const StringMap<unsigned> &Index,
                    std::vector<unsigned> &Refs,
-                    std::vector<unsigned> &Metadata) {
+                    std::vector<unsigned> &Metadata,
+                    std::vector<unsigned> *Calls = nullptr) {
  DenseSet<unsigned> Seen(Metadata.begin(), Metadata.end());
  bool Quoted = false;
  for (size_t I = 0; I < Text.size(); ++I) {
@@ -343,8 +414,12 @@ void scanReferences(StringRef Text, const StringMap<unsigned> &Index,
      size_t Length = nameLength(Text, I + 1);
      if (Length != 0) {
        auto It = Index.find(Text.substr(I + 1, Length));
-        if (It != Index.end())
+        if (It != Index.end()) {
          Refs.push_back(It->second);
+          if (Calls && I + 1 + Length < Text.size() &&
+              Text[I + 1 + Length] == '(')
+            Calls->push_back(It->second);
+        }
        I += Length;
      }
    } else if (C == '!' && I + 1 < Text.size() && isDigit(Text[I + 1])) {
@@ -359,6 +434,10 @@ void scanReferences(StringRef Text, const StringMap<unsigned> &Index,
  }
  llvm::sort(Refs);
  Refs.erase(std::unique(Refs.begin(), Refs.end()), Refs.end());
+  if (Calls) {
+    llvm::sort(*Calls);
+    Calls->erase(std::unique(Calls->begin(), Calls->end()), Calls->end());
+  }
 }

 // Appends Text with every unquoted metadata reference renumbered by Map.
@@ -539,7 +618,7 @@ std::optional<ProgramPartitions> partitionProgramText(StringRef Text,
    scanReferences(S.Header, Index, S.Refs, S.Metadata);
    if (!S.Body.empty()) {
      std::vector<unsigned> Refs;
-      scanReferences(S.Body, Index, Refs, S.Metadata);
+      scanReferences(S.Body, Index, Refs, S.Metadata, &S.Calls);
      S.Refs.insert(S.Refs.end(), Refs.begin(), Refs.end());
      llvm::sort(S.Refs);
      S.Refs.erase(std::unique(S.Refs.begin(), S.Refs.end()), S.Refs.end());
@@ -594,7 +673,15 @@ std::optional<ProgramPartitions> partitionProgramText(StringRef Text,
  // inlining candidate of it, or data only it uses; keep it with that
  // definition so its partition can inline it without importing, and so a
  // helper an edit adds lands in the partition of the code that uses it. A
-  // chain through a global continues to the function that uses the global.
+  // function that its referrer only takes the address of is not an inlining
+  // candidate: a closure body whose address its creator passes to the
+  // runtime, or a compiled regular expression's matcher that a mutable
+  // global points to, gives the referrer nothing to inline, and joining it
+  // would make every edit of the referrer optimize it again. Such a function
+  // is placed by its own name, unless its body is under AddressOnlyBytes (a
+  // trampoline or thread entry, which costs little to optimize again). A
+  // chain through a constant global, whose loads the optimizer can turn into
+  // direct calls, continues to the function that uses the global.
  // A function body larger than this is past LLVM's inline cost limit even
  // for the only call to a local function, so it is placed by its own name.
  // A chain stops where it would leave a function's source-module group for
@@ -608,9 +695,14 @@ std::optional<ProgramPartitions> partitionProgramText(StringRef Text,
  // partition, the program is optimized whole (see emitPartitions).
  constexpr size_t MaxInlinedBodyBytes = 256 * 1024;
  constexpr size_t GroupedChainBytes = 1024 * 1024;
+  constexpr size_t AddressOnlyBytes = 512;
  bool GroupedChains = Text.size() > GroupedChainBytes;
  constexpr unsigned Shared = ~1u;
  std::vector<unsigned> Referrer(Symbols.size(), NoPartition);
+  // Whether a definition may follow its first referrer: it is not a
+  // function, or the referrer calls it or is a constant global, or its body
+  // is under AddressOnlyBytes.
+  std::vector<char> MayFollow(Symbols.size(), 0);
  for (unsigned I = 0; I < Symbols.size(); ++I) {
    const Symbol &S = Symbols[I];
    if (S.Kind == SymbolKind::Declaration || !Live[I])
@@ -618,13 +710,23 @@ std::optional<ProgramPartitions> partitionProgramText(StringRef Text,
    for (unsigned Ref : S.Refs) {
      if (Ref == I)
        continue;
-      Referrer[Ref] = Referrer[Ref] == NoPartition ? I : Shared;
+      if (Referrer[Ref] == NoPartition) {
+        Referrer[Ref] = I;
+        MayFollow[Ref] =
+            Symbols[Ref].Kind != SymbolKind::Function ||
+            Symbols[Ref].Body.size() < AddressOnlyBytes ||
+            (S.Kind == SymbolKind::Function
+                 ? std::binary_search(S.Calls.begin(), S.Calls.end(), Ref)
+                 : S.Constant);
+      } else {
+        Referrer[Ref] = Shared;
+      }
    }
  }
  auto Follows = [&](unsigned Child, unsigned Parent) {
    const Symbol &S = Symbols[Child];
-    if (!S.Local || Parent == NoPartition || Parent == Shared ||
-        Symbols[Parent].Kind == SymbolKind::Declaration)
+    if (!S.Local || !MayFollow[Child] || Parent == NoPartition ||
+        Parent == Shared || Symbols[Parent].Kind == SymbolKind::Declaration)
      return false;
    if (S.Kind == SymbolKind::Function &&
        (S.Body.size() > MaxInlinedBodyBytes ||
@@ -656,6 +758,7 @@ std::optional<ProgramPartitions> partitionProgramText(StringRef Text,
      Symbols[I].Partition = stableHash(Symbols[Root].Name) % Count;
  }
  placeGroups(Symbols, Roots, Count);
+  balanceRoots(Symbols, Roots, Live, Count);
  // A global follows the function its chain ends at, then the first
  // function that references it, then its own name.
  for (unsigned I = 0; I < Symbols.size(); ++I)
@@ -667,6 +667,132 @@ describe.runIf(helper !== null)("LLVM program partitions", () => {
    expect(changed).toHaveLength(new Set([where.get("step7"), where.get("step9")]).size);
  });

+  test("a function that its referrer only takes the address of is placed by its own name", async () => {
+    const dir = await mkdtemp(join(tmpdir(), "scriptc-partitions-address-"));
+    dirs.push(dir);
+    // A body of `lines` additions, past the size under which an
+    // address-only function still follows its referrer.
+    const body = (name: string, lines: number) => [
+      `define internal void @${name}(i64 %value) noinline {`,
+      ...Array.from(
+        { length: lines },
+        (_, index) => `  %v${index} = add i64 ${index === 0 ? "%value" : `%v${index - 1}`}, 1`,
+      ),
+      "  call void @bump(i64 0)",
+      "  ret void",
+      "}",
+    ];
+    // step7 passes @closure and @trampoline to @spawn; step9 calls
+    // @matcher through a mutable global only it reads.
+    const source = program()
+      .split("\n")
+      .flatMap((line) => {
+        if (line === "define internal void @step7(i64 %value) noinline {")
+          return [
+            ...body("closure", 40),
+            ...body("trampoline", 1),
+            "define internal void @spawn(ptr %fn) noinline {",
+            "  call void %fn(i64 1)",
+            "  ret void",
+            "}",
+            line,
+            "  call void @spawn(ptr @closure)",
+            "  call void @spawn(ptr @trampoline)",
+          ];
+        if (line === "define internal void @step9(i64 %value) noinline {")
+          return [
+            "@descriptor = internal global ptr @matcher",
+            ...body("matcher", 40),
+            line,
+            "  %fn = load ptr, ptr @descriptor",
+            "  call void %fn(i64 0)",
+          ];
+        return [line];
+      })
+      .join("\n");
+    const objects = await emit(dir, source, { outputs: 4, optLevel: "0", name: "address" });
+    expect(await exitCode(dir, objects, "address")).toBe(expectedExit());
+    const where = new Map<string, number>();
+    for (const [index, object] of objects.entries()) {
+      const { stdout } = await execFileAsync("nm", [object]);
+      for (const line of stdout.split("\n")) {
+        const defined = /^[0-9a-fA-F]+ [A-Za-z] _?(step7|step9|closure|trampoline|matcher)$/.exec(
+          line.trim(),
+        );
+        if (defined !== null) where.set(defined[1]!, index);
+      }
+    }
+    expect(where.size).toBe(5);
+    expect(where.get("trampoline")).toBe(where.get("step7"));
+    expect(where.get("closure")).not.toBe(where.get("step7"));
+    expect(where.get("matcher")).not.toBe(where.get("step9"));
+  });
+
+  test("functions are balanced over a few partitions, and a small edit keeps their places", async () => {
+    const dir = await mkdtemp(join(tmpdir(), "scriptc-partitions-balance-"));
+    dirs.push(dir);
+    // Eight exported functions of about 7 KiB each, which main calls; each
+    // returns its argument plus its number of additions. By the hash of
+    // their names alone, six of them would share one partition.
+    const count = 8;
+    const additions = (index: number, edit: number) => 280 + (index === 3 ? edit : 0);
+    const source = (edit: number) => {
+      const lines: string[] = [];
+      for (let index = 0; index < count; index++) {
+        lines.push(`define i64 @big${index * 4}(i64 %value) noinline {`);
+        for (let step = 0; step < additions(index, edit); step++)
+          lines.push(`  %v${step} = add i64 ${step === 0 ? "%value" : `%v${step - 1}`}, 1`);
+        lines.push(`  ret i64 %v${additions(index, edit) - 1}`, "}");
+      }
+      lines.push("define i32 @main() {");
+      for (let index = 0; index < count; index++)
+        lines.push(
+          `  %r${index} = call i64 @big${index * 4}(i64 ${index})`,
+          `  %t${index} = add i64 ${index === 0 ? "0" : `%t${index - 1}`}, %r${index}`,
+        );
+      lines.push(
+        `  %low = urem i64 %t${count - 1}, 251`,
+        "  %code = trunc i64 %low to i32",
+        "  ret i32 %code",
+        "}",
+      );
+      return lines.join("\n") + "\n";
+    };
+    const expected = (edit: number) => {
+      let total = 0;
+      for (let index = 0; index < count; index++) total += index + additions(index, edit);
+      return total % 251;
+    };
+    const placed = async (objects: string[]) => {
+      const where = new Map<string, number>();
+      for (const [index, object] of objects.entries()) {
+        const { stdout } = await execFileAsync("nm", ["--defined-only", object]);
+        for (const line of stdout.split("\n")) {
+          const defined = /^[0-9a-fA-F]+ [A-Za-z] _?(big\d+)$/.exec(line.trim());
+          if (defined !== null) where.set(defined[1]!, index);
+        }
+      }
+      return where;
+    };
+    const cache = join(dir, "cache");
+    const options = { outputs: 4, optLevel: "0", cache } as const;
+    const objects = await emit(dir, source(0), { ...options, name: "balanced" });
+    expect(await exitCode(dir, objects, "balanced")).toBe(expected(0));
+    const where = await placed(objects);
+    expect(where.size).toBe(count);
+    for (let partition = 0; partition < 4; partition++)
+      expect([...where.values()].filter((home) => home === partition)).toHaveLength(2);
+
+    // Growing one function by a few hundred bytes recompiles only its
+    // partition and moves no function.
+    const before = await cacheEntries(cache, ".o");
+    const edited = await emit(dir, source(8), { ...options, name: "edited" });
+    expect(await exitCode(dir, edited, "edited")).toBe(expected(8));
+    expect(await placed(edited)).toEqual(where);
+    const changed = [...(await cacheEntries(cache, ".o"))].filter((entry) => !before.has(entry));
+    expect(changed).toHaveLength(1);
+  });
+
  test("a table written one element per line keeps the text partitioner and its reuse", async () => {
    const dir = await mkdtemp(join(tmpdir(), "scriptc-partitions-table-"));
    dirs.push(dir);