Defeat the Heap: Zero-Copy Data Movement in AXI4MLIR

Abstract: As custom hardware accelerators become increasingly central to machine learning workloads, efficient data transfer is critical for maximizing accelerator performance on linear algebra kernels. AXI4MLIR, an extension of the Multi-Level Intermediate Representation (MLIR) compiler framework for automated generation of host-accelerator driver code, incurs significant runtime overhead due to non-zero-copy CPU-accelerator data movement. During transfers from the host to the accelerator, data is copied from heap-allocated memory buffers into contiguous Direct Memory Access (DMA)-mapped buffers. This work identifies this copy as a redundant staging operation and eliminates it through zero-copy data movement. The optimization extends accel, an MLIR dialect introduced by AXI4MLIR, and implements lowering support that allocates buffers directly within DMA-mapped memory, thereby omitting the staging copy. We evaluate the proposed scheme using a configurable matrix-matrix multiplication accelerator and show that the zero-copy optimization reduces main memory data movement by up to 2x, increasing overall accelerator utilization.

• The paper identifies host-accelerator data movement as a major performance bottleneck and eliminates costly heap staging with a zero-copy technique.
• The authors integrate DMA-aware buffer allocation into the MLIR lowering process using extended memref descriptors and scatter-gather transfers to optimize tiling.
• Experimental evaluation on a Versal VP1902 SoC demonstrates up to 2x latency reduction and 1.7x average runtime improvement, validating the approach.

Zero-Copy Data Movement in AXI4MLIR: Eliminating Host-Accelerator Staging Overhead

Introduction

Custom hardware accelerators are pivotal for high-throughput linear algebra operations in ML workloads, yet their effectiveness is frequently throttled by inefficiencies in host-accelerator data movement. AXI4MLIR is an MLIR-based framework designed for automatic host-driver generation targeting AXI protocol-compliant accelerators, but its baseline implementation incurs significant overhead due to staged, heap-to-DMA buffer copies prior to data offload. This paper identifies this buffer-staging as a major performance bottleneck and proposes a zero-copy data movement technique that directly allocates MLIR buffers in the DMA-visible address space, eliminating intermediate heap copies and enhancing accelerator utilization.

Figure 1: Zero-copy DMA-aware optimization: System view.

AXI4MLIR Lowering and Data Movement Pipeline

The AXI4MLIR workflow transforms high-level MLIR linalg abstractions, which express algorithmic intent, through a multi-stage lowering process into representations suitable for custom accelerators. Tensors are bufferized into memref types and assigned custom accelerator attributes, then further lowered into operations in the accel dialect, with host-device data movement mediated via runtime library invocations. In the original pipeline, memrefs reside in standard heap-allocated memory, and prior to accelerator invocation, are copied to contiguous DMA-mapped buffers, leading to pronounced latency, particularly with fine-grained tiling strategies.

Zero-Copy Optimization: Technique and Pipeline Integration

The proposed optimization introduces custom attributes and lowering support in the accel dialect to allocate buffers in device-accessible DMA memory, obviating the need for explicit heap-to-DMA region staging. The modified lowering path for accel.send operations replaces the baseline two-step (heap staging, then DMA copy) with direct allocation via a runtime allocator (aximlir_dma_alloc), and all transfer metadata is encapsulated within extended memref descriptors. These descriptors fully encode offset, shape, strides, and physical memory addresses, enabling DMA routines to directly compute the requisite source and destination addresses for scatter-gather or contiguous transfers without manual memory management.

The approach exploits the MLIR memref.subview semantics, ensuring that even tiled or non-contiguous subregions are correctly described for the DMA engine. At runtime, filling and transferring a tile reduces to using the descriptor as input to a single DMA scatter-gather call, with completion events managed by blocking or as required by algorithmic dependencies.

Scatter-Gather Transfers and High-Rank Memory Layouts

Non-contiguous memory patterns, as encountered in higher-dimensional tiled operations, present additional complexity. To address this, the system leverages scatter-gather DMA capabilities, constructing a transfer list from memref descriptor metadata that enables logical tiles—regardless of internal contiguity—to be transferred in a single DMA operation. This strategy minimizes the DMA call overhead, eliminates the scaling cost of line-by-line transfers, and coalesces memory-intensive tile access patterns into bandwidth-efficient transactions.

Experimental Evaluation

The efficacy of the zero-copy strategy is modeled and evaluated on a simulated Versal VP1902 SoC. The experimentation targets tiled MatMul accelerators across varying matrix and tile sizes under multiple dataflow paradigms (Non-stationary, A-stationary, C-stationary). For each regime, the breakdown of execution time is analyzed with respect to heap-to-DMA movement, DMA-to-accelerator transfer, accelerator compute, and post-compute DMA/heap operations.

Figure 3: Normalized execution time breakdown across stationary data flows for various dimensions, tile sizes.

Results clearly indicate that, in the baseline ("naive") configuration, data movement dominates runtime, especially as tile dimensions decrease, with communication overhead overwhelming compute by an order of magnitude in some settings. The zero-copy and scatter-gather optimizations eliminate the heap-to-DMA staging cost, substantially increasing the compute fraction and improving end-to-end latency by up to $2\times$, with a mean improvement of $1.7\times$ across benchmarks. These gains are more pronounced with larger tiles and stationary dataflows that maximize temporal reuse, directly demonstrating the value of compiler and runtime co-design for memory placement.

Related Work and Context

Prior work on MLIR-based accelerator flows, notably mlir-aie [mliraie], focuses on architectures with pre-defined data movement dialects (e.g., aiex for AI Engines) but lacks generality for custom-AXI accelerators. EQueue (Li et al., 2022) introduces domain modeling for asynchronous control and data flows, enabling task queuing and synchronization in broader heterogeneous scenarios. The proposed zero-copy technique complements these by targeting non-AIE, generic accelerator pipelines and directly integrating memory allocation decisions into the MLIR lowering process—a necessary evolution for scalable, automated codesign in next-generation hardware-accelerated ML pipelines.

Implications and Future Directions

Eliminating heap-to-DMA staging in AXI4MLIR establishes a new efficiency baseline for host-accelerator communication in MLIR-centric accelerator development. The enhanced buffer placement and transfer strategies not only yield quantifiable improvements in accelerator utilization but also set a precedent for zero-copy semantics as a compile-time concern, not merely as a runtime optimization. The formalization of memory space-aware allocation in MLIR dialects unlocks further research directions:

• Automated selection of optimal data layouts based on target accelerator topology.
• Deeper integration of asynchronous DMA scheduling in the MLIR compilation flow.
• Generalization to distributed, multi-accelerator environments where direct, zero-copy Remote Direct Memory Access (RDMA) enables cross-device communication efficiency.

Conclusion

By identifying and excising the redundant host-side memory staging step in the AXI4MLIR toolchain, this paper delivers a pragmatic zero-copy optimization that substantially improves data movement efficiency for custom accelerator kernels. Achieving up to $2\times$ reduction in data transfer overhead and a $1.7\times$ average improvement in overall runtime establishes direct, DMA-region allocation as a critical design pattern for MLIR-based accelerator software. The work signals the increasing necessity of fine-grained, compiler-directed memory placement in the orchestration of future heterogeneous compute systems, and constitutes a migration path for other MLIR-based accelerator frameworks to adopt transparent zero-copy data movement.