Keep emission's workers busy (#2451)

packages/compiler/native/host.c

@@ -73,12 +73,14 @@ void scriptc_native_collect_cycles(void) {
 int32_t scriptc_native_fork(void) { return -1; }
 void scriptc_native_exit(int32_t code) { ExitProcess((UINT)code); }
 int32_t scriptc_native_wait(int32_t pid) { (void)pid; return -1; }
+int32_t scriptc_native_exited(void) { return -1; }
 uint8_t scriptc_native_lock_open(const uint8_t *bytes, size_t size) { (void)bytes; (void)size; return 0; }
 void scriptc_native_lock(void) {}
 void scriptc_native_unlock(void) {}
 #else
 #include <errno.h>
 #include <fcntl.h>
+#include <signal.h>
 #include <sys/file.h>
 #include <sys/wait.h>
 
@@ -119,6 +121,18 @@ int32_t scriptc_native_wait(int32_t pid) {
   return -1;
 }
 
+/* The pid of a child that has ended, waiting for one if none has, or -1
+ * when waitid fails. The child is left to scriptc_native_wait: a child
+ * another part of the process waits for is not reaped here. */
+int32_t scriptc_native_exited(void) {
+  siginfo_t info;
+  memset(&info, 0, sizeof info);
+  while (waitid(P_ALL, 0, &info, WEXITED | WNOWAIT) < 0) {
+    if (errno != EINTR) return -1;
+  }
+  return info.si_pid > 0 ? (int32_t)info.si_pid : -1;
+}
+
 static int scriptc_native_lock_fd = -1;
 
 uint8_t scriptc_native_lock_open(const uint8_t *bytes, size_t size) {

packages/compiler/native/host.ffi.json

@@ -6,6 +6,7 @@
     { "name": "compilerNativeFork", "symbol": "scriptc_native_fork", "params": [], "returns": "i32" },
     { "name": "compilerNativeExit", "symbol": "scriptc_native_exit", "params": ["i32"], "returns": "void" },
     { "name": "compilerNativeWait", "symbol": "scriptc_native_wait", "params": ["i32"], "returns": "i32" },
+    { "name": "compilerNativeExited", "symbol": "scriptc_native_exited", "params": [], "returns": "i32" },
     { "name": "compilerNativeLockOpen", "symbol": "scriptc_native_lock_open", "params": ["string"], "returns": "bool" },
     { "name": "compilerNativeLock", "symbol": "scriptc_native_lock", "params": [], "returns": "void" },
     { "name": "compilerNativeUnlock", "symbol": "scriptc_native_unlock", "params": [], "returns": "void" },

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

@@ -115,24 +115,56 @@ export function workerMinFunctions(): number {
   return Number.isFinite(value) && value > 0 ? value : 2000;
 }
 
+/** A function's estimated emission cost beyond its IR nodes, in nodes:
+ * fitted to the worker emission times of a large program's ranges. */
+const FUNCTION_EMIT_COST = 22;
+
+/** Where each of about `count` contiguous ranges of `functions` starts,
+ * followed by the end: ranges of about equal estimated emission cost, a
+ * function's node count plus FUNCTION_EMIT_COST. No range is empty. */
+export function emitRangeStarts(
+  functions: readonly IrFunction[],
+  count: number,
+  nodeCount: (fn: IrFunction) => number | undefined,
+): number[] {
+  const costs = functions.map((fn) => (nodeCount(fn) ?? 0) + FUNCTION_EMIT_COST);
+  const total = costs.reduce((sum, cost) => sum + cost, 0);
+  const starts = [0];
+  let sum = 0;
+  for (let i = 0; i < functions.length - 1 && starts.length < count; i++) {
+    sum += costs[i]!;
+    if (sum >= (total * starts.length) / count) starts.push(i + 1);
+  }
+  starts.push(functions.length);
+  return starts;
+}
+
 /** The function bodies emitted on forked worker processes, `workers` at a
- * time: the function list splits into contiguous ranges, each emitted in
- * its own process, which inherits this emitter as it is at the fork and
- * records a journal (emit-journal.ts). Ranges merge in function order, so
- * the text is the serial text. A range whose worker fails is emitted here
- * when its turn comes. Null without fork primitives (Node, Windows, inside
- * a worker) or when the module is small. */
+ * time: the function list splits into contiguous ranges of about equal
+ * estimated cost (emitRangeStarts), each emitted in its own process, which
+ * inherits this emitter as it is at the fork and records a journal
+ * (emit-journal.ts). A new range starts whenever a worker ends, and ranges
+ * merge in function order, so the text is the serial text. A range whose
+ * worker fails is emitted here when its turn comes. Null without fork
+ * primitives (Node, Windows, inside a worker) or when the module is small. */
 function emitFunctionsOnWorkers(emitter: LlEmitter, workers: number): string[] | null {
   const fork = emitter.targetOptions.forkPrimitives;
   const functions = emitter.mod.functions;
   if (fork === undefined || workers < 2) return null;
   if (functions.length < workerMinFunctions()) return null;
   // More ranges than workers: a range's merge overlaps the later ranges'
-  // emission, and a slow range holds back less.
-  const count = Math.min(functions.length, workers * 3);
-  const size = Math.ceil(functions.length / count);
+  // emission, and a slow range holds back less. Ranges of equal cost finish
+  // at about the same time.
+  const starts = emitRangeStarts(functions, Math.min(functions.length, workers * 3), (fn) =>
+    emitter.nodeIndex.nodeCount(fn),
+  );
+  const count = starts.length - 1;
+  const rangeFunctions = (range: number): readonly IrFunction[] =>
+    functions.slice(starts[range], starts[range + 1]);
   const dir = mkdtempSync(join(tmpdir(), "scriptc-emit-"));
   const pids: number[] = [];
+  // The workers not yet reaped, by pid, with their ranges.
+  const running = new Map<number, number>();
   const start = (range: number): void => {
     const pid = fork.fork();
     if (pid === 0) {
@@ -141,7 +173,7 @@ function emitFunctionsOnWorkers(emitter: LlEmitter, workers: number): string[] |
         emitter.journal = new EmitJournal();
         emitter.debug?.record();
         const texts: string[] = [];
-        for (const fn of functions.slice(range * size, (range + 1) * size)) {
+        for (const fn of rangeFunctions(range)) {
           texts.push(emitter.emitFunction(fn));
           relieveHeapPressure();
         }
@@ -155,17 +187,37 @@ function emitFunctionsOnWorkers(emitter: LlEmitter, workers: number): string[] |
       fork.exit(code);
     }
     pids.push(pid);
+    if (pid > 0) running.set(pid, range);
   };
   const fnDefs: string[] = [];
+  // Ended ranges by index, with what their workers wrote (null when a worker
+  // failed), waiting to merge in function order.
+  const ended = new Map<number, { result: ChunkResult; texts: string[] } | null>();
+  // Reaps a worker and starts the next range in its place. A worker's journal
+  // records requests, not table contents, so its range's text doesn't depend
+  // on what merged before it forked.
+  const reap = (pid: number, range: number): void => {
+    const status = pid > 0 ? fork.wait(pid) : -1;
+    running.delete(pid);
+    if (pids.length < count) start(pids.length);
+    ended.set(range, status === 0 ? readRange(dir, range) : null);
+  };
   try {
     for (let range = 0; range < Math.min(workers, count); range++) start(range);
     for (let range = 0; range < count; range++) {
-      const pid = pids[range]!;
-      const status = pid > 0 ? fork.wait(pid) : -1;
-      // The next worker starts from this emitter as merged so far, which
-      // the journal allows: it records requests, not table contents.
-      const merged = status === 0 ? readRange(dir, range) : null;
-      const slice = functions.slice(range * size, (range + 1) * size);
+      // Until this range's worker ends, reap whichever worker ends first, so
+      // that `workers` keep running while ranges merge in order. A child
+      // that isn't a worker (or no `exited`) waits for this range's worker.
+      while (!ended.has(range)) {
+        const pid = pids[range]!;
+        const first = pid > 0 && fork.exited !== undefined ? fork.exited() : pid;
+        const other = running.get(first);
+        if (first !== pid && other !== undefined) reap(first, other);
+        else reap(pid, range);
+      }
+      const merged = ended.get(range)!;
+      ended.delete(range);
+      const slice = rangeFunctions(range);
       if (merged !== null && merged.texts.length === slice.length) {
         const rename = emitter.mergeChunk(merged.result);
         slice.forEach((fn, index) => {
@@ -180,7 +232,6 @@ function emitFunctionsOnWorkers(emitter: LlEmitter, workers: number): string[] |
           fnDefs.push(def);
           relieveHeapPressure();
         }
-      if (pids.length < count) start(pids.length);
     }
   } finally {
     rmSync(dir, { recursive: true, force: true });

packages/compiler/src/backend/llvm/emitter.ts

@@ -508,6 +508,10 @@ export interface ForkPrimitives {
   exit(code: number): void;
   /** The worker's exit status, 128 plus a signal number, or -1. */
   wait(pid: number): number;
+  /** The pid of a child process that has ended, waiting for one if none
+   * has, or -1, without reaping it. Without it, workers are waited for in
+   * order. */
+  exited?(): number;
 }
 
 export interface LlvmTargetOptions {

packages/compiler/src/frontend/process-fork.ts

@@ -8,6 +8,9 @@ export interface ProcessFork {
   exit(code: number): void;
   /** The worker's exit status, 128 plus a signal number, or -1. */
   wait(pid: number): number;
+  /** The pid of a child process that has ended, waiting for one if none
+   * has, or -1. It stays unreaped: `wait` reaps it. */
+  exited(): number;
   /** Opens the lock file checker requests from workers take. */
   openLock(path: string): boolean;
   lock(): void;

packages/compiler/src/ir/node-index.ts

@@ -196,6 +196,13 @@ export class IrNodeIndex {
     return scan;
   }
 
+  /** The number of nodes recorded for `fn`, or undefined when `fn` is not
+   * recorded. */
+  nodeCount(fn: IrFunction): number | undefined {
+    const f = this.functionIndex.get(fn);
+    return f === undefined ? undefined : this.starts[f + 1]! - this.starts[f]!;
+  }
+
   /** Calls `note` with each recorded node of `fn` whose kind `scan` names,
    * in walk order, until it answers false. False (nothing called) when
    * `fn` is not recorded. */

packages/compiler/src/native/cli.ts

@@ -36,6 +36,7 @@ declare function compilerNativeCollectCycles(): void;
 declare function compilerNativeFork(): number;
 declare function compilerNativeExit(code: number): void;
 declare function compilerNativeWait(pid: number): number;
+declare function compilerNativeExited(): number;
 declare function compilerNativeLockOpen(path: string): boolean;
 declare function compilerNativeLock(): void;
 declare function compilerNativeUnlock(): void;
@@ -73,6 +74,7 @@ async function main(): Promise<number> {
       fork: () => compilerNativeFork(),
       exit: (code) => compilerNativeExit(code),
       wait: (pid) => compilerNativeWait(pid),
+      exited: () => compilerNativeExited(),
       openLock: (path) => compilerNativeLockOpen(path),
       lock: () => compilerNativeLock(),
       unlock: () => compilerNativeUnlock(),

packages/compiler/src/native/prepare.ts

@@ -603,5 +603,6 @@ function emitForkPrimitives(): ForkPrimitives | undefined {
     fork: () => primitives.fork(),
     exit: (code) => primitives.exit(code),
     wait: (pid) => primitives.wait(pid),
+    exited: () => primitives.exited(),
   };
 }

packages/compiler/test/node-index.test.ts

@@ -3,7 +3,8 @@ import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { beforeAll, expect, test } from "vitest";
 import { compile, deserializeModule } from "../src/index.js";
-import type { IrExpr, IrModule, IrStmt, IrType } from "../src/ir/ir.js";
+import type { IrExpr, IrFunction, IrModule, IrStmt, IrType } from "../src/ir/ir.js";
+import { emitRangeStarts } from "../src/backend/llvm/emit-module.js";
 import { eachNodeOfKinds, IrNodeIndex, type IrNodeKind } from "../src/ir/node-index.js";
 import { everyModuleNode, everyStmtList } from "../src/ir/traverse.js";
 
@@ -125,6 +126,31 @@ test("a false answer stops the visit, with or without the index", () => {
   }
 });
 
+test("a function's node count is the number of nodes its visit notes", () => {
+  const all = index.scanExcept(NO_KINDS);
+  for (const fn of mod.functions) expect(index.nodeCount(fn)).toBe(noted(fn, all).length);
+  expect(index.nodeCount({ ...mod.functions[0]! })).toBeUndefined();
+});
+
+test("emission ranges cover the functions in order, none empty, of about equal cost", () => {
+  const count = (fn: IrFunction): number | undefined => index.nodeCount(fn);
+  for (const ranges of [1, 2, 3, mod.functions.length, mod.functions.length + 4]) {
+    const starts = emitRangeStarts(mod.functions, ranges, count);
+    expect(starts[0]).toBe(0);
+    expect(starts.at(-1)).toBe(mod.functions.length);
+    expect(starts.length - 1).toBeLessThanOrEqual(ranges);
+    for (let i = 1; i < starts.length; i++) expect(starts[i]!).toBeGreaterThan(starts[i - 1]!);
+  }
+  // Equal costs split into equal ranges. A range ends after the function
+  // that reaches its share of the cost, so a large function ends its range.
+  const functions = Array.from({ length: 12 }, () => mod.functions[0]!);
+  expect(emitRangeStarts(functions, 4, () => 0)).toEqual([0, 3, 6, 9, 12]);
+  const sizes = [0, 0, 1000, 0, 0, 0];
+  const uneven = sizes.map((_, i) => ({ ...mod.functions[0]!, name: `f${i}` }));
+  const bySize = new Map(uneven.map((fn, i) => [fn, sizes[i]!]));
+  expect(emitRangeStarts(uneven, 3, (fn) => bySize.get(fn))).toEqual([0, 3, 4, 6]);
+});
+
 test("a function the index did not record is walked", () => {
   const all = index.scanExcept(NO_KINDS);
   for (const fn of mod.functions) {