Matrix Multiplication with Fusion Support

Extends MatMul to support kernel fusion with subsequent element-wise operations.

Fusion Strategy

MatMul itself is a complex operation requiring reduction across dimensions. However, the output of MatMul can be fused with subsequent element-wise ops:

-- This fuses the ReLU into the MatMul kernel's store phase
MatMul |> ReLU

Instead of:

1. MatMul: compute C[i,j], write to VRAM
2. ReLU: read C[i,j] from VRAM, apply ReLU, write back

We get:

1. MatMul: compute C[i,j], apply ReLU, write to VRAM

Implementation Notes

For now, this provides a placeholder kernel structure. A full implementation would:

1. Use the existing subgroup matmul from Examples/Compute/MainMatmul.lean
2. Allow fusion of element-wise operations into the output store phase
3. Generate optimized WGSL with fused operations

The key insight: MatMul produces values that can flow through a composable kernel before being written to memory.

MatMul Kernel Abstraction

def Hesper.Op.MatMulFusion.matmulKernel {_M _K _N : Nat} : WGSL.Kernel 256 1 1 Unit (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32))

Simplified MatMul kernel for demonstration.

In a full implementation, this would:

1. Load tiles of A and B into shared memory
2. Compute partial products using subgroup operations
3. Apply any fused operations to the result
4. Store to output buffer

For now, this is a placeholder showing the type structure.

Equations

• Hesper.Op.MatMulFusion.matmulKernel = { unKernel := fun (x : Unit) => pure (Hesper.WGSL.Exp.litF32 0.0) }

Fusion Example: MatMul + Activation

def Hesper.Op.MatMulFusion.matmulWithActivation {M K N : Nat} (activation : WGSL.Kernel 256 1 1 (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32)) (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32))) : WGSL.Kernel 256 1 1 Unit (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32))

Fuse MatMul with an element-wise activation function.

This demonstrates the key pattern:

let fused = matmulKernel |> activationKernel

The activation is applied to each output element before storing, eliminating a memory roundtrip.

Equations

• Hesper.Op.MatMulFusion.matmulWithActivation activation = activation |> Hesper.Op.MatMulFusion.matmulKernel

Common Fused Patterns

def Hesper.Op.MatMulFusion.matmulReLU {M K N : Nat} : WGSL.Kernel 256 1 1 Unit (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32))

MatMul + ReLU fusion

Equations

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

def Hesper.Op.MatMulFusion.matmulSigmoid {M K N : Nat} : WGSL.Kernel 256 1 1 Unit (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32))

MatMul + Sigmoid fusion

Equations

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

def Hesper.Op.MatMulFusion.matmulGELU {M K N : Nat} : WGSL.Kernel 256 1 1 Unit (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32))

MatMul + GELU fusion (Gaussian Error Linear Unit)

Equations

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

Documentation & Examples

def Hesper.Op.MatMulFusion.exampleFusedLayer : Unit

Example: How to use fused MatMul + ReLU in a neural network layer.

Without fusion:

let C ← matmul A B        -- GPU Kernel 1: compute matmul, write C
let R ← relu C            -- GPU Kernel 2: read C, apply ReLU, write R

With fusion:

let R ← matmulReLU A B    -- GPU Kernel 1: compute matmul, apply ReLU, write R

This saves one memory roundtrip and one kernel launch.

Equations

• Hesper.Op.MatMulFusion.exampleFusedLayer = ()