AXI4MLIR: Host-Accelerator Runtime Optimization
- AXI4MLIR is an extension of the MLIR framework that generates efficient host-driver code for custom AXI-based accelerators, focusing on linear algebra workloads.
- It employs cache-aware tiling, DMA-based optimizations, and zero-copy transfers to streamline communication between the host and accelerators.
- Benchmarks show up to 1.65× speedup in matrix multiplication, 56% CPU cache reduction, and 3.4× end-to-end improvement for applications like TinyBERT.
Searching arXiv for the cited AXI4MLIR and related MLIR-FPGA papers to ground the article.
AXI4MLIR is an extension of the Multi-Level Intermediate Representation (MLIR) compiler framework for automatic host-driver code generation for custom AXI-based accelerators, with a particular emphasis on linear algebra workloads such as matrix multiplication and convolution. Its central objective is not accelerator synthesis in isolation, but the generation of host-side communication and runtime logic that is aware of accelerator instructions, communication patterns, CPU cache hierarchy, dataflow, and tiling constraints. In the published line of work, AXI4MLIR evolves from a user-driven framework built around accelerator descriptors and an intermediate accel dialect (Agostini et al., 2023), to a system with explicit host-side data-transfer optimizations such as DMA-based allocation, coalescing of DMA transfers, and software pipelining (Haris et al., 2024), and then to zero-copy, descriptor-driven data movement that removes redundant heap-to-DMA staging copies (Cohavi et al., 9 Jun 2026).
1. Problem domain and design objective
AXI4MLIR was introduced to address a gap left by many accelerator-generation tools: they focus on designing accelerators, but devote comparatively little attention to the host-side driver code required to use those accelerators efficiently (Agostini et al., 2023). For custom AXI-based accelerators, the host must allocate and stage memory correctly, pack data into accelerator-specific layouts, send control opcodes and instructions, manage DMA transfers, synchronize send and receive phases, and exploit the CPU memory hierarchy to reduce cache misses and branch overhead. The framework treats these tasks as a compiler problem rather than as handwritten runtime engineering.
The motivating workloads are linear algebra algorithms used in machine learning and scientific computing. In this setting, communication can dominate runtime; different accelerators may require distinct dataflows such as nothing stationary, A-stationary, B-stationary, or C-stationary; and CPU cache-aware tiling and packing can materially affect end-to-end performance (Agostini et al., 2023). AXI4MLIR therefore targets the host-accelerator interaction layer, where data movement and loop structure determine whether a custom accelerator is actually beneficial.
The framework is explicitly user-driven. The user specifies accelerator capabilities, supported opcodes, supported opcode flows, tile sizes, data types, and tensor dimensions, while the compiler automates the generation of host code and DMA interactions. This division of labor is significant: AXI4MLIR does not infer arbitrary accelerator internals, but it does convert a structured accelerator description into executable host code and communication logic (Agostini et al., 2023).
2. Compilation flow and intermediate representation structure
The core AXI4MLIR workflow begins with a high-level MLIR program, typically expressed in the linalg dialect, and progressively lowers it into host code that drives an AXI DMA runtime (Agostini et al., 2023). For matrix multiplication, the paper illustrates a sequence from linalg.matmul to linalg.generic, then to tiled scf.for loops, and finally to accelerator runtime calls. The computational pattern is standard:
When the accelerator supports a tile such as MatMul_{4x4x4}, the loop body of the tiled software version is replaced by accelerator communication.
The end-to-end flow described in the original work is structured as follows. A user writes a high-level MLIR program, provides a configuration file describing CPU cache sizes and types, accelerator capabilities, supported opcodes, supported opcode flows, accelerator dimensions, data types, and tensor dimensions, and then the compiler parses that configuration, identifies compatible linalg.generic operations, applies tiling transformations using cache sizes and accelerator dimensions, introduces accel dialect operations, lowers those operations to DMA runtime calls, and emits final host code (Agostini et al., 2023).
Later work makes the lowering sequence more explicit. The 2026 zero-copy study describes the flow as: starting from high-level linalg ops such as linalg.matmul or linalg.conv_2d, converting named ops into linalg.generic based on traits such as indexing maps and iterator types, bufferizing tensors into memref form, annotating operations with custom accelerator-targeting attributes, lowering into the custom accel dialect plus supporting MLIR dialects, applying tiling transformations using memref and scf, and then lowering accel transfers into runtime library calls that control the AXI DMA engine (Cohavi et al., 9 Jun 2026). In that formulation, AXI4MLIR sits between high-level tensor and linear-algebra representation and the low-level runtime responsible for host-accelerator communication.
The principal architectural point is that AXI4MLIR introduces an intermediate representation for communication. Rather than lowering immediately to LLVM or C, it uses a custom accel dialect to represent accelerator initialization, sending of literals and operands, reception of results, and synchronization. The papers argue that this makes transformations easier, dependencies clearer, and communication placement in loop nests more explicit (Agostini et al., 2023).
3. Accelerator description model: attributes, opcode semantics, and the accel dialect
AXI4MLIR extends linalg.generic with a set of accelerator-oriented attributes and adds a new accel dialect that captures host-accelerator communication at a higher abstraction level than raw low-level code (Agostini et al., 2023). The added attributes are central to the framework’s configurability.
dma_init_config defines parameters needed to initialize the DMA engine for a specific accelerator, with example fields including id, inputAddress, inputBufferSize, outputAddress, and outputBufferSize. init_opcodes defines a flow of opcodes used to initialize or reset the accelerator before execution. accel_dim specifies the accelerator’s supported tile size or computation size; for example, map<(m, n, k) -> (4, 4, 4)>. permutation_map defines loop ordering and thus affects the implied dataflow. opcode_map encodes accelerator opcodes as sequences of communication actions. opcode_flow defines valid opcode sequences, including nested grouping used as a proxy for communication placement across loop nests (Agostini et al., 2023).
The syntax presented for opcode_map and opcode_flow is:
1 2 3 4 5 6 7 8 9 10 11 |
opcode_dict ::= `opcode_map` `<` opcode_entry (`,` opcode_entry)* `>`
opcode_entry ::= (bare_id | string_literal) `=` opcode_list
opcode_list ::= `[` opcode_expr (`,` opcode_expr)* `]`
opcode_expr ::= `send` `(` bare_id `)`
| `send_literal` `(` integer_literal `)`
| `send_dim` `(` bare_id `)`
| `send_idx` `(` bare_id `)`
| `recv` `(` bare_id `)`
opcode_flow_entry ::= `opcode_flow` `<` flow_expr >
flow_expr ::= `(` flow_expr `)` | bare_id (` ` bare_id)* |
The communication actions have fixed meanings. send(arg) sends data associated with an argument of the linalg.generic operation, send_literal(value) sends a fixed control literal, send_dim(arg) sends a dimension value, send_idx(arg) sends an index value, and recv(arg) receives a result associated with an argument (Agostini et al., 2023). In the matrix multiplication examples, opcodes such as sA, sB, cC, and rC correspond to specific sequences of literal and data-transfer actions.
These attributes allow the compiler to express dataflows directly. For matrix multiplication, an example trait set includes accel_dim = map<(m, n, k) -> (4, 4, 4)>, permutation_map = affine_map<(m, n, k) -> (m, k, n)>, and opcode_flow = opcode_flow < (sA (sBcCrC)) >, which encodes an A-stationary flow in which A is sent once at outer scope while B and C communication is repeated in an inner scope (Agostini et al., 2023). The paper also lists alternatives such as ((sA sB cC) rC) for C-stationary behavior and (sB sA cC rC) for nothing stationary or a different ordering. A convolution example similarly uses accel_dim, opcode_map, and opcode_flow <(sF (sIcO) rO)> to represent a filter plus output stationary flow.
The accel dialect then materializes these communication intents as explicit operations such as accel.dma_init(...), accel.sendLiteral(...), accel.send(...), accel.sendDim(...), and accel.recv(...) (Agostini et al., 2023). These operations are the substrate on which later host-side optimizations are applied.
4. Host memory hierarchy, DMA runtime, and data-transfer optimization
A major contribution of AXI4MLIR is that it does not generate merely correct communication; it generates communication that attempts to exploit the host memory hierarchy. The compiler uses CPU cache sizes from the configuration file, accelerator tile size, and problem dimensions to tile loops so that data is reused in cache and staged efficiently into DMA buffers (Agostini et al., 2023). At the runtime level, AXI4MLIR includes a lightweight AXI DMA runtime library that initializes DMA engines, maps memory for input and output staging buffers, copies from memref to DMA-accessible buffers, triggers DMA send, waits for completion, and copies received data back into memref buffers. The paper reports a size of 55 bytes for the target ARM SoC (Agostini et al., 2023).
The framework also specializes memref handling. Because MLIR memrefs carry an allocated pointer, aligned pointer, offset, sizes, and strides, generic MemRef-to-DMA copies can preserve arbitrary layouts. When strides[N-1] == 1, however, AXI4MLIR specializes the copy path and uses std::memcpy(src, dst, size) instead of recursive element-by-element copying, reducing branch instructions, cache references, and control-flow overhead (Agostini et al., 2023).
The 2024 extension, presented in "Data Transfer Optimizations for Host-CPU and Accelerators in AXI4MLIR" (Haris et al., 2024), studies a flexible stream-based MatMul accelerator with double buffering and identifies two bottlenecks: compute-core utilization below 10% on average, and heap-to-memory-mapped DMA-buffer copy latency as the critical host-side cost. It proposes three host-code data-movement optimizations. First, DMA-based data allocation adds a new attribute or tag, shown as #dma, so accelerator-bound memrefs are allocated directly in DMA-addressable memory. Second, coalescing of DMA transfers merges multiple accel.send operations into one variadic send such as accel.send([%op0,%sA,%sB,%op1]), reducing fixed DMA launch overhead and synchronization cost. Third, software pipelining rewrites the host loop structure to overlap the accelerator’s load, compute, and store stages and to exploit double buffering (Haris et al., 2024).
The 2026 work, "Defeat the Heap: Zero-Copy Data Movement in AXI4MLIR" (Cohavi et al., 9 Jun 2026), pushes this line further by removing the staging copy entirely. It reworks accel.alloc so that allocation is modeled using a memory-space-type attribute and an accompanying layout operand, and lowers allocation to aximlir_dma_alloc(), which returns memref descriptors allocated in the DMA-visible region. The conceptual descriptor shown is:
1 2 3 4 5 6 7 |
template<typename T, size_t N> struct MemRefDescriptor { T *allocated; T *aligned; intptr_t offset; intptr_t sizes[N]; intptr_t strides[N]; }; |
In the baseline flow, accel.send lowers through explicit copies and calls such as @copy_to_dma_region, @dma_start_send, and @dma_wait_send_completion. In the optimized flow, accel.send lowers directly from a memref or memref.subview descriptor to dma_send_descriptor() with no intermediate heap-to-DMA copy (Cohavi et al., 9 Jun 2026). Because memref.subview retains the original allocated and aligned pointers while updating offset, sizes, and strides, the runtime can compute the DMA transfer range directly from descriptor metadata. For non-contiguous tiles, the paper uses scatter-gather DMA, where multiple disjoint source regions are gathered into a contiguous on-device stream and disjoint outputs are scattered back.
5. Evaluation results and measured behavior
The original evaluation of AXI4MLIR uses a PYNQ-Z2 board with a Zynq-7000 SoC, a dual-core ARM Cortex-A9 CPU at 650 MHz, and a library of tile-based accelerators derived from SECDA-TFLite, synthesized with AMD/Xilinx Vitis at 200 MHz (Agostini et al., 2023). The accelerators are AXI-Stream connected and micro-ISA or opcode based. The framework is evaluated on matrix multiplication, all convolution layers of ResNet18, and an end-to-end TinyBERT workload compiled via Torch-MLIR.
For matrix multiplication, the paper reports that accelerator offload becomes beneficial when the problem size is at least dims >= 64, where dims = M = N = K, and when the accelerator tile size is at least accel_size >= 8 (Agostini et al., 2023). Compared with manually optimized driver code, AXI4MLIR-generated drivers achieve up to 1.65× speedup, an average speedup of 1.18×, up to 56% reduction in CPU cache references, and an average cache-reference reduction of 10%. For ResNet18 convolution layers, AXI4MLIR is faster on 10 out of 11 layers, with 1.28× average speedup and 1.54× max speedup; one layer slows down by about 10% because filter height is 1, input channels are small, and the specialized contiguous-copy optimization cannot help sufficiently. For TinyBERT, the reported gains are 3.4× end-to-end speedup, 18.4× speedup in accelerated MatMul layers, and the observation that MatMul layers constituted about 75% of the original CPU runtime (Agostini et al., 2023).
A concise summary of reported results is:
| Evaluation context | Reported result | Source |
|---|---|---|
| MatMul drivers vs manual code | up to 1.65× speedup | (Agostini et al., 2023) |
| CPU cache references | up to 56% reduction | (Agostini et al., 2023) |
| ResNet18 convolution layers | 1.28× average, 1.54× max | (Agostini et al., 2023) |
| TinyBERT end to end | 3.4× speedup | (Agostini et al., 2023) |
| Zero-copy main-memory movement | up to 2× reduction | (Cohavi et al., 9 Jun 2026) |
| Zero-copy end-to-end performance | 1.7× average improvement | (Cohavi et al., 9 Jun 2026) |
The 2024 MatMul case study adds cycle-level characterization. It reports average compute-core utilization below 10%, notes that Load A takes more cycles than load B, attributes this to higher startup latency on the first load, and mentions that later transfers can proceed with as little as a 1-cycle delay once the stream is active (Haris et al., 2024). The conclusion is that the data path, not the MAC array, is the principal throughput limiter in the studied configuration.
The 2026 zero-copy evaluation uses a configurable tiled matrix-matrix multiplication accelerator on a simulated Versal VP1902 SoC with 273 GB/s peak aggregate memory bandwidth over the NoC, AXI4-based hardened NoC implementation, 150 ns per-call DMA latency overhead, and an accelerator engine array capable of 4 TFLOPS (Cohavi et al., 9 Jun 2026). It considers non-stationary, A-stationary, and C-stationary dataflows, with and . The reported outcomes are a reduction in main-memory data movement by up to 2×, an average 1.7× improvement in end-to-end performance across all evaluated accelerator sizes and data flows, and 1.7× on average reductions in both load and store operation latency (Cohavi et al., 9 Jun 2026). The paper also states that in some configurations with small tiles, transfer cost exceeds compute time by more than 18×, underscoring that staging-copy overhead can dominate execution.
6. Scope, limitations, and relation to adjacent MLIR-based FPGA work
AXI4MLIR has a well-defined scope. It supports SoCs with ARM CPUs, accelerators connected via AXI-Stream, and systems in which DMA is used for communication (Agostini et al., 2023). It does not currently address host-accelerator cache coherence protocols, arbitrary interconnects without adapting the runtime, full scheduling optimization as a separate compiler component, or arbitrary accelerator internals. Most importantly, it does not model or optimize the internal accelerator architecture; the accelerator developer must provide supported operations, tile sizes, opcodes, flows, and dimensions (Agostini et al., 2023). In this sense, AXI4MLIR is a host-code generation and communication optimization framework, not a complete accelerator synthesis solution.
This boundary is useful when comparing AXI4MLIR to adjacent MLIR-based FPGA research. Olympus, introduced for platform-aware FPGA system architecture generation, targets the system architecture surrounding accelerators rather than host-driver code. It represents accelerator-system dataflow graphs using an Olympus dialect with olympus.kernel nodes and channel constructs such as olympus.make_channel, models channel semantics as stream, small, or complex, and performs platform-aware transformations including channel reassignment, replication, bus widening, Iris-based bus optimization, and PLM optimization (Soldavini et al., 2023). For Xilinx devices, kernels connected to complex channels have an AXI port that connects straight to global memory, and the backend uses a Vivado block diagram and a Vitis HLS bridge module. This suggests a complementary but distinct layer: Olympus focuses on generated hardware structure and memory-system fit, whereas AXI4MLIR focuses on host-side communication and driver generation.
By contrast, POM, "An Optimizing Framework on MLIR for Efficient FPGA-based Accelerator Generation" (Zhang et al., 2024), is explicitly not AXI4MLIR. It is an end-to-end FPGA HLS and compiler framework with a three-layer IR stack, polyhedral semantics via isl, a declarative DSL, and automatic design space exploration. The paper states that it does not mention AXI4 interfaces, does not describe AXI4 memory-mapped communication, and does not discuss host-device interaction over AXI4 (Zhang et al., 2024). The distinction clarifies a common misconception: not every MLIR-based FPGA compiler is addressing the host-driver and DMA problem that AXI4MLIR takes as its primary target.
Taken together, the AXI4MLIR papers define a coherent research trajectory. The framework begins with user-driven accelerator descriptors, cache-aware tiling, an accel dialect, and DMA runtime lowering (Agostini et al., 2023); it then identifies concrete host-side utilization bottlenecks and addresses them through DMA-aware allocation, transfer coalescing, and software pipelining (Haris et al., 2024); and finally it extends the dialect and lowering to support zero-copy, descriptor-driven DMA transfers and scatter-gather handling of non-contiguous subviews (Cohavi et al., 9 Jun 2026). The consistent theme is that efficient use of custom AXI-based accelerators depends not only on arithmetic throughput, but also on compiler-managed communication structure, host memory placement, and DMA semantics.