diff --git a/native/llvm-codegen/src/partition.cpp b/native/llvm-codegen/src/partition.cpp index d56839426..9ab1ce657 100644 --- a/native/llvm-codegen/src/partition.cpp +++ b/native/llvm-codegen/src/partition.cpp @@ -176,71 +176,108 @@ void balanceRoots(std::vector &Symbols, ArrayRef Roots, ArrayRef Live, unsigned Count) { if (Count >= MinGroupedPartitions) return; + // A root's size leaves out the digits of metadata references, which an + // edit elsewhere in the module can renumber. + auto TextBytes = [](StringRef Text) { + uint64_t Bytes = 0; + for (size_t I = 0; I < Text.size(); ++I) { + ++Bytes; + if (Text[I] == '!') + while (I + 1 < Text.size() && isDigit(Text[I + 1])) + ++I; + } + return Bytes; + }; 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; - } + RootBytes[Roots[I]] += + TextBytes(Symbols[I].Header) + TextBytes(Symbols[I].Body); 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; }); + DenseSet EntryRoots; + for (unsigned I = 0; I < Symbols.size(); ++I) + if (Symbols[I].Entry && Live[I]) + EntryRoots.insert(Roots[I]); + // A root of at least one unit counts in whole units, so that an edit that + // grows it by less than that leaves every comparison, and every other + // root's partition, unchanged. + constexpr uint64_t Unit = 4096; + uint64_t Total = 0; + for (unsigned Root : Order) + if (RootBytes[Root] >= Unit) { + RootBytes[Root] = (RootBytes[Root] + Unit - 1) / Unit * Unit; + Total += RootBytes[Root]; + } uint64_t Cap = Total * 5 / 4 / Count; std::vector Load(Count, 0); // Partitions that take no further root. std::vector Closed(Count, 0); DenseMap Placed; + auto FirstOpen = [&](unsigned Home) { + for (unsigned K = 0; K < Count; ++K) + if (!Closed[(Home + K) % Count]) + return (Home + K) % Count; + return Home; + }; // A root without room anywhere stays in the first open partition from its // own. auto Place = [&](unsigned Root) { uint64_t Bytes = RootBytes[Root]; - unsigned Home = stableHash(Symbols[Root].Name) % Count, Chosen = Home; - bool Found = false; - for (unsigned K = 0; K < Count && !Found; ++K) { + unsigned Home = stableHash(Symbols[Root].Name) % Count; + unsigned Chosen = FirstOpen(Home); + for (unsigned K = 0; K < Count; ++K) { unsigned P = (Home + K) % Count; if (!Closed[P] && Load[P] + Bytes <= Cap) { Chosen = P; - Found = true; - } - } - for (unsigned K = 0; K < Count && !Found; ++K) { - unsigned P = (Home + K) % Count; - if (!Closed[P]) { - Chosen = P; - Found = true; + break; } } Load[Chosen] += Bytes; Placed[Root] = Chosen; }; - DenseSet EntryRoots; - for (unsigned I = 0; I < Symbols.size(); ++I) - if (Symbols[I].Entry && Live[I]) - EntryRoots.insert(Roots[I]); - // Entry roots first, each closing its partition, then roots of more than - // half the cap, in name order, so that small roots fill around them; the - // others follow in name order. + // Entry roots first, each closing its partition; then roots of more than + // half the cap, in name order, so that the others fill around them; then + // the other roots of at least one unit, in name order. for (unsigned Root : Order) if (EntryRoots.contains(Root)) { Place(Root); Closed[Placed[Root]] = 1; } for (unsigned Root : Order) - if (!EntryRoots.contains(Root) && RootBytes[Root] * 2 > Cap) + if (!EntryRoots.contains(Root) && RootBytes[Root] >= Unit && + RootBytes[Root] * 2 > Cap) Place(Root); for (unsigned Root : Order) - if (!EntryRoots.contains(Root) && RootBytes[Root] * 2 <= Cap) + if (!EntryRoots.contains(Root) && RootBytes[Root] >= Unit && + RootBytes[Root] * 2 <= Cap) Place(Root); + // Roots under one unit (runtime helpers, mostly) count in no load of the + // larger roots, so a helper that an edit adds or removes moves no larger + // root. Each goes to the first open partition from its name's whose load, + // the helpers placed there included, is below the cap, or else to its + // name's: the helpers fill the partitions with room, and once those are + // full they spread by name instead of gathering in one partition (which + // could then hold most of the module and have it optimized whole). + std::vector SmallLoad(Count, 0); + for (unsigned Root : Order) + if (!EntryRoots.contains(Root) && RootBytes[Root] < Unit) { + unsigned Home = stableHash(Symbols[Root].Name) % Count; + unsigned Chosen = FirstOpen(Home); + for (unsigned K = 0; K < Count; ++K) { + unsigned P = (Home + K) % Count; + if (!Closed[P] && Load[P] + SmallLoad[P] < Cap) { + Chosen = P; + break; + } + } + SmallLoad[Chosen] += RootBytes[Root]; + Placed[Root] = Chosen; + } for (unsigned I = 0; I < Symbols.size(); ++I) if (Symbols[I].Kind == SymbolKind::Function && Live[I]) Symbols[I].Partition = Placed[Roots[I]]; diff --git a/packages/compiler/test/native-partitions.test.ts b/packages/compiler/test/native-partitions.test.ts index 52c6b42f1..836d1202b 100644 --- a/packages/compiler/test/native-partitions.test.ts +++ b/packages/compiler/test/native-partitions.test.ts @@ -873,6 +873,141 @@ describe.runIf(helper !== null)("LLVM program partitions", () => { } }); + test("a small helper that an edit adds moves no other function", async () => { + const dir = await mkdtemp(join(tmpdir(), "scriptc-partitions-helper-roots-")); + dirs.push(dir); + // Eight exported functions of these instruction counts and three small + // exported helpers, each a root of its own. The edit adds a fourth + // helper, which big1 calls. Counted in the loads, the added helper would + // raise the cap and move other functions. + const sizes = [426, 201, 333, 376, 152, 339, 327, 218]; + const helper = (name: string) => [ + `define i64 @${name}(i64 %value) noinline {`, + " %r = add i64 %value, 1", + " ret i64 %r", + "}", + ]; + const source = (edited: boolean) => { + const lines: string[] = []; + sizes.forEach((size, index) => { + lines.push(`define i64 @big${index}(i64 %value) noinline {`); + for (let step = 0; step < size; step++) + lines.push(` %v${step} = add i64 ${step === 0 ? "%value" : `%v${step - 1}`}, 1`); + if (edited && index === 1) + lines.push(` %h = call i64 @added(i64 %v${size - 1})`, " ret i64 %h"); + else lines.push(` ret i64 %v${size - 1}`); + lines.push("}"); + }); + for (let index = 0; index < 3; index++) lines.push(...helper(`helper${index}`)); + if (edited) lines.push(...helper("added")); + lines.push("define i32 @main() {"); + sizes.forEach((_, index) => + lines.push( + ` %r${index} = call i64 @big${index}(i64 ${index})`, + ` %t${index} = add i64 ${index === 0 ? "0" : `%t${index - 1}`}, %r${index}`, + ), + ); + for (let index = 0; index < 3; index++) { + const operand = index === 0 ? `%t${sizes.length - 1}` : `%h${index - 1}`; + lines.push(` %h${index} = call i64 @helper${index}(i64 ${operand})`); + } + lines.push( + " %low = urem i64 %h2, 251", + " %code = trunc i64 %low to i32", + " ret i32 %code", + "}", + ); + return lines.join("\n") + "\n"; + }; + const expected = (edited: boolean) => { + let total = 3 + (edited ? 1 : 0); + sizes.forEach((size, index) => (total += index + size)); + 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+|helper\d+|main)$/.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(false), { ...options, name: "before" }); + expect(await exitCode(dir, objects, "before")).toBe(expected(false)); + const before = await cacheEntries(cache, ".o"); + const edited = await emit(dir, source(true), { ...options, name: "after" }); + expect(await exitCode(dir, edited, "after")).toBe(expected(true)); + expect(await placed(edited)).toEqual(await placed(objects)); + // Only big1's partition and the new helper's change. + const changed = [...(await cacheEntries(cache, ".o"))].filter((entry) => !before.has(entry)); + expect(changed.length).toBeLessThanOrEqual(2); + }); + + test("small helpers spread when one partition is left below the cap", async () => { + const dir = await mkdtemp(join(tmpdir(), "scriptc-partitions-small-spread-")); + dirs.push(dir); + // Three exported functions of about 16 KiB, each over the cap of four + // partitions, whose names hash to partitions 3, 0 and 1, and forty + // exported helpers of about 2.5 KiB. Partition 2 is the only one left + // below the cap; counted only against the larger roots' loads, every + // helper would go there. + const chain = (header: string, lines: number) => [ + header, + ...Array.from( + { length: lines }, + (_, step) => ` %v${step} = add i64 ${step === 0 ? "%value" : `%v${step - 1}`}, 1`, + ), + ` ret i64 %v${lines - 1}`, + "}", + ]; + const helpers = 40; + const lines: string[] = []; + for (let index = 0; index < 3; index++) + lines.push(...chain(`define i64 @big${index}(i64 %value) noinline {`, 640)); + for (let index = 0; index < helpers; index++) + lines.push(...chain(`define i64 @helper${index}(i64 %value) noinline {`, 100)); + lines.push("define i32 @main() {"); + const calls = [ + ...[0, 1, 2].map((i) => `big${i}`), + ...Array.from({ length: helpers }, (_, i) => `helper${i}`), + ]; + calls.forEach((name, index) => + lines.push( + ` %r${index} = call i64 @${name}(i64 ${index})`, + ` %s${index} = add i64 ${index === 0 ? "0" : `%s${index - 1}`}, %r${index}`, + ), + ); + lines.push( + ` %low = urem i64 %s${calls.length - 1}, 251`, + " %code = trunc i64 %low to i32", + " ret i32 %code", + "}", + ); + let total = 3 * 640 + helpers * 100; + calls.forEach((_, index) => (total += index)); + const objects = await emit(dir, lines.join("\n") + "\n", { + outputs: 4, + optLevel: "0", + name: "spread", + }); + expect(await exitCode(dir, objects, "spread")).toBe(total % 251); + const homes = new Map(); + for (const [index, object] of objects.entries()) { + const { stdout } = await execFileAsync("nm", ["--defined-only", object]); + for (const line of stdout.split("\n")) + if (/^[0-9a-fA-F]+ [A-Za-z] _?helper\d+$/.test(line.trim())) + homes.set(index, (homes.get(index) ?? 0) + 1); + } + expect([...homes.values()].reduce((sum, count) => sum + count, 0)).toBe(helpers); + expect(Math.max(...homes.values())).toBeLessThan(helpers); + expect(homes.size).toBeGreaterThan(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);