Source-linked AI summary
Exo-GPU: Safe, Imperative, User-schedulable Programming for Tensor Cores
David Zhao Akeley, Yuka Ikarashi, Jonathan Ragan-Kelley
TL;DR
Modern GPU features expose asynchronous instruction ordering and synchronization, while existing approaches trade low-level control against easier analysis. Exo-GPU treats parallelism and synchronization as annotations on sequential code and checks their equivalence. Its GEMM kernels exceed 80% of theoretical peak on large problem sizes and sometimes outperform CUBLAS.
Problem
GPU programmers must manage explicit asynchronous instruction ordering, while high-level abstractions may hide details needed for tuning and low-level CUDA leaves synchronization exposed.
Source: §1 · p. 1Method
Exo-GPU treats parallelism and synchronization as annotations on sequential code and verifies that parallel execution preserves the sequential program’s behavior.
Source: §1 · p. 2, §2.3 · p. 6Results
Over 80% of theoretical peak on large GEMM problem sizes, with some kernels outperforming vendor-supplied CUBLAS.
Source: §1 · p. 3Takeaways & Limitations
The equivalence-checking approach preserves program behavior through sequential rewrites and final sequential-parallel verification.
Source: §2.3 · p. 6Generated code is not proven against formal CUDA semantics, and synchronization checking currently supports only concrete problem sizes.
Source: §8 · p. 21
Abstract
from arXiv · showhide
Modern GPUs require not only SIMT-style parallelism but also software-managed concurrency between compute and data movement to reach maximum performance. Performance engineers must reason about subdividing work into the hierarchy of computation resources (threads, warps, warpgroups, blocks, clusters), and, in many cases, also must use asynchronous tensor core and memcpy instructions on different levels of the memory hierarchy (registers, tensor core accumulators, shared memory, global memory). Unlike CPUs, where out-of-order execution is managed by hardware and hidden from programmers, GPUs expose explicit instruction reordering to software through these asynchronous instructions. Well-established GPU programming languages generally offer either direct low-level control without safety guarantees (e.g., CUDA C++ inline assembly or intrinsics) or easier-to-analyze, high-level abstractions (e.g., Triton's tile-based model) that hide asynchronous instructions in the compiler backend, which may prevent performance engineers from maximizing performance by tuning critical details. We propose Exo-GPU, an imperative, low-level language that creates minimal abstraction over CUDA. Our key idea is to treat parallelism and synchronization as mere annotations on sequential code rather than as fundamental control flow primitives, enabling verification that these constructs do not alter the program semantics. The benefit is twofold: programmers can reason about code without hidden control flow or mutation, while allowing the Exo-GPU compiler to verify sequential-parallel equivalence--guaranteeing that parallel execution is functionally equivalent to its sequential interpretation. We used Exo-GPU to author GEMM kernels for the H100 GPU, using wgmma, TMA, and split-k. Our kernels achieved over 80% of theoretical peak on large problem sizes, in some cases outperforming the vendor-provided CUBLAS library.
1 INTRODUCTION
Modern GPU features expose asynchronous instruction reordering and synchronization to programmers, while high-level abstractions can hide details needed for tuning. Exo-GPU adds low-level control and checks parallel execution against a sequential semantic baseline; its GEMM kernels exceed 80% of theoretical peak on large problems and sometimes outperform CUBLAS.
1 INTRODUCTION
- GPU tensor features expose explicit instruction reordering and shift synchronization work to programmers, complicating correct, high-performance CUDA programming.CUDA-like control offers access to asynchronous data movement, while higher-level abstractions can move responsibility for synchronization into the compiler.Source: §1 · p. 1
- CUDA’s low-level control exposes dataflow, instruction-format, inter-thread, and asynchronous intra-thread synchronization bugs.For example, wgmma requires a particular input format and uniform execution by a full warpgroup.Source: §1 · pp. 1–2
- Exo-GPU treats parallelism and synchronization as annotations on sequential code, so the sequential program defines the intended behavior.Programmers can reason without hidden control flow or side effects.Source: §1 · p. 2
- The compiler verifies that synchronization annotations preserve the sequential program’s behavior, while programmers explicitly control parallelism.The program’s parallel interpretation must be equivalent to its sequential interpretation.Source: §1 · p. 2
- Exo-GPU extends Exo and reuses its over 60 verified rewrites to transform sequential GEMM into a Hopper kernel through checked equivalence steps.Exo checks sequential rewrites; Exo-GPU checks equivalence between the final parallel and sequential interpretations.Source: §1 · pp. 2–3
- Over 80% of theoretical peak on large GEMM problem sizes, with some kernels outperforming vendor-supplied CUBLAS.The transformation chain checks sequential rewrite equivalence and final sequential-parallel equivalence.Source: §1 · p. 3
2 EXAMPLES AND SYSTEM OVERVIEW
Exo-GPU exposes CUDA-level parallelism, memory distribution, and synchronization while treating these constructs as annotations on sequential code. It illustrates synchronization hazards and checks that scheduling transformations preserve sequential behavior under parallel execution.
2.1 Exo-GPU Language Overview
- Parallel loops map iterations across clusters or subsets of threads, while collective tiling records which threads execute statement instances in each scope.Thread units specify participating set shapes, such as contiguous 32-thread warps or single threads.Source: §2.1 · p. 3
- Explicit allocations identify memory spaces, and distributed objects are sharded across threads; threads access only owned shards unless using explicit communication.Sharding is inferred from usage and checked for consistency by distributed-memory analysis.Source: §2.1 · p. 3
2.2 Exo-GPU Examples with Erroneous Synchronization
- A CTA-scoped fence cannot prevent one CTA’s TMA multicast from overwriting shared memory another CTA is still reading.The example demonstrates that multicast cooperation does not itself synchronize the two CTAs.Source: §2.2 · p. 5
- An in-order cluster barrier does not wait for asynchronous wgmma, so a multicast can overwrite its input before the tensor-core operation finishes.The examples contrast this unsafe ordering with Arrive/Await pairs that order the relevant accesses and completion.Source: §2.2 · pp. 5–6
2.3 Optimization Process and Equivalence Checking
- Scheduling transformations can be applied as reusable sequences of verified primitives, with compiler-managed index calculations for tiling.Programmers may instead write optimized object code directly; both approaches retain control over performance-critical details.Source: §2.3 · p. 6
- Exo-GPU’s scheduling flow combines sequentially checked rewrites with a final synchronization check to verify that parallel execution preserves the initial program’s behavior.The checker also guards against inter- and intra-thread synchronization bugs for selected concrete problem sizes.Source: §2.3 · p. 6
3 GPU IR
GPU IR is Exo-GPU’s frontend language, organizing its expressions, types, procedures, and statements. Its statement forms describe kernel launches, thread assignment, memory access, and synchronization.
3 GPU IR
- GPU IR is the frontend language of Exo-GPU and is used for the paper’s examples.It is a static control program with distinct control, data, barrier, and warp expressions.Source: §3 · p. 7, §3.1 · p. 7
3.1 Variables and Expressions
- GPU IR separates control, data, barrier, and warp expressions, with window coordinates specifying point or interval array accesses.Scalar reads are syntactic sugar for array accesses with an empty window-coordinate tuple.Source: §3.1 · p. 7
3.2 Types
- GPU IR types specify data precision, barrier configuration, collective thread domains, memory spaces, and qualitative timelines for buffer accesses.A newer Exo-GPU version removed the explicitly guarded barrier requirement noted for certain mbarrier configurations.Source: §3.2 · pp. 7–8
3.3 Programs
- Procedures model CPU functions or hardware instructions, with parameter annotations specifying data precision and memory type.Instruction annotations also encode hardware-specific properties and required CUDA collective units.Source: §3.3 · pp. 8–9
3.4 Statements
- Statements include loops, guards, indexed writes and reductions, windows, kernel launches, thread selection, and explicit synchronization.CUDA device functions launch nests of cuda_tasks loops, while cuda_threads and CudaWarps control intra-cluster thread execution.Source: §3.4 · pp. 9–10
4 COLLECTIVE ANALYSIS
Exo-GPU maps statements to thread collectives and statically checks whether CUDA-scoped data and barriers are used consistently with their declared memory types. Its code generation lowers those mappings and synchronization choices into CUDA C++ while enforcing safe shared-memory reuse.
4 COLLECTIVE ANALYSIS
- Collective analysis annotates scopes with thread tilings, deriving new tilings at device entry, thread loops, and warp-specialization blocks while other statements inherit their parent’s tiling.The detailed derivation of a new tiling is left for future work.Source: §4 · pp. 10–11
4.1 Distributed Memory Analysis
- Distributed-memory analysis checks allocation and use against memory ownership, consistent indexing across access sites, and the thread mappings induced by logical indices.Instruction calls are expanded into explicit cuda_threads loops before collective analysis and access triples are collected for checking.Source: §4.1 · pp. 11–12
- Passing distributed-memory analysis guarantees memory behavior is coherent with the collective structure and declared memory types.The analysis rejects mismatched allocation and accesses, inconsistent indexing across sites, and missing or extra CUDA-thread iterators.Source: §4.1 · p. 12
4.2 Code Generation to CUDA Code
- CUDA code generation inverts thread mappings to infer loop variables from thread indices and specializes warp paths to control register use.Warp-specific paths use setmaxnreg, and omit code guarded for other warps, allowing PTX assembly to check that register-heavy instructions stay on suitable paths.Source: §4.2 · p. 12
- Persistent kernels schedule lexicographically ordered tasks round-robin across clusters, while shared-memory reuse waits for all accesses to retire before freeing storage.The compiler may alias shared memory suboptimally, and its cluster-wide safety requirement may be conservative; explicit lifetime and aliasing control is future work.Source: §4.2 · pp. 12–13
5 ABSTRACT MACHINE
Exo-GPU checks synchronization by interpreting an Abstract Machine IR. The IR is used for validation, separating synchronization reasoning from value computation while keeping the frontend simple.
5 ABSTRACT MACHINE
- Abstract Machine interpretation checks synchronization to guarantee sequential–parallel equivalence, with validation factored into explicit operations orthogonal to value computation.The section introduces the AMIR grammar, its translation from GPU IR, and its execution semantics.Source: §5 · p. 13
5.1 Abstract Machine Grammar
AMIR retains GPU IR’s control, boolean, window-coordinate, and variable/type definitions, and adds visibility levels and synchronization-focused statements.
5.1 Abstract Machine Grammar
- AMIR reuses GPU IR definitions for control, booleans, window coordinates, and variables and types, while introducing windowed access expressions and visibility levels.The visibility levels are totally ordered from invisible through atomic_only, unordered, temporally_ordered, to fully_ordered.Source: §5.1 · p. 13
- AMIR statements include structural control, operations that record synchronization events, and checks that fail when required visibility obligations are unmet.The recording operations include allocation, reads, mutations, fences, arrivals, and awaits; checks include barrier, read, and mutation checks.Source: §5.1 · p. 13
5.2 Conversion from GPU IR to Abstract Machine IR
GPU IR is translated into AMIR through expression, statement, instruction-argument, and program conversions. The translation records synchronization effects and checks while preserving the relevant control structure.
5.2 Conversion from GPU IR to Abstract Machine IR
- Before conversion, window statements and subprocedure calls outside device functions are inlined; ordered variable sequences become left-to-right AMIR composition.An empty sequence converts to a no-op.Source: §5.2 · pp. 14–15
5.2.1 Expression Conversion.
- Expression conversion extracts windowed accesses from names, indexings, and arithmetic while leaving value expressions unchanged.Scalars contribute no accesses, and arithmetic combines accesses by set union.Source: §5.2.1 · p. 15
5.2.2 Statement Conversion.
- Statement conversion preserves sequencing and conditionals, normalizes loops, and lowers assignments into explicit read and write effects.Warp and block scopes are replaced with true guards; sync-exempt memory accesses are removed from checking.Source: §5.2.2 · p. 15
5.2.3 Instruction Argument Conversion.
- Instruction-argument conversion emits checks and recorded effects based on whether each parameter is read-only, write-only, read–write, or atomic.Access modes are determined by a static analysis of the instruction body.Source: §5.2.3 · p. 15
5.2.4 Program Conversion.
- Program conversion retains control arguments and variables but omits non-control arguments, then converts the program body with the statement translation.Source: §5.2.4 · p. 15
5.3 Abstract Machine Semantics
The abstract machine gives Exo-GPU sequential semantics for expressions and statements while tracking thread assignments and synchronization state. Its checks detect access-visibility and barrier-consistency violations without changing state when they succeed.
5.3 Abstract Machine Semantics
- The machine state tracks control values, global thread identities, and per-variable read and mutation visibility records plus barrier counts.Thread mappings determine which local threads execute each statement; CUDA and CPU statements use different scope mappings.Source: §5.3 · pp. 15–16
5.3.1 Expression semantics.
- Expression rules evaluate literals, reads, arithmetic, logic, comparisons, windows, and accesses into values or concrete indexed access sets.Windows represent points, inclusive ranges, or Cartesian products of index sets.Source: §5.3.1 · pp. 16–17, Figure 6 · p. 16
5.3.2 Statement semantics.
- Statement rules thread machine state through loops, sequencing, task-ID increments, guards, allocation, and synchronization-record updates.Loops iterate over a half-open interval; task-ID increments leave the store and synchronization state unchanged.Source: §5.3.2 · p. 17
5.3.3 Checks.
- Read checks reject insufficient visibility across required timelines, while barrier checks reject mismatched arrive and await counts; successful checks preserve state.Violations reduce the configuration to an error state.Source: §5.3.3 · pp. 17–18, Figure 7 · p. 17
5.3.4 Program semantics.
- The interpreter initializes control values from user inputs, starts task ID at zero with empty synchronization state, and evaluates the program statement.Source: §5.3.4 · p. 18
6 PERFORMANCE RESULTS
On H100, Exo-GPU's large square GEMM performance is very similar to CUBLAS, while GEMV generally trails it and unusual GEMM sizes show more variation. The synchronization-check interpreter's runtime scales with the number of memory operations interpreted.
6 PERFORMANCE RESULTS
- The evaluation used sm_90a H100 hardware, tf32 GEMM with wgmma and TMA, fp32 GEMV, and cuBLAS comparisons across problem sizes and kernel variants.Each problem-size and hardware combination used 100 timed samples after five warmups; reported kernel results include specialized and single-kernel choices.Source: §6 · pp. 18–19
- Large square GEMM performance is very similar between Exo-GPU and CUBLAS, while GEMV usually favors CUBLAS and unusual GEMM sizes vary more.Exo-GPU expresses warp specialization, cluster and TMA multicast, wgmma, and split-barrier synchronization for sm_90a GEMM.Source: §6 · p. 19
- 2/16/128 seconds for GEMM at 1024/2048/4096; synchronization-check runtime scales linearly with interpreted memory operations.The small-GEMM runtime bump is attributed to imperfect division by the per-cluster tile size; a 768-size development test can provide quick bug feedback.Source: §6 · p. 19, Figure 9 · p. 19
7 RELATED WORK
Related GPU languages trade abstraction and control differently, while Exo-GPU exposes low-level execution structure and adds synchronization checking. Its design extends Exo's sequential reasoning to GPUs, where parallel execution must be reasoned about explicitly.
7.1 GPU Kernel Authoring Languages
- Tile DSLs abstract thread mapping but can hide performance-critical details, while template libraries preserve low-level control without whole-program synchronization safety guarantees.Exo-GPU keeps a thin abstraction over CUDA and makes performance-critical control flow and data copies explicit.Source: §7.1 · p. 20
7.2 User-schedulable Languages
- Exo-GPU's WYSIWYG GPU IR exposes low-level control and checks sequential-parallel equivalence, unlike approaches that hide warp specialization or focus on lowering-based schedule generation.User-schedulable languages express optimization as programmer-directed rewrites to equivalent, faster programs.Source: §7.2 · p. 20
7.3 Exo vs. Exo-GPU
- Because GPUs expose data and instruction parallelism to software, Exo-GPU extends Exo with user-controlled parallel loops and synchronization, then checks equivalence between sequential and parallel interpretations.This completes an equivalence chain from the initial Exo program through its optimized parallel implementation.Source: §7.3 · p. 21
8 LIMITATIONS AND FUTURE WORK
The paper identifies limits in Exo's sequential checks for certain Hopper GEMM rewrites and in Exo-GPU's current GPU reasoning and synchronization-checking scope. It presents verified, explicit control as a promising direction for GPU programming.
8 LIMITATIONS AND FUTURE WORK
- Two Hopper GEMM rewrites involving TMA out-of-bounds behavior fall outside Exo's sequential safety checks and were marked unsafe despite being semantically sound.Dependency-based checks cannot prove that guard-removal writes are observationally redundant; the authors identify value-sensitive rewrites as needed.Source: §8 · p. 21
- The system lacks a proof against formal CUDA semantics, Blackwell tcgen05 support, non-CUDA platforms, and symbolic-size checking; interpretation also limits compiler performance.The authors suggest static analysis of the abstract-machine program as a possible improvement.Source: §8 · p. 21
- The authors characterize explicit programmer control paired with safety verification as a promising path for GPU programming.Source: §8 · p. 21
A GPU IR TYPE DEFINITIONS
This section’s GPU IR type reference covers variable and CUDA-specific types, collective units, and qualitative timelines. Collective units are parameterized by a domain and box, while SyncTL is defined as a composition of QualTL.
A GPU IR TYPE DEFINITIONS
- The reference lists variable types and qualitative timelines, and defines SyncTL as a composition of QualTL.The supplied figure captions identify these lists and the composition relation.Source: Figure 10 · p. 24, Figure 15 · p. 25, Figure 16 · p. 25
- Collective units are parameterized by δ, which consists of a domain and box.The figure identifies the domain and box as components of the collective type.Source: Figure 11 · p. 24
- The CUDA type reference includes memory allocations, barrier completion mechanisms, and special window types.The figures identify these three CUDA-related type categories.Source: Figure 12 · p. 24, Figure 13 · p. 25, Figure 14
- The Sm90 tensor-map special window is parameterized by swizzle mode and shared-memory box size.The window is denoted πw and uses a CUtensorMap.Source: A GPU IR TYPE DEFINITIONS · p. 25
B ABSTRACT MACHINE CONVERSION AND SEMANTICS DEFINITIONS
This section defines conversions and big-step semantics for the abstract machine, including how memory visibility records and synchronization operations are updated. Its rules cover expression reads, statement and argument conversion, and synchronization-aware execution.
B ABSTRACT MACHINE CONVERSION AND SEMANTICS DEFINITIONS
- Record reads and mutations update separate visibility-record sets while leaving the control state and store unchanged.The constructor R is used to add the corresponding record to the read or mutate set.Source: B ABSTRACT MACHINE CONVERSION AND… · pp. 25–26
- Fence, Arrive, and Await update visibility through lift, which applies a function pointwise to read and mutate visibility sets.The rules also update barrier arrive or wait counts; Await’s augmentation is conditional on an observed arrive count.Source: B ABSTRACT MACHINE CONVERSION AND… · p. 26
- The pointwise join of two visibility maps takes the maximum VL value at each indexed position.Source: B ABSTRACT MACHINE CONVERSION AND… · p. 26
- Expression conversion collects distinct reads in an expression and returns them as a set of AMIR reads.Source: Figure 17 · p. 27
- The conversion rules include helper functions for statement conversion and argument conversion.The figures identify helper functions for both conversion procedures.Source: Figure 18 · p. 27, Figure 19 · p. 28, Figure 20 · p. 28
- Big-step semantics are given for loops, sequencing, task-ID increments, guards, and abstract-machine statements with synchronization-environment updates.These cases are identified in the semantics figures.Source: Figure 21 · p. 28, Figure 22 · p. 29