Run paired row sums two columns per step (#2247)

packages/compiler/src/ir/row-pairs.test.ts

@@ -93,12 +93,39 @@ test("a row-sum nest runs two rows per pass behind its bounds and overlap test",
   // from the array the term reads.
   expect(JSON.stringify(go.init)).toContain('"fn":"bytes.disjoint"');
   expect(branch.else_).toBeNull();
-  const pairs = branch.then[0] as IrStmt & { kind: "for" };
+  // The four bounds, masked to the range the test proved, then the pairs.
+  expect(branch.then.slice(0, 4)).toMatchObject(
+    Array(4).fill({ kind: "varDecl", init: { op: "&", right: { value: 2 ** 25 - 1 } } }),
+  );
+  const pairs = branch.then[4] as IrStmt & { kind: "for" };
   expect(pairs.body.filter((s) => s.kind === "bytesSet")).toHaveLength(2);
-  const inner = pairs.body[2] as IrStmt & { kind: "for" };
-  // One read of u[j] for both rows, then the two sums.
-  expect(inner.body).toHaveLength(3);
+  const [, , inner, last] = pairs.body as [
+    IrStmt,
+    IrStmt,
+    IrStmt & { kind: "for" },
+    IrStmt & { kind: "if" },
+  ];
+  const column = inner.init as IrStmt & { kind: "varDecl" };
+  // Two columns per step: for each, one read of u[j] for both rows, then
+  // the two sums; the second column reads `column + 1`.
+  expect(inner.body).toHaveLength(6);
   expect(inner.body[0]).toMatchObject({ kind: "varDecl", init: { kind: "bytesIntrinsic" } });
+  expect(inner.body[3]).toMatchObject({
+    kind: "varDecl",
+    init: {
+      kind: "bytesIntrinsic",
+      args: [{ op: "+", left: { localId: column.localId }, right: { value: 1 } }],
+    },
+  });
+  expect(inner.cond).toMatchObject({ op: "<", right: { op: "-", right: { value: 1 } } });
+  expect(inner.update).toMatchObject({ value: { op: "+", right: { value: 2 } } });
+  // The last column of an odd count, `to - 1`.
+  expect(last.then).toHaveLength(3);
+  expect(last.then[0]).toMatchObject({
+    kind: "varDecl",
+    init: { args: [{ op: "-", right: { value: 1 } }] },
+  });
+  expect(last.else_).toBeNull();
   // The second row's term reads `row + 1`.
   const row = pairs.init?.kind === "varDecl" ? pairs.init.localId : "";
   expect((inner.body[2] as IrStmt & { kind: "assign" }).value).toMatchObject({
@@ -130,7 +157,7 @@ test("a row-sum nest runs two rows per pass behind its bounds and overlap test",
   expect(top.body).toHaveLength(4);
   expect(
     result.functions[0]!.locals.map((l) => l.id).filter((id) => id.startsWith("%pair.")),
-  ).toHaveLength(7);
+  ).toHaveLength(13);
 });
 
 test("nests whose term calls, writes or reads the sum keep their loops", () => {

packages/compiler/src/ir/row-pairs.ts

@@ -14,10 +14,17 @@ import {
   mapExprChildren,
   transformStmtList,
 } from "./traverse.js";
+import { PRODUCT_COUNTER_BOUND } from "./integer-loops.js";
 
 /** Inner-loop expressions above this many nodes keep the loop as written. */
 const MAX_TERM_NODES = 64;
 
+/** The largest bound a paired nest runs with. Counters below
+ * PRODUCT_COUNTER_BOUND keep products of their sums exact, so the nest's
+ * counted loops take integer counters without a run-time guard (and its
+ * copy of the body). */
+const SMALL_BOUND = PRODUCT_COUNTER_BOUND - 1;
+
 /** The number operators a term may use: they run no other code and cannot
  * throw. */
 const PURE_BINARY = new Set<string>([
@@ -202,7 +209,7 @@ function copy(
 }
 
 const bin = (
-  op: "+" | "-" | "&" | "|",
+  op: "+" | "-" | "/" | "&" | "|",
   left: IrExpr,
   right: IrExpr,
   loc: IrExpr["loc"],
@@ -216,7 +223,8 @@ const bin = (
 });
 
 /** Runs the row sums of `match` two rows per pass when the test holds and
- * as written otherwise. */
+ * as written otherwise. Each pass takes two columns per step, then the
+ * last column when their count is odd. */
 function pairRows(
   s: IrStmt & { kind: "for" },
   match: RowSum,
@@ -231,42 +239,93 @@ function pairRows(
     sum = fresh("sum"),
     next = fresh("next"),
     column = fresh("column");
-  const atColumn = (r: IrExpr & { kind: "varRef" }): IrExpr =>
-    r.localId === match.column ? { ...r, localId: column } : { ...r };
   // Element reads that do not depend on the row are read once per column
   // and shared by both rows.
-  const shared = new Map<IrExpr, string>();
-  const elements: IrStmt[] = [];
+  const sharedReads: IrExpr[] = [];
   const collect = (e: IrExpr): boolean => {
     if (e.kind === "bytesIntrinsic" && e.method === "get" && !readsLocal(e, match.row)) {
-      const id = fresh("element");
-      shared.set(e, id);
-      elements.push({ kind: "varDecl", localId: id, init: copy(e, atColumn), loc });
+      sharedReads.push(e);
       return true;
     }
     return everyExprChild(e, collect, () => true);
   };
   collect(match.term);
-  // The second row's term reads `row + 1`, so its counter sums stay visible
-  // to the counted-loop analysis.
-  const termAt = (rowExpr: () => IrExpr): IrExpr =>
-    copy(
-      match.term,
-      (r) => (r.localId === match.row ? rowExpr() : atColumn(r)),
-      (node) => {
-        const id = shared.get(node);
-        return id === undefined ? null : { kind: "varRef", localId: id, type: node.type, loc };
+  /** Both rows' additions at the column `columnAt()`: its shared element
+   * reads, then each row's sum. The second row's term reads `row + 1`, and
+   * the second column's `column + 1`, so their counter sums stay visible to
+   * the counted-loop analysis. */
+  const step = (columnAt: () => IrExpr): IrStmt[] => {
+    const at = (r: IrExpr & { kind: "varRef" }): IrExpr =>
+      r.localId === match.column ? columnAt() : { ...r };
+    const shared = new Map<IrExpr, string>();
+    const elements: IrStmt[] = sharedReads.map((read) => {
+      const id = fresh("element");
+      shared.set(read, id);
+      return { kind: "varDecl", localId: id, init: copy(read, at), loc };
+    });
+    const termAt = (rowExpr: () => IrExpr): IrExpr =>
+      copy(
+        match.term,
+        (r) => (r.localId === match.row ? rowExpr() : at(r)),
+        (node) => {
+          const id = shared.get(node);
+          return id === undefined ? null : { kind: "varRef", localId: id, type: node.type, loc };
+        },
+      );
+    return [
+      ...elements,
+      {
+        kind: "assign",
+        localId: sum,
+        value: bin(
+          "+",
+          ref(sum),
+          termAt(() => ref(row)),
+          loc,
+        ),
+        loc,
       },
-    );
+      {
+        kind: "assign",
+        localId: next,
+        value: bin(
+          "+",
+          ref(next),
+          termAt(() => bin("+", ref(row), num(1), loc)),
+          loc,
+        ),
+        loc,
+      },
+    ];
+  };
   const innerCond = match.inner.cond as IrExpr & { kind: "bin" };
+  const columnStart = (match.inner.init as IrStmt & { kind: "varDecl" }).init!;
+  // The bounds, as the paired nest reads them (see `small` below).
+  const low = fresh("low"),
+    high = fresh("high"),
+    first = fresh("first"),
+    last = fresh("last");
+  const bounds: IrStmt[] = (
+    [
+      [low, match.start],
+      [high, match.end],
+      [first, columnStart],
+      [last, innerCond.right],
+    ] as const
+  ).map(([id, bound]) => ({
+    kind: "varDecl",
+    localId: id,
+    init: bin("&", again(bound), num(SMALL_BOUND), loc),
+    loc,
+  }));
   const pairs: IrStmt = {
     kind: "for",
-    init: { kind: "varDecl", localId: row, init: again(match.start), loc },
+    init: { kind: "varDecl", localId: row, init: ref(low), loc },
     cond: {
       kind: "bin",
       op: "<",
       left: ref(row),
-      right: bin("-", again(match.end), num(1), loc),
+      right: bin("-", ref(high), num(1), loc),
       type: BOOL,
       loc,
     },
@@ -274,41 +333,40 @@ function pairRows(
     body: [
       { kind: "varDecl", localId: sum, init: { ...match.init }, loc },
       { kind: "varDecl", localId: next, init: { ...match.init }, loc },
+      // Both column bounds are int32, so `column < to - 1` is
+      // `column + 1 < to`. Each row adds its columns in order.
       {
         kind: "for",
-        init: {
-          kind: "varDecl",
-          localId: column,
-          init: again((match.inner.init as IrStmt & { kind: "varDecl" }).init!),
+        init: { kind: "varDecl", localId: column, init: ref(first), loc },
+        cond: {
+          ...innerCond,
+          left: ref(column),
+          right: bin("-", ref(last), num(1), loc),
+        },
+        update: { kind: "assign", localId: column, value: bin("+", ref(column), num(2), loc), loc },
+        body: [...step(() => ref(column)), ...step(() => bin("+", ref(column), num(1), loc))],
+        loc,
+      },
+      // An odd column count leaves the last column, `to - 1`.
+      {
+        kind: "if",
+        cond: {
+          kind: "logical",
+          op: "&&",
+          left: { ...innerCond, left: ref(first), right: ref(last) },
+          right: {
+            kind: "bin",
+            op: "===",
+            left: bin("&", bin("-", ref(last), ref(first), loc), num(1), loc),
+            right: num(1),
+            type: BOOL,
+            loc,
+          },
+          type: BOOL,
           loc,
         },
-        cond: { ...innerCond, left: ref(column), right: again(innerCond.right) },
-        update: { kind: "assign", localId: column, value: bin("+", ref(column), num(1), loc), loc },
-        body: [
-          ...elements,
-          {
-            kind: "assign",
-            localId: sum,
-            value: bin(
-              "+",
-              ref(sum),
-              termAt(() => ref(row)),
-              loc,
-            ),
-            loc,
-          },
-          {
-            kind: "assign",
-            localId: next,
-            value: bin(
-              "+",
-              ref(next),
-              termAt(() => bin("+", ref(row), num(1), loc)),
-              loc,
-            ),
-            loc,
-          },
-        ],
+        then: step(() => bin("-", ref(last), num(1), loc)),
+        else_: null,
         loc,
       },
       { kind: "bytesSet", arr: { ...match.out }, index: ref(row), value: ref(sum), loc },
@@ -322,9 +380,13 @@ function pairRows(
     ],
     loc,
   };
-  // Both bounds are int32, so `row < end - 1` is `row + 1 < end`, and the
-  // nest has at least two rows. The overlap tests run only then, so they never
-  // read an array the program assigns only when there are rows. Rows are
+  // Every bound is an integer in [0, SMALL_BOUND] and not -0, so `row <
+  // end - 1` is `row + 1 < end` (and the same for the columns), and the
+  // masks in `bounds` leave the bounds unchanged while proving both
+  // counters' ranges.
+  // The nest has at least two rows. The column bounds' tests and the
+  // overlap tests run only then, so they never read an array the program
+  // assigns only when there are rows. Rows are
   // independent when the written array shares no byte with an array the
   // term reads; an array that is not there answers false.
   const int32 = (e: IrExpr): IrExpr => ({
@@ -335,9 +397,24 @@ function pairRows(
     type: BOOL,
     loc,
   });
+  const compare = (op: "<=" | ">=", e: IrExpr, value: number): IrExpr => ({
+    kind: "bin",
+    op,
+    left: again(e),
+    right: num(value),
+    type: BOOL,
+    loc,
+  });
+  // `1 / e > 0` holds for +0 and every positive bound, but not for -0:
+  // masking would turn a -0 start into +0, which a term can tell apart.
+  const small = (e: IrExpr): IrExpr[] => [
+    int32(e),
+    { kind: "bin", op: ">", left: bin("/", num(1), again(e), loc), right: num(0), type: BOOL, loc },
+    compare("<=", e, SMALL_BOUND),
+  ];
   const tests: IrExpr[] = [
-    int32(match.start),
-    int32(match.end),
+    ...small(match.start),
+    ...small(match.end),
     {
       kind: "bin",
       op: "<",
@@ -346,6 +423,8 @@ function pairRows(
       type: BOOL,
       loc,
     },
+    ...small(columnStart),
+    ...small(innerCond.right),
     ...match.reads.map((read): IrExpr => ({
       kind: "libCall",
       fn: "bytes.disjoint",
@@ -400,7 +479,7 @@ function pairRows(
         },
         loc,
       },
-      { kind: "if", cond: { ...paired }, then: [pairs], else_: null, loc },
+      { kind: "if", cond: { ...paired }, then: [...bounds, pairs], else_: null, loc },
       rest,
     ],
     loc,
@@ -419,7 +498,15 @@ function pairRows(
  * store into shares no byte with any array the term reads, so no row reads
  * another row's result. An odd row count runs its last row as written.
  * Otherwise the nest runs as written. Two independent sums take turns in
- * the inner loop, so neither waits on the other's addition. */
+ * the inner loop, so neither waits on the other's addition.
+ *
+ * The inner loop takes two columns per step, each row adding the first
+ * column's term and then the second's, and the last column on its own
+ * when the count is odd, so the additions keep their order. Both column
+ * bounds must also be int32 (so `column < to - 1` is `column + 1 < to`).
+ * Neighboring rows and columns share counter sums (`row + 1 + column` and
+ * `row + (column + 1)`), and the step's divisions have no dependence
+ * between them. */
 export function pairRowSums(mod: IrModule): IrModule {
   let changed = false;
   const functions = mod.functions.map((fn): IrFunction => {

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/8355-row-pair-columns.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/8355-row-pair-columns.ts"
+ ],
+ "diags": []
+}

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/8356-shared-view-loop-versions.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/8356-shared-view-loop-versions.ts"
+ ],
+ "diags": []
+}

tests/corpus/8355-row-pair-columns.ts

@@ -0,0 +1,94 @@
+// Row sums whose inner loop runs two columns per step: even, odd, zero, one
+// and two columns, column ranges that start above zero or at or past their
+// end, fractional, negative, NaN and very large column bounds, a column
+// bound read from an array's length, terms that read two elements per
+// column, row bounds at and past 2^25, and row and column starts of -0
+// with terms that tell -0 from +0. Every sum adds its columns in order,
+// so each has the bits of the loop as written.
+function rows(u: Float64Array, v: Float64Array, start: number, end: number, from: number, to: number): void {
+  for (let i = start; i < end; i++) {
+    let t = 0;
+    for (let j = from; j < to; j++) t += u[j]! / ((((i + j) * (i + j + 1)) >>> 1) + i + 1);
+    v[i] = t;
+  }
+}
+
+function pairs(u: Float64Array, w: Float32Array, v: Float64Array, end: number): void {
+  for (let i = 0; i < end; i++) {
+    let t = -0.25;
+    for (let j = 0; j < u.length; j++) t += u[j]! * w[j + 1]! - (i + 1) / (j + 3);
+    v[i] = t;
+  }
+}
+
+// 1 / (i * 4 - j) is -Infinity at a -0 row and +0 column, and
+// 1 / (j - i * 4) at a -0 column and +0 row: the first pass sees the sign.
+function signs(v: Float64Array, start: number, end: number, from: number, to: number): void {
+  for (let i = start; i < end; i++) {
+    let t = 0;
+    for (let j = from; j < to; j++) t += 1 / (i * 4 - j) + 1 / (j - i * 4) + 1 / (i * 4 + j + 1);
+    v[i] = t;
+  }
+}
+
+const show = (label: string, v: Float64Array): void => {
+  const parts: string[] = [];
+  for (const x of v) parts.push(Object.is(x, -0) ? "-0" : x.toPrecision(17));
+  console.log(label, parts.join(" "));
+};
+
+const n = 9;
+const u = new Float64Array(n).map((_, k) => 1 / (k + 1) + k);
+for (const [from, to] of [
+  [0, n],
+  [0, n - 1],
+  [0, 0],
+  [0, 1],
+  [0, 2],
+  [3, n],
+  [4, n],
+  [5, 5],
+  [7, 2],
+  [0.5, 6],
+  [0, 6.5],
+  [-2, 5],
+  [-0, 4],
+  [NaN, 4],
+  [0, NaN],
+  [0, n + 3],
+  [2 ** 25, 2 ** 25 + 3],
+]) {
+  const v = new Float64Array(6);
+  rows(u, v, 0, 6, from!, to!);
+  show(`columns ${from} ${to}`, v);
+}
+for (const [start, end] of [
+  [2 ** 25 - 3, 2 ** 25 + 1],
+  [2 ** 25, 2 ** 25 + 2],
+  [1, 6],
+]) {
+  const v = new Float64Array(8);
+  rows(u, v, start!, end!, 0, n);
+  show(`rows ${start} ${end}`, v);
+}
+
+const w = new Float32Array(n + 1).map((_, k) => (k % 4) - 1.5);
+for (const end of [0, 1, 4, 5]) {
+  const v = new Float64Array(5);
+  pairs(u, w, v, end);
+  show(`two reads ${end}`, v);
+}
+const short = new Float64Array(4).map((_, k) => k * 2.5);
+const out = new Float64Array(5);
+pairs(short, w, out, 5);
+show("two reads short", out);
+for (const [start, from] of [
+  [0, 0],
+  [-0, 0],
+  [0, -0],
+  [-0, -0],
+]) {
+  const v = new Float64Array(4);
+  signs(v, start!, 4, from!, 5);
+  show(`signs ${Object.is(start, -0) ? "-0" : start} ${Object.is(from, -0) ? "-0" : from}`, v);
+}

tests/corpus/8356-shared-view-loop-versions.ts

@@ -0,0 +1,93 @@
+import { Worker, isMainThread, parentPort, workerData } from "node:worker_threads";
+
+// Loops over typed-array views of shared memory in a worker program: row
+// sums and element-wise passes whose accesses are in bounds, the same loops
+// running past a view's end, a loop that mixes a shared view with one the
+// thread owns, and views of one SharedArrayBuffer that overlap, written in
+// an order that shows whether each read sees the previous write. The worker
+// writes its results into shared memory, which the main thread reads after
+// the worker reports.
+
+interface Job {
+  sab: SharedArrayBuffer;
+  n: number;
+}
+
+function rowSums(u: Float64Array, v: Float64Array, start: number, end: number, n: number): void {
+  for (let i = start; i < end; i++) {
+    let t = 0;
+    for (let j = 0; j < n; j++) t = t + u[j]! / (((i + j) * (i + j + 1)) / 2 + i + 1);
+    v[i] = t;
+  }
+}
+
+function scaled(src: Float64Array, dst: Float64Array, k: number, count: number): void {
+  for (let i = 0; i < count; i++) dst[i] = src[i]! * k + 1;
+}
+
+function mixed(shared: Int32Array, owned: Int32Array, count: number): number {
+  let sum = 0;
+  for (let i = 0; i < count; i++) {
+    owned[i] = shared[i]! * 3;
+    shared[i] = owned[i]! - i;
+    sum = (sum + shared[i]! + owned[i]!) | 0;
+  }
+  return sum;
+}
+
+function shift(a: Float64Array, b: Float64Array, count: number): void {
+  for (let i = 0; i < count; i++) a[i + 1] = b[i]! + 1;
+}
+
+function checksum(values: Float64Array): string {
+  let h = 0;
+  for (let i = 0; i < values.length; i++) h = (Math.imul(h, 31) + Math.round(values[i]! * 1e6)) >>> 0;
+  return h.toString(16);
+}
+
+function work({ sab, n }: Job): string[] {
+  const out: string[] = [];
+  const u = new Float64Array(sab, 0, n);
+  const v = new Float64Array(sab, n * 8, n);
+  for (let i = 0; i < n; i++) u[i] = 1 + (i % 5) / 4;
+  rowSums(u, v, 0, n, n);
+  out.push(`rows ${checksum(v)}`);
+  // Past the end of u: the reads beyond it are undefined, so the sums are NaN.
+  rowSums(u, v, 0, 4, n + 3);
+  out.push(`past ${v[0]} ${v[3]} ${checksum(v.subarray(4))}`);
+  scaled(u, v, 0.5, n);
+  out.push(`scaled ${checksum(v)}`);
+  scaled(u, v, 2, n + 2);
+  out.push(`scaled past ${checksum(v)}`);
+  const ints = new Int32Array(sab, 2 * n * 8, 32);
+  for (let i = 0; i < 32; i++) ints[i] = i * 7 - 50;
+  const own = new Int32Array(32);
+  out.push(`mixed ${mixed(ints, own, 32)} ${mixed(ints, own, 34)} ${ints[5]} ${own[31]}`);
+  // a[i + 1] lands on b[i + 1] when b starts one element after a.
+  const a = new Float64Array(sab, 0, 24);
+  const b = new Float64Array(sab, 8, 24);
+  for (let i = 0; i < 25; i++) u[i] = i;
+  shift(a, b, 20);
+  out.push(`overlap ${Array.from(a.subarray(0, 8)).join(",")}`);
+  shift(b, a, 20);
+  out.push(`overlap back ${Array.from(b.subarray(0, 8)).join(",")}`);
+  return out;
+}
+
+if (isMainThread) {
+  const n = 40;
+  const sab = new SharedArrayBuffer(2 * n * 8 + 32 * 4);
+  const worker = new Worker(new URL(import.meta.url), { workerData: { sab, n } satisfies Job });
+  worker.on("message", (lines: string[]) => {
+    for (const line of lines) console.log("worker", line);
+    const v = new Float64Array(sab, n * 8, n);
+    const ints = new Int32Array(sab, 2 * n * 8, 32);
+    console.log("main sees", checksum(v), ints[0], ints[31]);
+    // The same passes on the main thread, over a fresh buffer.
+    const again = new SharedArrayBuffer(2 * n * 8 + 32 * 4);
+    for (const line of work({ sab: again, n })) console.log("main", line);
+  });
+  worker.on("exit", (code: number) => console.log("exit", code));
+} else {
+  parentPort!.postMessage(work(workerData as Job));
+}