perf(mesh-layer): memoize per-Device GLSL shader assembly

packages/deck.gl-raster/src/index.ts

@@ -1,6 +1,11 @@
 export type { RasterModule } from "./gpu-modules/types.js";
 // Not a public API; exported for use in COGLayer and ZarrLayer
 export { renderDebugTileOutline as _renderDebugTileOutline } from "./layer-utils.js";
+export type { MemoShaderAssemblerStats } from "./mesh-layer/shader-assembler-memo.js";
+export {
+  getMemoShaderAssemblerMissLog,
+  getMemoShaderAssemblerStats,
+} from "./mesh-layer/shader-assembler-memo.js";
 export type {
   MultiTilesetDescriptor,
   SecondaryTileIndex,

packages/deck.gl-raster/src/mesh-layer/mesh-layer.ts

@@ -6,6 +6,7 @@ import type { ShaderModule } from "@luma.gl/shadertools";
 import { CreateTexture } from "../gpu-modules/create-texture.js";
 import type { RasterModule } from "../gpu-modules/types.js";
 import fs from "./mesh-layer-fragment.glsl.js";
+import { getMemoizingShaderAssembler } from "./shader-assembler-memo.js";
 
 type _MeshTextureLayerProps =
   | { image: TextureSource; renderPipeline?: RasterModule[] }
@@ -69,6 +70,11 @@ export class MeshTextureLayer extends SimpleMeshLayer<
       // injection points
       fs,
       modules,
+      // Memoize shader assembly per (modules, vs, fs) across all sublayers of
+      // a tiled raster source. Within a `RasterTileLayer`, every tile passes
+      // identical inputs here — the per-Device cache collapses N regex-heavy
+      // assembly passes into one.
+      shaderAssembler: getMemoizingShaderAssembler(this.context.device),
     };
   }
 

packages/deck.gl-raster/src/mesh-layer/shader-assembler-memo.ts

@@ -0,0 +1,173 @@
+import type { Device } from "@luma.gl/core";
+import { ShaderAssembler } from "@luma.gl/shadertools";
+
+type AssembledPair = ReturnType<ShaderAssembler["assembleGLSLShaderPair"]>;
+
+type AssembleProps = Parameters<ShaderAssembler["assembleGLSLShaderPair"]>[0];
+
+const ASSEMBLE_GLSL = "assembleGLSLShaderPair";
+
+/**
+ * Cache statistics for a memoizing shader assembler. Useful for diagnosing
+ * whether the cache is actually catching the per-tile assembly load —
+ * `hits` should dominate `misses` once a tile layer has rendered a few tiles.
+ */
+export type MemoShaderAssemblerStats = {
+  /** Number of `assembleGLSLShaderPair` calls served from cache. */
+  hits: number;
+  /** Number of `assembleGLSLShaderPair` calls that produced a new cache entry. */
+  misses: number;
+  /** Distinct `(modules, vs, fs, defines)` tuples currently cached. */
+  entries: number;
+};
+
+type CacheRecord = {
+  cache: Map<string, AssembledPair>;
+  stats: MemoShaderAssemblerStats;
+  /** Cache keys for the first N misses, captured for debugging. */
+  missLog: string[];
+};
+
+const perDeviceAssembler = new WeakMap<Device, ShaderAssembler>();
+const perAssemblerRecord = new WeakMap<ShaderAssembler, CacheRecord>();
+
+/** How many cache-miss keys to retain in the in-memory miss log. */
+const MAX_LOGGED_MISSES = 20;
+
+/**
+ * Returns a `ShaderAssembler` whose `assembleGLSLShaderPair` results are
+ * memoized per `(modules, vs, fs, defines)` tuple.
+ *
+ * Instances are cached per `Device` — two Deck instances on the same page get
+ * independent caches, and the entries die with the device. All other
+ * `ShaderAssembler` methods (hook registration, default modules, WGSL
+ * assembly) delegate to {@link ShaderAssembler.getDefaultShaderAssembler}, so
+ * deck.gl's globally-registered hooks remain visible.
+ *
+ * Within a `RasterTileLayer`, every tile sublayer passes the same modules and
+ * shader source, so this collapses N regex-heavy assembly passes into one.
+ */
+export function getMemoizingShaderAssembler(device: Device): ShaderAssembler {
+  const existing = perDeviceAssembler.get(device);
+  if (existing) {
+    return existing;
+  }
+  const assembler = createMemoizingShaderAssembler(
+    ShaderAssembler.getDefaultShaderAssembler(),
+  );
+  perDeviceAssembler.set(device, assembler);
+  return assembler;
+}
+
+/**
+ * Reads the hit/miss counters for the assembler associated with a `Device`.
+ * Returns `null` if no memoizing assembler has been installed for the device.
+ *
+ * Intended for app-level diagnostics — call from devtools to confirm the cache
+ * is taking effect. A healthy mosaic-style workload should show
+ * `hits >> misses` after the first few tiles have rendered.
+ */
+export function getMemoShaderAssemblerStats(
+  device: Device,
+): MemoShaderAssemblerStats | null {
+  const assembler = perDeviceAssembler.get(device);
+  if (!assembler) {
+    return null;
+  }
+  const record = perAssemblerRecord.get(assembler);
+  if (!record) {
+    return null;
+  }
+  return { ...record.stats, entries: record.cache.size };
+}
+
+/**
+ * Reads the most recent cache-miss keys for a device's memoizing assembler.
+ * Use this to find out what's varying when the cache misses more than
+ * expected — the key encodes `modules|defines|vs|fs`, so a diff between two
+ * miss keys identifies the perturbing input.
+ */
+export function getMemoShaderAssemblerMissLog(
+  device: Device,
+): readonly string[] {
+  const assembler = perDeviceAssembler.get(device);
+  if (!assembler) {
+    return [];
+  }
+  const record = perAssemblerRecord.get(assembler);
+  if (!record) {
+    return [];
+  }
+  return record.missLog;
+}
+
+/**
+ * Reads stats from the assembler directly. Exposed for testing —
+ * production callers should use {@link getMemoShaderAssemblerStats}.
+ */
+export function readAssemblerStats(
+  assembler: ShaderAssembler,
+): MemoShaderAssemblerStats | null {
+  const record = perAssemblerRecord.get(assembler);
+  if (!record) {
+    return null;
+  }
+  return { ...record.stats, entries: record.cache.size };
+}
+
+/**
+ * Wraps a `ShaderAssembler` so that repeated calls to `assembleGLSLShaderPair`
+ * with the same inputs return the same cached result. Exposed for testing —
+ * production callers should use {@link getMemoizingShaderAssembler}.
+ */
+export function createMemoizingShaderAssembler(
+  inner: ShaderAssembler,
+): ShaderAssembler {
+  const cache = new Map<string, AssembledPair>();
+  const stats: MemoShaderAssemblerStats = { hits: 0, misses: 0, entries: 0 };
+  const missLog: string[] = [];
+
+  const assembler = new Proxy(inner, {
+    get(target, prop, receiver) {
+      if (prop !== ASSEMBLE_GLSL) {
+        return Reflect.get(target, prop, receiver);
+      }
+      return (props: AssembleProps): AssembledPair => {
+        const key = computeCacheKey(props);
+        const hit = cache.get(key);
+        if (hit) {
+          stats.hits++;
+          return hit;
+        }
+        stats.misses++;
+        if (missLog.length < MAX_LOGGED_MISSES) {
+          missLog.push(key);
+        }
+        const result = target.assembleGLSLShaderPair.call(target, props);
+        cache.set(key, result);
+        stats.entries = cache.size;
+        return result;
+      };
+    },
+  });
+
+  perAssemblerRecord.set(assembler, { cache, stats, missLog });
+  return assembler;
+}
+
+function computeCacheKey(props: AssembleProps): string {
+  const moduleKey = (props.modules ?? [])
+    .map((module) => module.name)
+    .join("|");
+  const definesKey = props.defines ? stableStringify(props.defines) : "";
+  return `${moduleKey}::${definesKey}::${props.vs ?? ""}::${props.fs ?? ""}`;
+}
+
+function stableStringify(value: Record<string, unknown>): string {
+  const keys = Object.keys(value).sort();
+  const parts: string[] = [];
+  for (const key of keys) {
+    parts.push(`${key}=${String(value[key])}`);
+  }
+  return parts.join(",");
+}

packages/deck.gl-raster/tests/mesh-layer/shader-assembler-memo.test.ts

@@ -0,0 +1,110 @@
+import type { ShaderAssembler, ShaderModule } from "@luma.gl/shadertools";
+import { describe, expect, it, vi } from "vitest";
+import {
+  createMemoizingShaderAssembler,
+  readAssemblerStats,
+} from "../../src/mesh-layer/shader-assembler-memo.js";
+
+type AssembleProps = Parameters<ShaderAssembler["assembleGLSLShaderPair"]>[0];
+
+function fakeModule(name: string): ShaderModule {
+  return { name } as ShaderModule;
+}
+
+function fakeAssembler() {
+  const assembleGLSLShaderPair = vi.fn(
+    (
+      props: AssembleProps,
+    ): ReturnType<ShaderAssembler["assembleGLSLShaderPair"]> => ({
+      vs: `assembled-vs:${props.vs ?? ""}`,
+      fs: `assembled-fs:${props.fs ?? ""}`,
+      getUniforms: () => ({}),
+      modules: props.modules ?? [],
+    }),
+  );
+  const addShaderHook = vi.fn();
+  const inner = {
+    assembleGLSLShaderPair,
+    addShaderHook,
+  } as unknown as ShaderAssembler;
+  return { inner, assembleGLSLShaderPair, addShaderHook };
+}
+
+function baseProps(overrides: Partial<AssembleProps> = {}): AssembleProps {
+  return {
+    platformInfo: { shaderLanguage: "glsl" } as AssembleProps["platformInfo"],
+    vs: "void main() {}",
+    fs: "void main() {}",
+    modules: [],
+    ...overrides,
+  };
+}
+
+describe("createMemoizingShaderAssembler", () => {
+  it("returns the same assembled result for identical inputs", () => {
+    const { inner, assembleGLSLShaderPair } = fakeAssembler();
+    const memo = createMemoizingShaderAssembler(inner);
+    const modules = [fakeModule("createTexture"), fakeModule("cutlineBbox")];
+
+    const first = memo.assembleGLSLShaderPair(baseProps({ modules }));
+    const second = memo.assembleGLSLShaderPair(baseProps({ modules }));
+
+    expect(second).toBe(first);
+    expect(assembleGLSLShaderPair).toHaveBeenCalledTimes(1);
+  });
+
+  it("does not collapse calls with different modules", () => {
+    const { inner, assembleGLSLShaderPair } = fakeAssembler();
+    const memo = createMemoizingShaderAssembler(inner);
+
+    memo.assembleGLSLShaderPair(baseProps({ modules: [fakeModule("a")] }));
+    memo.assembleGLSLShaderPair(baseProps({ modules: [fakeModule("b")] }));
+
+    expect(assembleGLSLShaderPair).toHaveBeenCalledTimes(2);
+  });
+
+  it("does not collapse calls with different shader source", () => {
+    const { inner, assembleGLSLShaderPair } = fakeAssembler();
+    const memo = createMemoizingShaderAssembler(inner);
+
+    memo.assembleGLSLShaderPair(baseProps({ fs: "// pipeline A" }));
+    memo.assembleGLSLShaderPair(baseProps({ fs: "// pipeline B" }));
+
+    expect(assembleGLSLShaderPair).toHaveBeenCalledTimes(2);
+  });
+
+  it("treats `defines` order as stable so reordered keys hit the cache", () => {
+    const { inner, assembleGLSLShaderPair } = fakeAssembler();
+    const memo = createMemoizingShaderAssembler(inner);
+
+    memo.assembleGLSLShaderPair(baseProps({ defines: { A: true, B: false } }));
+    memo.assembleGLSLShaderPair(baseProps({ defines: { B: false, A: true } }));
+
+    expect(assembleGLSLShaderPair).toHaveBeenCalledTimes(1);
+  });
+
+  it("delegates non-memoized methods to the inner assembler", () => {
+    const { inner, addShaderHook } = fakeAssembler();
+    const memo = createMemoizingShaderAssembler(inner);
+
+    memo.addShaderHook("vs:DECKGL_FILTER_GL_POSITION");
+
+    expect(addShaderHook).toHaveBeenCalledWith("vs:DECKGL_FILTER_GL_POSITION");
+  });
+
+  it("tracks hit/miss/entries counters", () => {
+    const { inner } = fakeAssembler();
+    const memo = createMemoizingShaderAssembler(inner);
+
+    memo.assembleGLSLShaderPair(baseProps({ fs: "// A" }));
+    memo.assembleGLSLShaderPair(baseProps({ fs: "// A" }));
+    memo.assembleGLSLShaderPair(baseProps({ fs: "// A" }));
+    memo.assembleGLSLShaderPair(baseProps({ fs: "// B" }));
+
+    expect(readAssemblerStats(memo)).toEqual({
+      hits: 2,
+      misses: 2,
+      entries: 2,
+    });
+  });
+});