diff --git a/packages/compiler/src/backend/llvm/bytes-loop-versions.test.ts b/packages/compiler/src/backend/llvm/bytes-loop-versions.test.ts index 34b75861c..de71a6ef0 100644 --- a/packages/compiler/src/backend/llvm/bytes-loop-versions.test.ts +++ b/packages/compiler/src/backend/llvm/bytes-loop-versions.test.ts @@ -11,7 +11,7 @@ import { type IrType, } from "../../ir/ir.js"; import { analyzeIntegerRanges } from "../../ir/integer-ranges.js"; -import { matchIntegerCountedForLoop } from "../../ir/integer-loops.js"; +import { matchIntegerCountedForLoop, matchIntegerPairedForLoop } from "../../ir/integer-loops.js"; import { findInitializedByteLoopBindings } from "./byte-windows.js"; import { matchBytesLoopVersion, type BytesLoopVersion } from "./bytes-loop-versions.js"; import { emitLlvmModule } from "./emitter.js"; @@ -623,3 +623,78 @@ test("while loop counters below a string's length version its reads", () => { expect(version.bounded.map((c) => c.slack)).toEqual([0]); expect(version.fixed).toEqual([]); }); + +test("a two-counter loop versions the typed arrays it indexes at either counter", () => { + // for (let i = 0, j = k; i < j; i++, j--) { const t = out[i]; out[i] = out[j]; out[j] = t; } + const loop: IrStmt & { kind: "for" } = { + kind: "for", + init: { + kind: "block", + body: [ + { kind: "varDecl", localId: "i", init: num(0), loc }, + { kind: "varDecl", localId: "j", init: ref("k"), loc }, + ], + loc, + }, + cond: { kind: "bin", op: "<", left: ref("i"), right: ref("j"), type: BOOL, loc }, + update: { + kind: "block", + body: [ + { kind: "assign", localId: "i", value: bin("+", ref("i"), num(1)), loc }, + { kind: "assign", localId: "j", value: bin("-", ref("j"), num(1)), loc }, + ], + loc, + }, + body: [ + { kind: "varDecl", localId: "t", init: get(ref("out", f64s), ref("i")), loc }, + set(ref("out", f64s), ref("i"), get(ref("out", f64s), ref("j"))), + set(ref("out", f64s), ref("j"), ref("t")), + // Another index stays checked. + set(ref("src", f64s), bin("+", ref("i"), num(1)), num(0)), + ], + loc, + }; + const kernel = fn([loop], [local("t", F64)]); + const ranges = analyzeIntegerRanges(kernel); + const locals = new Map(kernel.locals.map((l) => [l.id, l])); + const paired = matchIntegerPairedForLoop(loop, locals, ranges)!; + expect(paired).not.toBeNull(); + const version = matchBytesLoopVersion( + loop, + null, + { + locals, + invariantGlobals: new Set(), + captured: new Set(), + ranges, + initialized: findInitializedByteLoopBindings(kernel).get(loop) ?? new Set(), + proven: new Set(), + workers: false, + owned: new Set(), + scalarStrings: new Set(), + }, + paired, + )!; + expect(version.paired?.loop).toBe(paired); + expect(version.paired?.receivers.map((r) => version.receivers[r]!.ref)).toEqual([ + ref("out", f64s), + ]); + expect(version.bounds.size).toBe(4); + const mod: IrModule = { + irVersion: 17, + sourceFile: loc.file, + entry: "__main", + functions: [ + { name: "__main", params: [], returnType: VOID, locals: [], body: [], loc }, + kernel, + ], + }; + const body = /^define internal [^\n]*@sc_bf_kernel[^]*?^}/m.exec(emitLlvmModule(mod))![0]; + const guard = body.slice(0, body.search(/^bytes\.version\.(fast|slow)\d+:$/m)); + expect(guard).toMatch(/icmp sge i64 %\S+, 0/); + expect(guard).toMatch(/icmp slt i64 %\S+, %\S+/); + const fast = /^bytes\.version\.fast\d+:$[^]*?^bytes\.version\.done\d+:$/m.exec(body)![0]; + // The fast copy keeps only the `src[i + 1]` check, and no second guard. + expect(fast.match(/bytes\.index\.invalid\d+:/g)).toHaveLength(1); + expect(fast).not.toContain("counted.fast"); +}); diff --git a/packages/compiler/src/backend/llvm/bytes-loop-versions.ts b/packages/compiler/src/backend/llvm/bytes-loop-versions.ts index 704b6723a..2013d03e4 100644 --- a/packages/compiler/src/backend/llvm/bytes-loop-versions.ts +++ b/packages/compiler/src/backend/llvm/bytes-loop-versions.ts @@ -1,11 +1,15 @@ import { F64, type IrExpr, type IrLocal, type IrStmt } from "../../ir/ir.js"; import type { IntegerRange, IntegerRanges } from "../../ir/integer-ranges.js"; -import type { IntegerCountedForLoop } from "../../ir/integer-loops.js"; +import type { IntegerCountedForLoop, IntegerPairedForLoop } from "../../ir/integer-loops.js"; import type { LlvmEmitterContext } from "./expr-context.js"; import { byteNumberAccess } from "../../ir/byte-numbers.js"; import { everyExprChild, everyStmtChild } from "../../ir/traverse.js"; import { exactInteger, widenInteger } from "./integer-values.js"; -import { emitCountedLoopLimit, emitCountedLoopSafety } from "./counted-loops.js"; +import { + emitCountedLoopLimit, + emitCountedLoopSafety, + emitExactIntegerTest, +} from "./counted-loops.js"; import { emitBytesHeaderLength, emitBytesSharedStorage } from "./expr-bytes.js"; import { f64Lit } from "./common.js"; import { emitStringAsciiFlag, emitStringByteLength } from "./string-reads.js"; @@ -79,6 +83,10 @@ export interface BytesLoopVersion { guarded: ReadonlySet; /** The counted loop whose iteration space the affine checks cover. */ counted: IntegerCountedForLoop | null; + /** A two-counter loop (`for (let i = lo, j = hi; i < j; i++, j--)`) and + * the receivers it indexes at either counter. Both counters stay between + * the two starts, so the guard requires `lo >= 0` and `hi < length`. */ + paired: { loop: IntegerPairedForLoop; receivers: number[] } | null; /** The guard also requires every counter value within [-bound, bound]. */ counterBound: number | null; /** Integer ranges of affine index operands that hold in the unchecked @@ -153,6 +161,7 @@ export function matchBytesLoopVersion( loop: IrStmt & { kind: "for" | "while" }, counted: IntegerCountedForLoop | null, facts: BytesLoopFacts, + paired: IntegerPairedForLoop | null = null, ): BytesLoopVersion | null { const written = new Set(); const temps = new Map(); @@ -438,6 +447,9 @@ export function matchBytesLoopVersion( string, { check: BoundedCheck; nodes: { node: Node; access: WhileAccess }[] } >(); + const pairedCounters = + paired !== null && counted === null ? new Set([paired.low.localId, paired.high.localId]) : null; + const pairedNodes = new Map(); const access = (node: Node, array: IrExpr, index: IrExpr, width: number): void => { const ref = root(array); if (!ref) return; @@ -448,6 +460,20 @@ export function matchBytesLoopVersion( else owned.set(receiver, [node]); } if (facts.proven.has(node)) return; + // An element at either counter of a two-counter loop. + if ( + pairedCounters !== null && + width === 1 && + ref.type.kind === "bytes" && + index.kind === "varRef" && + pairedCounters.has(index.localId) + ) { + const receiver = receiverOf(ref); + const nodes = pairedNodes.get(receiver); + if (nodes) nodes.push(node); + else pairedNodes.set(receiver, [node]); + return; + } const shape = counter !== null ? affine(index) : null; if (shape && shape.counter !== 0) { const operands = new Set(); @@ -543,11 +569,16 @@ export function matchBytesLoopVersion( Math.max(0, MAX_CHECKS - fixedKept.length - affineKept.length), ); const ownedKept = [...owned.entries()].slice(0, MAX_CHECKS); + const pairedKept = [...pairedNodes.entries()].slice( + 0, + Math.max(0, MAX_CHECKS - fixedKept.length - affineKept.length - boundedKept.length), + ); if ( fixedKept.length === 0 && affineKept.length === 0 && boundedKept.length === 0 && - ownedKept.length === 0 + ownedKept.length === 0 && + pairedKept.length === 0 ) return null; const bounds = new Set(); @@ -616,6 +647,10 @@ export function matchBytesLoopVersion( if (access.index !== access.ref) narrow(access.index, { min: 0, max: SAFE }); } } + for (const [receiver, nodes] of pairedKept) { + used.add(receiver); + for (const node of nodes) bounds.add(node); + } const counterChecks: CounterCheck[] = []; for (const [id, min] of counterMins) { const info = whileAccesses!.counters.get(id)!; @@ -658,6 +693,10 @@ export function matchBytesLoopVersion( counters: counterChecks, guarded, counted: affineKept.length > 0 ? counted : null, + paired: + paired !== null && pairedKept.length > 0 + ? { loop: paired, receivers: pairedKept.map(([receiver]) => renumber.get(receiver)!) } + : null, counterBound: needsBound ? bound : null, ranges, }; @@ -1143,6 +1182,25 @@ export function emitBytesLoopVersionGuard( conditions.push(both); } } + if (version.paired !== null) { + // Both counters stay within [lo, hi] (exact integers below 2^53): every + // access at either one is in bounds when `lo >= 0` and `hi < length`. + const { loop } = version.paired; + const [low, high] = [loop.low.start, loop.high.start].map((expr) => { + const tested = emitExactIntegerTest(host, host.emitExpr(expr).name, -SAFE, SAFE); + conditions.push(tested.ok); + return tested.integer; + }); + const nonNegative = B.tmp(); + B.line(`${nonNegative} = icmp sge i64 ${low}, 0`); + conditions.push(nonNegative); + for (const receiver of version.paired.receivers) { + // Lengths are at most 2^53 elements: a signed comparison is exact. + const below = B.tmp(); + B.line(`${below} = icmp slt i64 ${high}, ${lengths[receiver]}`); + conditions.push(below); + } + } let valid = "true"; for (const cond of conditions) { if (valid === "true") valid = cond; @@ -1156,39 +1214,6 @@ export function emitBytesLoopVersionGuard( else B.condBr(valid, fast, slow); } -/** Tests that a double is an exact integer within [min, max] (not -0; NaN - * fails the ordered comparisons) and gives its i64. The integer is frozen, - * so it is a fixed value even when the test fails. */ -function emitExactIntegerTest( - host: LlvmEmitterContext, - value: string, - min: number, - max: number, -): { ok: string; integer: string } { - const B = host.B; - const low = B.tmp(), - high = B.tmp(), - inRange = B.tmp(), - raw = B.tmp(), - integer = B.tmp(), - back = B.tmp(), - bits = B.tmp(), - backBits = B.tmp(), - exact = B.tmp(), - ok = B.tmp(); - B.line(`${low} = fcmp oge double ${value}, ${f64Lit(min)}`); - B.line(`${high} = fcmp ole double ${value}, ${f64Lit(max)}`); - B.line(`${inRange} = and i1 ${low}, ${high}`); - B.line(`${raw} = fptosi double ${value} to i64`); - B.line(`${integer} = freeze i64 ${raw}`); - B.line(`${back} = sitofp i64 ${integer} to double`); - B.line(`${bits} = bitcast double ${value} to i64`); - B.line(`${backBits} = bitcast double ${back} to i64`); - B.line(`${exact} = icmp eq i64 ${bits}, ${backBits}`); - B.line(`${ok} = and i1 ${inRange}, ${exact}`); - return { ok, integer }; -} - function scaleCounter(host: LlvmEmitterContext, value: string, scale: number): string { if (scale === 1) return value; const scaled = host.B.tmp(); diff --git a/packages/compiler/src/backend/llvm/counted-loops.ts b/packages/compiler/src/backend/llvm/counted-loops.ts index ca2dd6b92..1ff4c6580 100644 --- a/packages/compiler/src/backend/llvm/counted-loops.ts +++ b/packages/compiler/src/backend/llvm/counted-loops.ts @@ -151,3 +151,36 @@ function emitExactStart(host: LlvmEmitterContext, expr: IrExpr): string { B.line(`${ok} = and i1 ${inRange}, ${exact}`); return ok; } + +/** Tests that a double is an exact integer within [min, max] (not -0; NaN + * fails the ordered comparisons) and gives its i64. The integer is frozen, + * so it is a fixed value even when the test fails. */ +export function emitExactIntegerTest( + host: LlvmEmitterContext, + value: string, + min: number, + max: number, +): { ok: string; integer: string } { + const B = host.B; + const low = B.tmp(), + high = B.tmp(), + inRange = B.tmp(), + raw = B.tmp(), + integer = B.tmp(), + back = B.tmp(), + bits = B.tmp(), + backBits = B.tmp(), + exact = B.tmp(), + ok = B.tmp(); + B.line(`${low} = fcmp oge double ${value}, ${f64Lit(min)}`); + B.line(`${high} = fcmp ole double ${value}, ${f64Lit(max)}`); + B.line(`${inRange} = and i1 ${low}, ${high}`); + B.line(`${raw} = fptosi double ${value} to i64`); + B.line(`${integer} = freeze i64 ${raw}`); + B.line(`${back} = sitofp i64 ${integer} to double`); + B.line(`${bits} = bitcast double ${value} to i64`); + B.line(`${backBits} = bitcast double ${back} to i64`); + B.line(`${exact} = icmp eq i64 ${bits}, ${backBits}`); + B.line(`${ok} = and i1 ${inRange}, ${exact}`); + return { ok, integer }; +} diff --git a/packages/compiler/src/backend/llvm/dense-array-loops.ts b/packages/compiler/src/backend/llvm/dense-array-loops.ts index 384c31b0c..b806939b5 100644 --- a/packages/compiler/src/backend/llvm/dense-array-loops.ts +++ b/packages/compiler/src/backend/llvm/dense-array-loops.ts @@ -36,8 +36,7 @@ import { everyExprChild, everyStmtChild } from "../../ir/traverse.js"; import type { IntegerCountedForLoop, IntegerPairedForLoop } from "../../ir/integer-loops.js"; import type { LlvmEmitterContext } from "./expr-context.js"; import { loadArrCount, SCR_ARR_FIELD } from "./common.js"; -import { emitCountedLoopLimit } from "./counted-loops.js"; -import { f64Lit } from "./common.js"; +import { emitCountedLoopLimit, emitExactIntegerTest } from "./counted-loops.js"; export const INT32_MAX = 2147483647; @@ -453,28 +452,10 @@ function emitPairedDenseGuard( header: (array: string, field: number) => string, ): string { const B = host.B; - const integer = (expr: IrExpr, min: number, max: number): { value: string; ok: string } => { - const start = host.emitExpr(expr); - const low = B.tmp(), - high = B.tmp(), - raw = B.tmp(), - value = B.tmp(), - back = B.tmp(), - exact = B.tmp(), - inRange = B.tmp(), - ok = B.tmp(); - B.line(`${low} = fcmp oge double ${start.name}, ${f64Lit(min)}`); - B.line(`${high} = fcmp ole double ${start.name}, ${f64Lit(max)}`); - B.line(`${inRange} = and i1 ${low}, ${high}`); - B.line(`${raw} = fptosi double ${start.name} to i64`); - B.line(`${value} = freeze i64 ${raw}`); - B.line(`${back} = sitofp i64 ${value} to double`); - B.line(`${exact} = fcmp oeq double ${back}, ${start.name}`); - B.line(`${ok} = and i1 ${inRange}, ${exact}`); - return { value, ok }; - }; - const low = integer(pair.lowStart, 0, INT32_MAX - 1); - const high = integer(pair.highStart, -1, INT32_MAX - 1); + const start = (expr: IrExpr, min: number, max: number) => + emitExactIntegerTest(host, host.emitExpr(expr).name, min, max); + const low = start(pair.lowStart, 0, INT32_MAX - 1); + const high = start(pair.highStart, -1, INT32_MAX - 1); let valid = B.tmp(); B.line(`${valid} = and i1 ${low.ok}, ${high.ok}`); for (const ref of loop.arrays.values()) { @@ -486,7 +467,7 @@ function emitPairedDenseGuard( both = B.tmp(); B.line(`${filled} = load i32, ptr ${filledPtr}${host.fieldAliasAttachment(filledPtr)}`); B.line(`${wide} = zext i32 ${filled} to i64`); - B.line(`${covered} = icmp slt i64 ${high.value}, ${wide}`); + B.line(`${covered} = icmp slt i64 ${high.integer}, ${wide}`); B.line(`${both} = and i1 ${valid}, ${covered}`); valid = both; } diff --git a/packages/compiler/src/backend/llvm/emit-statements.ts b/packages/compiler/src/backend/llvm/emit-statements.ts index aef5ccf30..55e31219c 100644 --- a/packages/compiler/src/backend/llvm/emit-statements.ts +++ b/packages/compiler/src/backend/llvm/emit-statements.ts @@ -1051,15 +1051,17 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void { LOOPS.of(emitter).unaliasedArrays, ) : countedLoop; - const version = versionable ? emitter.matchBytesLoopVersion(s, versionCounted) : null; + const version = versionable + ? emitter.matchBytesLoopVersion(s, versionCounted, versionCounted ? null : pairedLoop) + : null; if (version) { emitVersionedLoop(emitter, s, version); break; } let countedJoin: string | null = null; if ( - (countedLoop?.guarded && !LOOPS.of(emitter).provenCountedGuards.has(s)) || - pairedLoop?.guarded + (countedLoop?.guarded || pairedLoop?.guarded) && + !LOOPS.of(emitter).provenCountedGuards.has(s) ) { const safe = countedLoop ? emitCountedLoopSafety(emitter, countedLoop) @@ -2094,7 +2096,8 @@ function emitVersionedLoop( if (version.owned.size > 0) emitter.bytesOwned = new Set([...originalOwned, ...version.owned]); if (version.ranges.size > 0) emitter.integerRanges = new Map([...originalRanges, ...version.ranges]); - if (version.counted?.guarded) LOOPS.of(emitter).provenCountedGuards.add(s); + if (version.counted?.guarded || version.paired?.loop.guarded) + LOOPS.of(emitter).provenCountedGuards.add(s); emitter.emitStmt(s); B.br(done); } finally { diff --git a/packages/compiler/src/backend/llvm/emitter.ts b/packages/compiler/src/backend/llvm/emitter.ts index 4aa9b65ad..ad59c3b1b 100644 --- a/packages/compiler/src/backend/llvm/emitter.ts +++ b/packages/compiler/src/backend/llvm/emitter.ts @@ -60,7 +60,7 @@ import { captureFreeFunctions, observesStacks } from "../stack-observations.js"; import { findInitializedByteLoopBindings } from "./byte-windows.js"; import { matchBytesLoopVersion } from "./bytes-loop-versions.js"; import { globalSlotAlias } from "./bytes-alias.js"; -import type { IntegerCountedForLoop } from "../../ir/integer-loops.js"; +import type { IntegerCountedForLoop, IntegerPairedForLoop } from "../../ir/integer-loops.js"; import { materializeSplitPiece, type SplitSpan, @@ -4263,6 +4263,7 @@ export class LlEmitter { matchBytesLoopVersion( s: IrStmt & { kind: "for" | "while" }, counted: IntegerCountedForLoop | null, + paired: IntegerPairedForLoop | null = null, ): ReturnType { const fn = LOOPS.of(this).bytesLoopFunction; if (!fn) return null; @@ -4272,17 +4273,22 @@ export class LlEmitter { const invariantGlobals = new Set(this.stableGlobals); for (const global of this.initializerBindings.get(fn.name) ?? []) invariantGlobals.add(global.id); - return matchBytesLoopVersion(s, counted, { - locals: this.currentLocals, - invariantGlobals, - captured: this.captureIds, - ranges: this.integerRanges, - initialized, - proven: this.bytesBounds, - workers: this.mod.workers === true, - owned: this.bytesOwned, - scalarStrings: new Set(this.stringSlices.keys()), - }); + return matchBytesLoopVersion( + s, + counted, + { + locals: this.currentLocals, + invariantGlobals, + captured: this.captureIds, + ranges: this.integerRanges, + initialized, + proven: this.bytesBounds, + workers: this.mod.workers === true, + owned: this.bytesOwned, + scalarStrings: new Set(this.stringSlices.keys()), + }, + paired, + ); } /** The alias tag of a whole-slot load of a binding (module globals only). */ diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7981-typed-array-paired-loops.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7981-typed-array-paired-loops.ts.json new file mode 100644 index 000000000..71fa483a8 --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7981-typed-array-paired-loops.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/7981-typed-array-paired-loops.ts" + ], + "diags": [] +} diff --git a/tests/corpus/7981-typed-array-paired-loops.ts b/tests/corpus/7981-typed-array-paired-loops.ts new file mode 100644 index 000000000..d13eecbd6 --- /dev/null +++ b/tests/corpus/7981-typed-array-paired-loops.ts @@ -0,0 +1,57 @@ +// Two-counter loops over typed arrays (`for (let i = lo, j = hi; i < j; +// i++, j--)`), with starts inside the arrays and outside them: past the +// end, negative, fractional and NaN. Views over one buffer, other element +// types, inclusive conditions and steps other than one. + +function reverse(a: Int32Array, lo: number, hi: number): string { + for (let i = lo, j = hi; i < j; i++, j--) { + const t = a[i]!; + a[i] = a[j]!; + a[j] = t; + } + return Array.from(a).join(); +} + +function mirror(a: Float64Array, b: Float64Array, hi: number): number { + let sum = 0; + for (let i = 0, j = hi; i <= j; i++, j--) { + b[i] = a[j]! * 2; + sum += a[i]! - b[j]!; + } + return sum; +} + +function skipping(a: Uint8ClampedArray, hi: number): string { + for (let i = 0, j = hi; i < j; i += 2, j -= 3) { + a[i] = a[j]! + 200; + a[j] = i - 10; + } + return Array.from(a).join(); +} + +function palindrome(text: Uint16Array, n: number): boolean { + for (let i = 0, j = n - 1; i < j; i++, j--) if (text[i] !== text[j]) return false; + return true; +} + +console.log(reverse(new Int32Array([1, 2, 3, 4, 5, 6]), 0, 5)); +console.log(reverse(new Int32Array([1, 2, 3, 4, 5, 6]), 1, 3)); +console.log(reverse(new Int32Array([1, 2, 3, 4]), 0, 6)); +console.log(reverse(new Int32Array([1, 2, 3, 4]), -2, 2)); +console.log(reverse(new Int32Array([1, 2, 3, 4]), 0.5, 3.5)); +console.log(reverse(new Int32Array([1, 2, 3, 4]), 0, NaN)); +console.log(reverse(new Int32Array([1, 2, 3, 4]), 3, 0)); + +const buffer = new ArrayBuffer(48); +const whole = new Float64Array(buffer); +for (let k = 0; k < whole.length; k++) whole[k] = k + 0.5; +const tail = new Float64Array(buffer, 16, 3); +console.log(mirror(whole, tail, 4), Array.from(whole).join()); +console.log(mirror(tail, tail, 2), Array.from(whole).join()); +console.log(mirror(tail, whole, 7), Array.from(whole).join()); + +console.log(skipping(new Uint8ClampedArray([10, 20, 30, 40, 50, 60, 70, 80]), 7)); +console.log(skipping(new Uint8ClampedArray([10, 20, 30]), 9)); + +const word = new Uint16Array([114, 97, 99, 101, 99, 97, 114]); +console.log(palindrome(word, 7), palindrome(word, 6), palindrome(word, 9), palindrome(word, 0)); diff --git a/tests/harness/corpus-costs/7981-typed-array-paired-loops.ts.json b/tests/harness/corpus-costs/7981-typed-array-paired-loops.ts.json new file mode 100644 index 000000000..c93369880 --- /dev/null +++ b/tests/harness/corpus-costs/7981-typed-array-paired-loops.ts.json @@ -0,0 +1 @@ +{ "plain": 3500, "san": 4500 }