Test only the exponent of a masked value's shift and sum in speculative integers

packages/compiler/src/backend/llvm/emit-module.ts

@@ -1463,14 +1463,18 @@ function numberMapEntryHelpers(sz: string, loop: boolean): string[] {
     `miss:`,
     `  ret ptr null`,
     `slow:`,
-    `  %hashed = call ptr @sc_map_entry_f64_hashed_int${flavor}(ptr %m, i64 %ki)`,
+    `  %hashed = call ptr @sc_map_entry_f64_hashed_int${flavor}(ptr %m, i64 %ki, double %k)`,
     `  ret ptr %hashed`,
     `}`,
     ``,
     // Hashing (scr_map_find_f64): the SameValueZero key word and its
-    // scr_key_mix hash. sc_map_entry_f64_int's probe converts the integer
-    // view to the key word, which is never NaN or -0, and so skips the
-    // normalization.
+    // scr_key_mix hash. sc_map_entry_f64_int's probe takes the key's double
+    // with its integer view: its bits are the key word, never NaN, and +0
+    // for the integer 0. The caller converts the integer: cvtsi2sd writes
+    // only the low half of its register, so in a function entry, where
+    // LLVM cannot clear the register first, it waits for the register's
+    // last value, in a loop of lookups often a value the previous
+    // iteration loaded.
     ...hashedNumberProbe(
       `sc_map_entry_f64_hashed${flavor}`,
       loop,
@@ -1488,10 +1492,14 @@ function numberMapEntryHelpers(sz: string, loop: boolean): string[] {
     ...hashedNumberProbe(
       `sc_map_entry_f64_hashed_int${flavor}`,
       loop,
-      "i64 %ki",
+      "i64 %ki, double %k",
       sz,
-      [`  %kf = sitofp i64 %ki to double`, `  %kb = bitcast double %kf to i64`],
-      "%kf",
+      [
+        `  %kzero = icmp eq i64 %ki, 0`,
+        `  %kbits = bitcast double %k to i64`,
+        `  %kb = select i1 %kzero, i64 0, i64 %kbits`,
+      ],
+      "%k",
     ),
   ];
 }

packages/compiler/src/backend/llvm/map-number-lookup.test.ts

@@ -212,23 +212,27 @@ test("a key with an integer view tests the direct index on the integer", () => {
         integral.module,
       )![0];
     // An unsigned compare also rejects negative keys; the hashed probe takes
-    // the integer.
+    // the integer and its double.
     expect(helper).toContain("icmp ult i64 %ki, %nd");
     expect(helper).toContain(`getelementptr inbounds i32, ptr %dense, ${sz} %i`);
-    expect(helper).toContain("call ptr @sc_map_entry_f64_hashed_int(ptr %m, i64 %ki)");
+    expect(helper).toContain("call ptr @sc_map_entry_f64_hashed_int(ptr %m, i64 %ki, double %k)");
     expect(helper).toContain("br i1 %inside, label %slot, label %slow");
     expect(helper).not.toContain("fptosi");
-    // Its key word is the integer's double, which needs no NaN or -0
-    // normalization; the runtime lookup gets that double.
+    // Its key word is the bits of the integer's double, which needs no NaN
+    // normalization and is +0 for the integer 0; the runtime lookup gets
+    // that double. The probe converts nothing itself (cvtsi2sd would wait
+    // for its register's last value).
     const hashed =
-      /^define internal ptr @sc_map_entry_f64_hashed_int\(ptr %m, i64 %ki\) noinline [^]*?^}/m.exec(
+      /^define internal ptr @sc_map_entry_f64_hashed_int\(ptr %m, i64 %ki, double %k\) noinline [^]*?^}/m.exec(
         integral.module,
       )![0];
-    expect(hashed).toContain("sitofp i64 %ki to double");
+    expect(hashed).toContain("bitcast double %k to i64");
+    expect(hashed).toContain("select i1 %kzero, i64 0, i64 %kbits");
+    expect(hashed).not.toContain("sitofp");
     expect(hashed).not.toContain("fcmp");
     expect(hashed).toContain("mul i64 %x1, -4658895280553007687");
     expect(hashed).toContain(
-      "call ptr @scr_map_entry_f64(ptr %m, double %kf) memory(read) nounwind willreturn",
+      "call ptr @scr_map_entry_f64(ptr %m, double %k) memory(read) nounwind willreturn",
     );
     // A fractional key and a key without a proof keep the double lookup.
     for (const key of [{ kind: "numLit", value: 0.5, type: F64, loc } as IrExpr, ref("k", F64)]) {

packages/compiler/src/backend/llvm/speculative-integers.ts

@@ -234,6 +234,17 @@ function powerOfTwo(e: IrExpr): boolean {
 /** Typed-array element kinds whose elements are integers. */
 const INTEGER_ELEMENTS = new Set(["u8", "u8c", "i8", "u16", "i16", "u32", "i32"]);
 
+/** The values of each integer element kind. */
+const ELEMENT_RANGES: Record<string, { min: number; max: number }> = {
+  u8: { min: 0, max: 255 },
+  u8c: { min: 0, max: 255 },
+  i8: { min: -128, max: 127 },
+  u16: { min: 0, max: 65535 },
+  i16: { min: -32768, max: 32767 },
+  u32: { min: 0, max: 2 ** 32 - 1 },
+  i32: { min: -(2 ** 31), max: 2 ** 31 - 1 },
+};
+
 /** A typed-array element read of an integer kind with no effect: an
  * out-of-range read gives NaN, which fails its check. Float elements would
  * only bail. */
@@ -904,6 +915,17 @@ export function emitSafeIntegerCheck(host: LlvmEmitterContext, value: string): [
   return [integer, ok];
 }
 
+/** The condition that a bounded value's exponent is above `limit`. A
+ * value test it replaces can fail only past that limit; testing the
+ * exponent alone, which is usually the same in every iteration of a loop,
+ * keeps the value out of the test. Past the limit the write bails to the
+ * double copy, also for a value the replaced test would have passed. */
+function exponentAbove(host: LlvmEmitterContext, bound: MaskBound, limit: number): string {
+  const above = host.B.tmp();
+  host.B.line(`${above} = icmp sgt i64 ${bound.exponent}, ${limit}`);
+  return above;
+}
+
 /** A condition that is true when an i64 lies outside the safe range. */
 function unsafeInteger(host: LlvmEmitterContext, value: string): string {
   const B = host.B;
@@ -932,6 +954,29 @@ function outsideRange(host: LlvmEmitterContext, value: string, low: number, high
   return outside;
 }
 
+/** The range of an operand once its speculative evaluation continued: a
+ * literal's value, the range analysis' proof (well inside the safe
+ * integers), or for an integer element read, which bails on the NaN of an
+ * out-of-range index, the element kind's values. */
+function operandRange(
+  host: LlvmEmitterContext,
+  e: IrExpr,
+): { min: number; max: number } | undefined {
+  const range = staticRange(host, e);
+  if (range || !isSpeculativeByteRead(e) || e.kind !== "bytesIntrinsic") return range;
+  return e.receiver.type.kind === "bytes" ? ELEMENT_RANGES[e.receiver.type.elem] : undefined;
+}
+
+/** An i64 the speculative evaluation has bounded to [0, 2^(exponent +
+ * shift)): the floor-mod of a value by 2^exponent, shifted left by `shift`.
+ * `exponent` is an i64 already checked to be in [0, 52]; in a loop it is
+ * usually the same value on every iteration, so a test on it alone is
+ * loop-invariant. */
+interface MaskBound {
+  exponent: string;
+  shift: number;
+}
+
 /** A literal's value or the range analysis' proof for an expression whose
  * range lies well inside the safe integers. */
 function staticRange(
@@ -962,6 +1007,7 @@ export function emitSpeculativeInteger(
   bail: (condition: string) => void,
   reads: ReadonlyMap<IrExpr, LlValue> = new Map(),
   integral = false,
+  bounds = new Map<string, MaskBound>(),
 ): string {
   const B = host.B;
   // A value the range analysis already proves takes the ordinary integer
@@ -1000,12 +1046,13 @@ export function emitSpeculativeInteger(
       // 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 value = emitSpeculativeInteger(host, mod.value, bail, reads, integral, bounds);
         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}`);
+        bounds.set(low, { exponent: n, shift: 0 });
         return low;
       }
       // A product with 2^n shifts: it is safe iff the other factor's
@@ -1013,8 +1060,19 @@ export function emitSpeculativeInteger(
       // 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 value = emitSpeculativeInteger(host, factor, bail, reads, integral, bounds);
         const n = exponentOf(host, power, bail, reads, integral);
+        // A factor below 2^(m + s) times 2^k is below 2^(m + s + k), so
+        // safe when m + s + k <= 53.
+        const bound = bounds.get(value);
+        const k = power.kind === "numLit" ? exactPowerOfTwo(power.value) : null;
+        if (bound && k !== null) {
+          bail(exponentAbove(host, bound, 53 - bound.shift - k));
+          const product = B.tmp();
+          B.line(`${product} = shl nuw nsw i64 ${value}, ${n}`);
+          bounds.set(product, { exponent: bound.exponent, shift: bound.shift + k });
+          return product;
+        }
         const limit = B.tmp(),
           magnitude = B.tmp(),
           outside = B.tmp(),
@@ -1027,16 +1085,24 @@ export function emitSpeculativeInteger(
         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 a = emitSpeculativeInteger(host, e.left, bail, reads, integral, bounds);
+      const b = emitSpeculativeInteger(host, e.right, bail, reads, integral, bounds);
       const result = B.tmp();
       if (e.op === "+" || e.op === "-") {
         // With one operand's range known, the other is checked against the
         // narrowed bounds before the operation: the check then depends on
         // no new value, and LLVM keeps the operation itself simple.
-        const right = staticRange(host, e.right);
-        const left = e.op === "+" ? staticRange(host, e.left) : undefined;
-        if (right) {
+        const right = operandRange(host, e.right);
+        const left = e.op === "+" ? operandRange(host, e.left) : undefined;
+        // A bounded operand, in [0, 2^(m + s)), plus one in [lo, hi] with
+        // |lo|, hi <= 2^52 is safe when m + s <= 52.
+        const bound = (value: string): MaskBound | undefined =>
+          e.op === "+" ? bounds.get(value) : undefined;
+        const aBound = right ? bound(a) : undefined;
+        const bBound = !right && left ? bound(b) : undefined;
+        if (aBound) bail(exponentAbove(host, aBound, 52 - aBound.shift));
+        else if (bBound) bail(exponentAbove(host, bBound, 52 - bBound.shift));
+        else if (right) {
           // a + b with b in [lo, hi] is safe iff a in [-SAFE - lo, SAFE - hi];
           // a - b iff a in [-SAFE + hi, SAFE + lo].
           const low = e.op === "+" ? -SAFE - right.min : -SAFE + right.max;

packages/compiler/test/speculative-loop-emission.test.ts

@@ -177,13 +177,38 @@ console.log(key(new Uint8Array([1, 2, 3, 0]), 6));
   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/);
+  expect(fast).toMatch(/= shl (nuw )?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");
 });
 
+test("a masked value's shift and its sum with a byte test only the mask's exponent", 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));
+  // trim(k, bits) is below 2^bits: times 4 it is safe for bits <= 51, and
+  // plus a byte (0 to 255) for bits <= 50. Only the exponent is tested.
+  const exponent = /= icmp sgt i64 (%\S+), 51\n/.exec(fast)?.[1];
+  expect(exponent).toBeDefined();
+  expect(fast).toContain(`= icmp sgt i64 ${exponent}, 50\n`);
+  expect(fast).toMatch(/= shl nuw nsw i64 %\S+, 2\n/);
+  // No magnitude test of the masked value and no safe-range test of the sum.
+  expect(fast).not.toContain("@llvm.abs.i64");
+  expect(fast).not.toContain("18014398509481982");
+});
+
 test("a speculative loop reads its views' header fields once, before its copies", async () => {
   // A view held by a const local or by a parameter is the same view in
   // every iteration. Its data pointer is read before the branch to the fast

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/speculative-mask-bounds.ts.json

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

tests/corpus/speculative-mask-bounds.ts

@@ -0,0 +1,55 @@
+// Speculative integer loops whose writes shift and add a value masked to
+// its low bits: exponents up to 50 stay in the integer copy, and larger
+// ones (51 and 52) keep the value tests, which bail to the double copy
+// when a shift or a sum leaves the safe integers.
+function trim(value: number, bits: number): number {
+  return value - Math.floor(value / Math.pow(2, bits)) * Math.pow(2, bits);
+}
+
+function rolling(codes: Uint8Array, bits: number, scale: number): number {
+  let key = 2 ** 52 - 7;
+  let sum = 0;
+  for (let i = 0; i < codes.length; i++) {
+    key = trim(key, bits) * 4 + codes[i]!;
+    sum = (sum + (key % 1000003)) % 1000000007;
+  }
+  return sum * scale + (key % 999983);
+}
+
+function signedBytes(codes: Int8Array, bits: number): number {
+  let key = 12345;
+  let sum = 0;
+  for (let i = 0; i < codes.length; i++) {
+    key = trim(key, bits) * 2 + codes[i]!;
+    sum = (sum + (key % 65521)) % 1000000007;
+  }
+  return sum;
+}
+
+function wideWords(words: Uint32Array, bits: number): number {
+  let key = 1;
+  let sum = 0;
+  for (let i = 0; i < words.length; i++) {
+    key = trim(key, bits) * 8 + words[i]!;
+    sum = (sum + (key % 104729)) % 1000000007;
+  }
+  return sum;
+}
+
+// A read past the end gives undefined, so the sum is NaN from there on.
+function pastTheEnd(codes: Uint8Array, bits: number): number {
+  let key = 3;
+  for (let i = 0; i <= codes.length; i++) key = trim(key, bits) * 4 + codes[i]!;
+  return key;
+}
+
+const codes = new Uint8Array(4000);
+for (let i = 0; i < codes.length; i++) codes[i] = (i * 2654435761) % 256;
+const signed = new Int8Array(3000);
+for (let i = 0; i < signed.length; i++) signed[i] = ((i * 40503) % 256) - 128;
+const words = new Uint32Array(2000);
+for (let i = 0; i < words.length; i++) words[i] = (i * 2246822519) % 4294967296;
+
+for (const bits of [0, 1, 22, 34, 49, 50, 51, 52, 53])
+  console.log(bits, rolling(codes, bits, 3), signedBytes(signed, bits), wideWords(words, bits));
+console.log(pastTheEnd(codes.subarray(0, 10), 20), pastTheEnd(codes.subarray(0, 0), 50));

tests/harness/corpus-costs/speculative-mask-bounds.ts.json

@@ -0,0 +1 @@
+{ "plain": 6700, "san": 5700 }