Give the entry module's top-level code a partition of its own

The emitter writes "; scriptc-entry" before the entry module's top-level
function (its %init). In modules of fewer than 32 partitions, the helper
places the roots that hold a marked function first, and their partitions
take no other root. A root that finds no partition with room goes to the
first open partition from its own.

An edit of the program's top level then recompiles only the top-level code
and the chain that follows it, not the other roots that balancing placed
beside it.

native/llvm-codegen/src/partition.cpp

@@ -40,6 +40,8 @@ struct Symbol {
   std::vector<unsigned> Calls;
   // A variable defined with "constant", whose loads the optimizer can fold.
   bool Constant = false;
+  // A function the emitter marks as the entry module's top-level code.
+  bool Entry = false;
   std::vector<unsigned> Metadata;
   // The source-module group the emitter wrote before this definition, or
   // empty.
@@ -166,8 +168,10 @@ void placeGroups(std::vector<Symbol> &Symbols, ArrayRef<unsigned> Roots,
 // 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.
+// near the mean unless one root is larger. The roots of the entry module's
+// top-level code are placed first, and their partitions take no other root,
+// so an edit of the program's top level recompiles only that code. 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)
@@ -193,27 +197,49 @@ void balanceRoots(std::vector<Symbol> &Symbols, ArrayRef<unsigned> Roots,
   });
   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;
-  // 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.
+  // 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;
-    for (unsigned K = 0; K < Count; ++K) {
+    bool Found = false;
+    for (unsigned K = 0; K < Count && !Found; ++K) {
       unsigned P = (Home + K) % Count;
-      if (Load[P] + Bytes <= Cap) {
+      if (!Closed[P] && Load[P] + Bytes <= Cap) {
         Chosen = P;
-        break;
+        Found = true;
+      }
+    }
+    for (unsigned K = 0; K < Count && !Found; ++K) {
+      unsigned P = (Home + K) % Count;
+      if (!Closed[P]) {
+        Chosen = P;
+        Found = true;
       }
     }
     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.
   for (unsigned Root : Order)
-    if (RootBytes[Root] * 2 > Cap)
+    if (EntryRoots.contains(Root)) {
+      Place(Root);
+      Closed[Placed[Root]] = 1;
+    }
+  for (unsigned Root : Order)
+    if (!EntryRoots.contains(Root) && RootBytes[Root] * 2 > Cap)
       Place(Root);
   for (unsigned Root : Order)
-    if (RootBytes[Root] * 2 <= Cap)
+    if (!EntryRoots.contains(Root) && RootBytes[Root] * 2 <= Cap)
       Place(Root);
   for (unsigned I = 0; I < Symbols.size(); ++I)
     if (Symbols[I].Kind == SymbolKind::Function && Live[I])
@@ -500,9 +526,12 @@ std::optional<ProgramPartitions> partitionProgramText(StringRef Text,
 
   // The emitter writes "; scriptc-group <key>" before the definitions of each
   // source function; the group applies until the next such line, and an
-  // empty key ends it.
+  // empty key ends it. It writes "; scriptc-entry" before the definition of
+  // the entry module's top-level code.
   constexpr StringRef GroupLine = "; scriptc-group";
+  constexpr StringRef EntryLine = "; scriptc-entry";
   StringRef Group;
+  bool Entry = false;
   size_t Pos = 0;
   while (Pos < Text.size()) {
     size_t End = Text.find('\n', Pos);
@@ -516,11 +545,15 @@ std::optional<ProgramPartitions> partitionProgramText(StringRef Text,
                (Line.size() == GroupLine.size() ||
                 Line[GroupLine.size()] == ' ')) {
       Group = Line.drop_front(GroupLine.size()).trim();
+    } else if (Line.trim() == EntryLine) {
+      Entry = true;
     } else if (Line.starts_with("define ")) {
       Symbol S;
       S.Kind = SymbolKind::Function;
       S.Header = Line;
       S.Group = Group;
+      S.Entry = Entry;
+      Entry = false;
       if (!parseFunctionHeader(Line, S))
         return std::nullopt;
       // The emitter indents every body line, so the first line that is

packages/compiler/src/backend/llvm/emit-module.ts

@@ -291,10 +291,22 @@ function partitionGroupLines(functions: readonly IrFunction[]): string[] {
 /** Ends the group of the definitions that follow. */
 const PARTITION_GROUP_END = "; scriptc-group";
 
+/** Written before the entry module's top-level code (its `%init.<n>`). The
+ * LLVM helper gives that function and the chain that follows it a partition
+ * of their own, so an edit of the program's top level recompiles only that
+ * code (native/llvm-codegen/src/partition.cpp). */
+const PARTITION_ENTRY_LINE = "; scriptc-entry";
+
+/** Whether each function is the entry module's top-level code. */
+function partitionEntries(functions: readonly IrFunction[], sourceFile: string): boolean[] {
+  return functions.map((fn) => fn.name.startsWith("%init.") && fn.loc.file === sourceFile);
+}
+
 /** Writes the generated helpers, then the function definitions with their
- * group lines. With a layout, the definitions it lists, helpers included,
- * move after the unlisted helpers and before the other functions, by tier,
- * each in its tier's group (`layout-<tier>`). The helper compiles each
+ * group lines, and the entry line before the entry module's top-level code.
+ * With a layout, the definitions it lists, helpers included, move after the
+ * unlisted helpers and before the other functions, by tier, each in its
+ * tier's group (`layout-<tier>`). The helper compiles each
  * partition's definitions in text order, so listed ones come first in their
  * partitions, and it places a group's definitions in the same partitions.
  * Everything else keeps its order and its source-file group. */
@@ -304,6 +316,7 @@ function writeDefinitions(
   helperTexts: readonly (readonly string[])[],
   fnDefs: readonly string[],
   groups: readonly string[],
+  entries: readonly boolean[],
 ): void {
   const tierLine = (tier: number) => `${PARTITION_GROUP_END} layout-${tier}`;
   const listed = new Map<number, string[]>();
@@ -336,12 +349,13 @@ function writeDefinitions(
   fnDefs.forEach((text, i) => {
     const tier = layout === undefined || layout.size === 0 ? undefined : layoutTierOf(layout, text);
     if (tier === undefined) unlisted.push(i);
-    else listIn(tier, ["", tierLine(tier), text]);
+    else listIn(tier, ["", ...(entries[i] ? [PARTITION_ENTRY_LINE] : []), tierLine(tier), text]);
   });
   for (const tier of [...listed.keys()].sort((a, b) => a - b))
     for (const line of listed.get(tier)!) out.push(line);
   for (const [k, i] of unlisted.entries()) {
     if (k > 0 || listed.size > 0) out.push("");
+    if (entries[i]) out.push(PARTITION_ENTRY_LINE);
     out.push(groups[i] ?? PARTITION_GROUP_END, fnDefs[i]!);
   }
 }
@@ -1095,6 +1109,7 @@ export function emitModuleParts(emitter: LlEmitter): string[] {
     ],
     fnDefs,
     partitionGroupLines(emitter.mod.functions),
+    partitionEntries(emitter.mod.functions, emitter.mod.sourceFile),
   );
   if (fnDefs.length === 0) out.push("");
   out.push(PARTITION_GROUP_END, "");

packages/compiler/test/function-layout.test.ts

@@ -181,3 +181,48 @@ test("a changed function layout misses the executable cache", async () => {
     await rm(dir, { recursive: true, force: true });
   }
 });
+
+test("the entry line precedes only the entry module's top-level code", async () => {
+  const dir = await mkdtemp(join(tmpdir(), "scriptc-partition-entry-"));
+  try {
+    const entry = join(dir, "main.ts");
+    const outPath = join(dir, "main.ll");
+    await writeFile(
+      join(dir, "lib.ts"),
+      `function libOnly(n: number): number { return n * 2; }
+export function libValue(): number { return libOnly(Number(process.argv.length)); }
+console.log("lib");
+`,
+    );
+    await writeFile(
+      entry,
+      `import { libValue } from "./lib";
+function entryOnly(n: number): number { return n + 1; }
+console.log(entryOnly(libValue()));
+`,
+    );
+    const result = await compile(entry, { outDir: dir, outPath, outputKind: "llvm" });
+    expect(result.ok, JSON.stringify(result.diagnostics)).toBe(true);
+    const lines = (await readFile(outPath, "utf8")).split("\n");
+    const marked: string[] = [];
+    let pending = false;
+    for (const line of lines) {
+      if (line === "; scriptc-entry") pending = true;
+      const name = /^define [^@]*@([\w.$]+)\(/.exec(line)?.[1];
+      if (name !== undefined) {
+        if (pending) marked.push(name);
+        pending = false;
+      }
+    }
+    expect(marked).toHaveLength(1);
+    const defs = definitions(lines.join("\n"));
+    const groupOf = (pattern: RegExp) => defs.find((d) => pattern.test(d.name))?.group;
+    expect(marked[0]).toMatch(/init/);
+    expect(groupOf(new RegExp(`^${marked[0]!.replace(/[.$]/g, "\\$&")}$`))).toBe(
+      groupOf(/_entryOnly$/),
+    );
+    expect(groupOf(/_entryOnly$/)).not.toBe(groupOf(/_libOnly$/));
+  } finally {
+    await rm(dir, { recursive: true, force: true });
+  }
+});

packages/compiler/test/native-partitions.test.ts

@@ -793,6 +793,78 @@ describe.runIf(helper !== null)("LLVM program partitions", () => {
     expect(changed).toHaveLength(1);
   });
 
+  test("the entry line gives the top-level code and its chain a partition of their own", async () => {
+    const dir = await mkdtemp(join(tmpdir(), "scriptc-partitions-entry-"));
+    dirs.push(dir);
+    // Eight exported functions of about 7 KiB, and the marked top-level
+    // function @top, which only main calls, so it follows main.
+    const count = 8;
+    const body = (header: string, lines: number, ret: string) => [
+      header,
+      ...Array.from(
+        { length: lines },
+        (_, step) => `  %v${step} = add i64 ${step === 0 ? "%value" : `%v${step - 1}`}, 1`,
+      ),
+      ret.replace("LAST", `%v${lines - 1}`),
+      "}",
+    ];
+    const source = (marked: boolean) => {
+      const lines: string[] = [];
+      for (let index = 0; index < count; index++)
+        lines.push(
+          ...body(`define i64 @big${index * 4}(i64 %value) noinline {`, 280, "  ret i64 LAST"),
+        );
+      if (marked) lines.push("; scriptc-entry");
+      lines.push(
+        ...body("define internal i64 @top(i64 %value) noinline {", 280, "  ret i64 LAST"),
+        "define i32 @main() {",
+        "  %t = call i64 @top(i64 0)",
+      );
+      for (let index = 0; index < count; index++)
+        lines.push(
+          `  %r${index} = call i64 @big${index * 4}(i64 ${index})`,
+          `  %s${index} = add i64 ${index === 0 ? "%t" : `%s${index - 1}`}, %r${index}`,
+        );
+      lines.push(
+        `  %low = urem i64 %s${count - 1}, 251`,
+        "  %code = trunc i64 %low to i32",
+        "  ret i32 %code",
+        "}",
+      );
+      return lines.join("\n") + "\n";
+    };
+    let total = 280;
+    for (let index = 0; index < count; index++) total += index + 280;
+    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+|top|main)$/.exec(line.trim());
+          if (defined !== null) where.set(defined[1]!, index);
+        }
+      }
+      return where;
+    };
+    for (const marked of [false, true]) {
+      const objects = await emit(dir, source(marked), {
+        outputs: 4,
+        optLevel: "0",
+        name: marked ? "marked" : "plain",
+      });
+      expect(await exitCode(dir, objects, marked ? "marked" : "plain")).toBe(total % 251);
+      const where = await placed(objects);
+      expect(where.get("top")).toBe(where.get("main"));
+      const beside = [...where].filter(
+        ([name, home]) => name.startsWith("big") && home === where.get("top"),
+      );
+      // Without the line, balancing places other functions beside main's
+      // chain; with it, none.
+      if (marked) expect(beside).toEqual([]);
+      else expect(beside.length).toBeGreaterThan(0);
+    }
+  });
+
   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);