lib/accy/src/root.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 //! A compiler and runtime for programs over multidimensional arrays, written in Zig: a caller
 2 //! builds a program by calling array operations, and the package records each call and checks the
 3 //! shapes of its values. The package compiles the program into functions a device runs across many
 4 //! threads at once (*kernels*) and runs them on a CPU or a GPU.
 5 //!
 6 //! A shape mistake, such as adding two arrays whose axes line up by accident, has to fail when the
 7 //! program is built, before any device time is spent. One program has to run on NVIDIA GPUs through
 8 //! CUDA, on Vulkan, Metal and WebGPU devices, on the host CPU, and in WebAssembly. A long-running
 9 //! process recompiles programs as they change, so a compile has to reuse the work that did not
10 //! change and stay inside a budget the caller sets. Matrix products, factorizations, sorts and
11 //! neighbor searches run fast only as hand-written kernels whose schedules are tuned for each
12 //! device. A number that comes back from a GPU needs an independent reference before anyone can
13 //! trust it.
14 //!
15 //! When axes are named by position, a matrix product lists its batch axes and its summed axes in
16 //! separate lists, and a check that reads each list alone accepts one axis in both, after which the
17 //! output rank computed by subtraction disagrees with a count of the axes that remain. The math of
18 //! a program and the loops that carry it out on a device change for different reasons: the first
19 //! when the model changes, the second when the device, the memory layout or the tiling changes.
20 //! Devices take different code formats, such as PTX or cubin for CUDA, SPIR-V for Vulkan, MSL or
21 //! metallib for Metal and WGSL for WebGPU, and faster arithmetic such as TF32 on tensor cores
22 //! exists on some NVIDIA devices only. The fastest schedule for a kernel depends on the device and
23 //! on the size of the problem, so no fixed choice wins everywhere.
24 //!
25 //! XLA, StableHLO and PJRT, from Google and the OpenXLA community, faced these problems before: XLA
26 //! compiles typed graphs of array operations (https://openxla.org/xla), StableHLO fixes the
27 //! operation set and its meaning (https://openxla.org/stablehlo), and PJRT is the runtime interface
28 //! a compiler uses to reach any device (https://github.com/openxla/xla/tree/main/xla/pjrt). From
29 //! them the package takes typed graphs of array operations, operation meanings that stay fixed, and
30 //! one runtime boundary for every device: its operations carry StableHLO's names and fields, for
31 //! example a matrix product with separate batch and contracting axis lists, a broadcast that maps
32 //! each input axis to an output axis, gather, reduce and select, and every device is reached
33 //! through one interface.
34 //!
35 //! Each array axis carries a name and a length (a *named axis*), so operands line up by name, and a
36 //! missing name or an unequal length is an error while the program is recorded, before anything
37 //! runs. An operation multiplies and sums two arrays (the *generalized matrix product*,
38 //! `dot_general`) after splitting the axes of each operand into three disjoint groups, batch,
39 //! contracted and free, and paired batch and contracted axes must have equal lengths, so the output
40 //! rank is always the batch count plus the two free counts. A Lean proof in verification/accy shows
41 //! that accepted role lists partition each operand, and that letting one axis sit in two groups
42 //! breaks the rank count.
43 //!
44 //! Array operations lower into this repository's compiler infrastructure (*Choir*), a Zig library
45 //! modeled on MLIR. That infrastructure groups its operations by level into named groups of
46 //! compiler operations (*dialects*), and lowering the array operations into its structured loop
47 //! groups keeps the array math apart from device-specific loop scheduling and memory tiling.
48 //!
49 //! Compiling runs as a fixed chain of seven stages that takes a recorded program to kernels legal
50 //! on one device (*preparation*): semantic, contract, tensor, dispatch, memory, kernel and target.
51 //! The result of each stage is sealed as an immutable record (a *stage record*), and a later
52 //! compile reuses a record only when an exact check admits it for the new request, with every stage
53 //! charged against work limits the caller sets.
54 //!
55 //! Hand-written kernels for matrix products, factorizations, spatial grids, sorts and other
56 //! families are chosen by a typed request from one library (the *kernel library*), and measured
57 //! tables keyed by device and problem size pick their schedules. Every device sits behind one
58 //! interface, a pointer and a table of functions (a *backend*), through which a caller creates
59 //! buffers, loads compiled code and launches kernels. Compiled code records the arithmetic it
60 //! follows as part of its identity (its *math tier*): exact by default, or TF32 on CUDA tensor
61 //! cores, which fails to compile on any other device and never falls back to exact arithmetic.
62 //!
63 //! An evaluator of kernel programs that needs no device (an *oracle*) gives every backend result an
64 //! independent reference, and conformance tests compare each device output against it within a
65 //! stated tolerance.
66 //!
67 //! Programs travel between processes as versioned bytes (their *wire form*), which the receiver
68 //! replays through the same checked step that records a program. A small table inside the host
69 //! process compiles those programs for the CPU and runs them on caller memory (a *session*).
70 //!
71 //! The package's parts are element types (`dtype`), axis roles (`axis`), device interfaces
72 //! (`backend`), compiled code and its identity (`artifact`), loaded programs ready to launch
73 //! (`executable`), the compile chain (`preparation`), device targets (`target`), recorded programs
74 //! with their wire form and session (`tensor`), kernels and the kernel library (`kernel`), the
75 //! rules that decide when a device test runs or skips (`validation`), and the compiler dialects
76 //! (`choir`). The device-free evaluator lives in `eval` and is reached through the kernel part, and
77 //! benchmarks such as `accy-choir-bench` and `accy-versus-bench` are separate build steps outside
78 //! the library.
79 //!
80 //! - *tuning table*: measured schedule winners keyed by device, operation and problem size.
81 
82 pub const axis = @import("axis/root.zig");
83 pub const artifact = @import("artifact/root.zig");
84 pub const executable = @import("executable/root.zig");
85 pub const preparation = @import("preparation/root.zig");
86 pub const target = @import("target/root.zig");
87 pub const tensor = @import("tensor/root.zig");
88 pub const kernel = @import("kernel/root.zig");
89 pub const validation = struct {
90     pub const gating = @import("accy_validation_gating");
91 };
92 
93 pub const choir = @import("choir/root.zig");