Speculate floor-mods and products by powers of two as masks and shifts (#2199)

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();

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");
+});

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/8244-speculative-floor-mod-powers.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/8244-speculative-floor-mod-powers.ts"
+ ],
+ "diags": []
+}

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));