From 74a131bbcc3528a72fe1a2024cf7c9d876a777f9 Mon Sep 17 00:00:00 2001 From: Chris Tate Date: Wed, 28 Oct 2026 06:14:20 -0500 Subject: [PATCH] Run short runs of small innermost worker loops without a poll --- .../src/backend/llvm/emit-statements.ts | 177 ++++++++++++++---- .../test/counted-loop-emission.test.ts | 75 ++++++++ 2 files changed, 211 insertions(+), 41 deletions(-) diff --git a/packages/compiler/src/backend/llvm/emit-statements.ts b/packages/compiler/src/backend/llvm/emit-statements.ts index b910108b3..5eed52dce 100644 --- a/packages/compiler/src/backend/llvm/emit-statements.ts +++ b/packages/compiler/src/backend/llvm/emit-statements.ts @@ -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. */ diff --git a/packages/compiler/test/counted-loop-emission.test.ts b/packages/compiler/test/counted-loop-emission.test.ts index c42bda8b7..41f96e006 100644 --- a/packages/compiler/test/counted-loop-emission.test.ts +++ b/packages/compiler/test/counted-loop-emission.test.ts @@ -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 }); + } +});