insomni-plot 0.1.0-alpha.0

package/src/heatmap/gpu.ts

@@ -0,0 +1,565 @@
+// ---------------------------------------------------------------------------
+// GPU heatmap renderer — used by Renderer2D / WebGPU path.
+// ---------------------------------------------------------------------------
+
+import type { TgpuRoot } from "typegpu";
+import { createLayer, Layer, type Renderer2D, Texture } from "insomni";
+
+import type { ContinuousPalette } from "../colors.ts";
+import { LUT_SIZE, WG_1D, WG_2D, clamp01, nextPow2, type ResolvedSpec } from "./types.ts";
+
+interface GpuState {
+  root: TgpuRoot;
+  device: GPUDevice;
+  dataBuffer: GPUBuffer;
+  dataCapacity: number; // datums
+  binBuffer: GPUBuffer; // atomic<i32>[nx*ny]
+  maxBuffer: GPUBuffer; // atomic<i32>[1]
+  paramsBuffer: GPUBuffer; // uniform
+  lutTexture: GPUTexture;
+  lutView: GPUTextureView;
+  outputTexture: GPUTexture;
+  outputView: GPUTextureView;
+  outputWrapper: Texture;
+  clearPipeline: GPUComputePipeline;
+  splatPipeline: GPUComputePipeline;
+  reducePipeline: GPUComputePipeline;
+  colormapPipeline: GPUComputePipeline;
+  splatLayout: GPUBindGroupLayout;
+  clearBindGroup: GPUBindGroup;
+  splatBindGroup: GPUBindGroup;
+  reduceBindGroup: GPUBindGroup;
+  colormapBindGroup: GPUBindGroup;
+  spriteLayer: Layer;
+  dataCount: number;
+}
+
+export class HeatmapGpuRenderer<T> {
+  private readonly spec: ResolvedSpec<T>;
+  private data: readonly T[];
+  private state: GpuState | null = null;
+  private dirty = true;
+
+  constructor(spec: ResolvedSpec<T>, data: readonly T[]) {
+    this.spec = spec;
+    this.data = data;
+  }
+
+  setData(data: readonly T[]): void {
+    this.data = data;
+    this.dirty = true;
+  }
+
+  markDirty(): void {
+    this.dirty = true;
+  }
+
+  /**
+   * Build (or refresh) the GPU heatmap and return the drawable that renders it.
+   *
+   * `root` owns the GPU buffers/textures/pipelines; `renderer` is the v3
+   * {@link Renderer2D} whose `compute(cb)` seam the binning passes are recorded
+   * onto. The returned {@link Layer} is a `space:"ui"` layer with a single sprite
+   * pointing at the compute-written output texture, drawn in submission order
+   * (the same post-main slot as before). The sprite layer is cleared and
+   * re-pushed each frame so the pack content naturally changes when data is
+   * dirty — no separate `bumpVersion` needed.
+   *
+   * The compute callback is queued (not submitted) here; it runs when the
+   * caller invokes `renderer.render(...)`, immediately before the render pass,
+   * so the colormap output is visible to this layer's draw in the same frame.
+   */
+  build(root: TgpuRoot, renderer: Renderer2D): Layer {
+    const state = this.ensureState(root);
+    this.updateSprite(state);
+
+    if (this.dirty) {
+      this.uploadDataBuffer(state);
+      renderer.compute((encoder) => {
+        const pass = encoder.beginComputePass({ label: "heatmap-pass" });
+        this.runPasses(state, pass);
+        pass.end();
+      });
+      this.dirty = false;
+    }
+
+    return state.spriteLayer;
+  }
+
+  destroy(): void {
+    if (!this.state) return;
+    this.state.dataBuffer.destroy();
+    this.state.binBuffer.destroy();
+    this.state.maxBuffer.destroy();
+    this.state.paramsBuffer.destroy();
+    this.state.lutTexture.destroy();
+    this.state.outputWrapper.destroy();
+    this.state = null;
+  }
+
+  private ensureState(root: TgpuRoot): GpuState {
+    if (this.state) return this.state;
+    const spec = this.spec;
+    const device = root.device;
+    const nx = spec.nx;
+    const ny = spec.ny;
+
+    const binBuffer = device.createBuffer({
+      label: "heatmap-bins",
+      size: Math.max(nx * ny * 4, 16),
+      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+    const maxBuffer = device.createBuffer({
+      label: "heatmap-max",
+      size: 16,
+      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+    const paramsBuffer = device.createBuffer({
+      label: "heatmap-params",
+      size: 48, // padded — see params layout below
+      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
+    });
+
+    const initialCapacity = Math.max(1, this.data.length);
+    const dataBuffer = device.createBuffer({
+      label: "heatmap-data",
+      size: initialCapacity * 16, // vec4<f32> per datum
+      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+
+    const lutTexture = device.createTexture({
+      label: "heatmap-lut",
+      size: [LUT_SIZE, 1, 1],
+      format: "rgba8unorm",
+      usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST,
+    });
+    const lutView = lutTexture.createView();
+
+    const outputTexture = device.createTexture({
+      label: "heatmap-output",
+      size: [nx, ny, 1],
+      format: "rgba8unorm",
+      usage:
+        GPUTextureUsage.TEXTURE_BINDING |
+        GPUTextureUsage.STORAGE_BINDING |
+        GPUTextureUsage.COPY_DST,
+    });
+    const outputView = outputTexture.createView();
+    const outputWrapper = new Texture(outputTexture, nx, ny);
+
+    const clearModule = device.createShaderModule({ label: "heatmap-clear", code: CLEAR_WGSL });
+    const splatModule = device.createShaderModule({ label: "heatmap-splat", code: SPLAT_WGSL });
+    const reduceModule = device.createShaderModule({
+      label: "heatmap-reduce",
+      code: REDUCE_WGSL,
+    });
+    const colormapModule = device.createShaderModule({
+      label: "heatmap-colormap",
+      code: COLORMAP_WGSL,
+    });
+    const clearLayout = device.createBindGroupLayout({
+      label: "heatmap-clear",
+      entries: [
+        { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+        { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+        { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "uniform" } },
+      ],
+    });
+    const splatLayout = device.createBindGroupLayout({
+      label: "heatmap-splat",
+      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: "uniform" } },
+      ],
+    });
+    const reduceLayout = device.createBindGroupLayout({
+      label: "heatmap-reduce",
+      entries: [
+        { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+        { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+        { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "uniform" } },
+      ],
+    });
+    const colormapLayout = device.createBindGroupLayout({
+      label: "heatmap-colormap",
+      entries: [
+        { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+        { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
+        { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "uniform" } },
+        {
+          binding: 3,
+          visibility: GPUShaderStage.COMPUTE,
+          texture: { sampleType: "float", viewDimension: "2d" },
+        },
+        {
+          binding: 4,
+          visibility: GPUShaderStage.COMPUTE,
+          storageTexture: { access: "write-only", format: "rgba8unorm", viewDimension: "2d" },
+        },
+      ],
+    });
+
+    const clearPipeline = device.createComputePipeline({
+      label: "heatmap-clear-pipeline",
+      layout: device.createPipelineLayout({ bindGroupLayouts: [clearLayout] }),
+      compute: { module: clearModule, entryPoint: "main" },
+    });
+    const splatPipeline = device.createComputePipeline({
+      label: "heatmap-splat-pipeline",
+      layout: device.createPipelineLayout({ bindGroupLayouts: [splatLayout] }),
+      compute: { module: splatModule, entryPoint: "main" },
+    });
+    const reducePipeline = device.createComputePipeline({
+      label: "heatmap-reduce-pipeline",
+      layout: device.createPipelineLayout({ bindGroupLayouts: [reduceLayout] }),
+      compute: { module: reduceModule, entryPoint: "main" },
+    });
+    const colormapPipeline = device.createComputePipeline({
+      label: "heatmap-colormap-pipeline",
+      layout: device.createPipelineLayout({ bindGroupLayouts: [colormapLayout] }),
+      compute: { module: colormapModule, entryPoint: "main" },
+    });
+
+    const clearBindGroup = device.createBindGroup({
+      layout: clearLayout,
+      entries: [
+        { binding: 0, resource: { buffer: binBuffer } },
+        { binding: 1, resource: { buffer: maxBuffer } },
+        { binding: 2, resource: { buffer: paramsBuffer } },
+      ],
+    });
+    const splatBindGroup = device.createBindGroup({
+      layout: splatLayout,
+      entries: [
+        { binding: 0, resource: { buffer: dataBuffer } },
+        { binding: 1, resource: { buffer: binBuffer } },
+        { binding: 2, resource: { buffer: paramsBuffer } },
+      ],
+    });
+    const reduceBindGroup = device.createBindGroup({
+      layout: reduceLayout,
+      entries: [
+        { binding: 0, resource: { buffer: binBuffer } },
+        { binding: 1, resource: { buffer: maxBuffer } },
+        { binding: 2, resource: { buffer: paramsBuffer } },
+      ],
+    });
+    const colormapBindGroup = device.createBindGroup({
+      layout: colormapLayout,
+      entries: [
+        { binding: 0, resource: { buffer: binBuffer } },
+        { binding: 1, resource: { buffer: maxBuffer } },
+        { binding: 2, resource: { buffer: paramsBuffer } },
+        { binding: 3, resource: lutView },
+        { binding: 4, resource: outputView },
+      ],
+    });
+
+    // Pre-fill the LUT from the palette.
+    writeLut(device, lutTexture, spec.colorMap);
+    // Pre-fill the params uniform once; `updateSprite` keeps it in sync if the
+    // spec ever mutates (v1 treats spec as immutable, but params buffer still
+    // needs a first write).
+    writeParams(device, paramsBuffer, spec, this.data.length);
+
+    const spriteLayer = createLayer({ space: "ui" });
+
+    this.state = {
+      root,
+      device,
+      dataBuffer,
+      dataCapacity: initialCapacity,
+      binBuffer,
+      maxBuffer,
+      paramsBuffer,
+      lutTexture,
+      lutView,
+      outputTexture,
+      outputView,
+      outputWrapper,
+      clearPipeline,
+      splatPipeline,
+      reducePipeline,
+      colormapPipeline,
+      splatLayout,
+      clearBindGroup,
+      splatBindGroup,
+      reduceBindGroup,
+      colormapBindGroup,
+      spriteLayer,
+      dataCount: this.data.length,
+    };
+    return this.state;
+  }
+
+  private uploadDataBuffer(state: GpuState): void {
+    const spec = this.spec;
+    const n = this.data.length;
+    state.dataCount = n;
+
+    if (n > state.dataCapacity) {
+      state.dataBuffer.destroy();
+      const next = nextPow2(n);
+      state.dataBuffer = state.device.createBuffer({
+        label: "heatmap-data",
+        size: next * 16,
+        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+      });
+      state.dataCapacity = next;
+      // Rebuild the splat bind group since the buffer identity changed.
+      state.splatBindGroup = state.device.createBindGroup({
+        layout: state.splatLayout,
+        entries: [
+          { binding: 0, resource: { buffer: state.dataBuffer } },
+          { binding: 1, resource: { buffer: state.binBuffer } },
+          { binding: 2, resource: { buffer: state.paramsBuffer } },
+        ],
+      });
+    }
+
+    if (n > 0) {
+      const packed = new Float32Array(n * 4);
+      for (let i = 0; i < n; i++) {
+        const datum = this.data[i]!;
+        packed[i * 4 + 0] = spec.x(datum);
+        packed[i * 4 + 1] = spec.y(datum);
+        packed[i * 4 + 2] = spec.weight(datum);
+        packed[i * 4 + 3] = 0;
+      }
+      state.device.queue.writeBuffer(state.dataBuffer, 0, packed);
+    }
+
+    writeParams(state.device, state.paramsBuffer, spec, n);
+  }
+
+  private updateSprite(state: GpuState): void {
+    // Clear and re-push the single sprite every frame so the frame rect
+    // (caller may move it) is always current.
+    const spec = this.spec;
+    state.spriteLayer.clear();
+    state.spriteLayer.pushSprite({
+      pos: { x: spec.frame.x + spec.frame.width * 0.5, y: spec.frame.y + spec.frame.height * 0.5 },
+      size: { x: spec.frame.width, y: spec.frame.height },
+      texture: state.outputWrapper,
+      anchor: { x: 0.5, y: 0.5 },
+    });
+  }
+
+  private runPasses(state: GpuState, pass: GPUComputePassEncoder): void {
+    const spec = this.spec;
+    const cellCount = spec.nx * spec.ny;
+
+    // Clear: one thread per cell; also resets max to 0.
+    pass.setPipeline(state.clearPipeline);
+    pass.setBindGroup(0, state.clearBindGroup);
+    pass.dispatchWorkgroups(Math.ceil(cellCount / WG_1D));
+
+    // Splat: one thread per datum.
+    if (state.dataCount > 0) {
+      pass.setPipeline(state.splatPipeline);
+      pass.setBindGroup(0, state.splatBindGroup);
+      pass.dispatchWorkgroups(Math.ceil(state.dataCount / WG_1D));
+    }
+
+    // Reduce max across all cells. One workgroup, 256 threads, each thread
+    // strides through cells and atomicMax's into a single atomic<i32>.
+    pass.setPipeline(state.reducePipeline);
+    pass.setBindGroup(0, state.reduceBindGroup);
+    pass.dispatchWorkgroups(1);
+
+    // Colormap: one thread per output pixel.
+    pass.setPipeline(state.colormapPipeline);
+    pass.setBindGroup(0, state.colormapBindGroup);
+    pass.dispatchWorkgroups(Math.ceil(spec.nx / WG_2D), Math.ceil(spec.ny / WG_2D));
+  }
+}
+
+// ---------------------------------------------------------------------------
+// GPU helpers
+// ---------------------------------------------------------------------------
+
+function writeParams<T>(
+  device: GPUDevice,
+  buffer: GPUBuffer,
+  spec: ResolvedSpec<T>,
+  count: number,
+): void {
+  // Params layout (48 bytes):
+  //   bins: vec2u           @ 0   (8)
+  //   count: u32            @ 8   (4)
+  //   weightScale: f32      @ 12  (4)
+  //   xDomain: vec2f        @ 16  (8)
+  //   yDomain: vec2f        @ 24  (8)
+  //   cellCount: u32        @ 32  (4)
+  //   normalize: u32        @ 36  (4)  — 0=max, 1=log, 2=sqrt
+  //   _pad1: vec2f          @ 40  (8)
+  const bytes = new ArrayBuffer(48);
+  const u32 = new Uint32Array(bytes);
+  const f32 = new Float32Array(bytes);
+  u32[0] = spec.nx;
+  u32[1] = spec.ny;
+  u32[2] = count;
+  f32[3] = spec.weightScale;
+  f32[4] = spec.x0;
+  f32[5] = spec.x1;
+  f32[6] = spec.y0;
+  f32[7] = spec.y1;
+  u32[8] = spec.nx * spec.ny;
+  u32[9] = spec.normalize === "log" ? 1 : spec.normalize === "sqrt" ? 2 : 0;
+  device.queue.writeBuffer(buffer, 0, bytes);
+}
+
+function writeLut(device: GPUDevice, texture: GPUTexture, palette: ContinuousPalette): void {
+  const bytes = new Uint8Array(LUT_SIZE * 4);
+  for (let i = 0; i < LUT_SIZE; i++) {
+    const t = i / (LUT_SIZE - 1);
+    const c = palette(t);
+    // Premultiplied alpha (matches sprite sampler convention).
+    const a = clamp01(c.a);
+    bytes[i * 4 + 0] = Math.round(clamp01(c.r) * a * 255);
+    bytes[i * 4 + 1] = Math.round(clamp01(c.g) * a * 255);
+    bytes[i * 4 + 2] = Math.round(clamp01(c.b) * a * 255);
+    bytes[i * 4 + 3] = Math.round(a * 255);
+  }
+  device.queue.writeTexture({ texture }, bytes, { bytesPerRow: LUT_SIZE * 4, rowsPerImage: 1 }, [
+    LUT_SIZE,
+    1,
+    1,
+  ]);
+}
+
+// ---------------------------------------------------------------------------
+// WGSL
+// ---------------------------------------------------------------------------
+
+const PARAMS_WGSL = /* wgsl */ `
+struct Params {
+  bins: vec2u,
+  count: u32,
+  weightScale: f32,
+  xDomain: vec2f,
+  yDomain: vec2f,
+  cellCount: u32,
+  normalize: u32,  // 0=max, 1=log, 2=sqrt
+};
+`;
+
+const CLEAR_WGSL = /* wgsl */ `${PARAMS_WGSL}
+@group(0) @binding(0) var<storage, read_write> bins: array<atomic<i32>>;
+@group(0) @binding(1) var<storage, read_write> maxCell: array<atomic<i32>, 1>;
+@group(0) @binding(2) var<uniform> params: Params;
+
+@compute @workgroup_size(${WG_1D})
+fn main(@builtin(global_invocation_id) gid: vec3u) {
+  let i = gid.x;
+  if (i >= params.cellCount) {
+    return;
+  }
+  atomicStore(&bins[i], 0);
+  if (i == 0u) {
+    atomicStore(&maxCell[0], 0);
+  }
+}
+`;
+
+const SPLAT_WGSL = /* wgsl */ `${PARAMS_WGSL}
+@group(0) @binding(0) var<storage, read> data: array<vec4f>;
+@group(0) @binding(1) var<storage, read_write> bins: array<atomic<i32>>;
+@group(0) @binding(2) var<uniform> params: Params;
+
+@compute @workgroup_size(${WG_1D})
+fn main(@builtin(global_invocation_id) gid: vec3u) {
+  let i = gid.x;
+  if (i >= params.count) {
+    return;
+  }
+  let datum = data[i];
+  let xv = datum.x;
+  let yv = datum.y;
+  let w = datum.z;
+
+  let dx = params.xDomain.y - params.xDomain.x;
+  let dy = params.yDomain.y - params.yDomain.x;
+  if (dx == 0.0 || dy == 0.0) {
+    return;
+  }
+  let tx = (xv - params.xDomain.x) / dx;
+  let ty = (yv - params.yDomain.x) / dy;
+  if (tx < 0.0 || tx >= 1.0 || ty < 0.0 || ty >= 1.0) {
+    return;
+  }
+  let nx = params.bins.x;
+  let ny = params.bins.y;
+  let ix = min(nx - 1u, u32(tx * f32(nx)));
+  // Flip iy: row 0 represents the *top* of the frame (highest data-y).
+  let iy = ny - 1u - min(ny - 1u, u32(ty * f32(ny)));
+  let idx = iy * nx + ix;
+  let quant = i32(round(w * params.weightScale));
+  atomicAdd(&bins[idx], quant);
+}
+`;
+
+const REDUCE_WGSL = /* wgsl */ `${PARAMS_WGSL}
+@group(0) @binding(0) var<storage, read_write> bins: array<atomic<i32>>;
+@group(0) @binding(1) var<storage, read_write> maxCell: array<atomic<i32>, 1>;
+@group(0) @binding(2) var<uniform> params: Params;
+
+@compute @workgroup_size(${WG_1D})
+fn main(@builtin(local_invocation_id) lid: vec3u) {
+  let threads = ${WG_1D}u;
+  let total = params.cellCount;
+  var localMax: i32 = 0;
+  var i = lid.x;
+  loop {
+    if (i >= total) {
+      break;
+    }
+    let v = atomicLoad(&bins[i]);
+    if (v > localMax) {
+      localMax = v;
+    }
+    i = i + threads;
+  }
+  atomicMax(&maxCell[0], localMax);
+}
+`;
+
+const COLORMAP_WGSL = /* wgsl */ `${PARAMS_WGSL}
+@group(0) @binding(0) var<storage, read_write> bins: array<atomic<i32>>;
+@group(0) @binding(1) var<storage, read_write> maxCell: array<atomic<i32>, 1>;
+@group(0) @binding(2) var<uniform> params: Params;
+@group(0) @binding(3) var lut: texture_2d<f32>;
+@group(0) @binding(4) var outputTex: texture_storage_2d<rgba8unorm, write>;
+
+@compute @workgroup_size(${WG_2D}, ${WG_2D})
+fn main(@builtin(global_invocation_id) gid: vec3u) {
+  let nx = params.bins.x;
+  let ny = params.bins.y;
+  if (gid.x >= nx || gid.y >= ny) {
+    return;
+  }
+  let idx = gid.y * nx + gid.x;
+  let raw = atomicLoad(&bins[idx]);
+  let maxRaw = atomicLoad(&maxCell[0]);
+  var t: f32 = 0.0;
+  if (maxRaw > 0) {
+    let fraw = f32(raw);
+    let fmax = f32(maxRaw);
+    if (params.normalize == 1u) {
+      // log: log(1+v) / log(1+max)
+      t = clamp(log(1.0 + fraw) / log(1.0 + fmax), 0.0, 1.0);
+    } else if (params.normalize == 2u) {
+      // sqrt
+      t = clamp(sqrt(fraw / fmax), 0.0, 1.0);
+    } else {
+      // max (default)
+      t = clamp(fraw / fmax, 0.0, 1.0);
+    }
+  }
+  let lutWidth = ${LUT_SIZE};
+  let u = i32(round(t * f32(lutWidth - 1)));
+  let color = textureLoad(lut, vec2i(u, 0), 0);
+  textureStore(outputTex, vec2i(i32(gid.x), i32(gid.y)), color);
+}
+`;