Compute products of typed-array reads in 32 bits under bitwise operators (#2289)

packages/compiler/src/backend/llvm/expr-bytes.ts

@@ -303,7 +303,7 @@ function emitScopedBytesGet(
   const signed = elem === "i8" || elem === "i16" || elem === "i32";
   B.line(`${number} = ${signed ? "sitofp" : "uitofp"} i32 ${u32} to double`);
   B.line(`${out} = select i1 ${valid}, double ${number}, double ${f64Lit(NaN)}`);
-  return { name: out, type: F64, uint32: u32 };
+  return { name: out, type: F64, uint32: u32, ordered: valid };
 }
 
 function emitBytesLoad(

packages/compiler/src/backend/llvm/expr-context.ts

@@ -21,6 +21,10 @@ export interface LlValue {
    * its absolute value is below `below` (a zero may be -0). Index consumers
    * use it in place of a round trip (integer-values.ts `indexInteger`). */
   wideInteger?: WideInteger;
+  /** An i1 SSA value (or `true`) that is true exactly when this number is
+   * not NaN: set by NaN-coded integer element reads and the wrapping sums
+   * of them, so ToInt32 consumers need no floating compare. */
+  ordered?: string;
   slot?: boolean;
   /** A sequence local held in a capture box; release the box itself. */
   boxed?: boolean;

packages/compiler/src/backend/llvm/expr-primitives.test.ts

@@ -278,6 +278,18 @@ test("sums that only feed a bitwise conversion of int32 operands add in 32 bits"
     returning("float_read_add", bin("|", bin("+", length, read("floats")), zero)),
     returning("product", bin("|", bin("*", length, length), zero)),
     returning("number_add", bin("+", length, read("bytes"))),
+    returning(
+      "read_products",
+      bin(
+        ">>>",
+        bin(
+          "+",
+          bin("*", read("bytes"), { kind: "numLit", value: 77, type: F64, loc }),
+          bin("*", read("bytes"), { kind: "numLit", value: 150, type: F64, loc }),
+        ),
+        zero,
+      ),
+    ),
   ];
   const llvm = emitLlvmModule(mod);
   const body = (name: string): string =>
@@ -292,8 +304,10 @@ test("sums that only feed a bitwise conversion of int32 operands add in 32 bits"
     expect(body(name), name).not.toContain("uint32.coerce");
   }
   // A NaN-coded read converts to 0 through the sum: the result selects 0
-  // unless every NaN-coded operand is ordered.
-  expect(body("read_add")).toMatch(/ = fcmp ord double .*, 0\.0$/m);
+  // unless every NaN-coded operand is ordered. A read's own validity bit
+  // tells, so the number is not compared.
+  expect(body("read_add")).toMatch(/ = phi i1 \[ true, /);
+  expect(body("read_add")).not.toMatch(/fcmp ord/);
   expect(body("read_add")).toMatch(/ = add i32 /);
   expect(body("read_add")).toMatch(/ = select i1 .*, i32 .*, i32 0$/m);
   expect(body("read_chain")).toMatch(/ = and i1 /);
@@ -309,6 +323,12 @@ test("sums that only feed a bitwise conversion of int32 operands add in 32 bits"
     expect(body(name), name).not.toMatch(/ = (add|mul) i32 /);
   }
   expect(body("number_add")).not.toMatch(/ = add i32 |fcmp ord/);
+  // Products of NaN-coded reads and constants are exact integers or NaN, so
+  // their sum adds in 32 bits too, selected by the reads' validity bits.
+  expect(body("read_products")).toMatch(/ = mul i32 /);
+  expect(body("read_products")).toMatch(/ = add i32 /);
+  expect(body("read_products")).not.toContain("uint32.coerce");
+  expect(body("read_products")).not.toMatch(/fcmp ord/);
 });
 
 test.each([32, 64] as const)(

packages/compiler/src/backend/llvm/expr-primitives.ts

@@ -367,8 +367,8 @@ export function emitOperatorExpr(
               `${result} = ${e.op === "+" ? "add" : e.op === "-" ? "sub" : "mul"} nsw i64 ${left}, ${right}`,
             );
             const number = integerNumber(host, result, range, e.type);
-            const uint32 = emitWrappingSum(host, e, l, r, li, ri);
-            return uint32 === undefined ? number : { ...number, uint32 };
+            const wrapped = emitWrappingSum(host, e, l, r, li, ri);
+            return wrapped === undefined ? number : { ...number, ...wrapped };
           }
           if (e.op === "%" && li.range.min >= 0 && ri.range.min > 0) {
             const range = { min: 0, max: Math.min(li.range.max, ri.range.max - 1) };
@@ -411,12 +411,12 @@ export function emitOperatorExpr(
           return { name: emitPowerOfTwoQuotient(host, l.name, r.name), type: e.type };
         if (arith[e.op] !== undefined) {
           B.line(`${t} = ${arith[e.op]} double ${l.name}, ${r.name}`);
-          const uint32 = emitWrappingSum(host, e, l, r, li, ri);
+          const wrapped = emitWrappingSum(host, e, l, r, li, ri);
           // The exact integer beside the double, for index consumers.
           const wideInteger = wideArithmetic(host, e.op, l, li, r, ri);
-          if (uint32 !== undefined || wideInteger) {
+          if (wrapped !== undefined || wideInteger) {
             const value: LlValue = { name: t, type: e.type };
-            if (uint32 !== undefined) value.uint32 = uint32;
+            if (wrapped !== undefined) Object.assign(value, wrapped);
             if (wideInteger) value.wideInteger = wideInteger;
             return value;
           }

packages/compiler/src/backend/llvm/wrapping-sums.ts

@@ -22,8 +22,9 @@ export function markWrappingOperands(host: LlvmEmitterContext, ...operands: IrEx
 /** The values an f64 expression can have when each one is an exact integer
  * (never -0) or NaN: proven ranges, safe integer literals, NaN-coded reads
  * of integer typed-array elements and of Int32[] elements, `charCodeAt`,
- * and sums and differences of those whose every integer result is still
- * exact in a double. */
+ * and sums, differences and products of those whose every integer result
+ * is a safe integer, which the double holds exactly. A product that can be
+ * -0 (a zero times a negative) has no range. */
 export function integerOrNaNRange(host: LlvmEmitterContext, e: IrExpr): IntegerRange | null {
   if (e.type.kind !== "f64") return null;
   const range = host.integerRanges.get(e);
@@ -51,7 +52,7 @@ export function integerOrNaNRange(host: LlvmEmitterContext, e: IrExpr): IntegerR
   // A UTF-16 code unit, or NaN past either end.
   if (e.kind === "strIntrinsic" && e.method === "charCodeAt") return { min: 0, max: 65535 };
   if (e.kind === "varRef") return host.integerOrNaNLocals.get(e.localId) ?? null;
-  if (e.kind === "bin" && (e.op === "+" || e.op === "-")) {
+  if (e.kind === "bin" && (e.op === "+" || e.op === "-" || e.op === "*")) {
     const left = integerOrNaNRange(host, e.left);
     const right = left && integerOrNaNRange(host, e.right);
     return right ? integerArithmeticRange(e.op, left!, right) : null;
@@ -113,9 +114,14 @@ function wrappingOperand(
     return { low, valid: null };
   }
   // The number is an exact integer or NaN, and its ToUint32 view is the
-  // integer's low 32 bits or 0. An ordered compare tells the two apart.
-  const valid = B.tmp();
-  B.line(`${valid} = fcmp ord double ${value.name}, 0.0`);
+  // integer's low 32 bits or 0. Its ordered bit, or an ordered compare,
+  // tells the two apart.
+  let valid: string | null = value.ordered ?? null;
+  if (valid === "true") valid = null;
+  else if (valid === null) {
+    valid = B.tmp();
+    B.line(`${valid} = fcmp ord double ${value.name}, 0.0`);
+  }
   if (value.uint32 !== undefined) return { low: value.uint32, valid };
   // Without a view, convert once. The conversion of NaN is poison, but the
   // result only selects the sum when every operand is ordered.
@@ -137,8 +143,9 @@ function wrappingOperand(
  * low halves is the converted exact result; a NaN operand makes the result
  * NaN, which converts to 0. A product is exact when the operands' magnitude
  * bounds multiply to at most 2^53; it may be -0 (`0 * -1`), which converts
- * to 0 like the integer 0. Returns undefined, emitting nothing, for any
- * other expression. */
+ * to 0 like the integer 0. Returns the converted result with its ordered
+ * bit (`true` when no operand can be NaN), or undefined, emitting nothing,
+ * for any other expression. */
 export function emitWrappingSum(
   host: LlvmEmitterContext,
   e: IrExpr & { kind: "bin" },
@@ -146,7 +153,7 @@ export function emitWrappingSum(
   r: LlValue,
   li: LlInteger | null,
   ri: LlInteger | null,
-): string | undefined {
+): { uint32: string; ordered: string } | undefined {
   if ((e.op !== "+" && e.op !== "-" && e.op !== "*") || !host.wrappingSums.has(e)) return undefined;
   const leftRange = li ? li.range : integerOrNaNRange(host, e.left);
   const rightRange = leftRange && (ri ? ri.range : integerOrNaNRange(host, e.right));
@@ -165,14 +172,14 @@ export function emitWrappingSum(
     `${sum} = ${e.op === "+" ? "add" : e.op === "-" ? "sub" : "mul"} i32 ${left.low}, ${right.low}`,
   );
   let valid = left.valid ?? right.valid;
-  if (valid === null) return sum;
+  if (valid === null) return { uint32: sum, ordered: "true" };
   if (left.valid && right.valid) {
     valid = B.tmp();
     B.line(`${valid} = and i1 ${left.valid}, ${right.valid}`);
   }
   const out = B.tmp();
   B.line(`${out} = select i1 ${valid}, i32 ${sum}, i32 0`);
-  return out;
+  return { uint32: out, ordered: valid };
 }
 
 /** Whether every product of a value in `left` and one in `right` is an

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/6311-typed-array-integer-products.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/6311-typed-array-integer-products.ts"
+ ],
+ "diags": []
+}

tests/corpus/6311-typed-array-integer-products.ts

@@ -0,0 +1,83 @@
+// Products and sums of typed-array element reads that only reach bitwise
+// operators compute in 32-bit integers. Results must match Node for reads
+// past the end (NaN, which converts to 0), negative factors, products that
+// can be -0, sums that wrap past int32, and products near 2^53.
+
+const rgba = new Uint8Array([10, 20, 30, 255, 200, 100, 50, 0, 255, 255, 255, 128, 0, 0, 0, 1]);
+const pixels = rgba.length / 4;
+const gray = new Uint8Array(pixels);
+for (let i = 0; i < gray.length; i++) {
+  const j = i * 4;
+  gray[i] = (rgba[j]! * 77 + rgba[j + 1]! * 150 + rgba[j + 2]! * 29) >> 8;
+}
+console.log("gray", gray.join(","));
+
+// Reads past the end are NaN: the sum is NaN and converts to 0.
+const tail: number[] = [];
+for (let i = pixels - 2; i < pixels + 2; i++) {
+  const j = i * 4;
+  tail.push((rgba[j]! * 77 + rgba[j + 1]! * 150 + rgba[j + 2]! * 29) >> 8);
+  tail.push((rgba[j + 3]! * 3 - rgba[j]!) | 0);
+}
+console.log("tail", tail.join(","));
+
+// Signed elements and negative factors, including products that are -0.
+const signed = new Int8Array([0, -1, 1, -128, 127, 0, 5, -5]);
+const mixed: number[] = [];
+for (let i = 0; i < signed.length; i++) {
+  mixed.push((signed[i]! * -3) | 0);
+  mixed.push((signed[i]! * signed[(i + 1) % signed.length]! + signed[i]! * 1000) >> 1);
+  mixed.push(Object.is(signed[i]! * -1, -0) ? 1 : 0);
+}
+console.log("signed", mixed.join(","));
+
+// Unsigned 16-bit products whose sum passes int32 and wraps.
+const wide = new Uint16Array([65535, 65535, 1, 40000, 50000, 12345]);
+const wrapped: number[] = [];
+for (let i = 0; i + 1 < wide.length; i++) {
+  wrapped.push((wide[i]! * wide[i + 1]! + wide[i]! * wide[i]!) | 0);
+  wrapped.push((wide[i]! * wide[i + 1]!) >>> 0);
+  wrapped.push((wide[i]! * 65537 - wide[i + 1]!) ^ 0x5a5a);
+}
+console.log("uint16", wrapped.join(","));
+
+// Unsigned 32-bit elements times constants just below and above 2^53.
+const big = new Uint32Array([0xffffffff, 0x80000000, 3, 0]);
+const near: number[] = [];
+for (let i = 0; i < big.length; i++) {
+  near.push((big[i]! * 2097151) | 0);
+  near.push((big[i]! * 2097153) | 0);
+  near.push((big[i]! * 2097151 + big[(i + 1) % big.length]!) >>> 0);
+}
+console.log("uint32", near.join(","));
+
+// A float element is never an exact integer.
+const floats = new Float64Array([1.5, -2.25, 3]);
+const fractional: number[] = [];
+for (let i = 0; i < floats.length; i++) fractional.push((floats[i]! * 3 + floats[i]! * 2) | 0);
+console.log("float", fractional.join(","));
+
+// Clamped bytes and a sum used as a number keep their double values.
+const clamped = new Uint8ClampedArray([0, 128, 255]);
+let total = 0;
+for (let i = 0; i < clamped.length; i++) {
+  total += clamped[i]! * 77 + clamped[i]! * 0.5;
+  total += (clamped[i]! * 77 + clamped[(i + 1) % 3]! * 150) >> 8;
+}
+console.log("clamped", total);
+
+// A larger image, so the loop runs its vectorized body.
+const image = new Uint8Array(4 * 1000);
+for (let i = 0; i < image.length; i++) image[i] = (i * 37 + (i >> 3)) & 255;
+const luma = new Uint8Array(1000);
+for (let i = 0; i < luma.length; i++) {
+  const j = i * 4;
+  luma[i] = (image[j]! * 77 + image[j + 1]! * 150 + image[j + 2]! * 29) >> 8;
+}
+let sum = 0;
+let mix = 0;
+for (let i = 0; i < luma.length; i++) {
+  sum += luma[i]!;
+  mix = (mix * 31 + luma[i]!) | 0;
+}
+console.log("luma", sum, mix);