AXI4MLIR: User-Driven Automatic Host Code Generation for Custom AXI-Based Accelerators

TL;DR

AXI4MLIR is a compiler framework extension that automates host code generation for custom AXI accelerators, improving performance and reducing cache references in linear algebra workloads.

Contribution

It introduces new MLIR attributes and transformations enabling user-driven, automated host-accelerator driver code generation for AXI-based accelerators.

Findings

Up to 56% reduction in CPU cache references

Up to 1.65x speedup over manual driver code

Versatile across different accelerator types

Abstract

This paper addresses the need for automatic and efficient generation of host driver code for arbitrary custom AXI-based accelerators targeting linear algebra algorithms, an important workload in various applications, including machine learning and scientific computing. While existing tools have focused on automating accelerator prototyping, little attention has been paid to the host-accelerator interaction. This paper introduces AXI4MLIR, an extension of the MLIR compiler framework designed to facilitate the automated generation of host-accelerator driver code. With new MLIR attributes and transformations, AXI4MLIR empowers users to specify accelerator features (including their instructions) and communication patterns and exploit the host memory hierarchy. We demonstrate AXI4MLIR's versatility across different types of accelerators and problems, showcasing significant CPU cache reference reductions (up to 56%) and up to a 1.65x speedup compared to manually optimized driver code implementations. AXI4MLIR implementation is open-source and available at: https://github.com/AXI4MLIR/axi4mlir.