From 896432c3d7a78faaf72b43ae6bcfd442b31e8428 Mon Sep 17 00:00:00 2001
From: SamG-Coder <13449963+SamG-Coder@users.noreply.github.com>
Date: Wed, 23 Sep 2026 08:47:16 +1000
Subject: [PATCH] Add streaming GPU normal equations for small least-squares
 fits

---
 README.md                      | 36 +++++++++++++
 scripts/test-least-squares.mjs | 10 ++++
 src/runtime/least-squares.js   | 92 ++++++++++++++++++++++++++++++++++
 tests/least-squares-gpu.html   | 27 ++++++++++
 tests/least-squares.test.mjs   | 22 ++++++++
 5 files changed, 187 insertions(+)
 create mode 100644 scripts/test-least-squares.mjs
 create mode 100644 src/runtime/least-squares.js
 create mode 100644 tests/least-squares-gpu.html
 create mode 100644 tests/least-squares.test.mjs

diff --git a/README.md b/README.md
index 4e703f4..1ccc4ff 100644
--- a/README.md
+++ b/README.md
@@ -110,6 +110,42 @@ The default optimization changes the generated text. A cache keyed only by CUDA
 
 ## Run a kernel in the browser
 
+### Streaming least squares
+
+`src/runtime/least-squares.js` supplies `NormalEquations` for applications such as
+texture fitting that generate many observation rows but solve a small dense
+system. Each GPU row contains `[features..., targets...]`; include a constant
+feature when the model needs a bias. Append batches from a reusable GPU buffer,
+then call `solve()` to obtain output-major coefficients. Products and partial
+sums use f32 on the GPU, and the final reduction and pivoted dense solve use f64
+on the host. This is not a native f64 or QR solver; normal equations can lose
+accuracy for ill-conditioned inputs. The default relative ridge is `1e-6`.
+
+```js
+import {NormalEquations} from './src/runtime/least-squares.js';
+
+const fit = await NormalEquations.create(runtime, 3, 2);
+try {
+  // observationBuffer has rows [x0, x1, 1, target0, target1].
+  fit.append(observationBuffer, rowCount);
+  // More appends may follow, including GPU-generated rows in the same buffer.
+  const weights = await fit.solve(); // 2 output rows of 3 coefficients
+} finally {
+  await runtime.idle();
+  fit.dispose();
+}
+```
+
+Both feature and output counts are bounded to 1..64. `append` supports `stride`
+and `offset` in float elements, validates buffer ownership and bounds, and
+submits work immediately. Queue order protects a reused observation buffer.
+Singular systems without a ridge, empty fits and non-finite solved equations
+are rejected. Test the numerical GPU path with
+`node scripts/test-least-squares.mjs` (installed Edge on Windows; Playwright
+Chromium elsewhere; override with `CW_BROWSER`).
+
+### General kernel launch
+
 Serve the repository and run this as a module on that origin. `GpuRuntime.create()` requests a WebGPU device.
 
 ```js
diff --git a/scripts/test-least-squares.mjs b/scripts/test-least-squares.mjs
new file mode 100644
index 0000000..b78c968
--- /dev/null
+++ b/scripts/test-least-squares.mjs
@@ -0,0 +1,10 @@
+import {chromium} from 'playwright';
+import {createStaticServer} from './serve.mjs';
+const server=createStaticServer();await new Promise(resolve=>server.listen(0,'127.0.0.1',resolve));let browser;
+try {
+ browser=await chromium.launch({...(process.env.CW_BROWSER?{channel:process.env.CW_BROWSER}:process.platform==='win32'?{channel:'msedge'}:{}),headless:true,args:['--enable-unsafe-webgpu']});
+ const page=await browser.newPage();await page.goto(`http://127.0.0.1:${server.address().port}/tests/least-squares-gpu.html`);
+ await page.waitForFunction(()=>window.runLeastSquaresTests);
+ console.log(JSON.stringify(await page.evaluate(()=>window.runLeastSquaresTests()),null,2));
+}finally{await browser?.close();await new Promise(resolve=>server.close(resolve));}
+
diff --git a/src/runtime/least-squares.js b/src/runtime/least-squares.js
new file mode 100644
index 0000000..221c46e
--- /dev/null
+++ b/src/runtime/least-squares.js
@@ -0,0 +1,92 @@
+// SPDX-License-Identifier: MIT
+// Streaming normal equations: GPU accumulation, small dense solve on the host.
+export const NORMAL_EQUATIONS_SOURCE = `
+__global__ void normal_equations(const float* rows, float* partial,
+    int count, int features, int outputs, int stride, int offset) {
+  int entry = blockIdx.x * blockDim.x + threadIdx.x;
+  int entries = features * (features + outputs);
+  if (entry >= entries) return;
+  int a = entry / (features + outputs);
+  int b = entry % (features + outputs);
+  int lane = blockIdx.y;
+  float sum = 0.0f;
+  float correction = 0.0f;
+  for (int row = lane; row < count; row += 16) {
+    int start = offset + row * stride;
+    float value = rows[start + a] * rows[start + b] - correction;
+    float next = sum + value;
+    correction = (next - sum) - value;
+    sum = next;
+  }
+  partial[lane * entries + entry] += sum;
+}`;
+
+/** Partial-pivoted Float64 host solve. Rows of the result correspond to outputs. */
+export function solveNormalEquations(equations, features, outputs, ridge = 1e-6) {
+  if (!Number.isInteger(features) || features < 1 || features > 64 ||
+      !Number.isInteger(outputs) || outputs < 1 || outputs > 64 ||
+      equations.length !== features * (features + outputs) ||
+      !Number.isFinite(ridge) || ridge < 0 || Array.from(equations).some(v => !Number.isFinite(v))) {
+    throw new RangeError('Invalid normal equations or ridge.');
+  }
+  const stride = features + outputs, a = Float64Array.from(equations);
+  let trace = 0;
+  for (let i = 0; i < features; i++) trace += a[i * stride + i];
+  if (!(trace > 0)) throw new Error('Least-squares system has no nonzero observations.');
+  const lambda = ridge * trace / features;
+  for (let i = 0; i < features; i++) a[i * stride + i] += lambda;
+  for (let k = 0; k < features; k++) {
+    let pivot = k;
+    for (let i = k + 1; i < features; i++) if (Math.abs(a[i * stride + k]) > Math.abs(a[pivot * stride + k])) pivot = i;
+    if (!(Math.abs(a[pivot * stride + k]) > Number.EPSILON * trace)) throw new Error('Singular least-squares system; use a positive ridge.');
+    for (let j = k; j < stride; j++) {
+      const t = a[k * stride + j]; a[k * stride + j] = a[pivot * stride + j]; a[pivot * stride + j] = t;
+    }
+    const diagonal = a[k * stride + k];
+    for (let j = k; j < stride; j++) a[k * stride + j] /= diagonal;
+    for (let i = 0; i < features; i++) if (i !== k) {
+      const factor = a[i * stride + k];
+      for (let j = k; j < stride; j++) a[i * stride + j] -= factor * a[k * stride + j];
+    }
+  }
+  const weights = new Float32Array(features * outputs);
+  for (let o = 0; o < outputs; o++) for (let i = 0; i < features; i++) weights[o * features + i] = a[i * stride + features + o];
+  if (weights.some(v => !Number.isFinite(v))) throw new Error('Least-squares solution overflows float32.');
+  return weights;
+}
+
+/** Each GPU row is [features..., targets...]; include a constant feature for a bias.
+ * Rows may be streamed through a reusable buffer. All appends use one queue in order.
+ * f32 products/accumulation are not a substitute for native f64 normal equations.
+ */
+export class NormalEquations {
+  static async create(runtime, features, outputs) {
+    if (![features, outputs].every(n => Number.isInteger(n) && n >= 1 && n <= 64)) throw new RangeError('Feature and output counts must be 1..64.');
+    const kernel = await runtime.kernel(NORMAL_EQUATIONS_SOURCE, {entry:'normal_equations', workgroupSize:[64,1,1]});
+    return new NormalEquations(runtime, features, outputs, kernel);
+  }
+  constructor(runtime, features, outputs, kernel) {
+    this.runtime = runtime; this.features = features; this.outputs = outputs; this.kernel = kernel;
+    this.entries = features * (features + outputs);
+    this.partial = runtime.createBuffer(this.entries * 16 * 4);
+    this.rows = 0;
+  }
+  append(rows, count, {stride = this.features + this.outputs, offset = 0} = {}) {
+    this.runtime.checkResource(this.partial); this.runtime.checkResource(rows);
+    if (!Number.isInteger(count) || count < 1 || !Number.isInteger(stride) || stride < this.features + this.outputs ||
+        !Number.isInteger(offset) || offset < 0 || offset + (count - 1) * stride + this.features + this.outputs > rows.byteLength / 4 ||
+        offset + count * stride > 0x7fffffff) throw new RangeError('Observation rows do not fit the buffer or signed indexing range.');
+    this.runtime.batch().dispatch(this.kernel.bind({rows, partial:this.partial}, {
+      count, features:this.features, outputs:this.outputs, stride, offset
+    }), [Math.ceil(this.entries / 64),16,1]).submit();
+    this.rows += count;
+    return this;
+  }
+  async solve({ridge = 1e-6} = {}) {
+    if (!this.rows) throw new Error('No observations were appended.');
+    const partial = await this.runtime.read(this.partial), equations = new Float64Array(this.entries);
+    for (let lane = 0; lane < 16; lane++) for (let i = 0; i < this.entries; i++) equations[i] += partial[lane * this.entries + i];
+    return solveNormalEquations(equations, this.features, this.outputs, ridge);
+  }
+  dispose() { if (!this.partial.destroyed) this.runtime.destroyBuffer(this.partial); }
+}
diff --git a/tests/least-squares-gpu.html b/tests/least-squares-gpu.html
new file mode 100644
index 0000000..5d6b51d
--- /dev/null
+++ b/tests/least-squares-gpu.html
@@ -0,0 +1,27 @@
+<!doctype html><meta charset="utf-8"><link rel="icon" href="data:,"><title>Least squares GPU test</title>
+<script type="module">
+import {GpuRuntime} from '../src/runtime/runtime.js';
+import {NormalEquations,solveNormalEquations} from '../src/runtime/least-squares.js';
+window.runLeastSquaresTests=async()=>{
+ const runtime=await GpuRuntime.create(),errors=[];
+ runtime.device.addEventListener('uncapturederror',event=>errors.push(event.error.message));
+ const observations=new Float32Array(1031*5),reference=new Float64Array(3*5);
+ for(let i=0;i<1031;i++) {
+  const x=Math.sin(i*0.1),y=Math.cos(i*0.3),row=[x,y,1,2*x-3*y+4,-x+0.5*y-2];
+  observations.set(row,i*5);
+ }
+ for(let i=0;i<1031;i++)for(let a=0;a<3;a++)for(let b=0;b<5;b++)reference[a*5+b]+=observations[i*5+a]*observations[i*5+b];
+ const expected=solveNormalEquations(reference,3,2),fit=await NormalEquations.create(runtime,3,2);
+ const rows=runtime.createBuffer(observations);
+ try {
+  fit.append(rows,509);fit.append(rows,522,{offset:509*5});
+  const actual=await fit.solve(),difference=Math.max(...actual.map((v,i)=>Math.abs(v-expected[i])));
+  if(difference>2e-5)throw Error(`GPU / CPU difference ${difference}`);
+  for(const call of [()=>fit.append(rows,0),()=>fit.append(rows,1032),()=>fit.append(rows,1,{stride:4}),()=>fit.append(rows,1,{offset:-1})]) {
+   let refused=false;try{call();}catch{refused=true;}if(!refused)throw Error('Invalid row layout accepted');
+  }
+  await runtime.idle();if(errors.length)throw Error(errors.join('\n'));
+  return {actual:Array.from(actual),expected:Array.from(expected),maxDifference:difference,rows:fit.rows,adapter:runtime.adapter.info,errors};
+ }finally{await runtime.idle();fit.dispose();runtime.dispose();}
+};
+</script>
diff --git a/tests/least-squares.test.mjs b/tests/least-squares.test.mjs
new file mode 100644
index 0000000..329e767
--- /dev/null
+++ b/tests/least-squares.test.mjs
@@ -0,0 +1,22 @@
+import test from 'node:test';
+import assert from 'node:assert/strict';
+import {compile} from '../src/compiler/compiler.js';
+import {NORMAL_EQUATIONS_SOURCE, solveNormalEquations} from '../src/runtime/least-squares.js';
+
+test('normal-equation CUDA compiles with bounded bindings', () => {
+  const artifact = compile(NORMAL_EQUATIONS_SOURCE, {entry:'normal_equations',workgroupSize:[64,1,1]});
+  assert.equal(artifact.metadata.bindings.length, 2);
+});
+test('solves several outputs with partial pivoting', () => {
+  // x = -2..2, targets 2x+3 and -x+4; features [x,1].
+  const solution = solveNormalEquations([10,0,20,-10,0,5,15,20],2,2,0);
+  assert.deepEqual(Array.from(solution),[2,3,-1,4]);
+  assert.deepEqual(Array.from(solveNormalEquations([0,1,2,1,1,5],2,1,0)),[3,2]);
+});
+test('ridge regularizes duplicate features; invalid and empty systems fail', () => {
+  const solution = solveNormalEquations([2,2,4,2,2,4],2,1);
+  assert.ok(solution.every(v=>Math.abs(v-1)<1e-5));
+  assert.throws(()=>solveNormalEquations([2,2,4,2,2,4],2,1,0),/Singular/);
+  assert.throws(()=>solveNormalEquations([0,0],1,1),/no nonzero/);
+  assert.throws(()=>solveNormalEquations([NaN,1],1,1),/Invalid/);
+});