Design Space 06: Hardware Abstraction

Series: AI-Native Programming Language Research — Design Space Memos Document 6 of 6 (final memo).

1. Problem statement

Modern AI hardware is not one architecture, it is a zoo. NVIDIA GPUs (Volta through Blackwell, with each generation introducing incompatible new features). AMD CDNA. Google TPU v4 / v5. AWS Trainium and Inferentia. Intel Gaudi. Cerebras wafers. Graphcore IPUs. Tenstorrent dataflow chips. Groq's deterministic LPU. Apple's Neural Engine and AMX. ARM SVE and Intel AMX on CPUs. Each has its own memory hierarchy, concurrency model, instruction set, and cost model.

A user-facing ML language has to either:

1. Pick one target and hand-tune (the CUDA-only world). Best performance, no portability.
2. Abstract over targets via a graph IR (XLA, TVM, IREE). Good portability, performance gap on new hardware.
3. Provide a tile-level intermediate language (Triton, Pallas) that the user writes once and the compiler lowers. The current pragmatic answer.
4. Force the compiler to do all the work (TPU + XLA). Works only if one party owns the whole stack.

The state of the art is option 3 for kernel authors and option 2 for application code — a two-tier model. But the boundary between tiers is bespoke per-framework, the tile language is NVIDIA-first (Triton works best on NVIDIA hardware; AMD support exists; non-GPU support is patchy), and there is no clean way to describe hardware capabilities in a form the compiler can match against.

This memo's thesis: hardware capabilities are a refinement axis on the tensor type, just like shape, sparsity, and sharding. A tensor is "FP16, 2:4-sparse, sharded over model, and compatible with Hopper WGMMA acceleration." An operation type-checks if its inputs' refinements match a hardware capability the target supports. Lowering to hardware is matching capabilities; portability is parametricity over capability sets; verification is preservation of refinements through lowering.

This is the final memo, and the one that tests whether the unified-type-system framework from the previous five actually scales to a heterogeneous hardware reality.

2. The hardware zoo

A taxonomy by execution model, with the abstractions that fit each:

Cross-cutting features that complicate the picture further:

• Tensor cores / matrix units: NVIDIA's tensor cores have specific input shapes (16×16×16 originally, varying since), specific dtype combinations (FP16 in / FP32 acc, FP8 / BF16, etc.), and specific sparsity support (Ampere 2:4). AMD's matrix cores, TPU MXU, and Apple AMX have analogous-but-different shape and dtype constraints.
• Asynchronous memory operations: NVIDIA TMA (Hopper), AMD async copy, TPU DMA. These are effects in the memo 2 sense — operations that schedule work and complete later.
• Specialized dtypes: FP8 (E4M3, E5M2), FP4, INT4, bfloat16, NF4. Each has its own range, accumulation behavior, and hardware acceleration.
• Topology: NVLink, NVSwitch, infiniband, Inter-Chip Interconnect (ICI). Communication cost is non-uniform; sharding decisions interact with topology (memo 4).

The combinatorial observation: the number of valid (operation, dtype, layout, hardware) combinations is in the thousands. No single library can enumerate them, which is why CUTLASS exists as a generator, why Triton exists as a DSL, and why MLIR exists as an IR — three responses to the same combinatorial pressure.

3. Current abstraction layers

Five layers, ordered from highest to lowest:

3.1 Frameworks (PyTorch, JAX, TensorFlow)

The user writes Python; the framework dispatches to kernels. Eager mode is per-operation library calls; compiled mode (TorchDynamo + Inductor, JAX jit) traces and lowers to a graph IR.

Strength: widest user adoption; ergonomic. Weakness: the compiler sees only what tracing reveals; everything below the framework is bespoke per-target.

3.2 Graph compilers (XLA, IREE, TVM, TensorRT)

Take a graph IR (HLO, MLIR, Relax, ONNX) and produce target code. Apply graph-level optimizations (fusion, layout selection) and dispatch to lower-level kernels.

Strength: target portability via dialect/target selection; whole-graph optimization. Weakness: limited to the operations the dialect knows about; new operations (e.g., novel attention variants) require dialect work or fall back to library calls.

3.3 Tile-level DSLs (Triton, Pallas, ThunderKittens, CUTILE)

The user writes kernels in terms of tiles (fixed-size blocks of tensors) and operations on tiles. The compiler handles per-thread scheduling, memory placement, and instruction selection.

Strength: the right abstraction level for kernel authors. Triton-style tile programming has become the de facto standard for custom kernels (FlashAttention, Mamba kernels, custom MoE layers). Weakness: tile semantics are NVIDIA-first; portability to TPU/IPU/dataflow is limited because their execution models don't map to thread-based tiles. Pallas (JAX) extends this idea to TPU but with its own limitations.

3.4 Polyhedral and scheduling languages (Halide, TVM TensorIR, Exo, AKG, Tiramisu)

Separate the algorithm from the schedule. The algorithm specifies what to compute; the schedule specifies how to tile, vectorize, parallelize, and unroll.

Strength: principled separation; auto-tuning works well; foundation for compiler research. Weakness: the schedule language is the hard part — Halide schedules are notoriously tricky, TVM's auto-scheduler is a research artifact, Exo (Ikarashi et al., PLDI 2022) exposes the schedule explicitly to give the user control. None has been widely adopted by application programmers.

3.5 Vendor primitives (cuBLAS, cuDNN, oneDNN, MIOpen, vendor-libraries-of-choice)

Pre-tuned implementations of common operations. The compiler dispatches to these when shapes and dtypes match.

Strength: highest performance for supported cases. Weakness: opaque to the compiler; new operations (or unusual shapes) fall off the fast path; can't be composed or fused with surrounding code.

4. Triton as a case study

Triton (Tillet, Kung, Cox, MAPL 2019) is the most successful kernel-level abstraction in current ML. It deserves a closer look because what it gets right and what it gets wrong are both informative.

The model: a Triton kernel takes pointers to global memory and tile sizes; the user writes scalar code that operates on tiles (tl.load, tl.dot, tl.store); the compiler generates per-thread code with appropriate memory accesses, register allocation, and instruction selection. Block sizes are explicit; thread coordination is implicit.

What it gets right:

• The tile is the right abstraction unit for modern accelerators. Thinking in tiles matches the hardware (warps, tensor cores, async memory).
• Scalar-on-tile syntax hides the threading model without sacrificing performance.
• The compiler handles the boring parts (vectorization, register allocation, async memory scheduling).
• It's embedded in Python, which makes it accessible to ML researchers.

What it gets wrong (or hasn't solved):

• It's NVIDIA-first. AMD support exists but trails. TPU/IPU support is essentially impossible because their execution models don't have warps or shared memory.
• Hardware capabilities (which tensor core shapes, which dtypes, what sparsity) are implicit — the compiler picks the right tensor core variant based on tile shapes and dtypes, but the user has no way to declare "this kernel requires Hopper WGMMA" or "this kernel only runs on hardware with FP8 support."
• Composition is awkward. Calling one Triton kernel from another doesn't fuse; using Triton tiles with PyTorch tensors requires explicit boundary management.
• No cost model exposed. The user picks block sizes; the compiler doesn't tell them which sizes will be fastest on this target.

The diagnosis: Triton is the right abstraction level but not the right abstraction typing. Hardware capabilities are not in the type system; portability is by happy coincidence (the same code happens to lower OK on AMD); the user cannot reason about why their kernel is fast or slow.

A from-scratch design takes Triton's tile model and adds the typing layer.

5. Hardware capabilities as types

The proposal: a hardware target exposes a set of typed capabilities, and a tile-level program type-checks against a target if its operations match the target's capability set.

Sketch:

capability MatMul(M, N, K, A_dtype, B_dtype, C_dtype, sparsity_pattern) {
  -- A target supports this capability if it has a matrix unit
  -- accepting the given shapes, dtypes, and sparsity
}

target Hopper provides {
  MatMul(64, 128, 32, FP8 E4M3, FP8 E4M3, FP32, dense),
  MatMul(64, 128, 16, FP16, FP16, FP32, dense),
  MatMul(64, 128, 32, FP16, FP16, FP32, sparse_2_of_4),
  AsyncCopy(global, shared, ...),
  ...
}

target Ampere provides {
  MatMul(16, 16, 16, FP16, FP16, FP32, dense),
  MatMul(16, 16, 16, FP16, FP16, FP32, sparse_2_of_4),
  ...
}

A kernel type-checks against Hopper if every operation it performs has a matching capability. A kernel parameterized over a target capability set is portable across targets that supply those capabilities.

This is straightforward refinement typing — capabilities are predicates on hardware that operations require. The compiler's job is to match required capabilities against provided capabilities and emit the appropriate instruction. The user's job is to write kernels that compose capabilities correctly.

What this gives you:

1. Portability with explicit failure: a kernel that requires WGMMA doesn't compile for Ampere; you get a clear error rather than silently slow code.
2. Capability-polymorphic kernels: write once, parameterize over the matmul capability, lower to whatever the target provides.
3. Compositional reasoning: combining a Hopper-only matmul with an FP8 quantization op gives a kernel typed as requires Hopper ∩ FP8.
4. Cost model integration: capabilities can carry cost information (cycles per matmul, memory bandwidth); the compiler picks the cheapest valid lowering.
5. Verification: capability-preservation is the lowering-correctness theorem (memo 5).

Existing partial work in this direction:

• MLIR's target descriptions carry some capability information but are mostly used for code-generation flags, not type-checking.
• CUTLASS parameterizes over Arch template parameters, which is structurally similar but C++-template-based and not user-facing.
• Halide's target predicates allow conditional code generation.
• Hidet's task-mapping abstraction (Ding et al., CGO 2023) is the closest published work, separating logical work from hardware mapping.

Nobody has built this as a typed user-facing system.

6. The lowering story

Granted a capability-typed tile language, how does a high-level operation become a tile program?

The pipeline:

1. Algebraic IR (the refinement-typed tensor algebra from memos 1–5). Operations are typed by shape, effects, structure, and sharding.
2. Tile decomposition. Lower each algebraic operation to a tile program. The decomposition is parameterized: tile sizes, accumulator types, vectorization widths.
3. Capability matching. The tile program is type-checked against the target's capability set. Operations that have no matching capability either fall back to a less-efficient capability or fail.
4. Schedule selection. Among valid tilings, the compiler picks based on a cost model (memory traffic, register pressure, occupancy).
5. Code generation. Emit target-specific code (PTX, AMDGCN, vendor IR).

Each step is a refinement-preserving program transformation. Verification (memo 5) applies at each boundary.

The key claim: each step is typed. The algebraic IR has tensor refinements; the tile IR has tile refinements + capability requirements; the target IR has hardware-specific instructions. Refinements compose down the stack: a Tensor f16 [B, S, D] sharded sparse_2_of_4 lowers to tiles that require MatMul(..., FP16, ..., sparse_2_of_4) capability. A target that doesn't provide it can't run the program; a target that does provides a unique lowering.

This pulls in the previous five memos:

The lowering pipeline is, structurally, a verified compilation pipeline that propagates a single multi-axis refinement type from the user surface to the hardware instruction.

7. Comparison across systems

The pattern matches the previous memos: each system gets one thing right. Triton has the tile abstraction; Pallas extends it to TPU; TensorIR has explicit schedules; Halide/Exo have the algorithm-schedule separation; CUTLASS has compositional templates; MLIR has the multi-dialect lowering. None has all of them, and none has hardware capabilities as types.

A unified system isn't a research project from scratch — it's an integration of existing partial answers under a typed framework that doesn't currently exist.

8. The portability vs. performance trade

The eternal question. Three positions:

Portable-first (XLA/IREE/ONNX Runtime): the user code is hardware-agnostic; the compiler does everything. Works for the operations the compiler knows; fails open (slow) for novel work.

Performance-first (CUDA, Triton on NVIDIA): the user writes hardware-aware code. Best performance; no portability.

Two-tier (current pragmatic answer): high-level code is portable, kernel code is target-specific. Works in practice; bespoke per-framework.

The capability-typed proposal sits between these. It is parametrically portable — code parameterized over a capability set runs on any target providing those capabilities — but not implicitly portable — the user has to declare capability requirements rather than hoping the compiler figures it out. This is closer to how generic programming works in typed languages (Rust traits, Haskell type classes) than to either current ML extreme.

This trade is, I think, the correct one. Implicit portability has been tried and produces the "good for known ops, slow for novel ops" failure mode. Explicit capability requirements push the user to think about what they need, in exchange for predictable behavior across targets.

9. Where the unified framework breaks down

Honesty section. Things the unified type system doesn't handle cleanly:

Concurrency at the tile level. Tile operations run concurrently; reasoning about race freedom inside a tile program is a separation-logic problem (Iris, Perennial). The refinement type system doesn't address this; verification of concurrent tile code is the hardest part of the verification story (memo 5, §9).

Dataflow architectures. Cerebras, Graphcore, and Tenstorrent don't fit the "tile program over a SIMT execution model" mold cleanly. They have per-tile programs with explicit message passing. The capability-type approach probably extends — a "tile" becomes a unit of independent computation rather than a block of threads — but the surface programming model is more different than just renaming. This is a real research question, not just an integration exercise.

Specialized hardware that breaks the matmul/elementwise/reduce abstraction. Hardware with native attention units, native softmax, native FFT, etc. The abstraction has to either expose those or paper over them. Either way is a tradeoff: exposing them means user code has to know about them; papering over them means leaving performance on the floor.

Heterogeneous nodes. CPU + GPU + NPU + accelerator-of-the-quarter, with explicit data movement between them. The capability framing extends (each device has its own capability set; movement between them is a typed coercion) but the cost model gets hairy.

Numerical reproducibility. Different hardware produces different bit-patterns from the same algorithm (different reduction orders, different rounding). For reproducibility-critical workloads, this is a hard problem that the type system does not solve. It might expose the issue (a "bit-reproducible" capability that few targets actually have) but cannot fix it.

These are the places where the unified framework hits real complexity. None of them invalidate the framework, but all of them require additional work beyond what the previous memos sketched.

10. Open research problems

1. Capability-typed tile language as a paper-and-implementation. Take Triton's surface, add a capability type system, target multiple hardware. Likely 1-2 years of work; PLDI/POPL publishable. The most concrete next step from this memo.

2. Cross-architecture tile abstraction. Make tile programming work for systolic arrays (TPU) and dataflow architectures (IPU/Tenstorrent), not just SIMT. Pallas is the existing point of reference; full multi-target tile programming is open.

3. Capability-aware autotuning. Given a capability-typed kernel and a target, pick tile sizes and schedule. Combines existing autotuning work (Ansor, AutoTVM) with the capability framework.

4. Verified hardware lowering. Building on memo 5: prove that a refinement-typed algebraic operation lowers correctly to a capability-typed tile program. The structural piece is tractable; the numerical piece (FP semantics across hardware) is harder.

5. Heterogeneous-mesh lowering. Combine memo 4's mesh framework with memo 6's capability framework: a mesh of devices with different capabilities, sharding decisions that account for capability differences.

6. Specialized-hardware capability vocabulary. What's the right way to express "this hardware accelerates softmax"? An exhaustive op list doesn't scale; a small core capability set with composition rules is better. Open.

7. Cost model unification. Each layer (graph, tile, instruction) has its own cost model. A unified cost model that propagates through the lowering is a substantial research project, partially addressed by Welder, Roller, Hidet, Mirage.

8. Static reproducibility guarantees. A capability set ensuring bit-reproducibility (deterministic reduction order, IEEE-strict math) for verification-critical workloads. Currently nobody offers this.

11. Recommendation for an AI-native language design

Concrete positions:

• A two-tier surface language. Application code in the algebraic surface (refinement-typed tensors, automatic lowering); kernel code in a tile-level surface (capability-typed tiles, explicit schedule control). Boundary is principled, not bespoke.

• Hardware capabilities as a kind of type. Targets advertise capability sets; programs require capability sets; type-checking matches them. Portability is parametric over capabilities.

• Tiles as first-class. The user writes tile programs the way Triton allows. The compiler lowers tiles to threads/warps/whatever the target uses.

• Schedule explicit but optional. Default to compiler-chosen schedule; let the user pin schedule decisions when they know better. Halide-style.

• MLIR or equivalent as the IR substrate. Don't reinvent the multi-dialect lowering machinery. Build on existing infrastructure.

• Capability descriptions as first-class metadata. Targets are described in machine-readable form; the compiler uses them for type-checking, lowering, and cost modeling.

• Async memory and communication as effects. TMA, async copy, all-reduce — handlers from memo 2 apply.

• Verified lowering for the high-stakes path. AD, sharding, kernel selection — verify these. Other passes can use translation validation or testing.

The thesis: hardware abstraction is the place where a unified-type-system language pays off most concretely, because it converts a combinatorial integration problem ("which library, which target, which dtype, which sparsity") into a typed search ("find a capability assignment that satisfies the program's requirements on this target"). The combinatorics don't go away — they get pushed into the compiler's search, where they belong, instead of being the application programmer's problem.

12. References

Tile-level kernel languages

• Tillet, P., Kung, H. T. & Cox, D. (2019). Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations. MAPL 2019. — The reference for tile-level GPU kernels.
• JAX team. Pallas: A JAX kernel language. — Triton-style kernels for JAX, with TPU support.
• Spector, B. et al. (2024). ThunderKittens. — Simpler kernel abstractions for modern GPUs.
• NVIDIA. CUTLASS: CUDA Templates for Linear Algebra Subroutines. — Template-based kernel composition.

Polyhedral and scheduling languages

• Ragan-Kelley, J. et al. (2013). Halide. PLDI 2013. (Cited in memo 1.)
• Chen, T. et al. (2018). TVM. OSDI 2018. (Cited in memo 1.)
• Feng, S. et al. (2023). TensorIR: An Abstraction for Automatic Tensorized Program Optimization. ASPLOS 2023.
• Ikarashi, Y. et al. (2022). Exocompilation for Productive Programming of Hardware Accelerators. PLDI 2022. — Exo.
• Baghdadi, R. et al. (2019). Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code. CGO 2019.

Multi-target compiler infrastructure

• Lattner, C. et al. (2021). MLIR. CGO 2021. (Cited in memo 1.)
• IREE project. IREE: Intermediate Representation Execution Environment. Engineering project.
• TensorFlow XLA. XLA: Optimizing Compiler for Machine Learning.
• ONNX project. ONNX Runtime / ONNX MLIR.

Schedule search and autotuning

• Zheng, L. et al. (2020). Ansor: Generating High-Performance Tensor Programs for Deep Learning. OSDI 2020.
• Chen, T. et al. (2018). Learning to Optimize Tensor Programs. NeurIPS 2018. (AutoTVM.)
• Zhu, H. et al. (2022). ROLLER: Fast and Efficient Tensor Compilation for Deep Learning. OSDI 2022.
• Ding, Y. et al. (2023). Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs. CGO 2023.
• Wu, M. et al. (2024). Mirage: A Multi-Level Superoptimizer for Tensor Programs. OSDI 2024.
• Shi, Y. et al. (2023). Welder: Scheduling Deep Learning Memory Access via Tile-graph. OSDI 2023.

Vendor compilers and SDKs (engineering, not papers)

• NVIDIA: cuBLAS, cuDNN, CUTLASS, CUDA, PTX.
• AMD: ROCm, HIP, MIOpen, Composable Kernel.
• Google TPU: XLA, OpenXLA, Pallas-TPU.
• Intel: oneDNN, oneMKL, AMX intrinsics.
• Graphcore: Poplar, PopART.
• Cerebras: CSL (Cerebras Software Language).
• Tenstorrent: Metalium, TT-NN.
• Apple: Metal Performance Shaders, Accelerate, AMX intrinsics.

Adjacent — high-performance language design (HPC heritage)

• Chamberlain, B. et al. (2007). Chapel. IJHPCA. (Cited in memo 4.)
• Charles, P. et al. (2005). X10. OOPSLA 2005. (Cited in memo 4.)
• Fatahalian, K. et al. (2006). Sequoia. SC 2006. (Cited in memo 4.)
• Bauer, M. et al. (2012). Legion. SC 2012. (Cited in memo 4.)
• Stratton, J. et al. (2010). MCUDA: An Efficient Implementation of CUDA Kernels for Multi-Core CPUs. LCPC 2010.

Hardware reference (architecture papers)

• Jouppi, N. et al. (2017). In-Datacenter Performance Analysis of a Tensor Processing Unit. ISCA 2017. — The TPU paper.
• NVIDIA. NVIDIA H100 Architecture Whitepaper. Hopper details.
• NVIDIA. NVIDIA Ampere Architecture Whitepaper. TC and 2:4 sparsity details.
• Cerebras. The Cerebras Wafer-Scale Engine. Architecture overviews.

13. Final cross-memo synthesis

Six memos. Time to write down the synthesis as a single proposal, because I think the framework that emerged is more concrete than I expected when I started.

13.1 The unified type system

The "tensor" in an AI-native language is a type with five orthogonal refinement axes:

1. Shape — dimension sizes, possibly symbolic, with compile-time arithmetic and constraints (memo 1).
2. Effects — purity, mutation, randomness, divergence; the things AD and parallelization need to know (memo 2).
3. Structure — sparsity, low-rank factorization, quantization, hardware-imposed patterns (memo 3).
4. Sharding — distribution across a typed device mesh, with collective insertion as type-changing coercion (memo 4).
5. Hardware capability — the hardware features the tensor's operations require for efficient execution (memo 6).

Operations on tensors compose along all five axes. The compiler infers result-type refinements from operand-type refinements. Mismatches at boundaries are either type errors or compiler-inserted coercions: broadcast for shape, AD-handler insertion for effects, format conversion for structure, resharding for sharding, lowering selection for capability.

13.2 The verification story (memo 5)

The five-axis refinement framework is the contract for verified compilation. Each compiler pass preserves refinements (or has a documented coercion). Verification happens at three weights:

• Lightweight (SMT-discharged refinement types) for users.
• Medium (Verus-style Rust verification) for the runtime.
• Heavyweight (Coq/Lean) for the high-stakes passes (AD, sharding, kernel selection).
• Translation validation as a fallback during compiler development.

13.3 The architecture

A multi-tier system, not a single language:

• Surface language: refinement-typed Pythonic or ML-family syntax. Most users live here.
• Algebraic IR: typed tensor algebra. Most compiler passes operate at this level.
• Tile IR: capability-typed tiles for kernel authors.
• Target IR: hardware-specific dialects (PTX, AMDGCN, vendor IRs).
• All wired through MLIR-style dialect lowering, with verification annotations at each boundary.

13.4 What's tractable, what's hard, what's research

Tractable (1–2 years of work each):

• Refinement-typed shapes with gradual-typing fallback.
• Verified sharding propagation.
• Capability-typed tile language for SIMT GPUs.
• Effect-handler-based AD over a clean core calculus.

Hard (multi-year research programs):

• AD over algebraic effects with full soundness proof.
• Cross-architecture tile abstraction (SIMT + systolic + dataflow).
• Verified end-to-end lowering pipeline.
• Numerical-error bounds that compose across passes.

Research (open, no clear path):

• Compositional sparsity that includes dynamic routing.
• Verified concurrent training with semantic preservation.
• Heterogeneous-mesh sharding with cost-aware optimization.
• Static numerical reproducibility guarantees across hardware.

13.5 What this is, finally

The series did not produce a complete language design. What it produced is a research program: a set of related questions, with a unifying type-system framework, a tractability ranking, and a clear methodological path (existing verification techniques, existing IR infrastructure, existing tile-level surface languages, all integrated under a typed framework that doesn't currently exist).

The pieces are out there. CompCert showed verified compilation works. GSPMD showed sharding inference works. Triton showed tile-level kernels work. Brunel/Mazza/Pagani showed verified AD works. TACO showed sparse-format composition works. Refinement types in Liquid Haskell showed predicate-based type systems work.

Nobody has integrated them. The integration is the contribution.

For someone with the right background — compiler engineering, formal verification, OS-level systems work, and the patience for multi-year research projects — this is a credible direction. Not a hobby project; a research program. The first memo's question ("is there an AI-native programming language?") has an answer: not yet, but the components for one exist, and the integration is a tractable research target.

End of series.