diff --git a/packages/compiler/src/backend/llvm/speculative-integers.ts b/packages/compiler/src/backend/llvm/speculative-integers.ts index 07e5e3721..971bd00cb 100644 --- a/packages/compiler/src/backend/llvm/speculative-integers.ts +++ b/packages/compiler/src/backend/llvm/speculative-integers.ts @@ -187,6 +187,37 @@ function exactPowerOfTwo(value: number): number | null { return k >= 0 && k <= 52 && 2 ** k === value ? k : null; } +/** The same speculated operand written twice: a local, a literal, or a + * power of two of the same exponent. */ +function sameOperand(a: IrExpr, b: IrExpr): boolean { + if (a.kind === "varRef") return b.kind === "varRef" && a.localId === b.localId; + if (a.kind === "numLit") return b.kind === "numLit" && Object.is(a.value, b.value); + return ( + a.kind === "libCall" && + b.kind === "libCall" && + a.fn === "math.pow" && + b.fn === "math.pow" && + powerOfTwo(a) && + powerOfTwo(b) && + sameOperand(a.args[1]!, b.args[1]!) + ); +} + +/** `a - Math.floor(a / p) * p` (either factor order) with p a power of two + * and `a` a local or literal: the floor-mod of `a` by p. */ +function floorModPowerOfTwo(e: IrExpr): { value: IrExpr; power: IrExpr } | null { + if (e.kind !== "bin" || e.op !== "-" || e.right.kind !== "bin" || e.right.op !== "*") return null; + const { left, right } = e.right; + const [floor, power] = + left.kind === "libCall" && left.fn === "math.floor" ? [left, right] : [right, left]; + if (floor.kind !== "libCall" || floor.fn !== "math.floor" || floor.args.length !== 1) return null; + const quotient = floor.args[0]!; + if (quotient.kind !== "bin" || quotient.op !== "/") return null; + if (!powerOfTwo(power) || !sameOperand(quotient.right, power)) return null; + if (e.left.kind !== "varRef" && e.left.kind !== "numLit") return null; + return sameOperand(quotient.left, e.left) ? { value: e.left, power } : null; +} + /** Whether `e` is a power of two speculation can shift by: * `Math.pow(2, n)` or a literal 2^k with k in [0, 52]. */ function powerOfTwo(e: IrExpr): boolean { @@ -950,6 +981,37 @@ export function emitSpeculativeInteger( } case "bin": { if (e.op !== "+" && e.op !== "-" && e.op !== "*" && e.op !== "%") break; + // The floor-mod of a safe integer by 2^n (n in [0, 52]) is its low n + // bits, also for a negative value: in [0, 2^n), and never -0. + const mod = floorModPowerOfTwo(e); + if (mod) { + const value = emitSpeculativeInteger(host, mod.value, bail, reads, integral); + const n = exponentOf(host, mod.power, bail, reads, integral); + const mask = B.tmp(), + low = B.tmp(); + B.line(`${mask} = sub i64 ${shiftOne(host, n)}, 1`); + B.line(`${low} = and i64 ${value}, ${mask}`); + return low; + } + // A product with 2^n shifts: it is safe iff the other factor's + // magnitude is at most (2^53 - 1) >> n, and a zero product of a safe + // integer and a positive power is +0. + if (e.op === "*" && (powerOfTwo(e.right) || powerOfTwo(e.left))) { + const [factor, power] = powerOfTwo(e.right) ? [e.left, e.right] : [e.right, e.left]; + const value = emitSpeculativeInteger(host, factor, bail, reads, integral); + const n = exponentOf(host, power, bail, reads, integral); + const limit = B.tmp(), + magnitude = B.tmp(), + outside = B.tmp(), + product = B.tmp(); + B.line(`${limit} = ashr i64 ${SAFE}, ${n}`); + host.declare("declare i64 @llvm.abs.i64(i64, i1)"); + B.line(`${magnitude} = call i64 @llvm.abs.i64(i64 ${value}, i1 true)`); + B.line(`${outside} = icmp ugt i64 ${magnitude}, ${limit}`); + bail(outside); + B.line(`${product} = shl nsw i64 ${value}, ${n}`); + return product; + } const a = emitSpeculativeInteger(host, e.left, bail, reads, integral); const b = emitSpeculativeInteger(host, e.right, bail, reads, integral); const result = B.tmp(); diff --git a/packages/compiler/test/speculative-loop-emission.test.ts b/packages/compiler/test/speculative-loop-emission.test.ts index 748fd50c5..d26a0c96c 100644 --- a/packages/compiler/test/speculative-loop-emission.test.ts +++ b/packages/compiler/test/speculative-loop-emission.test.ts @@ -159,3 +159,27 @@ if (isMainThread) { expect(order.slice(slow, fast).some((l) => l.startsWith("spec.fast"))).toBe(true); expect(order.slice(fast).some((l) => l.startsWith("spec.fast"))).toBe(true); }); + +test("a floor-mod by a power of two speculates as a mask and a product by one as a shift", async () => { + const llvm = await emitLlvm(` +function trim(value: number, bits: number): number { + return value - Math.floor(value / Math.pow(2, bits)) * Math.pow(2, bits); +} +function key(bytes: Uint8Array, bits: number): number { + let k = 0; + for (let i = 0; i < bytes.length; i++) k = trim(k, bits) * 4 + bytes[i]!; + return k; +} +console.log(key(new Uint8Array([1, 2, 3, 0]), 6)); +`); + const fn = /^define internal double @sc_bf_\w*key[\w.]*\([^]*?^}/m.exec(llvm)?.[0] ?? ""; + const fast = fn.slice(fn.search(/^spec\.fast\d+:/m)); + expect(fast.length).toBeLessThan(fn.length); + // trim(k, bits): the low `bits` bits; * 4: a shift by two. + expect(fast).toMatch(/= sub i64 %\S+, 1\n\s+%\S+ = and i64 %\S+, %\S+\n/); + expect(fast).toMatch(/= shl nsw i64 %\S+, 2\n/); + // No floor division of k (an arithmetic shift of it) and no checked + // multiply remain. + expect(fast).not.toMatch(/ashr i64 %/); + expect(fast).not.toContain("@llvm.smul.with.overflow.i64"); +}); diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/8244-speculative-floor-mod-powers.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/8244-speculative-floor-mod-powers.ts.json new file mode 100644 index 000000000..65e7ccb38 --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/8244-speculative-floor-mod-powers.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/8244-speculative-floor-mod-powers.ts" + ], + "diags": [] +} diff --git a/tests/corpus/8244-speculative-floor-mod-powers.ts b/tests/corpus/8244-speculative-floor-mod-powers.ts new file mode 100644 index 000000000..db28dcde0 --- /dev/null +++ b/tests/corpus/8244-speculative-floor-mod-powers.ts @@ -0,0 +1,88 @@ +// Integer loops whose writes take a floor-mod by a power of two +// (`v - Math.floor(v / 2^n) * 2^n`) or multiply by one: negative and zero +// values, exponents 0, 52 and above, literal powers, both factor orders, +// fractional values and results past 2^53, which leave the integer copy. + +function trim(value: number, bits: number): number { + return value - Math.floor(value / Math.pow(2, bits)) * Math.pow(2, bits); +} + +function low3(value: number): number { + return value - Math.floor(value / 8) * 8; +} + +function rolling(bytes: Uint8Array, bits: number): number { + let key = 0; + let sum = 0; + for (let i = 0; i < bytes.length; i++) { + key = trim(key, bits) * 4 + bytes[i]!; + sum = (sum + key) % 1000003; + } + return key * 7 + sum; +} + +function stepped(start: number, step: number, bits: number, count: number): number { + let v = start; + let h = 0; + for (let i = 0; i < count; i++) { + v = trim(v, bits) * 3 - step; + h = (h * 31 + v) % 999983; + } + return h * 1000 + v; +} + +function literal(start: number, count: number): number { + let v = start; + let h = 0; + for (let i = 0; i < count; i++) { + v = low3(v) * 16 - i; + h = (h * 17 + v) % 65521; + } + return h * 100 + v; +} + +function scaled(start: number, bits: number, count: number): number { + let v = start; + let h = 0; + for (let i = 0; i < count; i++) { + v = Math.pow(2, bits) * v + 1; + h = (h * 7 + v) % 10007; + } + return h + v; +} + +function shifted(start: number, count: number): number { + let v = start; + let h = 0; + for (let i = 0; i < count; i++) { + v = v * 1024 - i; + h = (h * 3 + v) % 101; + } + return h + v; +} + +const bytes = new Uint8Array(4000); +for (let i = 0; i < bytes.length; i++) bytes[i] = (i * 7 + (i >> 3)) & 3; + +for (const bits of [0, 1, 2, 10, 22, 34, 50, 52, 53, 60]) console.log("rolling", bits, rolling(bytes, bits)); +for (const [start, step, bits] of [ + [0, 1, 3], + [-5, 3, 2], + [-1, -1, 0], + [123456789, 7, 20], + [-9007199254740991, 1, 52], + [2.5, 1, 4], + [7, 0.5, 3], +] as const) + console.log("stepped", start, step, bits, stepped(start, step, bits, 50)); +for (const start of [0, 13, -13, 4096, -7.25]) console.log("literal", start, literal(start, 40)); +for (const [start, bits] of [ + [1, 3], + [-2, 5], + [0, 52], + [3, 20], + [-1, 0], +] as const) + console.log("scaled", start, bits, scaled(start, bits, 6)); +for (const start of [0, 1, -1, 3, 1e6]) console.log("shifted", start, shifted(start, 7)); +console.log(Object.is(stepped(0, 0, 3, 5), 0), Object.is(scaled(0, 2, 3), 21 + 1));