mirror of
https://github.com/vercel-labs/scriptc.git
synced 2026-11-01 07:33:12 +08:00
Run short runs of small innermost worker loops without a poll (#2463)
This commit is contained in:
Chris Tate
@@ -74,7 +74,7 @@ import { bytesHeaderAlias } from "./bytes-alias.js";
|
||||
import { mangleLocal } from "../mangle.js";
|
||||
import { LlvmUnsupportedError } from "./unsupported.js";
|
||||
import type { LlValue } from "./expr-context.js";
|
||||
import { everyExprChild, everyStmtChild } from "../../ir/traverse.js";
|
||||
import { everyExprChild, everyStmtChild, everyStmtList } from "../../ir/traverse.js";
|
||||
import { widenInteger } from "./integer-values.js";
|
||||
import { SCALAR_UNION_RESULT } from "./int32-abi.js";
|
||||
import { emitLiteralSwitch } from "./switch-dispatch.js";
|
||||
@@ -1310,6 +1310,37 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
|
||||
emitter.workerLoopPolls(s) &&
|
||||
(loopAccessesTypedArrays(s) || (countedLoop !== null && innermost));
|
||||
const workerBudget = stripMined || bounded ? null : emitter.workerLoopBudget(s);
|
||||
// An innermost counted loop, or a two-counter loop closing in by one
|
||||
// step on each side, has a run whose length its counters fix on entry.
|
||||
// With a small body, a run of at most WORKER_UNPOLLED_RUN steps takes a
|
||||
// plain copy of the loop without a poll, which polls once as it leaves
|
||||
// when it ran more than WORKER_POLL_BLOCK steps, like a strip-mined run;
|
||||
// longer runs take the loop as it is emitted otherwise.
|
||||
const shortRuns: { stride: number; left: () => string } | null =
|
||||
!innermost ||
|
||||
bounded ||
|
||||
capturedLets.length > 0 ||
|
||||
!emitter.workerLoopPolls(s) ||
|
||||
!withinNodes(s.body, MAX_SHORT_RUN_NODES)
|
||||
? null
|
||||
: countedLoop !== null && countedLimit !== null && integerSlot !== null
|
||||
? {
|
||||
stride: Math.abs(countedLoop.step),
|
||||
left: () => {
|
||||
const start = B.tmp(),
|
||||
left = B.tmp();
|
||||
B.line(`${start} = load i64, ptr ${integerSlot}`);
|
||||
B.line(
|
||||
countedLoop.step > 0
|
||||
? `${left} = sub nsw i64 ${countedLimit}, ${start}`
|
||||
: `${left} = sub nsw i64 ${start}, ${countedLimit}`,
|
||||
);
|
||||
return left;
|
||||
},
|
||||
}
|
||||
: pairedLoop !== null && pairedLoop.low.step === 1 && pairedLoop.high.step === -1
|
||||
? { stride: 2, left: () => emitPairedLoopGap(emitter, pairedSlots) }
|
||||
: null;
|
||||
// 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;
|
||||
@@ -1325,48 +1356,12 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
|
||||
}
|
||||
}
|
||||
freshenBindings();
|
||||
emitter.emitLoopCopies(s, capturedLets.length === 0, le, (exit) => {
|
||||
if (stripMined) {
|
||||
const counter: StripMinedCounter = countedLoop
|
||||
? {
|
||||
localId: countedLoop.localId,
|
||||
slot: integerSlot!,
|
||||
type: "i64",
|
||||
signed: true,
|
||||
step: countedLoop.step,
|
||||
followers: countedLoop.followers.map((f, i) => ({
|
||||
slot: followerSlots[i]!,
|
||||
step: f.step,
|
||||
})),
|
||||
limit: () => countedLimit!,
|
||||
}
|
||||
: {
|
||||
localId: integerLoop!.localId,
|
||||
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, runSlot);
|
||||
return;
|
||||
}
|
||||
// The loop as a plain header, body and update, polling by `budget`.
|
||||
const emitPlainLoop = (exit: string, budget: string | null): void => {
|
||||
const lc = B.newLabel("loop.c");
|
||||
const lb = B.newLabel("loop.b");
|
||||
const lu =
|
||||
s.update || capturedLets.length > 0 || workerBudget !== null ? B.newLabel("loop.u") : lc;
|
||||
s.update || capturedLets.length > 0 || budget !== null ? B.newLabel("loop.u") : lc;
|
||||
// A periodic guard on the counter (periodic-guards.ts): the back
|
||||
// edge asks LLVM to unroll the loop by the guard's period.
|
||||
const period =
|
||||
@@ -1437,7 +1432,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
|
||||
B.br(lu, lu === lc ? backEdge : "");
|
||||
if (lu !== lc) {
|
||||
B.startBlock(lu);
|
||||
emitter.emitWorkerLoopCheck(workerBudget);
|
||||
emitter.emitWorkerLoopCheck(budget);
|
||||
freshenBindings();
|
||||
if (integerLoop && integerSlot) {
|
||||
const old = B.tmp();
|
||||
@@ -1469,6 +1464,78 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
|
||||
} else if (s.update) emitter.emitStmt(s.update);
|
||||
B.br(lc, backEdge);
|
||||
}
|
||||
};
|
||||
emitter.emitLoopCopies(s, capturedLets.length === 0, le, (exit) => {
|
||||
// Speculative copies resume iterations inside the copy they enter,
|
||||
// which must keep its polls.
|
||||
if (
|
||||
shortRuns &&
|
||||
emitter.speculation === null &&
|
||||
emitter.speculationResume === null &&
|
||||
emitter.speculationHeader === null &&
|
||||
emitter.int32Kinds === null
|
||||
) {
|
||||
// The run's length in counter steps (two counter steps an
|
||||
// iteration for a two-counter loop), signed: a run that does not
|
||||
// enter the body is short too.
|
||||
const left = shortRuns.left();
|
||||
const inRun = B.tmp(),
|
||||
pollsAfter = B.tmp();
|
||||
const shortLabel = B.newLabel("loop.short"),
|
||||
longLabel = B.newLabel("loop.long"),
|
||||
shortEnd = B.newLabel("loop.short.end"),
|
||||
shortPoll = B.newLabel("loop.short.poll");
|
||||
B.line(`${inRun} = icmp sle i64 ${left}, ${shortRuns.stride * WORKER_UNPOLLED_RUN}`);
|
||||
B.condBr(inRun, shortLabel, longLabel);
|
||||
B.startBlock(shortLabel);
|
||||
emitPlainLoop(shortEnd, null);
|
||||
// A run longer than one block polls once as it leaves.
|
||||
B.startBlock(shortEnd);
|
||||
B.line(`${pollsAfter} = icmp sgt i64 ${left}, ${shortRuns.stride * WORKER_POLL_BLOCK}`);
|
||||
B.condBr(pollsAfter, shortPoll, exit);
|
||||
B.startBlock(shortPoll);
|
||||
emitter.emitWorkerPoll();
|
||||
B.br(exit);
|
||||
B.startBlock(longLabel);
|
||||
}
|
||||
if (stripMined) {
|
||||
const counter: StripMinedCounter = countedLoop
|
||||
? {
|
||||
localId: countedLoop.localId,
|
||||
slot: integerSlot!,
|
||||
type: "i64",
|
||||
signed: true,
|
||||
step: countedLoop.step,
|
||||
followers: countedLoop.followers.map((f, i) => ({
|
||||
slot: followerSlots[i]!,
|
||||
step: f.step,
|
||||
})),
|
||||
limit: () => countedLimit!,
|
||||
}
|
||||
: {
|
||||
localId: integerLoop!.localId,
|
||||
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, runSlot);
|
||||
return;
|
||||
}
|
||||
emitPlainLoop(exit, workerBudget);
|
||||
});
|
||||
B.startBlock(le);
|
||||
emitter.releaseScope(emitter.scopes.pop()!);
|
||||
@@ -2801,6 +2868,34 @@ function boundedLoopRuns(emitter: LlEmitter, loop: IrStmt): number {
|
||||
return trips * (1 + nested);
|
||||
}
|
||||
|
||||
/** The most IR nodes the body of a loop with an unpolled copy for short
|
||||
* runs may have: the copy duplicates the loop. */
|
||||
const MAX_SHORT_RUN_NODES = 256;
|
||||
|
||||
/** Whether `body` has at most `limit` statement and expression nodes. */
|
||||
function withinNodes(body: IrStmt[], limit: number): boolean {
|
||||
let nodes = 0;
|
||||
return everyStmtList(body, {
|
||||
expr: () => ++nodes <= limit,
|
||||
stmt: () => ++nodes <= limit,
|
||||
});
|
||||
}
|
||||
|
||||
/** The gap between the counters of a two-counter loop stepping them by 1
|
||||
* and -1 (`for (let i = a, j = b; i < j; i++, j--)`), from the values in
|
||||
* `slots`: each iteration closes it by two, so a run takes about half the
|
||||
* gap in iterations, and none when the gap is closed. */
|
||||
function emitPairedLoopGap(emitter: LlEmitter, slots: readonly string[]): string {
|
||||
const B = emitter.B;
|
||||
const low = B.tmp(),
|
||||
high = B.tmp(),
|
||||
gap = B.tmp();
|
||||
B.line(`${low} = load i64, ptr ${slots[0]}`);
|
||||
B.line(`${high} = load i64, ptr ${slots[1]}`);
|
||||
B.line(`${gap} = sub nsw i64 ${high}, ${low}`);
|
||||
return gap;
|
||||
}
|
||||
|
||||
/** The counter of a strip-mined loop: its binding, 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. */
|
||||
|
||||
@@ -471,3 +471,78 @@ test("an exact integer loop limit becomes the endpoint without rounding", async
|
||||
await rm(dir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
test("worker programs run short runs of small innermost loops inside another loop without a poll", async () => {
|
||||
const dir = await mkdtemp(join(tmpdir(), "scriptc-short-run-polls-"));
|
||||
try {
|
||||
const entry = join(dir, "short.ts");
|
||||
await writeFile(
|
||||
entry,
|
||||
[
|
||||
'import { isMainThread, Worker } from "node:worker_threads";',
|
||||
"export function flips(p: number[], n: number): number {",
|
||||
" let total = 0;",
|
||||
" for (let r = 0; r < n; r++) {",
|
||||
" const k = p[r % p.length]!;",
|
||||
" for (let i = 1, j = k - 1; i < j; i++, j--) {",
|
||||
" const t = p[i]!;",
|
||||
" p[i] = p[j]!;",
|
||||
" p[j] = t;",
|
||||
" }",
|
||||
" total += p[0]!;",
|
||||
" }",
|
||||
" return total;",
|
||||
"}",
|
||||
"export function rows(u: Float64Array, v: Float64Array, n: number): void {",
|
||||
" for (let i = 0; i < n; i++) {",
|
||||
" let t = 0;",
|
||||
" for (let j = 0; j < n; j++) t += u[j]! / (i + j + 1);",
|
||||
" v[i] = t;",
|
||||
" }",
|
||||
"}",
|
||||
'const n = Number(process.argv[2] ?? "6");',
|
||||
"const u = new Float64Array(n).fill(1);",
|
||||
"const v = new Float64Array(n);",
|
||||
"rows(u, v, n);",
|
||||
"console.log(flips([3, 1, 2, 0, 5, 4], n), v[0]);",
|
||||
"if (!isMainThread) new Worker(new URL(import.meta.url));",
|
||||
].join("\n"),
|
||||
);
|
||||
const outPath = join(dir, "short.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"));
|
||||
const ll = emitLlvmModule(deserializeModule(await readFile(outPath, "utf8")), {
|
||||
targetTriple: "x86_64-unknown-linux-gnu",
|
||||
});
|
||||
const body = (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");
|
||||
for (const [name, limit] of [
|
||||
// Two counters close the gap between them by two steps an iteration.
|
||||
["flips", 2 * 16384],
|
||||
["rows", 16384],
|
||||
] as const) {
|
||||
const text = body(name);
|
||||
const blocks = text.split(/\n(?=[\w.]+:)/);
|
||||
// A test on entry compares the run with its limit and takes a copy
|
||||
// without a poll for a short run, which polls once as it leaves after
|
||||
// a run longer than one block; longer runs keep a copy that polls.
|
||||
expect(text).toContain(`, ${limit}\n`);
|
||||
expect(blocks.some((block) => /^loop\.short\d+:/.test(block))).toBe(true);
|
||||
expect(blocks.some((block) => /^loop\.long\d+:/.test(block))).toBe(true);
|
||||
const polls = blocks.filter((block) => /^loop\.short\.poll\d+:/.test(block));
|
||||
expect(polls.length).toBeGreaterThan(0);
|
||||
for (const poll of polls) expect(poll).toContain("@scr_exc_flag");
|
||||
}
|
||||
} finally {
|
||||
await rm(dir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user