Search arXiv⌕ Search
Home/Archive/arXiv

arXiv · 2404.15204

Towards a high-performance AI compiler with upstream MLIR

Renato Golin·Lorenzo Chelini·Adam Siemieniuk·Kavitha Madhu·Niranjan Hasabnis·Hans Pabst·Evangelos Georganas·Alexander Heinecke
View original source ↗Read source document ↗Export metadata ↓

Abstract

This work proposes a compilation flow using open-source compiler passes to build a framework to achieve ninja performance from a generic linear algebra high-level abstraction. We demonstrate this flow with a proof-of-concept MLIR project that uses input IR in Linalg-on-Tensor from TensorFlow and PyTorch, performs cache-level optimizations and lowering to micro-kernels for efficient vectorization, achieving over 90% of the performance of ninja-written equivalent programs. The contributions of this work include: (1) Packing primitives on the tensor dialect and passes for cache-aware distribution of tensors (single and multi-core) and type-aware instructions (VNNI, BFDOT, BFMMLA), including propagation of shapes across the entire function; (2) A linear algebra pipeline, including tile, fuse and bufferization strategies to get model-level IR into hardware friendly tile calls; (3) A mechanism for micro-kernel lowering to an open source library that supports various CPUs.

Explore related subjects

cs.PLcs.AIcs.ARcs.DCcs.LG

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Renato Golin, Lorenzo Chelini, Adam Siemieniuk, Kavitha Madhu, Niranjan Hasabnis, Hans Pabst, Evangelos Georganas, Alexander Heinecke. 2024-04-15. Towards a high-performance AI compiler with upstream MLIR. https://arxiv.org/abs/2404.15204

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

On the computational complexity of JavaScript regex matching

Despite widespread use, the complexity of the matching problem for modern regular expressions languages remains unclear. Previous work proved that an idealized regular expression language with backreferences and lookarounds had PSPACE-complete matching. We extend this work to a real-world regex language by proving that JavaScript regex matching with expanded lower-bounded quantifiers is PSPACE-complete. We then generalize the result: we show that PSPACE-hardness survives the removal of negative lookarounds, and that removing all lookarounds leads to an OptP-complete parsing problem. Our core arguments are formalized in Rocq.

cs.PL↗
arXiv2026

Opportunistic ZGC: Leveraging Idle Cores for More Effective Concurrent Garbage Collection

Managed language runtimes often provide concurrent garbage collectors so that latency-critical applications with large working sets can keep running while most collection work proceeds in the background. ZGC is a production-quality, generational, concurrent collector in OpenJDK with sub-millisecond pause times. While ZGC is designed to run concurrently, frequent and excessive collections with ZGC can still slow the mutators due to synchronization costs and interference with shared computing resources. Hence, the ZGC scheduler is conservative by default, and in most cases, will grow the heap toward the maximum allowed before scheduling a collection. While this approach minimizes collection effort, it can be wasteful, or even harmful, if the maximum heap size is not well tuned to the actual working set. We propose Opportunistic ZGC (OppZGC), a feedback-directed ZGC scheduling policy that constrains the heap dynamically and automatically, without per-application tuning. OppZGC identifies periods when CPU cores are underutilized and leverages them for concurrent collection with ZGC. We describe the design and implementation of OppZGC in OpenJDK's HotSpot Java VM and evaluate it with standard and latency-sensitive benchmarks from DaCapo Chopin and SPECjbb. OppZGC limits heap usage when there is CPU capacity sufficient for additional collections, and avoids scheduling extra collections when they would substantially degrade performance. Overall, it reduces maximum heap usage for our DaCapo benchmarks by between 61% and 90%, on average, depending on configuration, with minimal impact on throughput and request latency compared to default ZGC.

cs.PL↗
arXiv2026

From Rocq to Metal: A Pipeline for Formally Verified Microcontroller Firmware

Enforcing invariants in safety-critical firmware is increasingly urgent as generated code becomes widespread, but standard extraction targets for proof assistants require runtimes too large for many embedded devices. We present a pipeline for running formally verified Rocq firmware logic on Cortex-M microcontrollers. The pipeline extracts Gallina to Scheme, compiles it with Encore!, a bare-metal Continuation Passing Style (CPS) bytecode virtual machine, and embeds the result in no_std Rust firmware. We structure applications as pure state-transition functions, so the business logic is proved in Rocq while the event/effect boundary, host callbacks, compiler, and VM remain explicit trusted components. On ST33-class targets with a 50 KB RAM lower bound, Encore! executes Rocq-extracted code end-to-end, stays within the target memory budget on our benchmarks, and validates a transaction-signing application on physical Ledger Flex hardware.

cs.PL↗