GAIA-ML

GAIA-ML is an experimental ONNX compiler that generates AMD AI Engine v1 graphs, vector kernels, packed weights, and simulation projects for low-latency inference. Version 0.10.0 focuses on compact multilayer perceptrons, small convolutional networks, and bounded additive tree ensembles.

GAIA-ML is developed by the NGT WP2.1 Group and is distributed under the MIT license.

What the compiler does

The compiler processes a model through these representations:

ONNX model
  -> imported operator graph
  -> canonical Math IR
  -> mathematical rewrites
  -> Loop IR
  -> Schedule IR and implementation search
  -> AIE IR and target checks
  -> AMD ADF graph, AIE C++ kernels, packed weights, and test vectors

The implemented pipeline can:

• inspect ONNX graphs and inferred tensor shapes;
• simplify local algebraic identities;
• canonicalize MatMul + Add and Gemm into Dense operations;
• canonicalize ONNX convolution and layout helper patterns into Conv2D;
• fold BatchNorm constants into convolution weights and bias;
• fuse supported ReLU and sigmoid activations;
• recognize depthwise convolution with a depth multiplier and fuse supported one-pixel float32 depthwise-Conv-to-Dense classifier tails;
• propagate physical tensor layouts through convolution chains;
• rank scalar, vector-dot, streaming, native-MMUL, sliding, and line-buffer microkernels;
• pad streams and reorganize weights for legal AIE vector and MMUL shapes;
• split work across tiles by output, spatial, batch, reduction, or channel dimensions when the selected backend supports it;
• validate AIE v1 tile memory, stack guidance, stream-port limits, and model input transfer bounds;
• lower bounded ONNX-ML tree ensembles through a QuickScorer bit-vector backend;
• emit float32, int32, int16, or int8 neural-network kernels; and
• generate Vitis x86 simulation and AIE simulation projects.

Target

GAIA-ML 0.10.0 models an AIE v1 device with:

• 1.25 GHz AIE cores;
• 32 KiB memory per compute tile;
• no memory tiles;
• at most two input and two output streams per kernel; and
• 32-bit PLIO at either 312.5 MHz or 500 MHz.

Generated projects default to part xcvp2802-vsva5601-2MHP-e-S. Override it at build time when using another compatible AIE v1 part:

make aie AIE_PART=<part-name>

Installation

Install the release from PyPI:

python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install gaia-ml==0.10.0
gaia_ml version

For development from a checkout:

python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -e .
gaia_ml version

The Python package requires Python 3.9 or newer. Generating a project does not require Vitis. Compiling or simulating generated AIE sources requires a local AMD Vitis installation; the generated build flow has been exercised with Vitis 2024.2.

Inspecting a model

Inspect the complete compiler pipeline:

gaia_ml inspect model.onnx --stage all

Inspect an individual representation or pass:

gaia_ml inspect model.onnx --stage raw
gaia_ml inspect model.onnx --stage algebra
gaia_ml inspect model.onnx --stage canonicalize
gaia_ml inspect model.onnx --stage math
gaia_ml inspect model.onnx --stage loop
gaia_ml inspect model.onnx --stage schedule
gaia_ml inspect model.onnx --stage constraints

Use JSON output for scripts and regression tests:

gaia_ml inspect model.onnx --stage schedule --json > schedule.json

For a dynamic ONNX input, provide one static shape:

gaia_ml inspect model.onnx --input-shape 1,18,18,2 --stage constraints

List implemented compiler stages and scheduling controls:

gaia_ml show optimizers
gaia_ml show optimizers --json

Compiling a model

The recommended starting point lets GAIA-ML choose legal implementations:

gaia_ml compile model.onnx \
  --output gaia_project \
  --force \
  --schedule-objective latency \
  --data-file model_input.csv \
  --simulation-events 10

The output directory contains:

gaia_project/
├── graph/          ADF graph declaration and simulation host
├── kernels/        generated AIE C++ kernels
├── weights/        compile-time packed coefficient headers
├── testVectors/    scalar and AXI/PLIO input vectors
├── tools/          generated simulation-report helper
├── Makefile
├── CMakeLists.txt
├── manifest.json   selected IR, schedule, constraints, and estimates
└── compile.log     compiler decisions and diagnostics

compile.log is also printed to the terminal. manifest.json is the preferred machine-readable record of the chosen schedule.

Both outputs distinguish four latency quantities: the MAC-issue-only theoretical compute minimum, the physical generated-I/O lower bound, the emitted compiler-program estimate (first output and steady-state II), and the user-requested steady-state latency. These are deliberately separate because the theoretical MAC floor excludes stream transfer, reductions, activation, address generation, routing, and control overhead.

Objectives and explicit implementations

An objective ranks legal candidates; it is not itself a microkernel:

--schedule-objective latency
--schedule-objective throughput
--schedule-objective compact
--schedule-objective auto

Keep --schedule-implementation auto for normal compilation. An explicit implementation is primarily useful for controlled benchmarking or debugging:

gaia_ml compile model.onnx \
  --output gaia_project \
  --schedule-objective latency \
  --schedule-implementation conv2d_line_buffer

Not every explicit implementation is legal for every operation, shape, or datatype. GAIA-ML reports a structured compiler error instead of silently falling back when a forced implementation is illegal.

Target output interval

--target-latency-us means the requested steady-state pipeline initiation interval: after pipeline fill, the graph should be able to accept and produce events at that interval. It is not a promise about first-output latency.

gaia_ml compile model.onnx \
  --output gaia_project \
  --schedule-objective latency \
  --target-latency-us 3

GAIA-ML rejects a target below the calculated input-transfer bound. Meeting an estimated target still requires validation with AIE simulation or hardware, especially for shapes outside the calibrated microkernel set.

Datatypes and quantization

Select storage and arithmetic scheduling with --schedule-dtype or its --quantize alias:

gaia_ml compile model.onnx --output float_project --schedule-dtype float32
gaia_ml compile model.onnx --output int16_project --schedule-dtype int16
gaia_ml compile model.onnx --output int8_project --schedule-dtype int8

Supported code-generation types are float32, int32, int16, and int8. Integer kernels use wider accumulators and datatype-specific AIE primitives where legal. Reduced precision does not automatically guarantee lower latency: the available AIE instruction shape, padding, conversion work, routing, and model topology all affect the result.

int4 is accepted for schedule analysis but code generation deliberately rejects it because packed-nibble storage and arithmetic are not implemented.

Compile-time quantization converts weights and generated input vectors using the compiler's fixed-point policy. It is not a replacement for quantization- aware training or a model-accuracy study. Always compare generated output with an independently produced golden reference.

Multi-tile schedules

GAIA-ML exposes the following tiling controls:

none
output_split
spatial_split
reduction_split
batch_split
channel_unroll

With --multi-tile none, the automatic latency scheduler may still choose a multi-tile topology when a single tile is illegal or a calibrated parallel implementation ranks better. Explicit multi-tile modes constrain the search and can require additional broadcast, stitch, or reduction kernels.

Shape-driven float32 Conv selection

Schedule IR records a model-independent shape_analysis for each Conv2D. It classifies the layer as depthwise-multiplier, depthwise, pointwise, small-channel spatial, or channel-reduction; identifies the vector axis and estimated native-lane utilisation; records border and line-buffer properties; and lists independently tileable output dimensions and the required weight transform.

For float32 depthwise convolution, a depth multiplier of four through eight can be packed as [input channel][kernel tap][8 multiplier lanes]. GAIA then computes all multiplier filters for one input channel with one native AIE1 vector accumulator. Smaller multipliers remain on reduction/spatial vectorisation because filling fewer than half of the eight float lanes was not profitable in Vitis-generated code. Explicit multi-tile requests currently use the general depthwise backend so that tile boundaries cannot split a packed multiplier group.

Shape-driven Dense selection

Schedule IR also records a shape_analysis for every Dense operation. It classifies dynamic, tiny, quantized, convolution-fed, native-MMUL, and general vector Dense shapes; measures reduction- and output-lane utilisation; and describes the required weight transform. The analysis recommends scalar, vector-dot, streaming outer-product, buffered MMUL, or direct MMUL codegen and an appropriate single-tile, output-split, or batch-split topology.

These recommendations seed the implementation search rather than replacing it. GAIA's calibrated cost model remains the final ranking step. In particular, the direct MMUL accumulator is restricted to a single two-output block: previous AIE simulations showed that retaining several output blocks in one kernel can increase register pressure and perform worse than buffered MMUL. Streaming outer-product now uses the native AIE broadcast operation and packed output-major weights, without a scalar temporary broadcast array.

Use inspect --stage schedule to see the classification, vector utilisation, recommended microkernel and topology, and any cost-model override.

AIE v1 capability contracts

Version 0.9.8 gives every current Dense, Conv2D, depthwise Conv2D, activation, Softmax, and BDT implementation a formal capability declaration. A contract records supported dtypes, interfaces, parallel modes, layouts, vector widths, static-shape requirements, calibration provenance, fusion support, and the existing AIE lowering backend. Each scheduled candidate is evaluated through the contract for matching, target legality, code-generation readiness, cost, latency, and its concrete lowering layout.

The AIE v1 registry now owns the implementation universe and candidate admission. Dense and BDT builders enumerate their alternative tilings and microkernels through the registry; Conv2D and activation selection must also resolve to a registered strategy before lowering. Unsupported variants remain visible as rejected candidates but cannot displace a supported implementation.

Fusion is explicit as well: compact CNN, Dense-chain, depthwise-to-Dense, and streaming CNN-tail transformations are registered whole-region strategies. This preserves the measured fast paths while removing hidden lowering choices. Use JSON inspection to consume the complete report:

gaia_ml inspect model.onnx --stage schedule --json

The selected contracts are under schedule.analysis.capability_contracts, with an aggregate capability_summary. schedule.analysis.strategy_enumeration records declared, constructed, admitted, and rejected candidates. Human-readable inspection prints both views.

Conv2D scheduling additionally emits cross_validation for every operation. The compiler independently builds the applicable standard or depthwise strategies, compares their topology, memory, cycle, and steady-II estimates, and records whether each alternative has enough evidence to replace the current winner. MAC-only estimates remain visible as analytical shadow candidates. They cannot replace a Vitis/AIESim-validated implementation unless the alternative has independent same-geometry profile evidence and improves the confidence-adjusted objective by at least 20%. This guard prevents optimistic instruction-count estimates from silently regressing measured float32 latency.

Whole-graph planning

Version 0.10.0 includes an evidence-aware graph planner after local candidate enumeration. It performs a bounded search across the complete operator chain and compares steady-state pipeline II, first-output fill, tile count, stream barriers, and producer/consumer layouts. Layout mismatches receive an explicit adapter cost and cannot be selected until codegen can materialize the adapter. Flatten-to-Dense permutations are treated specially: when mathematically legal, GAIA reorganizes Dense weights offline instead of inserting a runtime transpose.

Candidate estimates are classified as analytical, geometry-scaled, profile-calibrated, or same-geometry measured. The planner reports optimistic challengers, but it changes a validated schedule only for a direct-layout plan with same-geometry evidence and at least a 20% whole-graph improvement. Fusion opportunities from the region registry are reported alongside their estimated kernel-boundary saving and are rechecked by AIE lowering for memory and semantic legality.

The JSON report is available under schedule.analysis.graph_planning. It contains the incumbent and challenger signatures, transition costs, evidence floor, evaluated-plan count, replacement decision, and eligible fusion regions.

Schedule-side Conv alternatives also use the emitted-backend calibration model. It charges scalar address generation, patch packing, padded stream traffic, fixed prologues, and backend-specific vector work before graph ranking. This corrects the canonical 18x18/K7/Cout5 direct-Conv shadow estimate from 1,102 cycles to 107,136 cycles (85.7088 us at 1.25 GHz), while leaving the selected spatial-vector estimate at 2.6576 us. Compile output includes latency_estimates.calibration_consistency, which compares Schedule IR and emitted-program steady-state II and reports their delta and ratio.

Empirical compiler feedback

Version 0.10.0 closes the validation loop without allowing unrelated profiles to influence scheduling. After a completed simulation, record its event timing with:

gaia_ml profile test_prj/gaia_compile \
  --database .gaia_ml/aie_v1_feedback.json \
  --source "Vitis 2024.2 AIESim"

GAIA derives a SHA-256 fingerprint from the complete AIE v1 schedule, tensor geometry, physical layouts, kernel topology, and PLIO policy. Repeated runs are aggregated with medians and assigned a confidence level. The fingerprint also includes the ONNX model, emitted weight/table, and optional input-data hashes. Simulation records capture the output/report artifact hash and detected Vitis version; importing the same artifacts again is deduplicated. A JSON copy of the observation is written to gaia_feedback.json in the generated project.

Pass the database to a later identical compilation to expose the measurement alongside the analytical and compiler estimates:

gaia_ml compile model.onnx \
  --output test_prj/gaia_compile \
  --calibration-db .gaia_ml/aie_v1_feedback.json \
  --force

Only an exact emitted-plan fingerprint matches. GAIA does not extrapolate a measurement across shapes, data types, tilings, layouts, interfaces, or kernel topologies, and empirical feedback does not silently replace the selected schedule. This makes discrepancies actionable while preserving the validated fast paths.

Bounded BDT support

GAIA can compile a single-output ONNX-ML TreeEnsembleRegressor. The scheduler ranks three representations: legacy QuickScorer, cumulative-prefix QuickScorer, and normalized direct traversal. Prefix QuickScorer replaces data-dependent split scans and scattered tree updates with a threshold-bucket lookup followed by contiguous bitwise masks. Its threshold, prefix-mask and leaf tables are placed in separate AIE data-memory banks. Direct traversal covers deeper trees without requiring one bit per leaf.

Generate a deterministic test model and 100 golden events with:

python tests/generate_bdt_models.py \
  --output-dir model \
  --features 32 \
  --depth 5 \
  --trees 8 \
  --events 100 \
  --force

gaia_ml compile model/test_bdt.onnx \
  --output test_prj/bdt \
  --data-file model/test_bdt_input.csv \
  --simulation-events 10 \
  --schedule-objective throughput \
  --multi-tile forest_split \
  --tile-output 2 \
  --force

This development backend supports float32 features, BRANCH_LEQ, SUM or AVERAGE, one regression target, at most 32 features, depth at most eight, and at most 1024 trees. Trees through depth five can use 32-bit QuickScorer masks; deeper trees use direct traversal. --multi-tile forest_split partitions the forest into independently scheduled workers, broadcasts each feature vector through two-port-safe binary fan-out, and combines partial scores with a binary reduction. --tile-output N bounds the trees assigned to each worker. Multi-target/class-label output, other split predicates and integer thresholds remain future work. Forest splitting expands ensemble capacity and throughput, but each individual tree and worker table must still fit one AIE tile.

The BDT cost model ranks emitted implementations rather than treating the split count as the latency by itself. It reports the arithmetic lower bound, physical PLIO floor, calibrated compiler estimate, and requested steady-state II. The current Vitis 2024.2 calibration includes feature ingest, dependent traversal, prefix-mask updates, two-port-safe fan-out, reduction fill and per-worker table memory. For the deterministic 32-feature/depth-5/eight-tree fixture, the measured single-tile prefix backend is 2.131 us/event; four direct-traversal workers measure 0.830 us/event with a 1.971 us first output. These are calibration points, not universal guarantees—tree shape and target routing can change the result, so publish AIESim or hardware measurements alongside the estimate.

Stream and window interfaces

Streaming is the default:

--interface stream

Window interfaces are available for supported kernels:

--interface window

A stream does not inherently mean lower latency. Backpressure, reconvergent paths, padding, and kernel production rates can dominate. Use the simulation report before drawing conclusions from the interface name alone.

Input data

--data-file accepts CSV rows where each row is one flattened input event:

x0,x1,x2,...,xN

The logical row width must match the compiled model input. GAIA-ML performs any required physical padding and emits both scalar test-vector files and Vitis AXI/PLIO CSV files with command, data, TLAST, and TKEEP fields.

--simulation-events N selects the first N rows; it does not duplicate one row to create additional events.

If a sibling <stem>_output.csv file exists, GAIA-ML can include a golden reference in the generated project. Only the input CSV is required.

Building the generated project

Load the Vitis environment, enter the generated directory, and use the generated Makefile:

source /opt/modules/Vitis/2024.2/settings64.sh
cd gaia_project

make x86
make x86sim
make aie
make aiesim

Useful additional targets include:

make aiesim_profile
make aiesim_report
make aie_fifo

The report helper summarizes output timestamps, first-output latency, steady-state event intervals, large bubbles, PLIO throughput, stream stalls, and available FIFO guidance.

If hardware compilation reports stack overflow, use the stack recommendation in compile.log or manifest.json:

make aie AIE_STACK_SIZE=<recommended-by-gaia>

Do not increase the stack beyond the reported safe tile-memory headroom.

Supported model patterns

Version 0.10.0 targets feed-forward inference graphs composed from:

• Dense layers represented by Gemm or canonicalizable MatMul + Add;
• standard 2D convolution;
• depthwise 2D convolution represented by supported grouped convolution;
• ReLU and sigmoid activations;
• terminal Softmax, including LUT-based implementations where selected;
• BatchNorm that can be folded into constant convolution parameters; and
• Flatten, Reshape, Transpose, and layout helpers that can be eliminated or represented as compiler layout metadata; and
• bounded additive TreeEnsembleRegressor models described above.

Supported mixtures include multilayer perceptrons, convolution chains, and convolution-to-Dense classifier tails. Static shapes and constant trained weights are strongly recommended.

Current limitations

GAIA-ML 0.10.0 is an experimental compiler, not a general ONNX backend.

• Only the operators and canonicalizable patterns listed above are compiled. Pooling, recurrent networks, attention, arbitrary elementwise graphs, and general ONNX control flow are not implemented.
• Batch size one has the most mature latency path. Batch tiling exists, but is not as extensively calibrated as single-event streaming.
• Dense output and batch splitting are supported. A Dense whose single output reduction cannot fit one tile is diagnosed as requiring reduction splitting; general Dense reduction-tree code generation is not implemented yet.
• Dynamic dimensions must be resolved with --input-shape; general dynamic- shape code generation is not supported.
• Grouped convolution support is primarily intended for standard or depthwise cases. Arbitrary group configurations may select conservative code or be rejected.
• Automatic layouts, fusion, and tiling are heuristic. The compiler does not exhaustively search every legal graph placement or routing solution.
• Performance estimates combine analytical bounds with a limited set of Vitis 2024.2 AIE simulation calibrations. Estimates are most reliable near tested shapes and can be inaccurate for substantially different networks.
• A legal schedule is not guaranteed to be globally latency-optimal. Vitis placement, routing, FIFO behavior, stack use, and stream backpressure can change observed performance.
• Integer conversion is a compiler storage transformation, not automatic accuracy-preserving quantization. Accuracy must be validated by the user.
• int4 code generation is not implemented.
• Generated C++ targets AIE v1 APIs and has not been validated as an AIE-ML or future-architecture backend.
• Hardware execution, platform integration, and host application generation are outside the current release; GAIA-ML emits the AIE graph project.

For a new architecture, begin with inspect --stage constraints, compile with automatic scheduling, run at least ten simulation events, and compare the measured output and timing against an independent reference before forcing a specific microkernel.

Python inspection API

The stable Python-facing inspection objects are available from gaia_ml:

from gaia_ml import GraphAnalyzer, OnnxModelLoader

graph = OnnxModelLoader("model.onnx").load()
report = GraphAnalyzer(graph).analyze()

print(report.operation_count)
print(report.operation_counts)

The IR and low-level scheduling classes remain available for compiler development, but the command-line interface and generated manifest are the recommended user interfaces for version 0.10.0.

Development and release checks

Run the unit suite:

python -m unittest discover -s tests -v

The v0.10.0 suite freezes the compiler decisions and performance ceilings of representative float32 CNN, Dense, mixed Conv-to-Dense, and BDT paths. The fast suite regenerates these deterministic models and checks selected implementations, tile topology, estimated first-output latency and steady II, tile memory, stack, and kernel count.

Completed AIESim projects can be checked against the separately recorded Vitis 2024.2 measurements:

python tools/check_performance_baseline.py --list
python tools/check_performance_baseline.py \
  --case bdt_forest4 \
  --project test_prj/bdt_forest4

Compiler baselines run without Vitis. Measured gates require timestamped output from a completed AIESim run; the checker never treats a missing measurement as a successful performance result.

Build release artifacts:

python -m pip install build
python -m build

The package version is defined in python/gaia_ml/_version.py.