Why do we need Auto-Perfect

MLX is possibly the go-to for local MLX inference given how much apple has been investing in it. Similar tools exist for Nvidia CUDA to write custom CUDA kernels.

How does this work?

MLX Code
  ↓
Graph Capture
  ↓
Fusion Detection
  ↓
Kernel IR
  ↓
LLM Kernel Generation
  ↓
Compile + Repair
  ↓
Validate
  ↓
Benchmark
  ↓
Cache
  ↓
Runtime Dispatch

Wait hol'up why do we need custom kernels?

We need custom GPU kernels because general-purpose kernels leave a lot of performance, memory efficiency, and control on the table. Once you push beyond “standard deep learning workloads,” the defaults become a bottleneck.

Kernel Fusion

Custom GPU kernels are how you turn mathematical insight into hardware-level speed. frameworks give correctness, kernels give dominance.

Every GPU kernel launch has overhead and memory traffic.

Typical PyTorch code:

load A → kernel
load B → kernel
add → kernel
relu → kernel
store

That’s:

• Multiple kernel launches
• Multiple global memory reads/writes

A custom kernel can do:

load once → compute everything → store once

Why this matters

GPU performance is usually memory-bound, not FLOP-bound.

Fusing ops:

• Reduces global memory traffic
• Keeps data in registers / shared memory
• Improves cache locality

FlashAttention exists almost entirely because of this.

Core Compiler (Deterministic)

• Identity-safe tensor tracing using id(tensor)
• Full op surface tracing (elementwise, reductions, linalg, shape, indexing, multi-output)
• SSA-style GraphIR with explicit producers, consumers, constants, and kwargs
• Deterministic input binding captured at trace time
• Faithful runtime interpreter for unfused GraphIR execution
• End-to-end numerical correctness vs eager MLX (mx.allclose)
• Explicit reduction semantics tests (axis, keepdims, broadcast correctness)
• Large-shape stress tests to surface memory-bound vs compute-bound behavior

---

Canonicalization & Semantics

• Op classification map
  (ELEMENTWISE, REDUCTION, GEMM, RESHAPE/VIEW, INDEXING, BARRIER)
• Conservative fusion barriers defined (IO, random, sort/select, unsafe multi-output)
• Numerically stable softmax canonicalization
• softmax eliminated from IR and lowered into primitive ops
• Structural + numerical regression tests for canonicalization
• Canonicalization coverage tests on mixed reduction + elementwise graphs
• Explicit broadcast-shape invariance tests

---

Fusion (Compiler-Owned)

Legality phase

• Legality-only fusion region discovery
• Deterministic toposort-based region construction
• SSA single-consumer enforcement
• Reduction boundary enforcement
• Barrier enforcement
• Adversarial fusion graphs (diamond, fan-out, fan-in) regression tests

Cost-gated fusion

• Greedy forward fusion
• Incremental cost gating (Δbenefit > Δpenalty)
• Peak live-byte estimation
• Hard footprint caps for pathological regions
• Cost model validated against large-tensor and over-fusion cases
• Roofline-style sanity checks (memory vs compute bound classification)

---

Kernel Interface Layer (Compiler ↔ Backend Boundary)

• Explicit FusionRegion abstraction

• Stable per-region input/output contract

• Region signature determinism
  (ops, order, shapes, dtypes)

• KernelIR defined
  (lower than GraphIR, higher than MLX)

• Explicit kernel ABI

  • inputs
  • outputs
  • temporaries

• Deterministic lowering: FusionRegion → KernelIR

• KernelIR structural validator

• Reference (correctness-first) KernelIR executor

• KernelIR → MLX codegen path

• Reference executor ≡ MLX codegen equivalence tests

• Multi-output kernel support (e.g. meshgrid)

• Strict numerical equivalence gate before any benchmarking

• Per-op KernelIR golden tests (hand-computed small tensors)

---

AI-Generated Kernel Layer (Future, Isolated)

This layer does not replace the compiler.
It plugs in strictly below KernelIR.

• Serialize KernelIR → deterministic, prompt-safe representation
• Prompt explicitly constrains:
  • pure MLX
  • no side effects
  • no graph mutation
  • explicit inputs/outputs only
• Prompt generation isolated from execution (LLM may be stubbed)
• Generate candidate kernels for a single region
• Enforce shape/dtype guards in generated code
• Static validation of generated kernel vs KernelIR ABI
• Run correctness check vs KernelIR reference executor
• Reject kernels that allocate, branch, or call high-level MLX ops

Optional optimization

• Generate multiple variants per region
• Score variants via:
  • cost model estimate
  • microbenchmark timing
• Deterministic winner selection
• Persist rejected variants for failure analysis

---

Caching & Reuse

• Cache generated kernels by:
  • (region signature, shapes, dtypes, device)
• Reuse kernels across graphs and runs
• Cache invalidation on compiler / KernelIR version change
• Cold vs warm cache behavior benchmarks

---

Execution & Integration

• Deterministic fallback execution for fused regions
• Dispatch AI-generated kernel when available
• Fallback to KernelIR interpreter on failure
• Preserve debuggability:
  • region → kernel → source ops mapping
• Per-kernel execution tracing (timing + correctness metadata)

---

Validation & Benchmarking (Robust)

• Fusion legality regression tests
• Cost-model regression tests
• KernelIR lowering / validation / execution tests
• Prompt determinism & safety tests
• KernelIR reference executor correctness gate (must pass before timing)
• AI-generated kernel correctness tests
• Numerical stability stress tests (large values, small eps, NaNs)
• Deterministic benchmark harness (sync, warmup, fixed iters)

Example Benchmark Scripts (Required)

• benchmark_elementwise.py
  (sanity check, expect no speedup)

• benchmark_reduction.py
  (mean / variance, memory-bound baseline)

• benchmark_layernorm.py
  (reduction + broadcast, fusion stress test)

• benchmark_softmax.py
  (canonicalization + reduction ordering)

• benchmark_fused_vs_eager.py
  (unfused vs fused KernelIR)

• benchmark_kernelir_vs_generated.py
  (KernelIR ref vs AI-generated kernel)

• benchmark_cache_effects.py
  (cold vs warm generated kernels)

• Automated benchmark validation:

  • correctness must pass
  • timing otherwise discarded