mirror of
https://github.com/vercel-labs/scriptc.git
synced 2026-11-01 07:33:12 +08:00
Poll worker cancellation between blocks of typed-array loop iterations
In a worker program, every loop without a pending test on each iteration decremented a cancellation budget in its body and polled the exception flag every 64 iterations. The poll's early return made each loop an early-exit loop, which LLVM's loop vectorizer does not vectorize. A counted for loop (integer counter, invariant limit), or a loop to a typed array's length, whose body reads or writes a typed-array element outside nested loops and closures, now runs as blocks of at most 64 iterations: the outer header tests the counter against the limit, the inner loop runs to the nearer of the limit and 64 steps on, and the poll follows each block. A typed array's length is fixed, so each block takes it as its limit. The loop polls as often as the budget did, and the inner loop's only exit besides break and return is its bound. The block bound lives in one slot shared by the loop's speculative copies, which resume iterations in each other's bodies. Other loops keep the budget, which costs less in loops of a few iterations. An elementwise Float64Array loop in a worker-enabled program vectorizes again: 949 ms to 691 ms. A byte-classifying loop over a 100M-byte Uint8Array takes 13 ms instead of 141 ms. New corpus program: native-worker-terminate-counted-loops.ts.
This commit is contained in:
Chris Tate
@@ -91,6 +91,7 @@ import {
|
||||
} from "./string-construction.js";
|
||||
import { emitStringInputs } from "./string-lifetimes.js";
|
||||
import { LOOPS } from "./loop-state.js";
|
||||
import { loopAccessesTypedArrays } from "../loop-polls.js";
|
||||
|
||||
/** A callback loop version copies the loop; keep the copies small. */
|
||||
const MAX_CALLBACK_VERSION_COST = 400;
|
||||
@@ -1211,9 +1212,67 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
|
||||
B.line(`store ptr ${fresh}, ptr ${slot}`);
|
||||
}
|
||||
};
|
||||
const workerBudget = emitter.workerLoopBudget(s);
|
||||
// A counted loop over typed arrays, or a loop to a typed array's
|
||||
// length, polls between blocks of WORKER_POLL_BLOCK iterations instead
|
||||
// of counting a budget in its body: the inner loop over a block then
|
||||
// has no exit besides its own bound, which LLVM can vectorize. A typed
|
||||
// array's length is fixed, so each block can take it as its limit.
|
||||
const stripMined =
|
||||
capturedLets.length === 0 &&
|
||||
integerSlot !== null &&
|
||||
((countedLoop !== null &&
|
||||
countedLimit !== null &&
|
||||
Math.abs(countedLoop.step) * WORKER_POLL_BLOCK <= 2 ** 52) ||
|
||||
(integerLoop !== null && integerArrayLoop === null)) &&
|
||||
emitter.workerLoopPolls(s) &&
|
||||
loopAccessesTypedArrays(s);
|
||||
const workerBudget = stripMined ? null : emitter.workerLoopBudget(s);
|
||||
// One slot for every copy of the loop: a speculative copy resumes an
|
||||
// iteration in another copy's body, which then reads this bound.
|
||||
let blockBound: string | null = null;
|
||||
if (stripMined) {
|
||||
blockBound = B.slot();
|
||||
B.entryAllocas.push(
|
||||
`${blockBound} = alloca ${countedLoop ? "i64" : emitter.sizeType} ; cancellation poll block bound`,
|
||||
);
|
||||
}
|
||||
freshenBindings();
|
||||
emitter.emitLoopCopies(s, capturedLets.length === 0, le, (exit) => {
|
||||
if (stripMined) {
|
||||
const counter: StripMinedCounter = countedLoop
|
||||
? {
|
||||
slot: integerSlot!,
|
||||
type: "i64",
|
||||
signed: true,
|
||||
step: countedLoop.step,
|
||||
followers: countedLoop.followers.map((f, i) => ({
|
||||
slot: followerSlots[i]!,
|
||||
step: f.step,
|
||||
})),
|
||||
limit: () => countedLimit!,
|
||||
}
|
||||
: {
|
||||
slot: integerSlot!,
|
||||
type: emitter.sizeType,
|
||||
signed: false,
|
||||
step: 1,
|
||||
followers: [],
|
||||
limit: () => {
|
||||
const receiver = emitter.emitStableReceiver(integerLoop!.limitReceiver, []);
|
||||
const lenPtr = B.tmp();
|
||||
const len = B.tmp();
|
||||
B.line(
|
||||
`${lenPtr} = getelementptr inbounds %ScrBytes, ptr ${receiver.name}, i64 0, i32 1`,
|
||||
);
|
||||
B.line(
|
||||
`${len} = load ${emitter.sizeType}, ptr ${lenPtr}${bytesHeaderAlias(emitter)}`,
|
||||
);
|
||||
return len;
|
||||
},
|
||||
};
|
||||
emitStripMinedCountedLoop(emitter, s, counter, blockBound!, exit);
|
||||
return;
|
||||
}
|
||||
const lc = B.newLabel("loop.c");
|
||||
const lb = B.newLabel("loop.b");
|
||||
const lu =
|
||||
@@ -2549,6 +2608,102 @@ function emitInt32KindLoop(
|
||||
|
||||
/** At the top of a loop copy's header: lets a speculative loop find its
|
||||
* ordinary copy's header, and keeps its third copy out of the loop. */
|
||||
/** Iterations between cancellation polls of a strip-mined counted loop. */
|
||||
const WORKER_POLL_BLOCK = 64;
|
||||
|
||||
/** The counter of a strip-mined loop: its slot, the LLVM integer type and
|
||||
* signedness of its values, its step, the counters stepped beside it, and
|
||||
* the limit, emitted in the outer header before each block. */
|
||||
interface StripMinedCounter {
|
||||
slot: string;
|
||||
type: string;
|
||||
signed: boolean;
|
||||
step: number;
|
||||
followers: { slot: string; step: number }[];
|
||||
limit: () => string;
|
||||
}
|
||||
|
||||
/** A counted `for` loop of a worker program as blocks of WORKER_POLL_BLOCK
|
||||
* iterations, with the cancellation poll between blocks: the outer header
|
||||
* tests the counter against the loop's limit, each block runs the body up
|
||||
* to the nearer of the limit and WORKER_POLL_BLOCK steps on, and the poll
|
||||
* follows the block. It polls as often as the budget it replaces, and the
|
||||
* inner loop's only exit besides `break` and `return` is its bound. */
|
||||
function emitStripMinedCountedLoop(
|
||||
emitter: LlEmitter,
|
||||
s: Extract<IrStmt, { kind: "for" }>,
|
||||
counter: StripMinedCounter,
|
||||
boundSlot: string,
|
||||
exit: string,
|
||||
): void {
|
||||
const B = emitter.B;
|
||||
const { slot, type, signed, step } = counter;
|
||||
const ascending = step > 0;
|
||||
const compare = !signed ? "ult" : ascending ? "slt" : "sgt";
|
||||
const saturating = `llvm.${signed ? "sadd" : "uadd"}.sat.${type}`;
|
||||
const nearer = `llvm.${!signed ? "umin" : ascending ? "smin" : "smax"}.${type}`;
|
||||
const wrap = signed ? "nsw" : "nuw";
|
||||
const lc = B.newLabel("loop.c");
|
||||
const block = B.newLabel("loop.block");
|
||||
const bc = B.newLabel("loop.bc");
|
||||
const lb = B.newLabel("loop.b");
|
||||
const lu = B.newLabel("loop.u");
|
||||
const poll = B.newLabel("loop.poll");
|
||||
emitter.declare(`declare ${type} @${saturating}(${type}, ${type})`);
|
||||
emitter.declare(`declare ${type} @${nearer}(${type}, ${type})`);
|
||||
B.br(lc);
|
||||
B.startBlock(lc);
|
||||
enterLoopHeader(emitter, lc);
|
||||
const start = B.tmp();
|
||||
const open = B.tmp();
|
||||
B.line(`${start} = load ${type}, ptr ${slot}`);
|
||||
const limit = counter.limit();
|
||||
B.line(`${open} = icmp ${compare} ${type} ${start}, ${limit}`);
|
||||
B.condBr(open, block, exit);
|
||||
B.startBlock(block);
|
||||
const stepped = B.tmp();
|
||||
const bound = B.tmp();
|
||||
B.line(
|
||||
`${stepped} = call ${type} @${saturating}(${type} ${start}, ${type} ${step * WORKER_POLL_BLOCK})`,
|
||||
);
|
||||
B.line(`${bound} = call ${type} @${nearer}(${type} ${stepped}, ${type} ${limit})`);
|
||||
B.line(`store ${type} ${bound}, ptr ${boundSlot}`);
|
||||
B.br(bc);
|
||||
B.startBlock(bc);
|
||||
const index = B.tmp();
|
||||
const end = B.tmp();
|
||||
const inBlock = B.tmp();
|
||||
B.line(`${index} = load ${type}, ptr ${slot}`);
|
||||
B.line(`${end} = load ${type}, ptr ${boundSlot}`);
|
||||
B.line(`${inBlock} = icmp ${compare} ${type} ${index}, ${end}`);
|
||||
B.condBr(inBlock, lb, poll);
|
||||
B.startBlock(lb);
|
||||
emitter.jumpTargets.push({
|
||||
kind: "loop",
|
||||
brkLabel: exit,
|
||||
contLabel: lu,
|
||||
labels: s.labels,
|
||||
frameDepth: emitter.frames.length,
|
||||
scopeDepth: emitter.scopes.length,
|
||||
finallyDepth: emitter.finallyStack.length,
|
||||
});
|
||||
emitter.emitBlock(s.body);
|
||||
emitter.jumpTargets.pop();
|
||||
B.br(lu);
|
||||
B.startBlock(lu);
|
||||
for (const { slot: stepSlot, step: by } of [{ slot, step }, ...counter.followers]) {
|
||||
const old = B.tmp();
|
||||
const next = B.tmp();
|
||||
B.line(`${old} = load ${type}, ptr ${stepSlot}`);
|
||||
B.line(`${next} = add ${wrap} ${type} ${old}, ${by}`);
|
||||
B.line(`store ${type} ${next}, ptr ${stepSlot}`);
|
||||
}
|
||||
B.br(bc);
|
||||
B.startBlock(poll);
|
||||
emitter.emitWorkerPoll();
|
||||
B.br(lc);
|
||||
}
|
||||
|
||||
function enterLoopHeader(emitter: LlEmitter, header: string): void {
|
||||
const dispatch = emitter.speculationHeader;
|
||||
emitter.speculationHeader = null;
|
||||
|
||||
@@ -3932,13 +3932,19 @@ export class LlEmitter {
|
||||
}
|
||||
}
|
||||
|
||||
/** Whether `loop` needs a cancellation poll: in a worker program, unless
|
||||
* the loop already runs a pending test on every iteration, which observes
|
||||
* a terminate request itself (a separate poll would only add work). */
|
||||
workerLoopPolls(loop: Extract<IrStmt, { kind: "while" | "doWhile" | "for" | "forOf" }>): boolean {
|
||||
return (
|
||||
this.mod.workers === true && !loopTestsEveryIteration(loop, (e) => this.callTestsPending(e))
|
||||
);
|
||||
}
|
||||
|
||||
workerLoopBudget(
|
||||
loop: Extract<IrStmt, { kind: "while" | "doWhile" | "for" | "forOf" }>,
|
||||
): string | null {
|
||||
if (!this.mod.workers) return null;
|
||||
// A loop that runs a pending test on every iteration observes a
|
||||
// terminate request through it; a separate poll would only add work.
|
||||
if (loopTestsEveryIteration(loop, (e) => this.callTestsPending(e))) return null;
|
||||
if (!this.workerLoopPolls(loop)) return null;
|
||||
const slot = this.B.slot();
|
||||
this.B.entryAllocas.push(`${slot} = alloca i32 ; worker cancellation budget`);
|
||||
this.B.line(`store i32 0, ptr ${slot}`);
|
||||
@@ -4056,14 +4062,19 @@ export class LlEmitter {
|
||||
B.condBr(due, poll, next);
|
||||
B.startBlock(poll);
|
||||
B.line(`store i32 63, ptr ${slot}`);
|
||||
// Calls and throwing operations retain their own checks. Amortize the
|
||||
// extra cancellation poll across bounded loop batches so scalar work
|
||||
// does not cross the runtime ABI on every iteration.
|
||||
this.emitPendingCheck(undefined, true);
|
||||
this.emitWorkerPoll();
|
||||
B.br(next);
|
||||
B.startBlock(next);
|
||||
}
|
||||
|
||||
/** The cancellation poll itself. Calls and throwing operations retain
|
||||
* their own checks. Loops amortize this extra poll across bounded
|
||||
* batches of iterations so scalar work does not cross the runtime ABI on
|
||||
* every iteration. */
|
||||
emitWorkerPoll(): void {
|
||||
this.emitPendingCheck(undefined, true);
|
||||
}
|
||||
|
||||
/** Moves an already-evaluated value into the runtime's exception cell —
|
||||
* the `throw` statement's kind dispatch (stmts.ts's), shared with
|
||||
* every synthetic thrower. Ownership of a refcounted payload must have
|
||||
|
||||
@@ -106,3 +106,34 @@ export function loopTestsEveryIteration(loop: Loop, tested: (call: IrExpr) => bo
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
/** True when the body of `loop`, outside its nested loops and closures,
|
||||
* reads or writes a typed-array element: the loops whose cancellation poll
|
||||
* the emitter moves between blocks of iterations, so that LLVM can
|
||||
* vectorize them. Other loops keep the budget in their body, which costs
|
||||
* less in a loop of a few iterations. */
|
||||
export function loopAccessesTypedArrays(loop: Loop): boolean {
|
||||
const inExpr = (e: IrExpr): boolean => {
|
||||
if (e.kind === "bytesIntrinsic" && e.method === "get") return true;
|
||||
let found = false;
|
||||
everyExprChild(
|
||||
e,
|
||||
(child) => !(found = inExpr(child)),
|
||||
(child) => !(found = inStmt(child)),
|
||||
);
|
||||
return found;
|
||||
};
|
||||
const inStmt = (s: IrStmt): boolean => {
|
||||
if (s.kind === "bytesSet") return true;
|
||||
if (s.kind === "while" || s.kind === "doWhile" || s.kind === "for" || s.kind === "forOf")
|
||||
return false;
|
||||
let found = false;
|
||||
everyStmtChild(
|
||||
s,
|
||||
(child) => !(found = inExpr(child)),
|
||||
(child) => !(found = inStmt(child)),
|
||||
);
|
||||
return found;
|
||||
};
|
||||
return loop.body.some(inStmt);
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import { mkdtemp, readFile, rm } from "node:fs/promises";
|
||||
import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { expect, test } from "vitest";
|
||||
@@ -45,3 +45,81 @@ test("counted loops guard exact induction and preserve wide remainders on both p
|
||||
await rm(dir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
test("worker programs poll counted loops over typed arrays between blocks of iterations", async () => {
|
||||
const dir = await mkdtemp(join(tmpdir(), "scriptc-counted-poll-"));
|
||||
try {
|
||||
const source = (workers: boolean): string =>
|
||||
[
|
||||
workers ? 'import { isMainThread, Worker } from "node:worker_threads";' : "",
|
||||
"export function scale(a: Float64Array, b: Float64Array, n: number): void {",
|
||||
" for (let i = 0; i < n; i++) b[i] = a[i]! * 2;",
|
||||
"}",
|
||||
"export function codes(a: Uint8Array): void {",
|
||||
" for (let i = 0; i < a.length; ++i) a[i] = a[i] === 97 ? 0 : 3;",
|
||||
"}",
|
||||
"export function mix(n: number): number {",
|
||||
" let t = 0;",
|
||||
" for (let i = 0; i < n; i++) t = (t * 31 + i) % 1000003;",
|
||||
" return t;",
|
||||
"}",
|
||||
"const a = new Float64Array(8).fill(1);",
|
||||
"const b = new Float64Array(8);",
|
||||
"scale(a, b, 8);",
|
||||
"const c = new Uint8Array(8).fill(97);",
|
||||
"codes(c);",
|
||||
"console.log(b[7], mix(8), c[0]);",
|
||||
workers ? "if (!isMainThread) new Worker(new URL(import.meta.url));" : "",
|
||||
].join("\n");
|
||||
const emit = async (workers: boolean): Promise<string> => {
|
||||
const entry = join(dir, `${workers ? "worker" : "plain"}.ts`);
|
||||
await writeFile(entry, source(workers));
|
||||
const outPath = join(dir, `${workers ? "worker" : "plain"}.ir.json`);
|
||||
const result = await compile(entry, { outPath, outDir: dir, outputKind: "ir" });
|
||||
if (!result.ok) throw new Error(result.diagnostics.map((d) => d.message).join("\n"));
|
||||
return emitLlvmModule(deserializeModule(await readFile(outPath, "utf8")));
|
||||
};
|
||||
// Every emitted version of a function (the number and int32 entries).
|
||||
const bodies = (ll: string, name: string): string =>
|
||||
[
|
||||
...ll.matchAll(
|
||||
new RegExp(
|
||||
`define internal [^\\n]+ @sc_b?f_${name}(?:\\.i32)?\\([^\\n]*\\n([\\s\\S]*?)\\n\\}`,
|
||||
"g",
|
||||
),
|
||||
),
|
||||
]
|
||||
.map((match) => match[1]!)
|
||||
.join("\n");
|
||||
const workerModule = await emit(true);
|
||||
const worker = bodies(workerModule, "scale");
|
||||
const blocks = worker.split(/\n(?=[\w.]+:)/);
|
||||
// No budget in the body: the inner loop's latch only steps the counter,
|
||||
// the poll runs after each block, and the block's bound is the nearer of
|
||||
// the limit and 64 steps on.
|
||||
expect(worker).not.toContain("worker cancellation budget");
|
||||
expect(worker).toContain("@llvm.sadd.sat.i64(i64 %");
|
||||
expect(worker).toMatch(/@llvm\.smin\.i64\(i64 %\w+, i64 %\w+\)/);
|
||||
for (const latch of blocks.filter((block) => block.startsWith("loop.u")))
|
||||
expect(latch).not.toContain("scr_exc_flag");
|
||||
const polls = blocks.filter((block) => block.startsWith("loop.poll"));
|
||||
expect(polls.length).toBeGreaterThan(0);
|
||||
for (const poll of polls) expect(poll).toContain("@scr_exc_flag");
|
||||
// A counted loop without typed-array accesses keeps the budget.
|
||||
// A loop to a typed array's length is strip-mined the same way, with the
|
||||
// length as each block's limit.
|
||||
const codes = bodies(workerModule, "codes");
|
||||
expect(codes).not.toContain("worker cancellation budget");
|
||||
expect(codes).toMatch(/@llvm\.uadd\.sat\.i64\(i64 %\w+, i64 64\)/);
|
||||
expect(codes).toMatch(/@llvm\.umin\.i64\(i64 %\w+, i64 %\w+\)/);
|
||||
expect(codes).toMatch(/^loop\.poll\d+:$/m);
|
||||
const mix = bodies(workerModule, "mix");
|
||||
expect(mix).toContain("worker cancellation budget");
|
||||
expect(mix).not.toContain("loop.poll");
|
||||
const plain = bodies(await emit(false), "scale");
|
||||
expect(plain).not.toContain("loop.poll");
|
||||
expect(plain).not.toContain("@scr_exc_flag");
|
||||
} finally {
|
||||
await rm(dir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"order": [
|
||||
"<repo>/tests/corpus/native-worker-terminate-counted-loops.ts"
|
||||
],
|
||||
"diags": []
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
// worker.terminate() stops a worker in counted loops over a typed array
|
||||
// with long bounds: ascending and descending counters, a step above one, a
|
||||
// second counter that steps with the first, and loops to a typed array's
|
||||
// length run again and again.
|
||||
import { Worker, isMainThread, parentPort, workerData } from "node:worker_threads";
|
||||
|
||||
function ascending(values: Float64Array, limit: number): number {
|
||||
let acc = 0;
|
||||
for (let i = 0; i < limit; i++) {
|
||||
acc = (acc + values[i & 1023]!) % 1000003;
|
||||
if (i === 5000) parentPort!.postMessage("ready");
|
||||
}
|
||||
return acc;
|
||||
}
|
||||
|
||||
function descending(values: Float64Array, limit: number): number {
|
||||
let acc = 0;
|
||||
for (let i = limit; i > 0; i -= 3) {
|
||||
acc = (acc * 31 + i + values[i & 1023]!) % 1000003;
|
||||
if (i === limit - 15000) parentPort!.postMessage("ready");
|
||||
}
|
||||
return acc;
|
||||
}
|
||||
|
||||
function paired(values: Float64Array, limit: number): number {
|
||||
let acc = 0;
|
||||
for (let i = 0, j = 7; i < limit; i++, j += 2) {
|
||||
acc = (acc + i * j + values[j & 1023]!) % 1000003;
|
||||
if (i === 5000) parentPort!.postMessage("ready");
|
||||
}
|
||||
return acc;
|
||||
}
|
||||
|
||||
function lengths(bytes: Uint8Array): number {
|
||||
let acc = 0;
|
||||
for (let round = 0; ; round++) {
|
||||
for (let i = 0; i < bytes.length; ++i) bytes[i] = bytes[i] === 97 ? 0 : (bytes[i]! + round) & 127;
|
||||
for (let i = 0; i < bytes.length; i++) acc = (acc + bytes[i]!) % 1000003;
|
||||
if (round === 50) parentPort!.postMessage("ready");
|
||||
}
|
||||
return acc;
|
||||
}
|
||||
|
||||
if (isMainThread) {
|
||||
for (const kind of ["ascending", "descending", "paired", "lengths"]) {
|
||||
await new Promise<void>((resolve) => {
|
||||
const worker = new Worker(new URL(import.meta.url), { workerData: kind });
|
||||
worker.once("message", () => {
|
||||
void worker.terminate().then((code: number) => console.log(kind, "terminate", code));
|
||||
});
|
||||
worker.once("exit", (code: number) => {
|
||||
console.log(kind, "exit", code);
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
}
|
||||
} else {
|
||||
const values = new Float64Array(1024).map((_, i) => i);
|
||||
const limit = 1e15;
|
||||
let result = 0;
|
||||
if (workerData === "ascending") result = ascending(values, limit);
|
||||
else if (workerData === "descending") result = descending(values, limit);
|
||||
else if (workerData === "paired") result = paired(values, limit);
|
||||
else result = lengths(new Uint8Array(1 << 16).fill(97));
|
||||
console.log("unexpected result", result);
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
{ "plain": 3594, "san": 23593 }
|
||||
Reference in New Issue
Block a user