Verified Operator Pattern with Kernel Fusion

This module extends the VerifiedOp pattern to support kernel fusion.

Key Design Evolution

Before (Immediate Execution - No Fusion):

let tmp ← MatMul.impl_forward A B   -- Writes to VRAM
let res ← ReLU.impl_forward tmp     -- Reads from VRAM, writes again
-- Problem: 2 memory roundtrips!

After (Lazy Composition - Fusion Enabled):

let fused = MatMul.impl_kernel |> ReLU.impl_kernel  -- Build recipe
let res ← run_kernel fused input                     -- Single GPU dispatch!
-- Benefit: 1 memory roundtrip, fused shader

Type Parameters

• I, O: High-level CPU types (e.g., MatMulInput, TensorData)
• WI, WO: WGSL expression types (e.g., Exp (.array .f32 256))

The separation allows:

1. CPU spec works on convenient high-level types
2. GPU impl works on low-level WGSL expressions
3. Kernels can be composed before execution

Workgroup Size

Kernels are parameterized by workgroup dimensions (wX, wY, wZ). This is tracked at the type level for safety.

Data Abstractions

structure Hesper.Core.TensorData : Type

CPU tensor data: wrapper around Array Float with shape information. Used for specification and testing.

• shape : Tensor.Shape

  Shape of the tensor

• data : Array Float

  Flattened data in row-major order

def Hesper.Core.instInhabitedTensorData.default : TensorData

Equations

• Hesper.Core.instInhabitedTensorData.default = { shape := default, data := default }

instance Hesper.Core.instInhabitedTensorData : Inhabited TensorData

Equations

• Hesper.Core.instInhabitedTensorData = { default := Hesper.Core.instInhabitedTensorData.default }

def Hesper.Core.instReprTensorData.repr : TensorData → Nat → Std.Format

Equations

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

instance Hesper.Core.instReprTensorData : Repr TensorData

Equations

• Hesper.Core.instReprTensorData = { reprPrec := Hesper.Core.instReprTensorData.repr }

def Hesper.Core.TensorData.zeros (shape : Tensor.Shape) : TensorData

Create a zero tensor with given shape

Equations

• Hesper.Core.TensorData.zeros shape = { shape := shape, data := { toList := List.replicate shape.size 0.0 } }

def Hesper.Core.TensorData.constant (shape : Tensor.Shape) (value : Float) : TensorData

Create a tensor filled with a constant value

Equations

• Hesper.Core.TensorData.constant shape value = { shape := shape, data := { toList := List.replicate shape.size value } }

def Hesper.Core.TensorData.get (t : TensorData) (idx : Nat) : Float

Get element at flattened index

Equations

• t.get idx = t.data[idx]!

def Hesper.Core.TensorData.set (t : TensorData) (idx : Nat) (value : Float) : TensorData

Set element at flattened index

Equations

• t.set idx value = { shape := t.shape, data := t.data.set! idx value }

def Hesper.Core.TensorData.size (t : TensorData) : Nat

Total number of elements

Equations

• t.size = t.shape.size

def Hesper.Core.TensorData.approxEq (a b : TensorData) (tolerance : Float := 1e-5) : Bool

Check if two tensors are approximately equal within tolerance

Equations

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

structure Hesper.Core.GPUHandle : Type

GPU tensor handle: wrapper around WebGPU Buffer with shape and device. Used for high-performance GPU operations.

• device : WebGPU.Device

  Device this buffer belongs to

• buffer : WebGPU.Buffer

  GPU buffer containing the data

• shape : Tensor.Shape

  Shape of the tensor

def Hesper.Core.GPUHandle.fromTensorData (device : WebGPU.Device) (data : TensorData) : IO GPUHandle

Upload CPU tensor data to GPU

Equations

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

def Hesper.Core.GPUHandle.toTensorData (handle : GPUHandle) : IO TensorData

Download GPU tensor data to CPU

Equations

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

def Hesper.Core.GPUHandle.zeros (device : WebGPU.Device) (shape : Tensor.Shape) : IO GPUHandle

Create a zero tensor on GPU

Equations

• Hesper.Core.GPUHandle.zeros device shape = Hesper.Core.GPUHandle.fromTensorData device (Hesper.Core.TensorData.zeros shape)

Fusable Verified Operator Type Class

class Hesper.Core.VerifiedOpFusion (wX wY wZ : Nat) (I O WI WO : Type) : Type

Type class for verified operators with kernel fusion support.

Type parameters:

• I, O: High-level CPU types (e.g., Matrix, TensorData)
• WI, WO: WGSL expression types (e.g., Exp (.array .f32 N))
• wX, wY, wZ: Workgroup dimensions (compile-time constants)

Design philosophy:

1. Specification (CPU): Pure mathematical definition for correctness
2. Implementation (GPU): Composable kernel for performance
3. Fusion: Kernels compose via |> before execution
4. Verification: Compare spec vs impl for correctness

Example usage:

-- Fuse two operators
let fused = MatMul.impl_kernel |> ReLU.impl_kernel

-- Execute fused kernel
let result ← execute_kernel device fused input

-- Verify correctness
let cpu_result := MatMul.spec_forward input |> ReLU.spec_forward
assert (cpu_result ≈ result)

• spec_forward : I → O

  Specification: Pure mathematical definition (CPU). Reference implementation for correctness and formal verification.

  This should be:

  • Easy to understand and verify
  • Provably correct
  • Used as the "ground truth" for testing
• impl_kernel : WGSL.Kernel wX wY wZ WI WO

  Implementation: Composable GPU kernel. Returns a Kernel that can be fused with other kernels.

  This should be:

  • Lazy (doesn't execute immediately)
  • Composable (can be chained with |>)
  • Optimized for GPU performance

  Example:

  impl_kernel : Kernel 256 1 1 (Exp (.array .f32 N)) (Exp (.array .f32 N))
  

• spec_backward : I → O → I

  Backward Specification: CPU gradient computation. Given (input, grad_output), compute grad_input.

  Used for training and automatic differentiation.

• impl_kernel_backward : WGSL.Kernel wX wY wZ (WI × WO) WI

  Backward Implementation: GPU gradient kernel. Composable kernel for backpropagation.

  Can be fused with other backward passes.

• run_forward : WI → IO WO

  Execution Helper: Optional wrapper for immediate execution. Compiles and runs the kernel immediately (for testing/debugging).

  Default implementation provided. Can be overridden per operator.

• verify_consistency : I → Float → IO Bool

  Verification: Check CPU spec matches GPU impl. Requires conversion between high-level (I/O) and low-level (WI/WO) types.

  Default implementation is a placeholder.

instance Hesper.Core.instDifferentiableVerifiedOpFusion {wX wY wZ : Nat} {I O WI WO : Type} [inst : VerifiedOpFusion wX wY wZ I O WI WO] : Differentiable (VerifiedOpFusion wX wY wZ I O WI WO) I O

Linking VerifiedOpFusion to the Unified Differentiable Interface

Equations

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

Helper Functions for Fusion

def Hesper.Core.composeOps {wX wY wZ : Nat} {I M O WI WM WO : Type} [inst1 : VerifiedOpFusion wX wY wZ I M WI WM] [inst2 : VerifiedOpFusion wX wY wZ M O WM WO] : WGSL.Kernel wX wY wZ WI WO

Compose two verified operators into a fused operation. The result is a new kernel that applies f then g in a single pass.

Example:

let matmul_relu = composeOps (I := MatMulInput) (M := MatMulOutput) (O := ReLUOutput)
-- Single GPU kernel that does matmul AND relu!

Equations

• Hesper.Core.composeOps = Hesper.Core.impl_kernel M O |> Hesper.Core.impl_kernel I M

def Hesper.Core.mkElementwiseOp {wX wY wZ : Nat} {ty : WGSL.WGSLType} (gpu_fn : WGSL.Exp ty → WGSL.Exp ty) : WGSL.Kernel wX wY wZ (WGSL.Exp ty) (WGSL.Exp ty)

Create a simple element-wise operation verified operator.

This is a helper for defining simple pointwise operations like ReLU, sigmoid, etc.

Parameters:

• gpu_fn: WGSL expression transformation (Exp ty → Exp ty)

Equations

• Hesper.Core.mkElementwiseOp gpu_fn = Hesper.WGSL.mapK gpu_fn