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.9.0 focuses on compact multilayer perceptrons and small convolutional networks built from supported Dense and Conv2D operations.

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;
• 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;
• emit float32, int32, int16, or int8 kernels; and
• generate Vitis x86 simulation and AIE simulation projects.

Target

GAIA-ML 0.9.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.9.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.

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.

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.9.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.

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.9.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.
• 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.9.0.

Development and release checks

Run the unit suite:

python -m unittest discover -s tests -v

Build release artifacts:

python -m pip install build
python -m build

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