Shader Monad for Imperative Shader Construction

The ShaderM monad provides an imperative interface for building WGSL compute shaders. It tracks:

• Accumulated statements
• Fresh variable generation
• Shared memory declarations
• Automatic buffer binding

Usage pattern:

def myShader : ShaderM Unit := do
  let gid ← globalId
  let idx := swizzleX gid

  -- Declare buffers
  input ← declareInputBuffer "input" (.array (.scalar .f32) 1024)
  output ← declareOutputBuffer "output" (.array (.scalar .f32) 1024)

  -- Read, compute, write
  val ← readBuffer input idx
  let result := val * litF32 2.0
  writeBuffer output idx result

inductive Hesper.WGSL.Monad.BufferAccessMode : Type

Buffer access mode for storage buffers

• read : BufferAccessMode
• readWrite : BufferAccessMode

instance Hesper.WGSL.Monad.instReprBufferAccessMode : Repr BufferAccessMode

Equations

• Hesper.WGSL.Monad.instReprBufferAccessMode = { reprPrec := Hesper.WGSL.Monad.instReprBufferAccessMode.repr }

def Hesper.WGSL.Monad.instReprBufferAccessMode.repr : BufferAccessMode → Nat → Std.Format

Equations

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

instance Hesper.WGSL.Monad.instBEqBufferAccessMode : BEq BufferAccessMode

Equations

• Hesper.WGSL.Monad.instBEqBufferAccessMode = { beq := Hesper.WGSL.Monad.instBEqBufferAccessMode.beq }

def Hesper.WGSL.Monad.instBEqBufferAccessMode.beq : BufferAccessMode → BufferAccessMode → Bool

Equations

• Hesper.WGSL.Monad.instBEqBufferAccessMode.beq x✝ y✝ = (x✝.ctorIdx == y✝.ctorIdx)

structure Hesper.WGSL.Monad.ShaderState : Type

Shader Construction State

• stmts : List Stmt
• varCounter : Nat
• sharedVars : List (String × WGSLType)
• declaredBuffers : List (String × WGSLType × BufferAccessMode)
• needsSubgroups : Bool
• maxnreg : Option Nat

  Optional __launch_bounds__ analogue: caps register allocation per thread. CUDA backend lowers to PTX .maxnreg. None = ptxas default (which may pick a higher count and reduce occupancy).

• minnctapersm : Option Nat

  Optional minimum CTAs per SM hint. Combined with maxntid (derived from workgroup size), this is the full __launch_bounds__(threads, minCtas) equivalent — ptxas uses it to size register allocation.

@[reducible, inline]

abbrev Hesper.WGSL.Monad.ShaderM (α : Type) : Type

The Shader Monad

Equations

• Hesper.WGSL.Monad.ShaderM α = StateM Hesper.WGSL.Monad.ShaderState α

def Hesper.WGSL.Monad.ShaderM.initialState : ShaderState

Initial state for shader construction

Equations

• Hesper.WGSL.Monad.ShaderM.initialState = { stmts := [], varCounter := 0, sharedVars := [], declaredBuffers := [] }

def Hesper.WGSL.Monad.ShaderM.run {α : Type} (m : ShaderM α) : α × ShaderState

Run a shader computation and extract the result and final state

Equations

• m.run = m Hesper.WGSL.Monad.ShaderM.initialState

def Hesper.WGSL.Monad.ShaderM.exec {α : Type} (m : ShaderM α) : ShaderState

Run a shader computation and extract only the final state

Equations

• m.exec = (m Hesper.WGSL.Monad.ShaderM.initialState).snd

def Hesper.WGSL.Monad.ShaderM.emitStmt (stmt : Stmt) : ShaderM Unit

Emit a statement to the shader body

Equations

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

def Hesper.WGSL.Monad.ShaderM.freshVar (pfx : String) : ShaderM String

Generate a fresh variable name with given prefix

Equations

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

def Hesper.WGSL.Monad.ShaderM.captureStmts {α : Type} (action : ShaderM α) : ShaderM (α × List Stmt)

Capture statements from a monadic action (for control flow)

Equations

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

def Hesper.WGSL.Monad.ShaderM.var (ty : WGSLType) (init : Exp ty) : ShaderM String

Declare a private variable with fresh name

Equations

• Hesper.WGSL.Monad.ShaderM.var ty init = do let name ← Hesper.WGSL.Monad.ShaderM.freshVar "v" Hesper.WGSL.Monad.ShaderM.emitStmt (Hesper.WGSL.Stmt.varDecl name ty (some ⟨ty, init⟩)) pure name

def Hesper.WGSL.Monad.ShaderM.varNamed (name : String) (ty : WGSLType) (init : Exp ty) : ShaderM Unit

Declare a named private variable

Equations

• Hesper.WGSL.Monad.ShaderM.varNamed name ty init = Hesper.WGSL.Monad.ShaderM.emitStmt (Hesper.WGSL.Stmt.varDecl name ty (some ⟨ty, init⟩))

def Hesper.WGSL.Monad.ShaderM.varRef (ty : WGSLType) (init : Exp ty) : ShaderM (String × Exp ty)

Declare a mutable variable with fresh name, returning both name and typed expression. This avoids manually constructing Exp.var name and passing raw string literals.

Equations

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

def Hesper.WGSL.Monad.ShaderM.let' (ty : WGSLType) (init : Exp ty) : ShaderM (Exp ty)

Bind a sub-expression to a named PTX register and return an Exp.var referring to it. Use this when an Exp is referenced multiple times inside an unrolled loop body or fan-out: without it, each reference is inlined as a fresh AST traversal and the CodeGen emits the same arithmetic at every use site (causing the 17× instruction inflation seen in V9 vs llama.cpp). Wrapping the value with let' materialises it as one PTX register that all uses share.

Example:

let kRowBase ← ShaderM.let' (.scalar .u32) (Exp.add base offset)
for pk in [0:4] do
  -- `kRowBase` is a single Exp.var, so the four iterations all
  -- read the same register instead of recomputing `base + offset`.
  let idx := Exp.add kRowBase (Exp.litU32 (pk * 32))
  ...

Functionally equivalent to do let n ← var ty e; pure (Exp.var n), but the intent is clearer at the call site.

Equations

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

def Hesper.WGSL.Monad.ShaderM.shared (ty : WGSLType) : ShaderM String

Declare shared memory (workgroup-scoped) with fresh name

Equations

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

def Hesper.WGSL.Monad.ShaderM.sharedNamed (name : String) (ty : WGSLType) : ShaderM Unit

Declare named shared memory

Equations

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

def Hesper.WGSL.Monad.ShaderM.assign {ty : WGSLType} (varName : String) (expr : Exp ty) : ShaderM Unit

Assign expression to variable

Equations

• Hesper.WGSL.Monad.ShaderM.assign varName expr = Hesper.WGSL.Monad.ShaderM.emitStmt (Hesper.WGSL.Stmt.assign varName ty expr)

def Hesper.WGSL.Monad.ShaderM.assignIndex {ty : WGSLType} (arrName : String) (idx : Exp (WGSLType.scalar ScalarType.u32)) (value : Exp ty) : ShaderM Unit

Assign to array index

Equations

• Hesper.WGSL.Monad.ShaderM.assignIndex arrName idx value = Hesper.WGSL.Monad.ShaderM.emitStmt (Hesper.WGSL.Stmt.assignIndex arrName idx ty value)

def Hesper.WGSL.Monad.ShaderM.if_ (cond : Exp (WGSLType.scalar ScalarType.bool)) (thenBranch elseBranch : ShaderM Unit) : ShaderM Unit

If-then-else statement

Equations

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

def Hesper.WGSL.Monad.ShaderM.while_ (_cond : Exp (WGSLType.scalar ScalarType.bool)) (body : ShaderM Unit) : ShaderM Unit

While loop - implemented as for loop with always-true condition

Equations

• Hesper.WGSL.Monad.ShaderM.while_ _cond body = do let __discr ← body.captureStmts match __discr with | (fst, bodyStmts) => Hesper.WGSL.Monad.ShaderM.emitStmt (Hesper.WGSL.Stmt.block bodyStmts)

def Hesper.WGSL.Monad.ShaderM.for_ (varName : String) (start end_ step : Exp (WGSLType.scalar ScalarType.u32)) (body : ShaderM Unit) : ShaderM Unit

For loop (start to end, incrementing by step) Builds proper WGSL for loop: for (var i: u32 = start; i < end; i = i + step)

Equations

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

@[inline]

def Hesper.WGSL.Monad.ShaderM.unrollFor (n : Nat) (body : Nat → ShaderM Unit) : ShaderM Unit

Compile-time unrolled loop: emits n copies of the body inline, with i bound to a Lean Nat (0, 1, ..., n-1). Use when iteration count is statically known and the unroll is desired (eg. dot-product over a fixed dim count, sub-warp partition iters).

The Nat i is a meta value — to use it in an Exp context, lift via Exp.litU32 i or simply (i : Exp _) (the Exp OfNat instance).

Equivalent to writing for i in [0:n] do body i, but the helper name documents intent: "this is meta-time unrolled, not a runtime loop". Mirrors CUDA C++ #pragma unroll for (int i = 0; i < N; ++i).

For runtime loops use ShaderM.loop (or its alias runtimeFor).

Equations

• Hesper.WGSL.Monad.ShaderM.unrollFor n body = (List.range n).forM body

LICM helpers (loop invariant code motion)

Approximate "free variable" check: serialise an Exp via Exp.toWGSL and look for the variable name as a whole token (name not preceded/followed by an alphanumeric or underscore). This is correct for ShaderM-emitted names because they're unique fresh names like v3 / i7 — there's no risk of accidental match against a substring of another identifier.

Why string-match rather than a structural traversal: Exp has 50+ constructors and many of them hold sub-Exps. A structural Exp.containsVar would be a lot of pattern matching that would have to stay in sync as new Exp cases are added. toWGSL is the existing single source of truth that walks every sub-Exp; reusing it for the free-var check costs O(n) extra string work per stmt but is robust to Exp evolution. In practice the runtime impact is negligible because LICM only runs at kernel-emission time, not on the GPU.

def Hesper.WGSL.Monad.ShaderM.loop (start end_ step : Exp (WGSLType.scalar ScalarType.u32)) (bodyFn : Exp (WGSLType.scalar ScalarType.u32) → ShaderM Unit) : ShaderM Unit

Higher-order loop: pass loop variable as Exp. Now applies LICM: varDecls in the body whose init expression doesn't reference the loop variable (or any other inner-loop binding) are hoisted before the for-loop. This automates the manual let'-outside-the-loop pattern that V9 / earlier kernels used by hand.

Usage: loop start end step fun i => do { ... use i ... }

Equations

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

def Hesper.WGSL.Monad.ShaderM.loopWithLICM (start end_ step : Exp (WGSLType.scalar ScalarType.u32)) (bodyFn : Exp (WGSLType.scalar ScalarType.u32) → ShaderM Unit) : ShaderM Unit

Variant of loop that applies LICM to hoist varDecls whose init is loop-invariant out of the for-loop. Use only when the body contains no shared-memory writes that other iterations read, no inter-iteration carries via externally-declared mutable vars beyond those captured by assign, and only varDecl (no if/inner for/barrier) sites that you want hoisted.

Semantics-changing in subtle ways — verify bit-parity before merging a kernel that uses this.

Equations

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

@[inline]

def Hesper.WGSL.Monad.ShaderM.runtimeFor (start end_ step : Exp (WGSLType.scalar ScalarType.u32)) (bodyFn : Exp (WGSLType.scalar ScalarType.u32) → ShaderM Unit) : ShaderM Unit

Runtime loop alias for loop, named to pair with unrollFor so the meta-vs-runtime distinction is explicit at the call site:

ShaderM.unrollFor 8 fun i => ...               -- meta-time, 8 inline copies
ShaderM.runtimeFor 0 cacheLen 1 fun i => ...   -- runtime for-loop

Mirrors CUDA C++ for (int i = 0; ...) (no #pragma unroll).

Equations

• Hesper.WGSL.Monad.ShaderM.runtimeFor start end_ step bodyFn = Hesper.WGSL.Monad.ShaderM.loop start end_ step bodyFn

def Hesper.WGSL.Monad.ShaderM.scope {α : Type} (body : ShaderM α) : ShaderM α

Block scope: emits { ... body ... } so var declared inside the body is block-scoped in WGSL, allowing Naga/Tint/Vulkan-driver register allocator to reuse the physical register once the scope exits.

This is the analog of CUDA { ... } block scope for register-pressure reduction. For example, when a temporary is needed only inside a hot inner loop iteration, putting it in a scope instead of at function level lets the compiler reuse its register slot afterwards.

Usage:

ShaderM.scope do
  let tmp ← ShaderM.var (.scalar .f32) (Exp.litF32 0.0)
  -- ... use tmp ...
  -- tmp's register is released at end of scope

Note: Lean-side let bindings to tmp outside the do-block can still refer to the var name string, but reading from it via Exp.var after scope-exit is undefined behavior — the WGSL register is gone. Discipline: only use tmp inside the scope body.

Equations

• body.scope = do let __discr ← body.captureStmts match __discr with | (result, bodyStmts) => do Hesper.WGSL.Monad.ShaderM.emitStmt (Hesper.WGSL.Stmt.block bodyStmts) pure result

@[inline]

def Hesper.WGSL.Monad.ShaderM.unrollForScoped (n : Nat) (body : Nat → ShaderM Unit) : ShaderM Unit

unrollFor + per-iter scope. Each of the n inline copies of the body lives in its own { ... } block, so any ShaderM.var declared inside is block-scoped and the WGSL→PTX backend can reuse the physical register across iters.

Collapses the V8/V9/V11 idiom

for i in [0:n] do ShaderM.scope do
  ...

into

ShaderM.unrollForScoped n fun i => do
  ...

Mirrors CUDA C++ #pragma unroll for (int i = ...) where each unrolled body is implicitly its own register-allocation scope.

Equations

• Hesper.WGSL.Monad.ShaderM.unrollForScoped n body = (List.range n).forM fun (i : Nat) => (body i).scope

def Hesper.WGSL.Monad.ShaderM.barrier : ShaderM Unit

Workgroup barrier (synchronization)

Equations

• Hesper.WGSL.Monad.ShaderM.barrier = Hesper.WGSL.Monad.ShaderM.emitStmt (Hesper.WGSL.Stmt.exprStmt Hesper.WGSL.Exp.workgroupBarrier)

def Hesper.WGSL.Monad.ShaderM.warpBarrier : ShaderM Unit

Warp-level barrier (CUDA __syncwarp()). PTX backend emits the cheap bar.warp.sync 0xFFFFFFFF; WGSL backend falls back to workgroupBarrier(). Use when only intra-warp ordering is needed — eg. between Phase 1 (KQ dot) and Phase 2a (cross-sub-warp max reduce) in the flash-attn vec kernel where all 32 lanes of a warp must finish writing before reading.

Equations

• Hesper.WGSL.Monad.ShaderM.warpBarrier = Hesper.WGSL.Monad.ShaderM.emitStmt (Hesper.WGSL.Stmt.exprStmt Hesper.WGSL.Exp.warpBarrier)

def Hesper.WGSL.Monad.ShaderM.setMaxnreg (n : Nat) : ShaderM Unit

Cap per-thread register usage at n. ptxas will spill to local memory if it would otherwise exceed this. Used for occupancy tuning — fewer registers per thread → more blocks per SM. CUDA-only (lowers to PTX .maxnreg). WGSL backend ignores.

Equations

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

def Hesper.WGSL.Monad.ShaderM.setMinnctapersm (n : Nat) : ShaderM Unit

Promise that ≥ n CTAs will run per SM. Combined with the workgroup size (which gives ptxas .maxntid), this is the __launch_bounds__ equivalent — lets ptxas pick a register budget that fits N blocks.

Equations

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

def Hesper.WGSL.Monad.ShaderM.bufferAddr (bufName : String) (elemSize : Nat) (idx : Exp (WGSLType.scalar ScalarType.u32)) : Exp (WGSLType.scalar ScalarType.u64)

Pure: raw u64 pointer to element idx of a global buffer. elemSize is the byte size of one element (typically 4 for u32/f32). Used as the global-address operand for cpAsync. CUDA-only.

Equations

• Hesper.WGSL.Monad.ShaderM.bufferAddr bufName elemSize idx = Hesper.WGSL.Exp.bufferAddr bufName elemSize idx

def Hesper.WGSL.Monad.ShaderM.sharedSymAddr (smemName : String) (elemSize : Nat) (idx : Exp (WGSLType.scalar ScalarType.u32)) : Exp (WGSLType.scalar ScalarType.u32)

Pure: raw u32 byte-address of element idx in a shared-memory array. elemSize must be 4, 8, or 16. Used as the smem-address operand for cpAsync. CUDA-only.

Equations

• Hesper.WGSL.Monad.ShaderM.sharedSymAddr smemName elemSize idx = Hesper.WGSL.Exp.sharedSymAddr smemName elemSize idx

def Hesper.WGSL.Monad.ShaderM.cpAsync (smemAddr : Exp (WGSLType.scalar ScalarType.u32)) (globalAddr : Exp (WGSLType.scalar ScalarType.u64)) (bytes : Nat) : ShaderM Unit

── cp.async (sm_80+) ── Issue an async global→shared memory copy of bytes bytes (must be 4, 8, or 16). Non-blocking — completion synchronised via cpAsyncCommit + cpAsyncWait. Used by llama.cpp's MMQ pipeline to overlap the next K-iteration's load with the current K-iteration's compute. WGSL backend has no equivalent.

Equations

• Hesper.WGSL.Monad.ShaderM.cpAsync smemAddr globalAddr bytes = Hesper.WGSL.Monad.ShaderM.emitStmt (Hesper.WGSL.Stmt.exprStmt (smemAddr.cpAsyncCgSharedGlobal globalAddr bytes))

def Hesper.WGSL.Monad.ShaderM.cpAsyncCa (smemAddr : Exp (WGSLType.scalar ScalarType.u32)) (globalAddr : Exp (WGSLType.scalar ScalarType.u64)) (bytes : Nat) : ShaderM Unit

cp.async.ca cache-all variant — supports bytes ∈ {4, 8, 16}. Slower per byte than .cg for 16-byte transfers but the only option for narrower transfers.

Equations

• Hesper.WGSL.Monad.ShaderM.cpAsyncCa smemAddr globalAddr bytes = Hesper.WGSL.Monad.ShaderM.emitStmt (Hesper.WGSL.Stmt.exprStmt (smemAddr.cpAsyncCaSharedGlobal globalAddr bytes))

def Hesper.WGSL.Monad.ShaderM.cpAsyncCommit : ShaderM Unit

cp.async.commit_group — bundle all preceding cpAsync issues by this thread into one group for later cpAsyncWait.

Equations

• Hesper.WGSL.Monad.ShaderM.cpAsyncCommit = Hesper.WGSL.Monad.ShaderM.emitStmt (Hesper.WGSL.Stmt.exprStmt Hesper.WGSL.Exp.cpAsyncCommitGroup)

def Hesper.WGSL.Monad.ShaderM.cpAsyncWait (n : Nat) : ShaderM Unit

cp.async.wait_group N — block this thread until all but the most recent N committed groups have completed. N=0 waits for all.

Equations

• Hesper.WGSL.Monad.ShaderM.cpAsyncWait n = Hesper.WGSL.Monad.ShaderM.emitStmt (Hesper.WGSL.Stmt.exprStmt (Hesper.WGSL.Exp.cpAsyncWaitGroup n))

def Hesper.WGSL.Monad.ShaderM.globalId : ShaderM (Exp (WGSLType.vec3 ScalarType.u32))

Global invocation ID (3D)

Equations

• Hesper.WGSL.Monad.ShaderM.globalId = pure (Hesper.WGSL.Exp.var "global_invocation_id")

def Hesper.WGSL.Monad.ShaderM.localId : ShaderM (Exp (WGSLType.vec3 ScalarType.u32))

Local invocation ID (3D)

Equations

• Hesper.WGSL.Monad.ShaderM.localId = pure (Hesper.WGSL.Exp.var "local_invocation_id")

def Hesper.WGSL.Monad.ShaderM.workgroupId : ShaderM (Exp (WGSLType.vec3 ScalarType.u32))

Workgroup ID (3D)

Equations

• Hesper.WGSL.Monad.ShaderM.workgroupId = pure (Hesper.WGSL.Exp.var "workgroup_id")

def Hesper.WGSL.Monad.ShaderM.numWorkgroups : ShaderM (Exp (WGSLType.vec3 ScalarType.u32))

Number of workgroups (3D)

Equations

• Hesper.WGSL.Monad.ShaderM.numWorkgroups = pure (Hesper.WGSL.Exp.var "num_workgroups")

Thread index helpers (CUDA-style shortcuts)

These mirror CUDA's threadIdx.x, blockIdx.x, etc. for the common 1D case so kernel code reads close to the CUDA original. The vec3 versions (localId, workgroupId) are still available when 2D/3D index components are needed.

@[inline]

def Hesper.WGSL.Monad.ShaderM.tidX : ShaderM (Exp (WGSLType.scalar ScalarType.u32))

threadIdx.x — local invocation X coordinate (= "tid" in CUDA kernels).

Equations

• Hesper.WGSL.Monad.ShaderM.tidX = do let __do_lift ← Hesper.WGSL.Monad.ShaderM.localId pure __do_lift.vec3X

@[inline]

def Hesper.WGSL.Monad.ShaderM.tidY : ShaderM (Exp (WGSLType.scalar ScalarType.u32))

threadIdx.y

Equations

• Hesper.WGSL.Monad.ShaderM.tidY = do let __do_lift ← Hesper.WGSL.Monad.ShaderM.localId pure __do_lift.vec3Y

@[inline]

def Hesper.WGSL.Monad.ShaderM.tidZ : ShaderM (Exp (WGSLType.scalar ScalarType.u32))

threadIdx.z

Equations

• Hesper.WGSL.Monad.ShaderM.tidZ = do let __do_lift ← Hesper.WGSL.Monad.ShaderM.localId pure __do_lift.vecZ

@[inline]

def Hesper.WGSL.Monad.ShaderM.bidX : ShaderM (Exp (WGSLType.scalar ScalarType.u32))

blockIdx.x — workgroup X coordinate.

Equations

• Hesper.WGSL.Monad.ShaderM.bidX = do let __do_lift ← Hesper.WGSL.Monad.ShaderM.workgroupId pure __do_lift.vec3X

@[inline]

def Hesper.WGSL.Monad.ShaderM.bidY : ShaderM (Exp (WGSLType.scalar ScalarType.u32))

blockIdx.y

Equations

• Hesper.WGSL.Monad.ShaderM.bidY = do let __do_lift ← Hesper.WGSL.Monad.ShaderM.workgroupId pure __do_lift.vec3Y

@[inline]

def Hesper.WGSL.Monad.ShaderM.bidZ : ShaderM (Exp (WGSLType.scalar ScalarType.u32))

blockIdx.z

Equations

• Hesper.WGSL.Monad.ShaderM.bidZ = do let __do_lift ← Hesper.WGSL.Monad.ShaderM.workgroupId pure __do_lift.vecZ

Warp / sub-warp index decomposition

NVIDIA warps are 32 lanes. These helpers compute laneId (tid & 31) and warpId (tid >>> 5) once and return them as Exp values, so kernel code that uses them many times shares one PTX register. The inner expressions are simple enough that lean codegen folds the extraction efficiently — no let' needed at the use site.

Sub-warp partition (used in V8/V11 attention vec kernels) splits a warp into N-lane groups via subWarpSplit n.

@[inline]

def Hesper.WGSL.Monad.ShaderM.laneId : ShaderM (Exp (WGSLType.scalar ScalarType.u32))

Lane index within the warp (0..31). Equivalent to CUDA's threadIdx.x & 31 when threads are 1D, or threadIdx.x % WARP_SIZE.

Equations

• Hesper.WGSL.Monad.ShaderM.laneId = do let __do_lift ← Hesper.WGSL.Monad.ShaderM.tidX pure (__do_lift &&& Hesper.WGSL.Exp.litU32 31)

@[inline]

def Hesper.WGSL.Monad.ShaderM.warpId : ShaderM (Exp (WGSLType.scalar ScalarType.u32))

Warp index within the workgroup. Equivalent to CUDA's threadIdx.x / WARP_SIZE (= >> 5). For 2D blocks where threadIdx.y selects warps, use tidY directly.

Equations

• Hesper.WGSL.Monad.ShaderM.warpId = do let __do_lift ← Hesper.WGSL.Monad.ShaderM.tidX pure (__do_lift >>> Hesper.WGSL.Exp.litU32 5)

@[inline]

def Hesper.WGSL.Monad.ShaderM.subWarpSplit (n : Nat) : ShaderM (Exp (WGSLType.scalar ScalarType.u32) × Exp (WGSLType.scalar ScalarType.u32))

Split a warp's lane into (subWarp, subLane) for nthreads_KQ-style sub-warp partitioning where each sub-warp handles a different K position in parallel. Returns (laneId / n, laneId % n).

Example (V11 sub-warp partition with n = 8):

let (sw, sl) ← ShaderM.subWarpSplit 8
-- sw ∈ [0, 4), sl ∈ [0, 8)
let kPos := warpBase + sw * 8 + iKQ0    -- this sub-warp's K
let dimBase := sl * 32                  -- this lane's D slice

Mirrors llama.cpp's (threadIdx.x & ~(nthreads_KQ-1), threadIdx.x % nthreads_KQ).

Equations

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

Warp-level reductions

warpReduceSum n e reduces e across the n-lane group within the warp:

• n = 32 (full warp) → emits a single subgroupAdd (5-shfl on PTX, hardware wide reduction on WGSL).
• n = 8 (sub-warp, eg. nthreads_KQ from llama.cpp) → emits 3 shfl-xor by 1, 2, 4. All 8 lanes end with the sum of their 8-lane group.
• Other n (must be power of 2, ≤ 32) → log2(n) shfl-xor butterflies.

After the call, every lane in the n-lane group holds the same reduced value. Mirrors llama.cpp's warp_reduce_sum<nthreads> template.

Example (V11-like sub-warp dot product):

let partialVar ← ShaderM.mutVar (.scalar .f32) 0
ShaderM.unrollFor dimsPerLane fun k => partialVar +↦ q[k]! * kVec[k]!
let sum ← ShaderM.warpReduceSum 8 partialVar.read   -- 3 shfl

def Hesper.WGSL.Monad.ShaderM.warpReduceSum {ty : WGSLType} (n : Nat) (e : Exp ty) : ShaderM (Exp ty)

Sum-reduce e across an n-lane group within the warp. Uses Exp.subgroupAdd when n = 32, otherwise emits log2(n) butterfly shuffles via Exp.subgroupShuffleXor.

Pre: n is a power of 2 and 1 ≤ n ≤ 32.

Each lane in the n-lane group receives the same sum. Cost: log2(n) shfl + n−1 add per lane (so a sub-warp reduce of 8 lanes is 3 shfl

• 3 add — half the cost of the full 32-lane reduce, useful when only 8-lane groups need to coordinate).

Equations

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

def Hesper.WGSL.Monad.ShaderM.warpReduceMax (n : Nat) (e : Exp (WGSLType.scalar ScalarType.f32)) : ShaderM (Exp (WGSLType.scalar ScalarType.f32))

Max-reduce variant: warpReduceMax n e. Same shfl-xor butterfly, using Exp.max instead of +. After the call, every lane in the group holds the same max.

Equations

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

def Hesper.WGSL.Monad.ShaderM.softmaxOnlineUpdate (kqMaxName kqSumName : String) (kqScore : Exp (WGSLType.scalar ScalarType.f32)) : ShaderM (Exp (WGSLType.scalar ScalarType.f32) × Exp (WGSLType.scalar ScalarType.f32) × Exp (WGSLType.scalar ScalarType.f32))

One step of online softmax (Milakov & Gimelshein 2018 / FlashAttention). Given the running max kqMax, running sum kqSum, and a new score kqScore, computes:

kqMaxNew  = max(kqMax, kqScore)
scale     = exp(kqMax - kqMaxNew)        -- multiplier for old VKQ accumulator
kqExp     = exp(kqScore - kqMaxNew)      -- per-thread weight for current K
kqSumNew  = kqSum * scale + kqExp
kqMax := kqMaxNew                        -- updated in-place
kqSum := kqSumNew                        -- updated in-place

Returns (kqMaxNew, scale, kqExp). The caller rescales their VKQ accumulators by scale and accumulates kqExp * V for this K position.

kqMaxName/kqSumName are the names of the running registers so we can assign them in place. The new max is also returned as an Exp so the caller can use it without an extra Exp.var lookup.

Mirrors the inner-loop accumulation in llama.cpp's flash_attn_ext_vec (lines 273, 287-291 in fattn-vec.cuh).

Equations

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

def Hesper.WGSL.Monad.ShaderM.readBuffer {ty : WGSLType} {n : Nat} (bufferName : String) (idx : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM (Exp ty)

Read from a global storage buffer at index Note: You need to provide the element type explicitly

Equations

• Hesper.WGSL.Monad.ShaderM.readBuffer bufferName idx = pure ((Hesper.WGSL.Exp.var bufferName).index idx)

def Hesper.WGSL.Monad.ShaderM.readBufferByte {n : Nat} (bufferName : String) (byteIdx : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM (Exp (WGSLType.scalar ScalarType.u32))

Read a single byte (zero-extended to u32) from a buffer declared as array<u32, n>, addressed by byte index. Lowers to one ld.global.u8 on CUDA; emulated via u32-load+shift+mask on WGSL. Essential for Q6_K scale reads (avoids issuing a full u32 load per byte).

Equations

• Hesper.WGSL.Monad.ShaderM.readBufferByte bufferName byteIdx = pure (Hesper.WGSL.Exp.loadByteFromU32Buf bufferName byteIdx)

def Hesper.WGSL.Monad.ShaderM.readBufferU16 {n : Nat} (bufferName : String) (byteIdx : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM (Exp (WGSLType.scalar ScalarType.u32))

Read a halfword (16 bits, zero-extended to u32) from a buffer declared as array<u32, n>, addressed by byte index. Lowers to ld.global.u16 on CUDA. Use for fp16 block scales.

Equations

• Hesper.WGSL.Monad.ShaderM.readBufferU16 bufferName byteIdx = pure (Hesper.WGSL.Exp.loadU16FromU32Buf bufferName byteIdx)

def Hesper.WGSL.Monad.ShaderM.readBufferU32x4 (bufferName : String) (u32Idx : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM (Exp (WGSLType.scalar ScalarType.u32) × Exp (WGSLType.scalar ScalarType.u32) × Exp (WGSLType.scalar ScalarType.u32) × Exp (WGSLType.scalar ScalarType.u32))

128-bit (4× u32) wide read from a buffer declared array<u32, n>, addressed by u32 index. The starting index must be 4-aligned and the buffer pointer 16-byte aligned (caller's responsibility — typical callers iterate by pk * 4).

On CUDA this lowers to a single ld.global.nc.v4.u32 instruction (one MIO op delivering 16 bytes), 4× more efficient than four scalar ld.global.u32 reads in MIO-pipe-saturated kernels (FlashAttn V11). On WGSL it emulates as four scalar reads.

Returns four Exp (.scalar .u32) referring to fresh declared vars. Each value can be used independently with Exp.unpack2x16float etc.

Equations

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

def Hesper.WGSL.Monad.ShaderM.readWorkgroupF32x4 (sharedName : String) (f32Idx : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM (Exp (WGSLType.scalar ScalarType.f32) × Exp (WGSLType.scalar ScalarType.f32) × Exp (WGSLType.scalar ScalarType.f32) × Exp (WGSLType.scalar ScalarType.f32))

128-bit (4× f32) wide read from shared memory. f32Idx must be 4-aligned. CUDA lowers to one ld.shared.v4.f32; WGSL emulates as four scalar reads.

Used to relieve MIO-pipe saturation on shared-memory traffic (FlashAttn V11 partial aggregation: 4× LDS instruction count reduction).

Equations

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

def Hesper.WGSL.Monad.ShaderM.writeBuffer {ty : WGSLType} (bufferName : String) (idx : Exp (WGSLType.scalar ScalarType.u32)) (value : Exp ty) : ShaderM Unit

Write to a global storage buffer at index

Equations

• Hesper.WGSL.Monad.ShaderM.writeBuffer bufferName idx value = Hesper.WGSL.Monad.ShaderM.assignIndex bufferName idx value

def Hesper.WGSL.Monad.ShaderM.readWorkgroup {ty : WGSLType} {n : Nat} (sharedName : String) (idx : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM (Exp ty)

Read from workgroup shared memory at index

Equations

• Hesper.WGSL.Monad.ShaderM.readWorkgroup sharedName idx = pure ((Hesper.WGSL.Exp.var sharedName).index idx)

def Hesper.WGSL.Monad.ShaderM.writeWorkgroup {ty : WGSLType} (sharedName : String) (idx : Exp (WGSLType.scalar ScalarType.u32)) (value : Exp ty) : ShaderM Unit

Write to workgroup shared memory at index

Equations

• Hesper.WGSL.Monad.ShaderM.writeWorkgroup sharedName idx value = Hesper.WGSL.Monad.ShaderM.assignIndex sharedName idx value

def Hesper.WGSL.Monad.ShaderM.declareInputBuffer (name : String) (ty : WGSLType) : ShaderM String

Declare an input buffer (read-only) with automatic binding assignment

Equations

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

def Hesper.WGSL.Monad.ShaderM.declareOutputBuffer (name : String) (ty : WGSLType) : ShaderM String

Declare an output buffer (read-write) with automatic binding assignment

Equations

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

def Hesper.WGSL.Monad.ShaderM.declareStorageBuffer (name : String) (ty : WGSLType) (mode : BufferAccessMode := BufferAccessMode.readWrite) : ShaderM String

Declare a storage buffer with explicit access mode

Equations

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

def Hesper.WGSL.Monad.ShaderM.declareReadOnlyBuffer (name : String) (ty : WGSLType) : ShaderM String

Declare a read-only storage buffer.

Semantically equivalent to declareInputBuffer on WGSL, but on CUDA this emits ld.global.nc.* (the read-only L1/tex cache hint, aka __ldg). Use for weight matrices and other buffers the kernel never writes to. It is a logical error to call writeBuffer on the returned name.

Equations

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

def Hesper.WGSL.Monad.ShaderM.declareInputBufferArray (name : String) (elemTy : WGSLType) (n : Nat) : ShaderM String

Declare a read-only buffer array binding: N separate runtime-sized storage buffers addressable by runtime layer index. The host side must pass a device-pointer-table — N × 8 bytes holding each layer's CUdeviceptr — as the kernel argument bound to name. Inside the kernel, read elements with readBufferArray name layerIdx elemIdx.

This is the primitive that lets a single kernel iterate over all 42 transformer layers, collapsing 42 dispatches into 1. The existing per-layer kernels remain; use bufferArray only when a kernel actually wants to fuse across layers.

Equations

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

def Hesper.WGSL.Monad.ShaderM.readBufferArray {elemTy : WGSLType} {n : Nat} (name : String) (bufIdx elemIdx : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM (Exp elemTy)

Read element elemIdx from the bufIdx-th buffer of a bufferArray. Emits a single-indirection load via the pointer table.

Equations

• Hesper.WGSL.Monad.ShaderM.readBufferArray name bufIdx elemIdx = pure ((Hesper.WGSL.Exp.var name).indexBuf bufIdx elemIdx)

def Hesper.WGSL.Monad.ShaderM.writeBufferArray {ty : WGSLType} (name : String) (bufIdx elemIdx : Exp (WGSLType.scalar ScalarType.u32)) (value : Exp ty) : ShaderM Unit

Write value to arr[bufIdx][elemIdx] where arr is a bufferArray.

Equations

• Hesper.WGSL.Monad.ShaderM.writeBufferArray name bufIdx elemIdx value = Hesper.WGSL.Monad.ShaderM.emitStmt (Hesper.WGSL.Stmt.assignIndexBuf name bufIdx elemIdx ty value)

def Hesper.WGSL.Monad.ShaderM.staticFor {α β : Type} (xs : List α) (f : α → ShaderM β) : ShaderM Unit

Compile-time loop for unrolling (Haskell-side loop, not WGSL loop)

Equations

• Hesper.WGSL.Monad.ShaderM.staticFor xs f = xs.forM fun (x : α) => f x *> pure ()

def Hesper.WGSL.Monad.ShaderM.litF (x : Float) : Exp (WGSLType.scalar ScalarType.f32)

Float32 literal

Equations

• Hesper.WGSL.Monad.ShaderM.litF x = Hesper.WGSL.Exp.litF32 x

def Hesper.WGSL.Monad.ShaderM.litI (x : Int) : Exp (WGSLType.scalar ScalarType.i32)

Int32 literal

Equations

• Hesper.WGSL.Monad.ShaderM.litI x = Hesper.WGSL.Exp.litI32 x

def Hesper.WGSL.Monad.ShaderM.litU (x : Nat) : Exp (WGSLType.scalar ScalarType.u32)

UInt32 literal

Equations

• Hesper.WGSL.Monad.ShaderM.litU x = Hesper.WGSL.Exp.litU32 x

def Hesper.WGSL.Monad.ShaderM.declareMatrixLeftArray (name : String) (st : ScalarType) (m k count : Nat) (init : Exp (WGSLType.subgroupMatrixLeft st m k)) : ShaderM Unit

Declare an array of subgroup_matrix_left matrices with initialization

Equations

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

def Hesper.WGSL.Monad.ShaderM.declareMatrixRightArray (name : String) (st : ScalarType) (k n count : Nat) (init : Exp (WGSLType.subgroupMatrixRight st k n)) : ShaderM Unit

Declare an array of subgroup_matrix_right matrices with initialization

Equations

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

def Hesper.WGSL.Monad.ShaderM.declareMatrixResultArray (name : String) (st : ScalarType) (m n count : Nat) (init : Exp (WGSLType.subgroupMatrixResult st m n)) : ShaderM Unit

Declare an array of subgroup_matrix_result matrices with initialization

Equations

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

def Hesper.WGSL.Monad.ShaderM.loadMatrixLeft {st : ScalarType} {m k : Nat} (arrayName : String) (index : Nat) (bufferName : String) (offset stride : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM Unit

Load subgroup_matrix_left from buffer

Example: Ax[i] = subgroupMatrixLoad<subgroup_matrix_left<f32,8,8>>(&A, offset, false, stride)

Equations

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

def Hesper.WGSL.Monad.ShaderM.loadMatrixRight {st : ScalarType} {k n : Nat} (arrayName : String) (index : Nat) (bufferName : String) (offset stride : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM Unit

Load subgroup_matrix_right from buffer

Equations

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

def Hesper.WGSL.Monad.ShaderM.loadMatrixRightT {st : ScalarType} {k n : Nat} (arrayName : String) (index : Nat) (bufferName : String) (offset : Exp (WGSLType.scalar ScalarType.u32)) (colMajor : Bool) (stride : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM Unit

loadMatrixRight with an explicit column-major flag — for loading a right (k×n) fragment DIRECTLY from a storage buffer laid out n-major (e.g. Bᵀ [N,K] row-major: element (k,n) at b[n·K+k] ⇒ colMajor=true, stride=K). Buffer-space subgroupMatrixLoad is supported by the pinned May Tint (verified; lowers to device simdgroup_load).

Equations

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

def Hesper.WGSL.Monad.ShaderM.matrixMultiplyAccumulate {st : ScalarType} {m k n : Nat} (resultArrayName : String) (resultIndex : Nat) (leftArrayName : String) (leftIndex : Nat) (rightArrayName : String) (rightIndex : Nat) : ShaderM Unit

Perform matrix multiply-accumulate: acc = a * b + acc

Example: accxx[idx] = subgroupMatrixMultiplyAccumulate(Ax[i], Bx[j], accxx[idx])

Equations

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

def Hesper.WGSL.Monad.ShaderM.matrixMultiplyAccumulateMixed {inSt outSt : ScalarType} {m k n : Nat} (resultArrayName : String) (resultIndex : Nat) (leftArrayName : String) (leftIndex : Nat) (rightArrayName : String) (rightIndex : Nat) : ShaderM Unit

Mixed-precision variant of matrixMultiplyAccumulate: A and B use inSt, C / D use outSt. This matches NVIDIA cooperative matrix's native (f16, f16) → f32 config.

Equations

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

def Hesper.WGSL.Monad.ShaderM.storeMatrixResult {st : ScalarType} {m n : Nat} (arrayName : String) (index : Nat) (bufferName : String) (offset stride : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM Unit

Store subgroup_matrix_result to buffer

Example: subgroupMatrixStore(&C, offset, accxx[idx], false, stride)

Equations

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

Fragment-as-local-variable helpers (real-CUDA-style WMMA)

The array-of-fragments helpers above (loadMatrixLeft, matrixMultiplyAccumulate, storeMatrixResult) match WGSL's array<subgroup_matrix_left, N> model — nice for blocked algorithms but awkward to lower to PTX since fragments live in registers, not memory. The helpers below match the more direct nvcuda::wmma::fragment a_frag; load_matrix_sync(a_frag, ...); style: each fragment is a single named variable, with no array indirection. This is the recommended path for new WMMA kernels — the array helpers remain for backward-compat and WGSL-targeted code.

def Hesper.WGSL.Monad.ShaderM.loadFragmentLeft (st : ScalarType) (m k : Nat) (name bufferName : String) (offset stride : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM Unit

Declare a left fragment (subgroup_matrix_left<st, m, k>) bound to a name, initialized via a load from bufferName at byte-offset offset with row-major stride (in elements).

Equations

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

def Hesper.WGSL.Monad.ShaderM.loadFragmentRight (st : ScalarType) (k n : Nat) (name bufferName : String) (offset stride : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM Unit

Declare a right fragment, loaded col-major (the WGSL subgroupMatrixLoadRight instruction, which lowers to PTX wmma.load.b...col.f16).

Equations

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

def Hesper.WGSL.Monad.ShaderM.declareFragmentResultZero (st : ScalarType) (m n : Nat) (name : String) : ShaderM Unit

Declare a result fragment initialized to zero. Maps to PTX mov.f32 %f0..7, 0; for the f32 case (8 regs).

Equations

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

def Hesper.WGSL.Monad.ShaderM.fragmentMultiplyAccumulate (st : ScalarType) (m k n : Nat) (cName aName bName : String) : ShaderM Unit

cName ← cName * (a × b). Uses the existing Exp.subgroupMatrix- MultiplyAccumulate; lowers to one wmma.mma.sync instruction and rebinds cName to the result fragment.

Equations

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

def Hesper.WGSL.Monad.ShaderM.storeFragmentResult (st : ScalarType) (m n : Nat) (cName bufferName : String) (offset stride : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM Unit

Store a result fragment to a buffer at byte-offset offset with row-major stride (in elements).

Equations

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

def Hesper.WGSL.Monad.ShaderM.staticLoop (count : Nat) (body : Nat → ShaderM Unit) : ShaderM Unit

Static loop unrolling for matrix operations

Executes an action for each index in the range [0, count)

Equations

• Hesper.WGSL.Monad.ShaderM.staticLoop count body = do forIn [:count] PUnit.unit fun (i : Nat) (r : PUnit) => do body i pure (ForInStep.yield PUnit.unit) pure PUnit.unit

def Hesper.WGSL.Monad.ShaderM.staticLoop2D (rows cols : Nat) (body : Nat → Nat → ShaderM Unit) : ShaderM Unit

Static nested loop for 2D iteration

Equations

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

Typed mutable variables

MutVar ty wraps a private (PTX-register / WGSL var) variable with its type carried in Lean. This closes the second cognitive gap from porting CUDA C++ — operator overloading in Step 1 still needed (Exp.var name : Exp _) type annotations everywhere. With MutVar, v.read + x and v +↦ x type-check directly because read returns a typed Exp ty.

Idiomatic use:

let acc ← ShaderM.mutVar (.scalar .f32) 0
for k in [0:8] do
  acc +↦ qVec[k]! * kVec[k]!     -- in-place +=
let final := acc.read

structure Hesper.WGSL.Monad.MutVar (ty : WGSLType) : Type

Mutable shader-side variable carrying its type in Lean. Construct via ShaderM.mutVar — never by hand (the contained name must be one emitted by ShaderM.var so PTX/WGSL codegen sees a declaration).

• name : String

instance Hesper.WGSL.Monad.instReprMutVar {ty✝ : WGSLType} : Repr (MutVar ty✝)

Equations

• Hesper.WGSL.Monad.instReprMutVar = { reprPrec := Hesper.WGSL.Monad.instReprMutVar.repr }

def Hesper.WGSL.Monad.instReprMutVar.repr {ty✝ : WGSLType} : MutVar ty✝ → Nat → Std.Format

Equations

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

@[inline]

def Hesper.WGSL.Monad.MutVar.read {ty : WGSLType} (v : MutVar ty) : Exp ty

Read the current value as a typed Exp ty.

Equations

• v.read = Hesper.WGSL.Exp.var v.name

@[inline]

def Hesper.WGSL.Monad.MutVar.write {ty : WGSLType} (v : MutVar ty) (x : Exp ty) : ShaderM Unit

Overwrite the variable with x.

Equations

• v ↦= x = Hesper.WGSL.Monad.ShaderM.assign v.name x

@[inline]

def Hesper.WGSL.Monad.MutVar.addAssign {ty : WGSLType} (v : MutVar ty) [HAdd (Exp ty) (Exp ty) (Exp ty)] (x : Exp ty) : ShaderM Unit

In-place add: v += x. Sugar for v.write (v.read + x).

Equations

• v +↦ x = Hesper.WGSL.Monad.ShaderM.assign v.name (v.read + x)

@[inline]

def Hesper.WGSL.Monad.MutVar.mulAssign {ty : WGSLType} (v : MutVar ty) [HMul (Exp ty) (Exp ty) (Exp ty)] (x : Exp ty) : ShaderM Unit

In-place multiply: v *= x.

Equations

• v *↦ x = Hesper.WGSL.Monad.ShaderM.assign v.name (v.read * x)

@[inline]

def Hesper.WGSL.Monad.MutVar.subAssign {ty : WGSLType} (v : MutVar ty) [HSub (Exp ty) (Exp ty) (Exp ty)] (x : Exp ty) : ShaderM Unit

In-place subtract: v -= x.

Equations

• v.subAssign x = Hesper.WGSL.Monad.ShaderM.assign v.name (v.read - x)

def Hesper.WGSL.Monad.«term_+↦_» : Lean.TrailingParserDescr

Mutating-add operator. v +↦ x ≡ MutVar.addAssign v x.

Equations

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

def Hesper.WGSL.Monad.«term_*↦_» : Lean.TrailingParserDescr

Mutating-multiply operator.

Equations

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

def Hesper.WGSL.Monad.«term_↦=_» : Lean.TrailingParserDescr

Mutating-write operator (≡ MutVar.write). v ↦= x.

Equations

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

def Hesper.WGSL.Monad.ShaderM.mutVar (ty : WGSLType) (init : Exp ty) : ShaderM (MutVar ty)

Declare a typed mutable variable. Use v.read / v.write / v +↦ x / v ↦= x. The successor to ShaderM.var.

Example:

let acc ← ShaderM.mutVar (.scalar .f32) 0
for k in [0:dimsPerLane] do
  acc +↦ qVec[k]! * kVec[k]!

Equations

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

Buffer pointers (Ptr ty)

Ptr ty is a (bufferName, offset) pair carrying its element type so the load / store calls don't need explicit (ty := ...) annotations each time. Mirrors CUDA C++ T *p arithmetic — particularly the common attention/matmul pattern:

const float *K = K_base + kvHead * maxSeq * D;
for (int k = 0; k < cacheLen; ++k) {
    sum += *K * q;
    K += D;
}

In ShaderM:

let K ← ShaderM.ptr (.scalar .f32) "k_cache" (kvHead * maxSeqLen * D)
ShaderM.runtimeFor 0 cacheLen 1 fun _ => do
  let v ← K.load
  acc +↦ v * q
  K := K.advance D    -- value-level advance (Ptr is immutable)

For mutating advance inside a runtime loop, use MutPtr (declared below). For pure pointer arithmetic at meta-time, use Ptr.atOffset.

Ptr is intentionally pure (no ShaderM effect). The constructor ShaderM.ptr only computes a let'-bound base offset for safety, which is the common case.

structure Hesper.WGSL.Monad.Ptr (ty : WGSLType) : Type

Buffer pointer: typed (buffer name, current u32 offset, declared array size) triple.

• buf: buffer name registered via declareInputBuffer etc.
• offset: current element index (NOT byte) into that buffer.
• bufLen: the buffer's declared array length (passed to readBuffer's {n} for the WGSL array<ty, n> type). Set this to the same value used at declareInputBuffer time.

Ptr is value-level (no ShaderM effect needed for arithmetic). Use ShaderM.ptr to construct one with a let'-stored base.

• buf : String
• offset : Exp (WGSLType.scalar ScalarType.u32)
• bufLen : Nat

@[inline]

def Hesper.WGSL.Monad.Ptr.load {ty : WGSLType} (p : Ptr ty) : ShaderM (Exp ty)

Element-wise read at the current pointer offset. Mirrors CUDA *p.

Equations

• p.load = Hesper.WGSL.Monad.ShaderM.readBuffer p.buf p.offset

@[inline]

def Hesper.WGSL.Monad.Ptr.store {ty : WGSLType} (p : Ptr ty) (x : Exp ty) : ShaderM Unit

Element-wise write at the current pointer offset. Mirrors CUDA *p = x.

Equations

• p.store x = Hesper.WGSL.Monad.ShaderM.writeBuffer p.buf p.offset x

@[inline]

def Hesper.WGSL.Monad.Ptr.advance {ty : WGSLType} (p : Ptr ty) (delta : Exp (WGSLType.scalar ScalarType.u32)) : Ptr ty

Advance the pointer by delta elements, returning a new Ptr. Mirrors CUDA p + delta (or p += delta if you reassign). Pure — no shader-side effect.

Equations

• p.advance delta = { buf := p.buf, offset := p.offset + delta, bufLen := p.bufLen }

@[inline]

def Hesper.WGSL.Monad.Ptr.atOffset {ty : WGSLType} (p : Ptr ty) (extra : Exp (WGSLType.scalar ScalarType.u32)) : Ptr ty

Read at p[k] without permanently advancing. Mirrors CUDA p[k]. Composes well with unrollFor:

ShaderM.unrollFor n fun k =>
  let v ← (p.atOffset (Exp.litU32 k)).load
  ...

Equations

• p.atOffset extra = { buf := p.buf, offset := p.offset + extra, bufLen := p.bufLen }

def Hesper.WGSL.Monad.ShaderM.ptr (ty : WGSLType) (buf : String) (bufLen : Nat) (baseOffset : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM (Ptr ty)

Construct a Ptr with a base offset materialised in a register (via let'). Use when the base offset is computed once (eg. kvHead * maxSeqLen * D) and the resulting pointer is then walked many times in inner loops.

bufLen must match the array size declared via declareInputBuffer for buf.

Example:

let K ← ShaderM.ptr (.scalar .u32) "k_cache_f16" kvWords
          (kvHead * maxSeqLen * (D/2))
-- K.offset is a single PTX register; subsequent advance/atOffset
-- only emit `add` against this register.

Equations

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

structure Hesper.WGSL.Monad.MutPtr (ty : WGSLType) : Type

Mutable pointer — like Ptr but offset is a var (read/write register) instead of an Exp, so advance actually updates the register in-place rather than constructing a new value.

Mirrors the CUDA outer-loop idiom

K += blockIdx.y * nthreads * nb11;          // initial
for (int k = ...; k < kmax; k += step,
     K += step * nb11, V += step * nb21) {  // advance per iter
    sum += vec_dot_KQ(K + i_KQ * nb11, ...);// inner per-iter offset
}

Step 6's Ptr only handles the inner per-iter offset; MutPtr handles the outer per-iter pointer advance.

• buf : String
• offsetVar : String

  Name of the u32 register that holds the current offset. Mutable — gets reassigned by MutPtr.advance.

• bufLen : Nat

def Hesper.WGSL.Monad.instReprMutPtr.repr {ty✝ : WGSLType} : MutPtr ty✝ → Nat → Std.Format

Equations

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

instance Hesper.WGSL.Monad.instReprMutPtr {ty✝ : WGSLType} : Repr (MutPtr ty✝)

Equations

• Hesper.WGSL.Monad.instReprMutPtr = { reprPrec := Hesper.WGSL.Monad.instReprMutPtr.repr }

@[inline]

def Hesper.WGSL.Monad.MutPtr.offset {ty : WGSLType} (p : MutPtr ty) : Exp (WGSLType.scalar ScalarType.u32)

Current offset as an Exp.var so it can be used in index arithmetic without materialising another register.

Equations

• p.offset = Hesper.WGSL.Exp.var p.offsetVar

@[inline]

def Hesper.WGSL.Monad.MutPtr.load {ty : WGSLType} (p : MutPtr ty) : ShaderM (Exp ty)

Read at the current offset.

Equations

• p.load = Hesper.WGSL.Monad.ShaderM.readBuffer p.buf p.offset

@[inline]

def Hesper.WGSL.Monad.MutPtr.loadAt {ty : WGSLType} (p : MutPtr ty) (extra : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM (Exp ty)

Read at p.offset + extra without advancing. Mirrors CUDA p[k] inside an unrolled inner loop.

Equations

• p.loadAt extra = Hesper.WGSL.Monad.ShaderM.readBuffer p.buf (p.offset + extra)

@[inline]

def Hesper.WGSL.Monad.MutPtr.advance {ty : WGSLType} (p : MutPtr ty) (delta : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM Unit

Advance the pointer by delta. The next load / loadAt reads from the new offset. Mirrors CUDA p += delta.

Equations

• p.advance delta = Hesper.WGSL.Monad.ShaderM.assign p.offsetVar (p.offset + delta)

@[inline]

def Hesper.WGSL.Monad.MutPtr.store {ty : WGSLType} (p : MutPtr ty) (value : Exp ty) : ShaderM Unit

Store at the current offset.

Equations

• p.store value = Hesper.WGSL.Monad.ShaderM.writeBuffer p.buf p.offset value

@[inline]

def Hesper.WGSL.Monad.MutPtr.freeze {ty : WGSLType} (p : MutPtr ty) : Ptr ty

Snapshot to an immutable Ptr at the current offset. Useful when handing the pointer to a helper that doesn't need advance.

Equations

• p.freeze = { buf := p.buf, offset := p.offset, bufLen := p.bufLen }

def Hesper.WGSL.Monad.ShaderM.mutPtr (ty : WGSLType) (buf : String) (bufLen : Nat) (baseOffset : Exp (WGSLType.scalar ScalarType.u32)) : ShaderM (MutPtr ty)

Construct a MutPtr whose offset register is initialised from baseOffset. The resulting pointer can be advanced in-place inside an outer loop, matching CUDA's K += stride idiom.

Example (V11 outer K loop):

let K ← ShaderM.mutPtr (.scalar .u32) "k_cache_f16" kvWords
          (kvHead * maxSeqLen * (D/2) + laneId)
ShaderM.runtimeFor splitStart splitEnd (Exp.litU32 wgSize) fun _ => do
  ...
  let kPacked ← K.loadAt (Exp.litU32 (pk * 32))
  ...
  K.advance (Exp.litU32 (wgSize * (D/2)))

Equations

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

structure Hesper.WGSL.Monad.RegArray (ty : WGSLType) (n : Nat) : Type

Typed register array — n named ShaderM vars that share a type and can be indexed by a meta-time Nat. Replaces the V8/V11 idiom

let mut q0Vars : Array String := #[]
for pk in [0:n] do
  let v ← ShaderM.var ty (init pk)
  q0Vars := q0Vars.push v
-- later:
let q0 : Exp _ := Exp.var q0Vars[pk]!

with

let q0 ← RegArray.mk ty n init
-- later:
let q0Exp := q0.get pk
q0.set pk newVal

The Array String field stays in Lean meta land — at codegen time it materialises as n separate varDecl stmts (one per slot). Mirrors CUDA's T arr[N]; register array.

• names : Array String

def Hesper.WGSL.Monad.instReprRegArray.repr {ty✝ : WGSLType} {n✝ : Nat} : RegArray ty✝ n✝ → Nat → Std.Format

Equations

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

instance Hesper.WGSL.Monad.instReprRegArray {ty✝ : WGSLType} {n✝ : Nat} : Repr (RegArray ty✝ n✝)

Equations

• Hesper.WGSL.Monad.instReprRegArray = { reprPrec := Hesper.WGSL.Monad.instReprRegArray.repr }

@[inline]

def Hesper.WGSL.Monad.RegArray.get {ty : WGSLType} {n : Nat} (a : RegArray ty n) (i : Nat) : Exp ty

Read slot i as an Exp. Returns Exp.litU32 0 if out-of-range (shouldn't happen in well-typed code; the bounds check is a safety net against meta-time index typos).

Equations

• a.get i = Hesper.WGSL.Exp.var (a.names[i]?.getD "RegArray.get/oob")

@[inline]

def Hesper.WGSL.Monad.RegArray.set {ty : WGSLType} {n : Nat} (a : RegArray ty n) (i : Nat) (v : Exp ty) : ShaderM Unit

Assign a new value to slot i.

Equations

• a.set i v = Hesper.WGSL.Monad.ShaderM.assign (a.names[i]?.getD "RegArray.set/oob") v

def Hesper.WGSL.Monad.ShaderM.regArray (ty : WGSLType) (n : Nat) (init : Nat → Exp ty) : ShaderM (RegArray ty n)

Construct a RegArray ty n by emitting n var declarations. init i produces the initial value for slot i.

Equations

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