GPU Backend Abstraction

Typeclass-based backend abstraction. Same ShaderM kernels run on WebGPU or CUDA — selected by HESPER_BACKEND env var at runtime.

Design

GPUBackend β provides all GPU operations parameterized by context type β.

• WebGPU instance: β = Device, Buf β = Buffer
• CUDA instance: β = CUDAContext, Buf β = CUDABuffer

Layer/model code is written as:

def myLayer [GPUBackend β] (ctx : β) (input output : GPUBackend.Buf β) : IO Unit := do
  GPUBackend.execute ctx myKernel [("input", input), ("output", output)]
    (GPUBackend.dispatch1D 1024)

structure Hesper.ExecConfig : Type

Execution config — backend-agnostic. extensions/diagnostics are used by WebGPU (e.g., chromium_experimental_subgroup_matrix) and ignored by CUDA.

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

def Hesper.ExecConfig.dispatch1D (n : Nat) (wgSize : Nat := 256) : ExecConfig

Equations

• Hesper.ExecConfig.dispatch1D n wgSize = { workgroupSize := { x := wgSize }, numWorkgroups := ((n + wgSize - 1) / wgSize, 1, 1) }

def Hesper.ExecConfig.dispatch2D (nx ny : Nat) (bx by_ : Nat := 16) : ExecConfig

Equations

• Hesper.ExecConfig.dispatch2D nx ny bx by_ = { workgroupSize := { x := bx, y := by_ }, numWorkgroups := ((nx + bx - 1) / bx, (ny + by_ - 1) / by_, 1) }

class Hesper.GPUBackend (β : Type) : Type 1

Typeclass for GPU compute backends.

• Buf : Type
• CachedDispatch : Type
• CompiledKernel : Type
• executeWithConfig : β → WGSL.Monad.ShaderM Unit → List (String × Buf β) → ExecConfig → IO Unit
• executeWithConfigCached : β → WGSL.Monad.ShaderM Unit → List (String × Buf β) → ExecConfig → UInt64 → IO.Ref (Option (CachedDispatch β)) → IO Unit
• executeWithConfigCachedArrays : β → WGSL.Monad.ShaderM Unit → List (String × Buf β) → List (String × List (Buf β)) → ExecConfig → UInt64 → IO.Ref (Option (CachedDispatch β)) → IO Unit

  Variant that accepts both single-buffer bindings and buffer-array bindings (for cross-layer kernel fusion). Each entry in namedBufferArrays maps a declared bufferArray name to the N underlying buffers; the backend materialises a device-side pointer table and passes its pointer as the kernel argument. Default impl: falls back to the single-buffer path if there are no array bindings.

• replayCached : β → CachedDispatch β → Nat × Nat × Nat → IO Unit
• allocBuffer : β → USize → IO (Buf β)
• allocBufferUsage : β → USize → List String → IO (Buf β)
• bufferUid : β → Buf β → IO UInt64

  Stable identity of a buffer for tracing (JS-replay provenance); 0 = backend has no notion of identity.

• freeBuffer : β → Buf β → IO Unit
• writeBuffer : β → Buf β → ByteArray → IO Unit
• writeBufferOffset : β → Buf β → USize → ByteArray → IO Unit
• readBuffer : β → Buf β → USize → IO ByteArray
• buildKernel : β → WGSL.Monad.ShaderM Unit → ExecConfig → IO (CompiledKernel β)
• dispatchCompiledKernel : β → CompiledKernel β → Array (Buf β) → Nat × Nat × Nat → Option (IO.Ref (Option (CachedDispatch β))) → IO Unit
• hasSubgroupSupport : β → IO Bool
• hasShaderF16Support : β → IO Bool
• newCacheRef : IO (IO.Ref (Option (CachedDispatch β)))
• rawDevicePtr : β → Buf β → IO (Option USize)

  Raw device pointer for a buffer, if the backend has a native address space exposable as a USize. CUDA returns some (CUdeviceptr), WebGPU returns none. Used by Phase-0 hybrid path that calls externally-JIT'd PTX directly. Must not be used by portable code.

• bufFromRawDevicePtr : β → USize → USize → IO (Option (Buf β))

  Wrap an already-device-resident pointer (e.g. from cuMemHostGetDevicePointer or an externally-allocated CUdeviceptr) as a Buf without copying. CUDA returns some buf; WebGPU returns none (no native pointer abstraction). The pointer must remain valid for the buffer's lifetime; the backend will NOT free it (callers handle deregistration externally).

• beginBatch : β → IO Unit

  Begin recording dispatches for batch submission. CUDA: no-op.

• endBatch : β → IO Unit

  Submit all recorded dispatches and wait. CUDA: no-op (sync is per-launch).

@[inline]

def Hesper.GPUBackend.execute {β : Type} [GPUBackend β] (ctx : β) (computation : WGSL.Monad.ShaderM Unit) (namedBuffers : List (String × Buf β)) (config : ExecConfig) : IO Unit

Convenience: execute with ExecConfig

Equations

• Hesper.GPUBackend.execute ctx computation namedBuffers config = Hesper.GPUBackend.executeWithConfig ctx computation namedBuffers config

structure Hesper.CachedKernel (β : Type) [GPUBackend β] : Type

A kernel bundled with its dispatch cache. Production code should use this instead of bare GPUBackend.execute to ensure dispatch caching. Create with CachedKernel.create, dispatch with CachedKernel.exec.

• computation : WGSL.Monad.ShaderM Unit
• cacheKey : UInt64
• cacheRef : IO.Ref (Option (GPUBackend.CachedDispatch β))

def Hesper.CachedKernel.create {β : Type} [GPUBackend β] (computation : WGSL.Monad.ShaderM Unit) (name : String) : IO (CachedKernel β)

Create a cached kernel. Call once at init time (e.g., in createInferenceState).

Equations

• Hesper.CachedKernel.create computation name = do let __do_lift ← liftM (IO.mkRef none) pure { computation := computation, cacheKey := hash name, cacheRef := __do_lift }

@[inline]

def Hesper.CachedKernel.exec {β : Type} [GPUBackend β] (k : CachedKernel β) (ctx : β) (namedBuffers : List (String × GPUBackend.Buf β)) (config : ExecConfig) : IO Unit

Dispatch with caching. First call compiles; subsequent calls skip PTX generation.

Equations

• k.exec ctx namedBuffers config = Hesper.GPUBackend.executeWithConfigCached ctx k.computation namedBuffers config k.cacheKey k.cacheRef

def Hesper.GPUBackend.debugExecuteOnce {β : Type} [GPUBackend β] (ctx : β) (computation : WGSL.Monad.ShaderM Unit) (namedBuffers : List (String × Buf β)) (config : ExecConfig) : IO Unit

For debug/test use only. Creates an ephemeral cache (not reused).

Equations

• Hesper.GPUBackend.debugExecuteOnce ctx computation namedBuffers config = do let ref ← liftM (IO.mkRef none) Hesper.GPUBackend.executeWithConfigCached ctx computation namedBuffers config 0 ref

def Hesper.smartDispatch (totalThreads : Nat) (wgSize : Nat := 256) : ExecConfig × Nat

Smart dispatch: 1D if fits, 2D otherwise. Returns (config, gridDimX) — same signature as TTT.Kernels.smartDispatch.

Equations

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

inductive Hesper.BackendChoice : Type

Which backend to use

• WebGPU : BackendChoice
• CUDA : BackendChoice

instance Hesper.instBEqBackendChoice : BEq BackendChoice

Equations

• Hesper.instBEqBackendChoice = { beq := Hesper.instBEqBackendChoice.beq }

def Hesper.instBEqBackendChoice.beq : BackendChoice → BackendChoice → Bool

Equations

• Hesper.instBEqBackendChoice.beq x✝ y✝ = (x✝.ctorIdx == y✝.ctorIdx)

def Hesper.instReprBackendChoice.repr : BackendChoice → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.
• Hesper.instReprBackendChoice.repr Hesper.BackendChoice.CUDA prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Hesper.BackendChoice.CUDA")).group prec✝

instance Hesper.instReprBackendChoice : Repr BackendChoice

Equations

• Hesper.instReprBackendChoice = { reprPrec := Hesper.instReprBackendChoice.repr }

def Hesper.detectBackend : IO BackendChoice

Detect backend from HESPER_BACKEND env var. Default: WebGPU.

Equations

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