@aardworx/wombat.rendering 0.9.10 → 0.9.12

package/src/runtime/derivedUniforms/dispatch.ts

@@ -1,407 +1,201 @@
-// §7 single-dispatch leaf dispatcher.
+// §7 v2 — uber-kernel dispatcher (GPU plumbing).
 //
-// Each bucket holds a static-ish records buffer (only edited on RO
-// add/remove). Per frame the caller:
-//   1. Drains dirty constituent avals and uploads the changed byte
-//      range to the shared constituents buffer.
-//   2. Builds a dirty bitmask (one bit per constituent slot) and
-//      uploads it to the shared dirty buffer.
-//   3. For each bucket, calls `dispatcher.encode(enc)` — one
-//      compute pass with one dispatch per bucket. The kernel
-//      early-returns on records whose inputs are all clean.
-//
-// No worklist building, no atomicAdd compaction, no per-frame record
-// encoding. The records buffer changes only when ROs come and go.
-
-import {
-  ConstituentSlots, DF32_MAT4_BYTES,
-  type DerivationRecord, type SlotIndex,
-} from "./slots.js";
-import {
-  RECORD_STRIDE_BYTES, RECORD_STRIDE_U32, RECORD_FIELD,
-  DERIVED_UBER_KERNEL_WGSL,
-} from "./uberKernel.wgsl.js";
+// Owns: the GPU mirror of the CPU records buffer, a tiny `Count` uniform, and a
+// recompilable compute pipeline. The kernel WGSL is regenerated whenever the rule
+// registry gains a new rule (`registry.version`) or the record stride grows
+// (`records.layoutVersion`). The records GPU buffer is re-uploaded whenever the CPU
+// records buffer mutates (`records.generation`) — records change only on RO add/remove,
+// never per frame. One `dispatchWorkgroups` per frame, no levels, no barriers.
+// See docs/derived-uniforms-extensible.md § "Dispatcher".
+
+import type { DerivedUniformRegistry } from "./registry.js";
+import type { RecordsBuffer } from "./records.js";
+import { buildUberKernel } from "./codegen.js";
+import { DF32_MAT4_BYTES, type SlotIndex } from "./slots.js";
+
+/** Upload the changed-constituent value range to the constituents GPU buffer. */
+export function uploadConstituentsRange(
+  device: GPUDevice,
+  buf: GPUBuffer,
+  mirror: Float32Array,
+  dirty: ReadonlySet<SlotIndex>,
+): void {
+  let minSlot = Infinity;
+  let maxSlot = -1;
+  for (const s of dirty) {
+    if (s < minSlot) minSlot = s;
+    if (s > maxSlot) maxSlot = s;
+  }
+  if (maxSlot < 0) return;
+  const startByte = minSlot * DF32_MAT4_BYTES;
+  const endByte = (maxSlot + 1) * DF32_MAT4_BYTES;
+  device.queue.writeBuffer(buf, startByte, mirror.buffer, mirror.byteOffset + startByte, endByte - startByte);
+}
 
-/** Live getters for the GPU resources the dispatcher needs. The scene
- *  may replace these buffers in-place when capacity grows; the
- *  dispatcher rebinds each frame. */
+/** Live getters for the GPU buffers the kernel binds (scene may grow/replace them). */
 export interface DerivedUniformsResources {
+  /** `array<vec2<f32>>` — df32 trafo halves (the constituents heap). */
   readonly constituentsBuf: () => GPUBuffer;
-  readonly mainHeapBuf:     () => GPUBuffer;
+  /** `array<f32>` — the main heap (drawHeaders); written by the kernel. */
+  readonly mainHeapBuf: () => GPUBuffer;
 }
 
-/** Per-bucket records buffer. Records are added/removed when ROs
- *  enter/leave the bucket. Uses swap-remove to keep the buffer dense.
- *  CPU mirror is kept up-to-date; the dirty range is uploaded once
- *  per `flush()`. */
-export class RecordsBuffer {
-  private readonly device: GPUDevice;
-  private cpu: Uint32Array;
-  private cap = 0;             // record slots (capacity, in records)
-  private count = 0;           // live records
-  private nextId = 0;
-  private readonly slotById = new Map<number, number>();
-  private readonly idBySlot: number[] = [];
-
-  /** Dirty CPU-mirror byte range that must be uploaded next flush. */
-  private dirtyMinByte = Infinity;
-  private dirtyMaxByte = -1;
+function bglEntries(): GPUBindGroupLayoutEntry[] {
+  return [
+    { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
+    { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+    { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
+    { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: "uniform" } },
+  ];
+}
 
-  private gpu: GPUBuffer | undefined;
-  private gpuCap = 0;          // record slots in `gpu`
+/** GPU mirror of a CPU `RecordsBuffer` + the `Count` uniform. */
+class RecordsGpu {
+  private readonly device: GPUDevice;
+  private buf: GPUBuffer | undefined;
+  private bufCapBytes = 0;
+  private uploadedGen = -1;
+  readonly countBuf: GPUBuffer;
+  private readonly countCpu = new Uint32Array(1);
+  private uploadedCount = -1;
 
-  constructor(device: GPUDevice, initialCapacity = 64) {
+  constructor(device: GPUDevice) {
     this.device = device;
-    this.cap = initialCapacity;
-    this.cpu = new Uint32Array(this.cap * RECORD_STRIDE_U32);
-  }
-
-  /** Add a record. Returns a stable handle for later removal. */
-  add(rec: DerivationRecord): number {
-    if (this.count === this.cap) this.growCpu();
-    const slot = this.count++;
-    const id = this.nextId++;
-    this.slotById.set(id, slot);
-    this.idBySlot[slot] = id;
-    this.writeAt(slot, rec);
-    return id;
-  }
-
-  /** Swap-remove. The last record moves into the freed slot. */
-  remove(id: number): void {
-    const slot = this.slotById.get(id);
-    if (slot === undefined) {
-      throw new Error("RecordsBuffer.remove: unknown id");
-    }
-    const last = this.count - 1;
-    if (slot !== last) {
-      // Move last → slot.
-      const srcOff = last * RECORD_STRIDE_U32;
-      const dstOff = slot * RECORD_STRIDE_U32;
-      for (let i = 0; i < RECORD_STRIDE_U32; i++) {
-        this.cpu[dstOff + i] = this.cpu[srcOff + i]!;
-      }
-      const movedId = this.idBySlot[last]!;
-      this.slotById.set(movedId, slot);
-      this.idBySlot[slot] = movedId;
-      this.markDirty(slot, slot + 1);
-    }
-    this.slotById.delete(id);
-    this.idBySlot.length = last;
-    this.count = last;
+    this.countBuf = device.createBuffer({
+      label: "derivedUniforms.recordCount",
+      size: 16, // uniform min binding size
+      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
+    });
   }
 
-  /** Live record count. */
-  get recordCount(): number { return this.count; }
-
-  /** Ensure GPU buffer matches cap and upload dirty CPU range. Call
-   *  before `encode`. */
-  flush(): void {
-    if (this.gpu === undefined || this.gpuCap < this.cap) {
-      if (this.gpu !== undefined) this.gpu.destroy();
-      this.gpu = this.device.createBuffer({
+  /** Re-upload the packed record data if it changed since last sync, and refresh `Count`. */
+  sync(records: RecordsBuffer): void {
+    const data = records.data;
+    const needBytes = Math.max(16, data.byteLength);
+    if (this.buf === undefined || this.bufCapBytes < needBytes) {
+      if (this.buf !== undefined) this.buf.destroy();
+      let cap = Math.max(16, this.bufCapBytes || 16);
+      while (cap < needBytes) cap *= 2;
+      this.buf = this.device.createBuffer({
         label: "derivedUniforms.records",
-        size:  this.cap * RECORD_STRIDE_BYTES,
+        size: cap,
         usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
       });
-      this.gpuCap = this.cap;
-      // Whole CPU mirror needs re-upload.
-      this.dirtyMinByte = 0;
-      this.dirtyMaxByte = this.count * RECORD_STRIDE_BYTES;
+      this.bufCapBytes = cap;
+      this.uploadedGen = -1; // force re-upload into the fresh buffer
+    }
+    if (records.generation !== this.uploadedGen && data.byteLength > 0) {
+      this.device.queue.writeBuffer(this.buf!, 0, data.buffer, data.byteOffset, data.byteLength);
+      this.uploadedGen = records.generation;
     }
-    if (this.dirtyMaxByte > this.dirtyMinByte) {
-      const start = this.dirtyMinByte;
-      const end   = Math.min(this.dirtyMaxByte, this.count * RECORD_STRIDE_BYTES);
-      if (end > start) {
-        this.device.queue.writeBuffer(
-          this.gpu!, start,
-          this.cpu.buffer,
-          this.cpu.byteOffset + start,
-          end - start,
-        );
-      }
-      this.dirtyMinByte = Infinity;
-      this.dirtyMaxByte = -1;
+    if (records.recordCount !== this.uploadedCount) {
+      this.countCpu[0] = records.recordCount;
+      this.device.queue.writeBuffer(this.countBuf, 0, this.countCpu.buffer, this.countCpu.byteOffset, 4);
+      this.uploadedCount = records.recordCount;
     }
   }
 
   buffer(): GPUBuffer {
-    if (this.gpu === undefined) {
-      // Empty bucket — make a tiny stub so binding always has something.
-      this.gpu = this.device.createBuffer({
+    if (this.buf === undefined) {
+      this.buf = this.device.createBuffer({
         label: "derivedUniforms.records.stub",
-        size:  RECORD_STRIDE_BYTES,
+        size: 16,
         usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
       });
-      this.gpuCap = 1;
+      this.bufCapBytes = 16;
     }
-    return this.gpu;
+    return this.buf;
   }
 
   dispose(): void {
-    if (this.gpu !== undefined) {
-      this.gpu.destroy();
-      this.gpu = undefined;
-    }
-  }
-
-  // ── Internals ─────────────────────────────────────────────────────
-
-  private writeAt(slot: number, rec: DerivationRecord): void {
-    const off = slot * RECORD_STRIDE_U32;
-    this.cpu[off + RECORD_FIELD.recipe_id] = rec.recipe;
-    this.cpu[off + RECORD_FIELD.in0]       = rec.in0;
-    this.cpu[off + RECORD_FIELD.in1]       = rec.in1;
-    this.cpu[off + RECORD_FIELD.in2]       = rec.in2;
-    this.cpu[off + RECORD_FIELD.out_byte]  = rec.outByte;
-    this.markDirty(slot, slot + 1);
-  }
-
-  private markDirty(startSlot: number, endSlot: number): void {
-    const startByte = startSlot * RECORD_STRIDE_BYTES;
-    const endByte   = endSlot   * RECORD_STRIDE_BYTES;
-    if (startByte < this.dirtyMinByte) this.dirtyMinByte = startByte;
-    if (endByte   > this.dirtyMaxByte) this.dirtyMaxByte = endByte;
-  }
-
-  private growCpu(): void {
-    const newCap = this.cap * 2;
-    const grown = new Uint32Array(newCap * RECORD_STRIDE_U32);
-    grown.set(this.cpu);
-    this.cpu = grown;
-    this.cap = newCap;
+    this.buf?.destroy();
+    this.buf = undefined;
+    this.countBuf.destroy();
   }
 }
 
-/** One compute pipeline shared across buckets — created once per scene. */
-export class DerivedUniformsPipeline {
-  readonly pipeline: GPUComputePipeline;
-  readonly bgl:      GPUBindGroupLayout;
+/** Recompilable uber-kernel pipeline. Rebuilt on `(registry.version, strideU32)` change. */
+class UberPipeline {
+  private readonly device: GPUDevice;
+  readonly bgl: GPUBindGroupLayout;
+  private pipe: GPUComputePipeline | undefined;
+  private builtVersion = -1;
+  private builtStride = -1;
 
   constructor(device: GPUDevice) {
-    const module = device.createShaderModule({ code: DERIVED_UBER_KERNEL_WGSL });
-    this.bgl = device.createBindGroupLayout({
-      entries: [
-        { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
-        { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
-        { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
-        { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: "uniform" } },
-      ],
-    });
-    this.pipeline = device.createComputePipeline({
-      layout: device.createPipelineLayout({ bindGroupLayouts: [this.bgl] }),
+    this.device = device;
+    this.bgl = device.createBindGroupLayout({ entries: bglEntries() });
+  }
+
+  pipeline(registry: DerivedUniformRegistry, strideU32: number): GPUComputePipeline | undefined {
+    if (registry.size === 0) return undefined;
+    if (this.pipe !== undefined && this.builtVersion === registry.version && this.builtStride === strideU32) {
+      return this.pipe;
+    }
+    const { wgsl } = buildUberKernel(registry, strideU32);
+    const module = this.device.createShaderModule({ code: wgsl, label: "derivedUniforms.uber" });
+    this.pipe = this.device.createComputePipeline({
+      label: "derivedUniforms.uber",
+      layout: this.device.createPipelineLayout({ bindGroupLayouts: [this.bgl] }),
       compute: { module, entryPoint: "main" },
     });
+    this.builtVersion = registry.version;
+    this.builtStride = strideU32;
+    return this.pipe;
   }
 }
 
-/** Per-bucket dispatcher. One bind group cached, rebuilt only when a
- *  bound buffer reference changes (constituents or main heap grew, or
- *  records buffer was re-created). */
 export class DerivedUniformsDispatcher {
   private readonly device: GPUDevice;
-  private readonly pipe:   DerivedUniformsPipeline;
   private readonly resources: DerivedUniformsResources;
-  readonly records: RecordsBuffer;
-
+  private readonly gpu: RecordsGpu;
+  private readonly pipe: UberPipeline;
   private cachedBg: GPUBindGroup | undefined;
-  private bgKeyConsts: GPUBuffer | undefined;
-  private bgKeyHeap:   GPUBuffer | undefined;
-  private bgKeyRecs:   GPUBuffer | undefined;
+  private bgKey: [GPUBuffer | undefined, GPUBuffer | undefined, GPUBuffer | undefined] = [undefined, undefined, undefined];
 
-  /** Tiny uniform holding the live record count — gates threads past
-   *  the live range from reading stale swap-removed records. */
-  private readonly countBuf: GPUBuffer;
-  private readonly countCpu = new Uint32Array(1);
-  private lastCountUploaded = -1;
-
-  constructor(
-    device: GPUDevice,
-    pipe:   DerivedUniformsPipeline,
-    resources: DerivedUniformsResources,
-    initialRecordCapacity = 64,
-  ) {
-    this.device    = device;
-    this.pipe      = pipe;
+  constructor(device: GPUDevice, resources: DerivedUniformsResources) {
+    this.device = device;
     this.resources = resources;
-    this.records   = new RecordsBuffer(device, initialRecordCapacity);
-    this.countBuf = device.createBuffer({
-      label: "derivedUniforms.recordCount",
-      size:  16,                             // uniform min size
-      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
-    });
+    this.gpu = new RecordsGpu(device);
+    this.pipe = new UberPipeline(device);
   }
 
-  /** Dispatch into an open compute pass. Caller is responsible for
-   *  beginning/ending the pass and setting the pipeline once across
-   *  all dispatchers (the pipeline is shared). */
-  dispatchInto(pass: GPUComputePassEncoder): boolean {
-    const count = this.records.recordCount;
-    if (count === 0) return false;
-    this.records.flush();
-    if (count !== this.lastCountUploaded) {
-      this.countCpu[0] = count;
-      this.device.queue.writeBuffer(
-        this.countBuf, 0, this.countCpu.buffer, this.countCpu.byteOffset, 4,
-      );
-      this.lastCountUploaded = count;
-    }
-    pass.setBindGroup(0, this.bindGroup());
-    pass.dispatchWorkgroups(Math.ceil(count / 64));
-    return true;
-  }
-
-  /** Convenience: open/close a pass and dispatch. Use only when this
-   *  is the sole dispatcher in the pass; otherwise use `dispatchInto`
-   *  inside a caller-managed pass. */
-  encode(enc: GPUCommandEncoder): boolean {
-    if (this.records.recordCount === 0) return false;
+  /** Open a compute pass and run the uber-kernel. No-ops if there are no records / no rules. */
+  encode(enc: GPUCommandEncoder, registry: DerivedUniformRegistry, records: RecordsBuffer): boolean {
+    if (records.recordCount === 0) return false;
+    const pipeline = this.pipe.pipeline(registry, records.strideWords);
+    if (pipeline === undefined) return false;
+    this.gpu.sync(records);
     const pass = enc.beginComputePass({ label: "derivedUniforms.uber" });
-    pass.setPipeline(this.pipe.pipeline);
-    const ran = this.dispatchInto(pass);
+    pass.setPipeline(pipeline);
+    pass.setBindGroup(0, this.bindGroup());
+    pass.dispatchWorkgroups(Math.ceil(records.recordCount / 64));
     pass.end();
-    return ran;
+    return true;
   }
 
   dispose(): void {
-    this.records.dispose();
-    this.countBuf.destroy();
+    this.gpu.dispose();
   }
 
   private bindGroup(): GPUBindGroup {
     const consts = this.resources.constituentsBuf();
-    const heap   = this.resources.mainHeapBuf();
-    const recs   = this.records.buffer();
-    if (
-      this.cachedBg !== undefined &&
-      this.bgKeyConsts === consts && this.bgKeyHeap === heap &&
-      this.bgKeyRecs === recs
-    ) {
+    const heap = this.resources.mainHeapBuf();
+    const recs = this.gpu.buffer();
+    if (this.cachedBg !== undefined && this.bgKey[0] === consts && this.bgKey[1] === heap && this.bgKey[2] === recs) {
       return this.cachedBg;
     }
     this.cachedBg = this.device.createBindGroup({
+      label: "derivedUniforms.uber",
       layout: this.pipe.bgl,
       entries: [
         { binding: 0, resource: { buffer: consts } },
         { binding: 1, resource: { buffer: heap } },
         { binding: 2, resource: { buffer: recs } },
-        { binding: 3, resource: { buffer: this.countBuf } },
+        { binding: 3, resource: { buffer: this.gpu.countBuf } },
       ],
     });
-    this.bgKeyConsts = consts; this.bgKeyHeap = heap;
-    this.bgKeyRecs = recs;
+    this.bgKey = [consts, heap, recs];
     return this.cachedBg;
   }
 }
-
-/** Scene-side helper: O(changed) maintenance of a GPU dirty bitmask.
- *
- *  Steady state (clean frame): zero work. The CPU mirror reflects the
- *  GPU buffer, so we only ever touch bits that actually flip.
- *
- *  Each `upload(dirty)` call:
- *    1. Clears the bits set last frame (O(prev_dirty)).
- *    2. Sets the bits dirty this frame (O(dirty)).
- *    3. Writes only the affected u32 range to the GPU (one writeBuffer).
- *  If both sets are empty, returns immediately. */
-export class DirtyMaskUploader {
-  private readonly device: GPUDevice;
-  private cpu: Uint32Array;
-  private gpu: GPUBuffer | undefined;
-  private gpuCapU32 = 0;
-  /** Bits set on the previous upload — must be cleared next time. */
-  private prevDirty: SlotIndex[] = [];
-
-  constructor(device: GPUDevice, initialU32 = 8) {
-    this.device = device;
-    this.cpu = new Uint32Array(Math.max(1, initialU32));
-  }
-
-  /** Resize CPU + GPU mirrors so they cover `slotCount` bits. */
-  ensureCapacity(slotCount: number): void {
-    const need = Math.max(1, Math.ceil(slotCount / 32));
-    if (need > this.cpu.length) {
-      let cap = Math.max(this.cpu.length, 1);
-      while (cap < need) cap *= 2;
-      const grown = new Uint32Array(cap);
-      grown.set(this.cpu);
-      this.cpu = grown;
-    }
-    if (this.gpu === undefined || this.gpuCapU32 < need) {
-      if (this.gpu !== undefined) this.gpu.destroy();
-      const cap = this.cpu.length;
-      this.gpu = this.device.createBuffer({
-        label: "derivedUniforms.dirtyMask",
-        size:  cap * 4,
-        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
-      });
-      this.gpuCapU32 = cap;
-      // GPU buffer is fresh-zero; CPU mirror should match. Wipe any
-      // prevDirty (those bits live in the destroyed buffer).
-      this.prevDirty.length = 0;
-    }
-  }
-
-  /** Apply the new dirty set. O(prev_dirty + dirty), no O(slotCount). */
-  upload(dirty: ReadonlySet<SlotIndex>): void {
-    let minU32 =  Infinity;
-    let maxU32 = -Infinity;
-    for (const s of this.prevDirty) {
-      const i = s >>> 5; const b = s & 31;
-      this.cpu[i] = (this.cpu[i]! & ~(1 << b)) >>> 0;
-      if (i < minU32) minU32 = i;
-      if (i > maxU32) maxU32 = i;
-    }
-    this.prevDirty.length = 0;
-    for (const s of dirty) {
-      const i = s >>> 5; const b = s & 31;
-      this.cpu[i] = (this.cpu[i]! | (1 << b)) >>> 0;
-      if (i < minU32) minU32 = i;
-      if (i > maxU32) maxU32 = i;
-      this.prevDirty.push(s);
-    }
-    if (maxU32 < 0 || maxU32 === -Infinity) return;
-    const startByte = minU32 * 4;
-    const endByte   = (maxU32 + 1) * 4;
-    this.device.queue.writeBuffer(
-      this.gpu!, startByte,
-      this.cpu.buffer, this.cpu.byteOffset + startByte,
-      endByte - startByte,
-    );
-  }
-
-  buffer(): GPUBuffer {
-    if (this.gpu === undefined) this.ensureCapacity(1);
-    return this.gpu!;
-  }
-
-  dispose(): void {
-    if (this.gpu !== undefined) {
-      this.gpu.destroy();
-      this.gpu = undefined;
-    }
-  }
-}
-
-/** Upload the dirty constituent value range to the constituents GPU
- *  buffer. Helper used by the scene each frame after draining. */
-export function uploadConstituentsRange(
-  device: GPUDevice,
-  buf:    GPUBuffer,
-  mirror: Float32Array,
-  dirty:  ReadonlySet<SlotIndex>,
-): void {
-  let minSlot = Infinity, maxSlot = -1;
-  for (const s of dirty) {
-    if (s < minSlot) minSlot = s;
-    if (s > maxSlot) maxSlot = s;
-  }
-  if (maxSlot < 0) return;
-  const startByte = minSlot * DF32_MAT4_BYTES;
-  const endByte   = (maxSlot + 1) * DF32_MAT4_BYTES;
-  device.queue.writeBuffer(
-    buf, startByte,
-    mirror.buffer, mirror.byteOffset + startByte,
-    endByte - startByte,
-  );
-}

package/src/runtime/derivedUniforms/flatten.ts

@@ -0,0 +1,89 @@
+// §7 v2 — chain flattening.
+//
+// A rule whose IR mentions another derived-uniform name has that producer's
+// IR substituted in place (recursively) until the rule reads only non-derived
+// inputs (constituents / host uniforms). After flattening every
+// record is independent ⇒ one compute dispatch, no levels, no barriers.
+// See docs/derived-uniforms-extensible.md § "Chain flattening".
+
+import type { Expr, Type } from "@aardworx/wombat.shader/ir";
+import { visitExprChildren, stableStringify } from "@aardworx/wombat.shader/ir";
+import { mapExpr } from "@aardworx/wombat.shader/passes";
+import type { DerivedRule } from "./rule.js";
+
+export interface RuleInput {
+  readonly name: string;
+  readonly type: Type;
+  /** `true` ⇒ the leaf appeared as `Inverse(ReadInput(name))` — for a constituent
+   *  trafo this selects the stored backward half (free); the forward half otherwise. */
+  readonly inverse: boolean;
+}
+
+/**
+ * The rule's input leaves in **first-appearance order** (DFS pre-order of `ir`),
+ * deduplicated by `(name, inverse)`. A leaf is either `ReadInput("Uniform", name)`
+ * or `Inverse(ReadInput("Uniform", name))`. Order matters: it is the order the
+ * generated kernel arm consumes its arguments (e.g. the multiplication order of a
+ * matmul chain). Throws if a name is used at two structurally-different types.
+ */
+export function inputsOf(ir: Expr): RuleInput[] {
+  const out: RuleInput[] = [];
+  const seen = new Set<string>(); // key = `${inverse ? "i" : "v"}:${name}`
+  const typeByName = new Map<string, Type>();
+  const record = (name: string, type: Type, inverse: boolean): void => {
+    const prev = typeByName.get(name);
+    if (prev !== undefined && !typesEqual(prev, type)) {
+      throw new Error(`derived rule: uniform '${name}' used at conflicting types`);
+    }
+    typeByName.set(name, type);
+    const key = `${inverse ? "i" : "v"}:${name}`;
+    if (!seen.has(key)) {
+      seen.add(key);
+      out.push({ name, type, inverse });
+    }
+  };
+  const walk = (e: Expr): void => {
+    if (e.kind === "ReadInput" && e.scope === "Uniform") {
+      record(e.name, e.type, false);
+      return;
+    }
+    if (e.kind === "Inverse" && e.value.kind === "ReadInput" && e.value.scope === "Uniform") {
+      record(e.value.name, e.value.type, true);
+      return; // do not also count the inner ReadInput as a separate forward leaf
+    }
+    visitExprChildren(e, walk);
+  };
+  walk(ir);
+  return out;
+}
+
+/**
+ * Substitute every derived-uniform leaf in `ir` with its (recursively
+ * flattened) producer. `name` is the uniform name `ir` is being registered
+ * under; it is pushed on the expansion stack so a self-reference — or any
+ * cycle through other rules — is caught.
+ */
+export function flatten(
+  name: string,
+  ir: Expr,
+  derived: ReadonlyMap<string, DerivedRule>,
+  _stack: readonly string[] = [],
+): Expr {
+  if (_stack.includes(name)) {
+    throw new Error(`derived uniform cycle: ${[..._stack, name].join(" → ")}`);
+  }
+  const stack = [..._stack, name];
+  return mapExpr(ir, (e) => {
+    if (e.kind === "ReadInput" && e.scope === "Uniform") {
+      const producer = derived.get(e.name);
+      if (producer !== undefined) {
+        return flatten(e.name, producer.ir, derived, stack);
+      }
+    }
+    return e;
+  });
+}
+
+function typesEqual(a: Type, b: Type): boolean {
+  return a === b || stableStringify(a) === stableStringify(b);
+}