Convert sums built on unsigned shifts from 32 bits

An integer `+`, `-` or `*` whose operands and result are proven within
the uint32 range, and whose operand is a `>>>` result or arithmetic on
one, such as `((i + j) * (i + j + 1) >>> 1) + i + 1`, is now 32-bit
arithmetic without wrap, and its double converts from the 32 bits
(`uitofp i32`). Vectorized, that conversion takes three instructions on
x86-64 without AVX-512, where the 64-bit sum's unsigned conversion took
six. Sums of counters and literals, such as `i + 1`, keep the 64-bit
form, whose conversion LLVM folds with their uses as indexes.

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

@@ -26,6 +26,7 @@ import {
   int32Number,
   integerNumber,
   narrowInteger,
+  uint32Number,
   wideArithmetic,
   widenInteger,
   type LlInteger,
@@ -129,6 +130,16 @@ function emitIntegerOrNaNCompare(
   return result;
 }
 
+/** A `>>>` result, or a sum, difference or product with one as an operand. */
+function unsignedShiftArithmetic(e: IrExpr): boolean {
+  if (e.kind !== "bin") return false;
+  if (e.op === ">>>") return true;
+  return (
+    (e.op === "+" || e.op === "-" || e.op === "*") &&
+    (unsignedShiftArithmetic(e.left) || unsignedShiftArithmetic(e.right))
+  );
+}
+
 export function emitLiteralExpr(
   host: LlvmEmitterContext,
   e: ExprOf<"numLit" | "boolLit" | "strLit" | "moduleNsRef" | "unitLit" | "varRef">,
@@ -383,6 +394,28 @@ export function emitOperatorExpr(
             );
             return int32Number(host, result, range, e.type);
           }
+          if (
+            range &&
+            withinUint32(range) &&
+            withinUint32(li.range) &&
+            withinUint32(ri.range) &&
+            (unsignedShiftArithmetic(e.left) || unsignedShiftArithmetic(e.right))
+          ) {
+            // A uint32 result built on a `>>>` result, such as
+            // `(x >>> 1) + i`: 32-bit arithmetic that cannot wrap. Its double
+            // converts from the 32 bits, which vectorizes to three
+            // instructions on x86-64 without AVX-512, where the unsigned
+            // 64-bit conversion takes six. Sums of counters and literals,
+            // such as `i + 1`, keep the 64-bit form, whose conversion LLVM
+            // folds with their uses as indexes.
+            const result = B.tmp();
+            B.line(
+              `${result} = ${e.op === "+" ? "add" : e.op === "-" ? "sub" : "mul"} nuw i32 ${narrowInteger(host, li)}, ${narrowInteger(host, ri)}`,
+            );
+            const number = uint32Number(host, result, range, e.type);
+            const wrapped = emitWrappingSum(host, e, l, r, li, ri);
+            return wrapped === undefined ? number : { ...number, ...wrapped };
+          }
           const left = widenInteger(host, li),
             right = widenInteger(host, ri);
           if (range) {

packages/compiler/src/backend/llvm/integer-values.ts

@@ -238,6 +238,18 @@ export function int32Number(
   return { name: number, type, uint32: name, integer: { name, type: "i32", signed: true, range } };
 }
 
+/** `integerNumber` for an i32 whose range is within uint32. */
+export function uint32Number(
+  host: LlvmEmitterContext,
+  name: string,
+  range: IntegerRange,
+  type: IrType,
+): LlValue {
+  const number = host.B.tmp();
+  host.B.line(`${number} = uitofp i32 ${name} to double`);
+  return { name: number, type, uint32: name, integer: { name, type: "i32", signed: false, range } };
+}
+
 export function integerNumber(
   host: LlvmEmitterContext,
   name: string,

packages/compiler/test/integer-induction-conversion.test.ts

@@ -47,3 +47,57 @@ console.log(weigh([1, 2, 3, 4]));
     await rm(dir, { recursive: true, force: true });
   }
 });
+
+test("an integer sum with a uint32 operand converts from 32 bits; counter sums keep 64 bits", async () => {
+  const dir = await mkdtemp(join(tmpdir(), "scriptc-integer-uint32-"));
+  try {
+    const entry = join(dir, "main.ts");
+    const outPath = join(dir, "main.ir.json");
+    const source = `
+function au(u: Float64Array, v: Float64Array): void {
+  for (let i = 0; i < v.length; i++) {
+    let t = 0;
+    for (let j = 0; j < u.length; j++) t += u[j]! / ((((i + j) * (i + j + 1)) >>> 1) + i + 1);
+    v[i] = t;
+  }
+}
+function next(values: number[]): number {
+  let total = 0;
+  for (let i = 0; i + 1 < values.length; i++) total += values[i + 1]! - (i + 1);
+  return total;
+}
+const u = new Float64Array(8).fill(1);
+const v = new Float64Array(8);
+au(u, v);
+console.log(v[3], next([1, 2, 3, 4]));
+`;
+    await writeFile(entry, source);
+    const result = await compile(entry, { outDir: dir, outPath, outputKind: "ir" });
+    if (!result.ok)
+      throw new Error(result.diagnostics.map((d) => `${d.code}: ${d.message}`).join("\n"));
+    const ll = emitLlvmModule(deserializeModule(await readFile(outPath, "utf8")));
+    const body = (name: string): string =>
+      ll.match(
+        new RegExp(
+          `define internal [^\\n]+ @sc_(?:b)?f_${name}(?:\\.i32)?\\([^\\n]*\\) #0 \\{([\\s\\S]*?)\\n\\}`,
+        ),
+      )![1]!;
+    // The divisor is at most 2^31 - 1 + 2^25 + 1: above int32, within uint32,
+    // and built on the uint32 \`>>>\` result. It is a 32-bit sum, and its double
+    // converts from the 32 bits.
+    const au = body("au");
+    const divisor = /(%\w+) = fdiv double %\w+, (%\w+)/.exec(au)?.[2];
+    expect(divisor).toBeDefined();
+    expect(au).toMatch(new RegExp(`${divisor} = uitofp i32 %\\w+ to double`));
+    expect(au).toMatch(/ = add nuw i32 /);
+    // \`i + 1\` sums a counter and a literal: it keeps the 64-bit sum and its
+    // signed conversion.
+    const next = body("next");
+    expect(next).not.toMatch(/ = add nuw i32 /);
+    const sum = /(%\w+) = add nsw i64 %\w+, 1\n/.exec(next)?.[1];
+    expect(sum).toBeDefined();
+    expect(next).toMatch(new RegExp(` = sitofp i64 ${sum} to double`));
+  } finally {
+    await rm(dir, { recursive: true, force: true });
+  }
+});