WGSL DSL Helpers for Common Patterns

Provides high-level helpers for common GPU operations:

• Matrix-vector multiplication
• Element-wise operations with activation functions
• Reductions (sum, max, softmax)
• Fused operations (MatMul + Bias + Activation)

Matrix Operations

def Hesper.WGSL.Helpers.matVecMulBiasActivation (inputName weightsName biasName outputName : String) (inputSize outputSize : Nat) (activation : Exp (WGSLType.scalar ScalarType.f32) → Exp (WGSLType.scalar ScalarType.f32)) (outputIdx : Exp (WGSLType.scalar ScalarType.u32)) : Monad.ShaderM Unit

Matrix-vector multiplication with bias and activation (fused Layer operation)

Computes: output[i] = activation(sum_j(weights[j, i] * input[j]) + bias[i])

Parameters:

• inputName: name of input buffer
• weightsName: name of weights buffer
• biasName: name of bias buffer
• outputName: name of output buffer
• inputSize: number of input features
• outputSize: number of output features
• activation: activation function to apply (or identity)
• outputIdx: expression for output index (usually from global_invocation_id)

Equations

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

def Hesper.WGSL.Helpers.reluActivation (x : Exp (WGSLType.scalar ScalarType.f32)) : Exp (WGSLType.scalar ScalarType.f32)

ReLU activation: max(0, x)

Equations

• Hesper.WGSL.Helpers.reluActivation x = x.max (Hesper.WGSL.Exp.litF32 0.0)

def Hesper.WGSL.Helpers.identityActivation (x : Exp (WGSLType.scalar ScalarType.f32)) : Exp (WGSLType.scalar ScalarType.f32)

Identity activation: x

Equations

• Hesper.WGSL.Helpers.identityActivation x = x

def Hesper.WGSL.Helpers.tanhActivation (x : Exp (WGSLType.scalar ScalarType.f32)) : Exp (WGSLType.scalar ScalarType.f32)

Tanh activation (using WGSL built-in)

Equations

• Hesper.WGSL.Helpers.tanhActivation x = x.tanh

Softmax Operation

def Hesper.WGSL.Helpers.softmaxInPlace (dataName : String) (size : Nat) : Monad.ShaderM Unit

Softmax normalization (single-workgroup version for small arrays)

Computes: output[i] = exp(input[i] - max) / sum(exp(input - max))

Note: This is a simplified version for small arrays that fit in one workgroup. For production use with large arrays, use multi-pass reduction.

Parameters:

• dataName: name of buffer (input/output in-place)
• size: number of elements

Equations

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

Shader Generation Helpers

def Hesper.WGSL.Helpers.generateMatVecBiasActivationShader (inputSize outputSize : Nat) (activation : Exp (WGSLType.scalar ScalarType.f32) → Exp (WGSLType.scalar ScalarType.f32)) (workgroupSizeX : Nat := 256) (workgroupSizeY workgroupSizeZ : Nat := 1) : String

Generate complete compute shader for matrix-vector multiply + bias + activation

Generates a WGSL compute shader with proper bindings for:

• @binding(0): input array [inputSize]
• @binding(1): weights array [inputSize * outputSize]
• @binding(2): bias array [outputSize]
• @binding(3): output array [outputSize]

The shader fuses MatMul + Bias + Activation into a single kernel.

Equations

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

def Hesper.WGSL.Helpers.generateSoftmaxShader (size : Nat) (workgroupSizeX workgroupSizeY workgroupSizeZ : Nat := 1) : String

Generate complete compute shader for softmax (in-place)

Generates a WGSL compute shader for softmax normalization:

• @binding(0): data array [size] (read_write)

Note: Uses single workgroup, suitable for small arrays (size ≤ 1024)

Equations

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