Binding WebGPU Render Passes to deck.gl Custom Layers

import { Layer, type LayerContext, type UpdateParameters } from "@deck.gl/core";

// The native handles deck.gl exposes when the WebGPU backend is active.
interface WebGPUContext extends LayerContext {
  device: GPUDevice;
  queue: GPUQueue;
  // luma.gl surfaces the canvas-backed context; getCurrentTexture() returns
  // the SAME swap-chain texture deck.gl rasterizes into this frame.
  gpuCanvasContext: GPUCanvasContext;
  depthTexture: GPUTexture; // deck.gl's shared depth attachment
}

interface RenderState {
  pipeline: GPURenderPipeline;
  bindGroup: GPUBindGroup;
  vertexBuffer: GPUBuffer;     // packed position data, 16-byte aligned
  uniformBuffer: GPUBuffer;    // view/projection matrices, 256-byte sized
  vertexCount: number;
  frameInFlight: Promise<undefined> | null; // the previous frame's fence
}

export class WebGPULayer extends Layer {
  static layerName = "WebGPULayer";
  declare state: RenderState;

  initializeState(context: LayerContext): void {
    const { device } = context as WebGPUContext;

    // Pre-validate the limits a large spatial tile grid will hit, BEFORE any
    // pipeline creation — a failed pipeline throws an opaque error otherwise.
    const required = this.props.tileBufferBytes as number;
    if (device.limits.maxStorageBufferBindingSize < required) {
      throw new Error("maxStorageBufferBindingSize too small for tile grid");
    }

    const pipeline = this.buildPipeline(device);

    // Allocate once. Per-frame createBuffer() is the top cause of allocator jank.
    const vertexBuffer = device.createBuffer({
      size: required,
      usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
    });
    // 256 bytes covers two mat4x4<f32> (view + projection) at the dynamic-uniform
    // offset granularity WebGPU enforces; see memory-alignment reference.
    const uniformBuffer = device.createBuffer({
      size: 256,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });

    const bindGroup = device.createBindGroup({
      layout: pipeline.getBindGroupLayout(0),
      entries: [{ binding: 0, resource: { buffer: uniformBuffer } }],
    });

    this.setState({
      pipeline,
      bindGroup,
      vertexBuffer,
      uniformBuffer,
      vertexCount: 0,
      frameInFlight: null,
    });
  }

  // deck.gl calls draw() once per rAF tick with the live viewport state.
  draw({ uniforms, context }: { uniforms: Record<string, unknown>; context: LayerContext }): void {
    const { device, queue, gpuCanvasContext, depthTexture } = context as WebGPUContext;
    const { pipeline, bindGroup, vertexBuffer, uniformBuffer, vertexCount } = this.state;
    if (vertexCount === 0) return;

    // Push the frame's view/projection matrices. writeBuffer is non-blocking and
    // the right channel for a <256 KB per-frame uniform on a moving viewport.
    const viewProjection = uniforms.viewProjectionMatrix as Float32Array;
    queue.writeBuffer(uniformBuffer, 0, viewProjection.buffer, viewProjection.byteOffset, 64);

    // Scope the encoder to THIS frame only — never reuse across rAF ticks.
    const encoder = device.createCommandEncoder({ label: "spatial-layer-frame" });

    // Reuse deck.gl's own swap-chain + depth attachment. Allocating a private
    // colour target here would force a full-screen blit and a context switch.
    const pass = encoder.beginRenderPass({
      colorAttachments: [{
        view: gpuCanvasContext.getCurrentTexture().createView(),
        loadOp: "load",   // 'load', not 'clear' — deck.gl already cleared the frame
        storeOp: "store",
      }],
      depthStencilAttachment: {
        view: depthTexture.createView(),
        depthLoadOp: "load",
        depthStoreOp: "store",
      },
    });

    pass.setPipeline(pipeline);              // cached — no recompile on pan/zoom
    pass.setBindGroup(0, bindGroup);         // cached — static layout
    pass.setVertexBuffer(0, vertexBuffer);
    pass.draw(vertexCount);
    pass.end();

    queue.submit([encoder.finish()]);

    // Fence: hold a handle to this frame's completion so the next buffer mutation
    // can wait on it rather than racing the GPU still reading these vertices.
    this.state.frameInFlight = queue.onSubmittedWorkDone();
  }

  // Stream new geometry only after the in-flight frame has drained.
  async updateGeometry(data: Float32Array): Promise<void> {
    const { queue, frameInFlight, vertexBuffer } = this.state;
    if (frameInFlight) await frameInFlight; // avoid write-while-read corruption
    queue.writeBuffer(vertexBuffer, 0, data.buffer, data.byteOffset, data.byteLength);
    this.setState({ vertexCount: data.length / 4 }); // vec4<f32> stride
  }

  finalizeState(): void {
    // Release native handles on unmount / hot-reload to prevent VRAM leaks.
    this.state.vertexBuffer?.destroy();
    this.state.uniformBuffer?.destroy();
  }

  private buildPipeline(device: GPUDevice): GPURenderPipeline {
    const module = device.createShaderModule({ code: this.props.wgsl as string });
    return device.createRenderPipeline({
      layout: "auto",
      vertex: {
        module,
        entryPoint: "vs_main",
        buffers: [{
          arrayStride: 16, // vec4<f32>: lon, lat, elevation, attribute
          attributes: [{ shaderLocation: 0, offset: 0, format: "float32x4" }],
        }],
      },
      fragment: { module, entryPoint: "fs_main", targets: [{ format: navigator.gpu.getPreferredCanvasFormat() }] },
      primitive: { topology: "point-list" },
      depthStencil: { format: "depth24plus", depthWriteEnabled: true, depthCompare: "less-equal" },
    });
  }
}

const querySet = device.createQuerySet({ type: "timestamp", count: 2 });
const resolveBuffer = device.createBuffer({
  size: 16, // 2 timestamps × 8 bytes (uint64)
  usage: GPUBufferUsage.QUERY_RESOLVE | GPUBufferUsage.COPY_SRC,
});
const readBuffer = device.createBuffer({
  size: 16,
  usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
});

// Declare timestamp writes on the pass descriptor (modern path).
const pass = encoder.beginRenderPass({
  ...renderPassDescriptor,
  timestampWrites: { querySet, beginningOfPassWriteIndex: 0, endOfPassWriteIndex: 1 },
});
// ... setPipeline / draw / end ...

// resolveQuerySet lives on the command encoder, NOT the pass encoder.
encoder.resolveQuerySet(querySet, 0, 2, resolveBuffer, 0);
encoder.copyBufferToBuffer(resolveBuffer, 0, readBuffer, 0, 16);
queue.submit([encoder.finish()]);

await queue.onSubmittedWorkDone();
await readBuffer.mapAsync(GPUMapMode.READ);
const ts = new BigUint64Array(readBuffer.getMappedRange());
const durationMs = Number(ts[1] - ts[0]) / 1e6; // nanoseconds -> ms
readBuffer.unmap();