WGSL Shader Execution Layer

Integration between ShaderM monad, code generation, and WebGPU execution.

This module provides high-level functions to:

1. Compile ShaderM computations to WGSL
2. Create GPU pipelines
3. Execute shaders with buffer management
4. Handle synchronization

Usage Pattern:

def myKernel : ShaderM Unit := do
  let gid ← globalId
  let idx := Exp.vecZ gid
  let _input ← declareInputBuffer "input" (.array (.scalar .f32) 1024)
  let _output ← declareOutputBuffer "output" (.array (.scalar .f32) 1024)
  let val ← readBuffer (ty := .scalar .f32) (n := 1024) "input" idx
  writeBuffer (ty := .scalar .f32) "output" idx (Exp.mul val (Exp.litF32 2.0))

-- Execute on GPU
executeShader device myKernel
  [("input", inputBuffer), ("output", outputBuffer)]
  {x := 256, y := 1, z := 1}
  (256, 1, 1)

structure Hesper.WGSL.Execute.NamedBuffer : Type

GPU Buffer with name for binding

• name : String
• buffer : WebGPU.Buffer

structure Hesper.WGSL.Execute.ExecutionConfig : Type

Execution configuration for compute shaders

• funcName : String
• workgroupSize : WorkgroupSize
• numWorkgroups : Nat × Nat × Nat
• extensions : List String
• diagnostics : List (String × String)

instance Hesper.WGSL.Execute.instInhabitedExecutionConfig : Inhabited ExecutionConfig

Equations

• Hesper.WGSL.Execute.instInhabitedExecutionConfig = { default := { numWorkgroups := (1, 1, 1) } }

def Hesper.WGSL.Execute.ExecutionConfig.default (numWorkgroups : Nat × Nat × Nat) : ExecutionConfig

Create default execution config with specified workgroup count

Equations

• Hesper.WGSL.Execute.ExecutionConfig.default numWorkgroups = { numWorkgroups := numWorkgroups }

def Hesper.WGSL.Execute.ExecutionConfig.dispatch1D (totalThreads : Nat) (workgroupSize : Nat := 256) : ExecutionConfig

Create config for 1D dispatch

Equations

• One or more equations did not get rendered due to their size.

Subgroup Feature Detection

Cached runtime check for subgroup support. Queried once per session, used to select between subgroup-based kernels and shared-memory fallback kernels.

opaque Hesper.WGSL.Execute.subgroupSupportRef : IO.Ref (Option Bool)

Cached subgroup support flag (queried once, then reused)

def Hesper.WGSL.Execute.hasSubgroupSupport (device : WebGPU.Device) : IO Bool

Check if the device supports subgroup operations (subgroupAdd, etc.). Result is cached after the first call.

Equations

• One or more equations did not get rendered due to their size.

opaque Hesper.WGSL.Execute.subgroupMatrixSupportRef : IO.Ref (Option Bool)

Cached SubgroupMatrix support flag

def Hesper.WGSL.Execute.hasSubgroupMatrixSupport (device : WebGPU.Device) : IO Bool

Check if the device supports Chromium experimental subgroup matrix operations (subgroup_matrix_{left,right,result}, subgroupMatrixLoad, subgroupMatrixMultiplyAccumulate, subgroupMatrixStore). Cached.

Equations

• One or more equations did not get rendered due to their size.

opaque Hesper.WGSL.Execute.shaderF16SupportRef : IO.Ref (Option Bool)

Cached ShaderF16 support flag

def Hesper.WGSL.Execute.hasShaderF16Support (device : WebGPU.Device) : IO Bool

Check if the device supports ShaderF16 (f16 values, f16 arithmetic).

Equations

• One or more equations did not get rendered due to their size.

Pipeline Cache

Caches compiled GPU pipelines keyed by WGSL source hash. Pipeline compilation is expensive (~1-5ms per shader). With ~270 dispatches per forward pass, caching eliminates 270-1350ms of per-token overhead.

structure Hesper.WGSL.Execute.CachedPipeline : Type

Cached GPU pipeline components

• shaderModule : WebGPU.ShaderModule
• bindGroupLayout : WebGPU.BindGroupLayout
• pipeline : WebGPU.ComputePipeline
• declaredNames : List String
• declaredModes : List Monad.BufferAccessMode

opaque Hesper.WGSL.Execute.pipelineCacheRef : IO.Ref (Std.HashMap UInt64 CachedPipeline)

Global pipeline cache: maps WGSL source hash to cached pipeline. HashMap — a linear Array scan cost ~ms/token at ~600 dispatches.

opaque Hesper.WGSL.Execute.cacheHitsRef : IO.Ref Nat

Pipeline cache hit counter

opaque Hesper.WGSL.Execute.cacheMissesRef : IO.Ref Nat

Pipeline cache miss counter

def Hesper.WGSL.Execute.getPipelineCacheStats : IO (Nat × Nat)

Get pipeline cache statistics: (hits, misses)

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.resetPipelineCache : IO Unit

Reset pipeline cache (call when device is destroyed or for benchmarking)

Equations

• Hesper.WGSL.Execute.resetPipelineCache = do ST.Ref.set Hesper.WGSL.Execute.pipelineCacheRef ∅ ST.Ref.set Hesper.WGSL.Execute.cacheHitsRef 0 ST.Ref.set Hesper.WGSL.Execute.cacheMissesRef 0

Bind Group Cache

Caches WebGPU BindGroups keyed by (pipeline hash, buffer IDs). BindGroup creation involves internal validation and allocation in the WebGPU runtime, costing ~20-30µs per call. With 572 dispatches per token, this adds ~12-17ms overhead.

Since inference reuses the same pipelines with the same pre-allocated buffers, bind groups are almost always identical across tokens. Caching eliminates 572 redundant createBindGroup calls per token.

opaque Hesper.WGSL.Execute.bindGroupCacheRef : IO.Ref (Std.HashMap UInt64 WebGPU.BindGroup)

Global bind group cache: maps (pipeline + buffer IDs) hash to cached BindGroup

opaque Hesper.WGSL.Execute.bgCacheHitsRef : IO.Ref Nat

Bind group cache hit/miss counters

opaque Hesper.WGSL.Execute.bgCacheMissesRef : IO.Ref Nat

def Hesper.WGSL.Execute.getBindGroupCacheStats : IO (Nat × Nat)

Get bind group cache statistics: (hits, misses)

Equations

• One or more equations did not get rendered due to their size.

PreparedDispatch (Graph Capture)

Pre-computed dispatch state for instant replay. Stores the pipeline and bind group so that subsequent tokens skip ALL Lean-side processing:

• No WGSL generation/lookup
• No buffer name matching
• No bind group key computation
• No cache lookups
• Just one FFI call: recordDispatch

Usage:

-- In layer struct:
structure BitLinear where
  ...
  prepared : IO.Ref (Option PreparedDispatch)

-- In forward function:
def forward (device : Device) (layer : BitLinear) ... := do
  let (wx, wy, wz) := computeWorkgroups ...
  -- Fast path: replay if prepared
  if let some p ← layer.prepared.get then
    replayPreparedDispatch device p wx wy wz
    return
  -- Slow path: full execution (first token only)
  executeShaderNamed device shader namedBuffers config cacheKey (some layer.prepared)

structure Hesper.WGSL.Execute.PreparedDispatch : Type

Pre-computed dispatch: pipeline + bind group, ready for instant replay. The trace* fields carry JS-trace identity through replay (DG_TRACE_JS); zero-cost defaults otherwise.

• pipeline : WebGPU.ComputePipeline
• bindGroup : WebGPU.BindGroup
• traceKey : UInt64
• traceName : String
• traceBufs : Array (String × UInt64)
• traceBufObjs : Array WebGPU.Buffer

Command Buffer Batching

Global batching mode: when enabled, executeShaderNamed records dispatches into a shared command encoder instead of creating individual command buffers. This eliminates per-dispatch overhead (encoder creation + submit + wait).

Usage:

beginBatch device           -- Create shared encoder
-- All executeShaderNamed calls now record instead of submit
layer1.forward device ...
layer2.forward device ...
endBatch device             -- Submit all + wait once

opaque Hesper.WGSL.Execute.batchEncoderRef : IO.Ref (Option WebGPU.CommandEncoder)

Global batch encoder: when some, executeShaderNamed records into it

opaque Hesper.WGSL.Execute.batchDispatchCountRef : IO.Ref Nat

Dispatch counter for current batch (for diagnostics)

def Hesper.WGSL.Execute.beginBatch (device : WebGPU.Device) : IO Unit

Begin command buffer batching. All subsequent executeShaderNamed calls will record into a shared encoder instead of submitting individually.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.endBatch (device : WebGPU.Device) : IO Unit

End command buffer batching. Submits all recorded dispatches and waits.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.flushBatch (device : WebGPU.Device) : IO Unit

Split the current batch WITHOUT a CPU sync: submit the recorded encoder (no wait) and start a fresh one. Keeps Dawn-on-Metal's per-encoder inter-pass barriers correct (it drops them in a very large single encoder) at a fraction of endBatch's cost. Cross-encoder buffer hazards are handled by queue ordering + the driver.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.isBatching : IO Bool

Check if currently in batch mode

Equations

• Hesper.WGSL.Execute.isBatching = do let __do_lift ← ST.Ref.get Hesper.WGSL.Execute.batchEncoderRef pure __do_lift.isSome

Section profiling

Lightweight per-section wall-clock timing for the "everything else" bucket breakdown. When sectionProfilingRef is true, withSection name act wraps act with monoNanos timestamps and accumulates total ns + call count into sectionTotalsRef keyed by name. Nested sections are allowed (the outer section's total includes nested ones). Only meaningful in unbatched mode where each dispatch auto-syncs.

opaque Hesper.WGSL.Execute.sectionProfilingRef : IO.Ref Bool

opaque Hesper.WGSL.Execute.sectionTotalsRef : IO.Ref (Array (String × UInt64 × Nat))

opaque Hesper.WGSL.Execute.sectionDispatchesRef : IO.Ref (Array (String × Nat × Nat))

Per-section kernel-dispatch counter. Populated by withSection when sectionProfilingRef is on. The caller must wire preDispatch / postDispatch (readers of the global dispatch counter) so withSection can compute the dispatches attributed to one section.

opaque Hesper.WGSL.Execute.sectionDispatchReader : IO.Ref (Option (IO Nat))

Callback to read the global dispatch counter. Set by the CUDA backend via registerDispatchCounter. When none, per-section dispatch counting is disabled.

def Hesper.WGSL.Execute.registerDispatchCounter (reader : IO Nat) : IO Unit

Equations

• Hesper.WGSL.Execute.registerDispatchCounter reader = ST.Ref.set Hesper.WGSL.Execute.sectionDispatchReader (some reader)

def Hesper.WGSL.Execute.resetSectionTotals : IO Unit

Equations

• Hesper.WGSL.Execute.resetSectionTotals = do ST.Ref.set Hesper.WGSL.Execute.sectionTotalsRef #[] ST.Ref.set Hesper.WGSL.Execute.sectionDispatchesRef #[]

def Hesper.WGSL.Execute.getSectionTotals : IO (Array (String × UInt64 × Nat))

Equations

• Hesper.WGSL.Execute.getSectionTotals = ST.Ref.get Hesper.WGSL.Execute.sectionTotalsRef

def Hesper.WGSL.Execute.getSectionDispatches : IO (Array (String × Nat × Nat))

Equations

• Hesper.WGSL.Execute.getSectionDispatches = ST.Ref.get Hesper.WGSL.Execute.sectionDispatchesRef

opaque Hesper.WGSL.Execute.sectionTraceRef : IO.Ref Bool

Chrome-trace event log for withSection calls. When sectionTraceRef is on, every withSection invocation appends an entry (name, start_ns, dur_ns, depth). Dumped to a JSON file by dumpSectionTrace at end-of-run. Enabled by env HESPER_TRACE_OUT=<path>.

opaque Hesper.WGSL.Execute.sectionTraceEvents : IO.Ref (List (String × UInt64 × UInt64 × Nat))

opaque Hesper.WGSL.Execute.sectionTraceDepth : IO.Ref Nat

def Hesper.WGSL.Execute.dumpSectionTrace (path : String) : IO Unit

Dump accumulated section-trace events to a Chrome-trace-format JSON fragment file (just the traceEvents array body — the merger in scripts/nsys_to_chrome_trace.py reads it).

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.withSection {α : Type} (name : String) (act : IO α) : IO α

Time an IO action under a section name; no-op when profiling disabled.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.jsTraceCksum (device : WebGPU.Device) (uids : Array (String × UInt64)) (bufObjs : Array WebGPU.Buffer) : IO Unit

DG_TRACE_JS_CKSUM: XOR of per-buffer FNV32 over each bound buffer's first 4KB, emitted after every armed dispatch — the replayer mirrors it to find the FIRST diverging dispatch exactly. Forces a flush + sync per dispatch (trace runs only).

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.replayPreparedDispatch (device : WebGPU.Device) (prepared : PreparedDispatch) (wx wy wz : Nat) : IO Unit

Replay a prepared dispatch directly. Skips ALL Lean-side processing. Works in both batch mode (record into shared encoder) and standalone mode.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.compileToWGSL (computation : Monad.ShaderM Unit) (funcName : String := "main") (workgroupSize : WorkgroupSize := { x := 256 }) (extensions : List String := []) (diagnostics : List (String × String) := []) : String

Compile a ShaderM computation to WGSL source code

Equations

• Hesper.WGSL.Execute.compileToWGSL computation funcName workgroupSize extensions diagnostics = Hesper.WGSL.CodeGen.generateWGSL funcName workgroupSize extensions diagnostics computation

CompiledKernel (Zero-Overhead Dispatch API)

Separates shader compilation from dispatch. A CompiledKernel holds the compiled pipeline and binding layout, ready for buffer binding and dispatch without any string matching, WGSL regeneration, or cache lookups.

Usage:

-- At initialization (once):
let kernel ← buildKernel device myShaderM config
let bg ← bindKernel device kernel [("input", inBuf), ("output", outBuf)]

-- At dispatch time (hot loop, zero overhead):
dispatchKernel device kernel bg (numWorkgroups, 1, 1)

-- Or combine into PreparedDispatch for even fewer indirections:
let prepared := kernel.prepare bg
replayPreparedDispatch device prepared wx wy wz

structure Hesper.WGSL.Execute.CompiledKernel : Type

Pre-compiled kernel: pipeline + layout + binding order. Created once via buildKernel, reused across dispatches.

• pipeline : WebGPU.ComputePipeline
• bindGroupLayout : WebGPU.BindGroupLayout
• declaredNames : Array String
• sourceHash : UInt64

def Hesper.WGSL.Execute.CompiledKernel.prepare (kernel : CompiledKernel) (bindGroup : WebGPU.BindGroup) : PreparedDispatch

Create a PreparedDispatch from this kernel and a bind group

Equations

• kernel.prepare bindGroup = { pipeline := kernel.pipeline, bindGroup := bindGroup }

def Hesper.WGSL.Execute.buildKernel (device : WebGPU.Device) (computation : Monad.ShaderM Unit) (config : ExecutionConfig) : IO CompiledKernel

Compile a ShaderM computation into a reusable CompiledKernel. Uses the global pipeline cache. Thread-safe for repeated calls.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.bindKernel (device : WebGPU.Device) (kernel : CompiledKernel) (namedBuffers : List (String × WebGPU.Buffer)) : IO WebGPU.BindGroup

Create a BindGroup by matching named buffers to a CompiledKernel's bindings. Uses the global bind group cache.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.bindKernelDirect (device : WebGPU.Device) (kernel : CompiledKernel) (buffers : Array WebGPU.Buffer) : IO WebGPU.BindGroup

Create a BindGroup from pre-sorted buffer array (no name matching). Buffers must be in binding order (matching kernel.declaredNames).

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.dispatchKernel (device : WebGPU.Device) (kernel : CompiledKernel) (bindGroup : WebGPU.BindGroup) (numWorkgroups : Nat × Nat × Nat) : IO Unit

Dispatch a compiled kernel with a pre-built BindGroup. Zero string matching. Works in both batch mode and standalone mode.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.createShaderFromComputation (device : WebGPU.Device) (computation : Monad.ShaderM Unit) (config : ExecutionConfig) : IO WebGPU.ShaderModule

Create shader module from ShaderM computation

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.executeShaderNamed (device : WebGPU.Device) (computation : Monad.ShaderM Unit) (namedBuffers : List (String × WebGPU.Buffer)) (config : ExecutionConfig) (cacheKey : Option UInt64 := none) (preparedRef : Option (IO.Ref (Option PreparedDispatch)) := none) : IO Unit

Execute a ShaderM computation on the GPU with named buffers.

This is the main high-level execution function. It:

1. Compiles the ShaderM computation to WGSL
2. Looks up or creates the GPU pipeline (cached)
3. Binds buffers by name
4. Dispatches the compute shader
5. Waits for completion

Parameters:

• device: GPU device
• computation: ShaderM monad defining the shader
• namedBuffers: List of (name, buffer) pairs for binding
• config: Execution configuration (workgroup size, dispatch size)

Example:

let kernel : ShaderM Unit := do
  let gid ← globalId
  let idx := Exp.vecZ gid
  let _input ← declareInputBuffer "input" (.array (.scalar .f32) 1024)
  let _output ← declareOutputBuffer "output" (.array (.scalar .f32) 1024)
  let val ← readBuffer (ty := .scalar .f32) (n := 1024) "input" idx
  writeBuffer (ty := .scalar .f32) "output" idx (Exp.mul val (Exp.litF32 2.0))

executeShaderNamed device kernel
  [("input", inputBuf), ("output", outputBuf)]
  (ExecutionConfig.dispatch1D 1024)

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.executeShaderRecorded (device : WebGPU.Device) (encoder : WebGPU.CommandEncoder) (computation : Monad.ShaderM Unit) (namedBuffers : List (String × WebGPU.Buffer)) (config : ExecutionConfig) : IO Unit

Record a ShaderM computation into a command encoder (no submit, no wait).

This is the batched variant of executeShaderNamed. Instead of creating its own command encoder and waiting, it records the dispatch into a pre-existing encoder. The caller must call submitAndWait after recording all dispatches.

Pipeline caching is shared with executeShaderNamed.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.executeShaderSimple (device : WebGPU.Device) (computation : Monad.ShaderM Unit) (inputBuffer outputBuffer : WebGPU.Buffer) (numThreads : Nat) : IO Unit

Execute a simple ShaderM computation with a single input/output buffer.

Convenience wrapper for the common case of one input buffer and one output buffer.

Example:

let kernel : ShaderM Unit := do
  let gid ← globalId
  let idx := Exp.vecZ gid
  let _input ← declareInputBuffer "input" (.array (.scalar .f32) 1024)
  let _output ← declareOutputBuffer "output" (.array (.scalar .f32) 1024)
  let val ← readBuffer (ty := .scalar .f32) (n := 1024) "input" idx
  writeBuffer (ty := .scalar .f32) "output" idx (Exp.mul val (Exp.litF32 2.0))

executeShaderSimple device kernel inputBuf outputBuf 1024

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Execute.debugPrintWGSL (computation : Monad.ShaderM Unit) (config : ExecutionConfig := ExecutionConfig.default (1, 1, 1)) : IO Unit

Print generated WGSL for debugging

Equations

• One or more equations did not get rendered due to their size.