Composable kernel abstraction for kernel fusion.

This module provides a high-level abstraction for composable GPU kernels. Instead of writing imperative procedures, we treat kernels as composable functions (Input → Output), enabling kernel fusion optimizations.

Key benefits:

• Compose operations with monadic bind (>>=) and sequencing (>>)
• Fuse multiple operations into single shader pass
• Reduce global memory traffic
• Type-safe workgroup size tracking

Example: -- Fuse: Load → Multiply → Add → ReLU → Store into single pass fusedKernel : Kernel 256 1 1 (Exp (.scalar .u32)) Unit fusedKernel = loadVec inputPtr |> andThen (mapK (· * Exp.litF32 2.0)) |> andThen (mapK (· + Exp.litF32 1.0)) |> andThen (mapK relu) |> andThen (storeVec outputPtr)

@[reducible, inline]

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

ShaderM monad - builds up statement lists while computing a value. This is a state monad that accumulates statements.

Equations

• Hesper.WGSL.ShaderM α = StateM (List Hesper.WGSL.Stmt) α

structure Hesper.WGSL.Kernel (wX wY wZ : Nat) (i o : Type) : Type

A composable kernel function running on a specific workgroup size.

Type parameters: wX, wY, wZ: Workgroup dimensions (natural numbers) i: Input type (e.g., Exp (.scalar .f32), or tuples) o: Output type

This abstraction allows us to compose operations and perform kernel fusion. Multiple operations can be fused into a single shader pass, reducing global memory roundtrips.

• unKernel : i → ShaderM o

def Hesper.WGSL.Kernel.id {wX wY wZ : Nat} {α : Type} : Kernel wX wY wZ α α

Identity kernel: passes input through unchanged

Equations

• Hesper.WGSL.Kernel.id = { unKernel := fun (x : α) => pure x }

def Hesper.WGSL.Kernel.comp {wX wY wZ : Nat} {a b c : Type} (g : Kernel wX wY wZ b c) (f : Kernel wX wY wZ a b) : Kernel wX wY wZ a c

Compose two kernels: g ∘ f means "f first, then g" The composition is performed in the ShaderM monad, so side effects (like memory operations) are properly sequenced.

Equations

• g |> f = { unKernel := fun (x : a) => do let y ← f.unKernel x g.unKernel y }

def Hesper.WGSL.Kernel.«term_|>_» : Lean.TrailingParserDescr

Infix operator for kernel composition (like >>> in Haskell Category)

Equations

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

def Hesper.WGSL.Kernel.andThen {wX wY wZ : Nat} {a b : Type} (k : Kernel wX wY wZ a b) (f : b → Kernel wX wY wZ b b) : Kernel wX wY wZ a b

Monadic bind for kernels

Equations

• k.andThen f = { unKernel := fun (x : a) => do let y ← k.unKernel x have k' : Hesper.WGSL.Kernel wX wY wZ b b := f y k'.unKernel y }

def Hesper.WGSL.mapK {wX wY wZ : Nat} {ty : WGSLType} (f : Exp ty → Exp ty) : Kernel wX wY wZ (Exp ty) (Exp ty)

Lift a pure DSL expression transformation into a Kernel.

This allows you to turn any pure expression transformation (like (· * 2.0) or (· + 1.0)) into a composable kernel.

Example: mapK (· * Exp.litF32 2.0) -- Multiply by 2 mapK (· + Exp.litF32 1.0) -- Add 1 mapK relu -- Apply ReLU

These can be composed: mapK (· * 2.0) |> mapK (· + 1.0) |> mapK relu

Equations

• Hesper.WGSL.mapK f = { unKernel := fun (x : Hesper.WGSL.Exp ty) => pure (f x) }

def Hesper.WGSL.emit {wX wY wZ : Nat} (s : Stmt) : Kernel wX wY wZ Unit Unit

Emit a statement (side effect) in a kernel. This adds the statement to the accumulated list.

Equations

• Hesper.WGSL.emit s = { unKernel := fun (x : Unit) => do modify fun (x : List Hesper.WGSL.Stmt) => x ++ [s] pure () }

def Hesper.WGSL.emitMany {wX wY wZ : Nat} (stmts : List Stmt) : Kernel wX wY wZ Unit Unit

Emit multiple statements

Equations

• Hesper.WGSL.emitMany stmts = { unKernel := fun (x : Unit) => do modify fun (x : List Hesper.WGSL.Stmt) => x ++ stmts pure () }

def Hesper.WGSL.loadBuffer {wX wY wZ : Nat} {ty : WGSLType} {n : Nat} (bufferName : String) : Kernel wX wY wZ (Exp (WGSLType.scalar ScalarType.u32)) (Exp ty)

Load operation: reads from a buffer at given index.

Input: index expression Output: loaded value

Note: The buffer must be an array type. The size parameter n is for type checking.

Equations

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

def Hesper.WGSL.storeBuffer {wX wY wZ : Nat} {ty : WGSLType} (bufferName : String) : Kernel wX wY wZ (Exp (WGSLType.scalar ScalarType.u32) × Exp ty) Unit

Store operation: writes to a buffer at given index.

Input: (index, value) pair Output: unit (side effect only)

Equations

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

def Hesper.WGSL.pairWithIndex {wX wY wZ : Nat} {ty : WGSLType} (idx : Exp (WGSLType.scalar ScalarType.u32)) : Kernel wX wY wZ (Exp ty) (Exp (WGSLType.scalar ScalarType.u32) × Exp ty)

Pair an expression with an index for storage

Equations

• Hesper.WGSL.pairWithIndex idx = { unKernel := fun (val : Hesper.WGSL.Exp ty) => pure (idx, val) }

def Hesper.WGSL.runKernel {wX wY wZ : Nat} {i o : Type} (k : Kernel wX wY wZ i o) (input : i) : o × List Stmt

Execute a kernel and extract the generated statements.

This runs the kernel computation and returns both the result value and the list of statements that were generated.

Equations

• Hesper.WGSL.runKernel k input = match k.unKernel input [] with | (result, stmts) => (result, stmts)

def Hesper.WGSL.execKernel {wX wY wZ : Nat} {i o : Type} (k : Kernel wX wY wZ i o) (input : i) : List Stmt

Execute a kernel and return only the statements (for side-effect kernels)

Equations

• Hesper.WGSL.execKernel k input = match Hesper.WGSL.runKernel k input with | (fst, stmts) => stmts