From 007a144cb68fc5c2582faa2a4569f1a785c693c3 Mon Sep 17 00:00:00 2001 From: Chris Tate Date: Wed, 28 Oct 2026 05:12:02 -0500 Subject: [PATCH] Balance program partitions by size --- native/llvm-codegen/src/partition.cpp | 125 +++++++++++++++-- .../compiler/test/native-partitions.test.ts | 126 ++++++++++++++++++ 2 files changed, 240 insertions(+), 11 deletions(-) diff --git a/native/llvm-codegen/src/partition.cpp b/native/llvm-codegen/src/partition.cpp index 9fdaeba9a..7174e2b24 100644 --- a/native/llvm-codegen/src/partition.cpp +++ b/native/llvm-codegen/src/partition.cpp @@ -36,6 +36,10 @@ struct Symbol { // Local definitions are rewritten with hidden visibility. std::string PromotedHeader; std::vector Refs; + // The functions its body calls directly (sorted). + std::vector Calls; + // A variable defined with "constant", whose loads the optimizer can fold. + bool Constant = false; std::vector 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 &Symbols, ArrayRef 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 &Symbols, ArrayRef 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 &Symbols, ArrayRef Roots, + ArrayRef Live, unsigned Count) { + if (Count >= MinGroupedPartitions) + return; + DenseMap 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 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 Load(Count, 0); + DenseMap 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 &Index, std::vector &Refs, - std::vector &Metadata) { + std::vector &Metadata, + std::vector *Calls = nullptr) { DenseSet 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 &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 &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 partitionProgramText(StringRef Text, scanReferences(S.Header, Index, S.Refs, S.Metadata); if (!S.Body.empty()) { std::vector 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 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 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 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 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 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 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) diff --git a/packages/compiler/test/native-partitions.test.ts b/packages/compiler/test/native-partitions.test.ts index e85802db9..9a4a5c755 100644 --- a/packages/compiler/test/native-partitions.test.ts +++ b/packages/compiler/test/native-partitions.test.ts @@ -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(); + 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(); + 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);