Keep partition placement stable when an edit adds a small helper (#2561)

native/llvm-codegen/src/partition.cpp

@@ -176,71 +176,108 @@ void balanceRoots(std::vector<Symbol> &Symbols, ArrayRef<unsigned> Roots,
                   ArrayRef<char> 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<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;
-  }
+      RootBytes[Roots[I]] +=
+          TextBytes(Symbols[I].Header) + TextBytes(Symbols[I].Body);
   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;
   });
+  DenseSet<unsigned> 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<uint64_t> Load(Count, 0);
   // Partitions that take no further root.
   std::vector<char> Closed(Count, 0);
   DenseMap<unsigned, unsigned> 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<unsigned> 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<uint64_t> 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]];

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<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+|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<number, number>();
+    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);