diff --git a/packages/compiler/src/backend/llvm/emit-statements.ts b/packages/compiler/src/backend/llvm/emit-statements.ts index 7239db7d8..a22de3445 100644 --- a/packages/compiler/src/backend/llvm/emit-statements.ts +++ b/packages/compiler/src/backend/llvm/emit-statements.ts @@ -803,6 +803,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void { const lp = workerBudget !== null ? B.newLabel("loop.p") : lc; B.br(lc); B.startBlock(lc); + enterLoopHeader(emitter, lc); B.condBr(emitter.emitCondition(s.cond), lb, exit); B.startBlock(lb); emitter.jumpTargets.push({ @@ -835,6 +836,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void { const lc = B.newLabel("loop.c"); B.br(lb); B.startBlock(lb); + enterLoopHeader(emitter, lb); emitter.jumpTargets.push({ kind: "loop", brkLabel: exit, @@ -1135,6 +1137,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void { s.update || capturedLets.length > 0 || workerBudget !== null ? B.newLabel("loop.u") : lc; B.br(lc); B.startBlock(lc); + enterLoopHeader(emitter, lc); if (integerLoop && integerSlot) { const receiver = emitter.emitStableReceiver(integerLoop.limitReceiver, []); const lenPtr = B.tmp(); @@ -1885,6 +1888,10 @@ export function emitSwitch(emitter: LlEmitter, s: IrStmt & { kind: "switch" }): B.startBlock(end); } +/** Entries in a row whose fast copy bailed before a speculative loop site + * stops trying it. */ +const SPECULATION_BAILS = 8; + /** Emit a loop's condition, body and update through `emitCopy`, twice * when speculative-integers.ts plans a fast integer copy. Both copies * share the loop's setup and exit; the ordinary copy comes first so the @@ -1939,23 +1946,66 @@ export function emitLoopCopies( ok = both; } } + // Each entry whose fast copy bails counts at the loop site, and a fast + // copy that completes a speculated write and leaves the loop clears the + // count. After SPECULATION_BAILS bailing entries in a row, the site runs + // the ordinary copy only. The count is a hint: both copies give the same + // results, and a race between workers loses at most an update. + const bails = `@scr.spec.bails.${emitter.speculativeLoopSites++}`; + emitter.declare(`${bails} = internal global i8 0`); + const count = B.tmp(), + trying = B.tmp(); + B.line(`${count} = load atomic i8, ptr ${bails} unordered, align 1`); + B.line(`${trying} = icmp ult i8 ${count}, ${SPECULATION_BAILS}`); + if (ok === "true") ok = trying; + else { + const both = B.tmp(); + B.line(`${both} = and i1 ${ok}, ${trying}`); + ok = both; + } const fast = B.newLabel("spec.fast"), slow = B.newLabel("spec.slow"); B.condBr(ok, fast, slow); const resume = new Map(); for (const s of plan.writes) resume.set(s, { label: B.newLabel("spec.resume"), placed: false }); + // A failed check finishes its iteration in a third copy, entered only + // there: it re-executes the statement and the rest of the body and the + // update, then enters the ordinary copy through its header. Jumping + // into the middle of the ordinary copy instead would give that loop a + // second entry, and LLVM does not optimize an irreducible loop. + const points = [...resume.values()]; + let header: string | null = null; B.startBlock(slow); + emitter.speculationHeader = (label) => { + header = label; + }; + try { + emitCopy(exit); + } finally { + emitter.speculationHeader = null; + } + if (header === null) + throw new InternalCompilerError("llvm emitter bug: speculative loop copy without a header"); + B.startBlock(B.newLabel("spec.tail")); emitter.speculationResume = resume; + emitter.speculationHeader = () => { + B.br(header!); + B.startBlock(B.newLabel("spec.unused")); + }; try { emitCopy(exit); } finally { emitter.speculationResume = null; + emitter.speculationHeader = null; } B.startBlock(fast); - if (![...resume.values()].every((r) => r.placed)) { + if (!points.every((r) => r.placed)) { B.br(slow); return; } + const ran = B.slot(); + B.entryAllocas.push(`${ran} = alloca i1 ; speculated write completed`); + B.line(`store i1 false, ptr ${ran}`); for (const id of plan.locals.keys()) { const shadow = B.slot(); B.entryAllocas.push( @@ -1992,22 +2042,43 @@ export function emitLoopCopies( writes: plan.writes, resume: new Map([...resume].map(([s, r]) => [s, r.label])), deferred, + bails, + ran, }; try { - if (deferred.size === 0) emitCopy(exit); - else { - const leave = B.newLabel("spec.leave"); - emitCopy(leave); - B.startBlock(leave); - emitter.writeBackDeferred(deferred); - B.br(exit); - } + const leave = B.newLabel("spec.leave"); + emitCopy(leave); + B.startBlock(leave); + emitter.writeBackDeferred(deferred); + // An entry that never ran a speculated write (a loop of no iterations) + // says nothing about the site. + const wrote = B.tmp(), + left = B.tmp(), + counted = B.tmp(), + both = B.tmp(), + clear = B.newLabel("spec.clear"); + B.line(`${wrote} = load i1, ptr ${ran}`); + B.line(`${left} = load atomic i8, ptr ${bails} unordered, align 1`); + B.line(`${counted} = icmp ne i8 ${left}, 0`); + B.line(`${both} = and i1 ${wrote}, ${counted}`); + B.condBr(both, clear, exit); + B.startBlock(clear); + B.line(`store atomic i8 0, ptr ${bails} unordered, align 1`); + B.br(exit); } finally { emitter.speculation = null; for (const id of plan.locals.keys()) emitter.integerLoopBindings.delete(id); } } +/** 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. */ +function enterLoopHeader(emitter: LlEmitter, header: string): void { + const dispatch = emitter.speculationHeader; + emitter.speculationHeader = null; + dispatch?.(header); +} + /** A proven integer loop (speculative-integers.ts) for this loop, or null. * A shadowed local's double is stale until the loop is left, which only * the loop's exit, a return and an exception do: the loop is planned only @@ -2134,7 +2205,8 @@ export function emitSpeculativeWrite( if (emitter.frames.pop()!.length > 0) throw new InternalCompilerError("llvm emitter bug: speculative write owns a temporary"); B.line(`store i64 ${result}, ptr ${binding.slot}`); - const deferred = emitter.speculation!.deferred; + const { deferred, bails: site, ran } = emitter.speculation!; + if (ran !== undefined) B.line(`store i1 true, ptr ${ran}`); if (!deferred.has(target.localId)) { const number = B.tmp(); B.line(`${number} = sitofp i64 ${result} to double`); @@ -2145,6 +2217,14 @@ export function emitSpeculativeWrite( B.br(done); B.startBlock(failed); emitter.writeBackDeferred(deferred); + if (site !== undefined) { + const count = B.tmp(), + next = B.tmp(); + emitter.declare("declare i8 @llvm.uadd.sat.i8(i8, i8)"); + B.line(`${count} = load atomic i8, ptr ${site} unordered, align 1`); + B.line(`${next} = call i8 @llvm.uadd.sat.i8(i8 ${count}, i8 1)`); + B.line(`store atomic i8 ${next}, ptr ${site} unordered, align 1`); + } B.br(resume!); B.startBlock(done); } diff --git a/packages/compiler/src/backend/llvm/emitter.ts b/packages/compiler/src/backend/llvm/emitter.ts index ed1785912..c1cf66af4 100644 --- a/packages/compiler/src/backend/llvm/emitter.ts +++ b/packages/compiler/src/backend/llvm/emitter.ts @@ -1102,9 +1102,20 @@ export class LlEmitter { resume: Map; /** Globals whose double is written back only when the copy is left. */ deferred: Set; + /** The loop site's count of bailing entries in a row (a global), and + * whether this entry completed a speculated write (a slot). A proven + * integer loop has neither. */ + bails?: string; + ran?: string; } | null = null; + /** Speculative loop sites numbered so far in the module. */ + speculativeLoopSites = 0; /** The ordinary copy being emitted: resume labels still to place. */ speculationResume: Map | null = null; + /** Run once at the top of the next loop copy's header + * (`enterLoopHeader`): records the ordinary copy's header, or sends the + * third copy's header there. */ + speculationHeader: ((header: string) => void) | null = null; /** The function being emitted, for speculative globals. */ speculativeFunction = ""; /** Every function that writes each number global (computed on demand). */ diff --git a/packages/compiler/test/speculative-loop-emission.test.ts b/packages/compiler/test/speculative-loop-emission.test.ts new file mode 100644 index 000000000..78a7de295 --- /dev/null +++ b/packages/compiler/test/speculative-loop-emission.test.ts @@ -0,0 +1,85 @@ +import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { expect, test } from "vitest"; +import { compile } from "../src/index.js"; + +async function emitLlvm(source: string): Promise { + const dir = await mkdtemp(join(tmpdir(), "scriptc-speculative-loop-")); + try { + const entry = join(dir, "main.ts"); + const outPath = join(dir, "main.ll"); + await writeFile(entry, source); + const result = await compile(entry, { outDir: dir, outPath, outputKind: "llvm" }); + expect(result.ok, JSON.stringify(result.diagnostics)).toBe(true); + return await readFile(outPath, "utf8"); + } finally { + await rm(dir, { recursive: true, force: true }); + } +} + +const source = ` +class Point { + x: number; + w = 0; + constructor(x: number) { + this.x = x; + } +} +function scaled(points: Point[], scale: number): number { + let s = 0; + for (let i = 0; i < points.length; i++) { + const p = points[i]!; + p.w += 1; + const k = p.x * scale; + s += k * k; + } + return s; +} +console.log(scaled([new Point(1), new Point(2)], 2), scaled([new Point(1)], 0.5)); +`; + +function labels(fn: string): string[] { + return [...fn.matchAll(/^([\w.]+):/gm)].map((m) => m[1]!); +} + +test("a failed check finishes its iteration outside the ordinary copy", async () => { + const llvm = await emitLlvm(source); + const fn = /^define internal double @sc_bf_scaled\([^]*?^}/m.exec(llvm)?.[0] ?? ""; + const order = labels(fn); + const slow = order.findIndex((l) => l.startsWith("spec.slow")); + const tail = order.findIndex((l) => l.startsWith("spec.tail")); + const fast = order.findIndex((l) => l.startsWith("spec.fast")); + expect(slow).toBeGreaterThanOrEqual(0); + expect(tail).toBeGreaterThan(slow); + expect(fast).toBeGreaterThan(tail); + // Every resume point is in the third copy, between its start and the + // fast copy, so the ordinary copy has one entry. + const resumes = order.flatMap((l, i) => (l.startsWith("spec.resume") ? [i] : [])); + expect(resumes.length).toBeGreaterThan(0); + for (const i of resumes) { + expect(i).toBeGreaterThan(tail); + expect(i).toBeLessThan(fast); + } + // The third copy's header leaves for the ordinary copy's header. + const header = order[slow + 1]!; + const third = fn.slice(fn.indexOf(`\n${order[tail]}:`), fn.indexOf(`\n${order[fast]}:`)); + expect(third).toContain(`br label %${header}`); +}); + +test("a loop site stops entering its fast copy after repeated bails", async () => { + const llvm = await emitLlvm(source); + const site = /^(@scr\.spec\.bails\.\d+) = internal global i8 0$/m.exec(llvm)?.[1]; + expect(site).toBeDefined(); + const fn = /^define internal double @sc_bf_scaled\([^]*?^}/m.exec(llvm)?.[0] ?? ""; + const escaped = site!.replace(/\./g, "\\."); + // The entry test reads the count, a bail adds one, and a completed fast + // copy clears it. + expect(fn).toMatch( + new RegExp( + `(%\\S+) = load atomic i8, ptr ${escaped} unordered, align 1\\n\\s*%\\S+ = icmp ult i8 \\1, 8`, + ), + ); + expect(fn).toMatch(new RegExp(`call i8 @llvm\\.uadd\\.sat\\.i8\\(i8 %\\S+, i8 1\\)`)); + expect(fn).toContain(`store atomic i8 0, ptr ${site} unordered, align 1`); +}); diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7321-speculative-loop-bail-sites.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7321-speculative-loop-bail-sites.ts.json new file mode 100644 index 000000000..113eeedb7 --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7321-speculative-loop-bail-sites.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/7321-speculative-loop-bail-sites.ts" + ], + "diags": [] +} diff --git a/tests/corpus/7321-speculative-loop-bail-sites.ts b/tests/corpus/7321-speculative-loop-bail-sites.ts new file mode 100644 index 000000000..460595feb --- /dev/null +++ b/tests/corpus/7321-speculative-loop-bail-sites.ts @@ -0,0 +1,111 @@ +// Speculative integer loops whose checks fail: on the first statement of +// every entry (fractional field data), after an effect earlier in the same +// iteration, before a `continue`, `break` or `return`, and on entries +// that alternate between integer and fractional data. A site that keeps +// failing stops trying its integer copy; every result is Node's. + +class Point { + x: number; + y: number; + w = 0; + + constructor(x: number, y: number) { + this.x = x; + this.y = y; + } +} + +function show(label: string, value: number): void { + console.log(label, Object.is(value, -0) ? "-0" : String(value)); +} + +// Each inner entry fails at its first statement when the fields hold +// fractions, and runs no iteration at all for the last point. +function pairs(points: Point[]): number { + let total = 0; + for (let i = 0; i < points.length; i++) { + const p = points[i]!; + let acc = p.w; + for (let j = i + 1; j < points.length; j++) { + const q = points[j]!; + const dx = p.x - q.x; + const dy = p.y - q.y; + const d2 = dx * dx + dy * dy; + acc += d2; + q.w += d2 * 0.5; + } + p.w = acc; + total += acc; + } + return total; +} + +const fractional = [ + new Point(0.5, 1), + new Point(1.25, -2), + new Point(3, 4.75), + new Point(-1.5, 0), +]; +const integral = [new Point(0, 1), new Point(2, -2), new Point(3, 4), new Point(-1, 0)]; +for (let round = 0; round < 12; round++) { + show(`pairs fractional ${round}`, pairs(fractional)); +} +// The same site, now with integer data after it stopped speculating. +show("pairs integral", pairs(integral)); +show("pairs integral again", pairs(integral)); + +// A field write comes before the failing statement: the iteration resumes +// after it and must not repeat it. +function afterEffect(points: Point[], scale: number): number { + let s = 0; + for (let i = 0; i < points.length; i++) { + const p = points[i]!; + p.w += 1; + const k = p.x * scale; + s += k * k; + } + return s; +} +const marks = [new Point(1, 0), new Point(2, 0), new Point(3, 0)]; +show("after effect 2", afterEffect(marks, 2)); +show("after effect 0.5", afterEffect(marks, 0.5)); +show("after effect 3", afterEffect(marks, 3)); +console.log(marks.map((p) => p.w).join(" ")); + +// The failing statement is followed by `continue`, `break` and `return`. +function control(n: number, step: number): number { + let acc = 0; + let i = 0; + while (i < n) { + const v = i * step; + i++; + if (v * v > 400) break; + if ((i & 1) === 0) continue; + acc += v * 3; + if (acc > 1000) return -acc; + } + return acc; +} +for (const step of [1, 0.5, 2, 0.75, 7, 1.5, 3, 0.25, 5, 1.75, 4]) { + show(`control ${step}`, control(40, step)); +} + +// do-while: the check fails part-way through, on some entries only. +function steps(start: number, factor: number): number { + let v = start; + let k = 0; + do { + v = v * factor - k; + k++; + } while (k < 6); + return v; +} +for (let r = 0; r < 20; r++) show(`steps ${r}`, steps(r % 3 === 0 ? r + 0.5 : r, 3)); + +// Overflow past 2^53 on later entries only. +function powers(base: number, n: number): number { + let p = 1; + for (let i = 0; i < n; i++) p = p * base; + return p; +} +for (const n of [3, 10, 40, 80, 5, 60, 2, 70, 1, 90, 0, 4]) show(`powers ${n}`, powers(3, n)); diff --git a/tests/harness/corpus-costs/7321-speculative-loop-bail-sites.ts.json b/tests/harness/corpus-costs/7321-speculative-loop-bail-sites.ts.json new file mode 100644 index 000000000..1e5609a66 --- /dev/null +++ b/tests/harness/corpus-costs/7321-speculative-loop-bail-sites.ts.json @@ -0,0 +1 @@ +{ "plain": 2000, "san": 2600 }