Building my first MLIR-based Tensor Compiler

I'm writing this article to document my learning journey. This is my first compiler. I first discovered tensor compilers last year while digging through the PyTorch codebase. At the time, I wanted to transition into low-level programming, and since PyTorch was the only framework I was familiar with, I decided to explore their codebase. There, I discovered their compiler toolchain, TorchDynamo and TorchInductor, which, in simple terms, serve as PyTorch's machine learning compiler frontend and backend.

After investigating further, I learned that these tools primarily optimize training and inference performance across diverse hardware architectures.[1] This led me to look deeper into ML compiler workloads. My current project is the result of the learning process in the field so far.

For this project, I built a minimal, end-to-end tensor compiler designed to lower and optimize core ML operations. At a high-level, the compiler loads an ONNX graph, translates it into a custom MLIR dialect, lowers it through progressive pipelines down to LLVM IR, and executes it via Just-In-Time (JIT) compilation.

To maximize hardware efficiency, I implemented custom optimization passes, including RHS matrix transposition and a tiling pass to dramatically improve cache locality, alongside built-in MLIR transformation passes. Combined, these optimizations resulted in a 9.7x speedup on a 2048 x 2048 matrix multiplication workload over naive baselines.

Note: I focused strictly on an inference pipeline (forward pass) at this point, as it allowed me to master the lowering of operations before tackling the complexity of autograd and training.

I have broken this series into multiple parts. This first article focuses on the high-level rationale and architectural choices for the project. The subsequent articles will provide a deep dive into the technical implementation:

• End-to-End Skeleton: Focuses on the lowering pipeline

• Linear Layers and Optimizations: Focuses on the implementation of the optimization methods used and custom builders to support GeMM ops

• Performance Analysis and Testing: Focuses on analyzing the performance and correctness of the program.

Link to Github repo: Tensor Compiler repo

The end-to-end skeleton

Before adding complex operations like MatMul, I wanted to make sure that the pipeline worked and started with the simplest possible operations, Add and Constant, with the goal of moving a simple calculation all the way from an ONNX file to a JIT-executed binary

I chose ONNX as my starting point because almost any modern ML model (from PyTorch to TensorFlow) can be easily exported to this format, making the compiler framework-agnostic. To handle the translation, I used MLIR (Multi-Level Intermediate Representation)[2]. I went this route because it allows for progressive lowering, instead of trying to translate high-level Python math directly to machine code in one giant, opaque leap, I can break it down into small verifiable steps.

Think of a dialect as a specialized vocabulary for a specific level of abstraction. By using a chain of dialects, we can stop at several checkpoints along the way to perform optimizations where they make the most sense.

The full lowering pipeline looks like this:

ONNX -> dlc (custom MLIR dialect) -> tensor/linalg -> memref -> llvm (mlir) -> llvm -> JIT

ONNX -> Custom Dialect

I first lower to a custom dialect on MLIR for the following reasons:

• Semantic Preservation: ONNX operations are coarse-grained (e.g., onnx.MatMul). If I lower directly to low-level loops, I lose the high-level attributes of the operation. Keeping it in a custom dialect allows me to see the intent of the code before it gets broken down into fragments.

  
  

Custom Dialect -> Linalg on Tensors

Once the intent is preserved, I lower to Linalg on Tensors. I chose Linalg as the primary optimization layer because it provides a specialized infrastructure for structured transformations, such as the 'TilingInterface'[3], which allows me to experiment with structural changes without having to rewrite the lowering logic for every operation.

While other dialects like TOSA are excellent for hardware-agnostic specification, they lack the built-in concepts of "tiles". Conversely, dialects like Affine[4] or SCF are great for loop-level tuning but are better suited for the MemRef space, where they can operate directly on physical memory offsets.

By staying in the Value space (Tensors) rather than Memory space (MemRef) at this stage, I have the following advantages:

• Easier optimization: Because there are no pointers or memory locations yet, the compiler can track data flow perfectly using SSA logic

• Tiling and Fusion: I can tile an operation while it is still a "Value". This allows the compiler to plan how the hardware will process data before I decide where those chunks live in RAM

Tensors -> MemRef

Once the high-level schedule is determined, I perform One-Shot Bufferization. This moves the compiler from "Values" to actual memory addresses and buffers

• CPUs don't understand "tensors"; they understand memory addresses. To generate machine code, we must eventually define exactly where data lives in the heap or stack.

• In the MemRef world, optimization is difficult. The compiler has to worry about aliasing (the possibility that two different pointers hit the same spot in memory), which makes moving code around much riskier

MemRef -> LLVM IR -> JIT

Finally, I lower the IR to the LLVM dialect and then to raw LLVM IR. I chose LLVM because it acts as the final translator, taking the high-level logic and handles the details of specific CPU instructions (like x86 or ARM). Since my host machine uses an ARM-based architecture, using LLVM simplifies the process of generating machine code that runs natively within my Docker environment.

To see it work, I use a JIT (Just-In-Time) engine, which allows for the execution of the compiled code immediately. It turns all the abstract compiler theory into actual results running on the machine.

Supporting Linear Layers and Cache-Friendly Memory Access

Once the initial skeleton was done, I moved on to supporting GeMM operations. This meant adding MatMul and ReLu ops.

The initial naive run on a 2048 x 2048 matrix multiplication took 58 seconds on my M1 Pro. I wanted to see if I could make it faster by changing how the hardware accesses memory and used Valgrind to get a better picture of what was happening under the hood.

I focused on two main optimizations:

• Transpose one matrix (Spatial Locality): In a basic MatMul, the CPU grabs data in a way that misses the cache constantly. By transposing the second matrix during the lowering phase, I aligned the data with how the hardware actually reads memory

• Tiling (Temporal Locality): I broke the matrices into 64 x 64 blocks (tiling). This keeps the data fresh in the L1 cache so the CPU doesn't have to keep going back to the main RAM

Beyond these optimizations, many other optimization exist to improve the performance of matrix multiplications workloads, such as vectorization, multi-level tiling, loop unrolling, and parallelization, among others. I largely only focused on the two optimizations for this first step of my project as it's the largest bottleneck (memory latency), by ensuring data fits cleanly into the cache rather than constantly fetching from the main RAM. Furthermore, these are high-level algorithmic loop transformations that can be managed within MLIR's standard dialects, whereas low-level SIMD vectorization requires complex, target-specific hardware generation that is more suited for a next step.

Checking the result with Valgrind

To see if these changes actually mattered, I compared the execution time and cache misses:

Metric	Naive	Transposed B matrix	Transposed + Tiled
Execution Time	56s	9.7s	5.8s
L1d Misses	8.59B	273M	426M
LLd Miss	8.59B	273M	9M

The most significant improvement came from transposeB, which reduced L1 misses from 8.59B down to 273M and cut the execution time by nearly 50 seconds. This confirms that aligning memory access with the hardware's row-major layout is the highest-impact optimization for GeMM

Adding Tiling dropped the execution time by another 40%. Looking at the LLd Misses, while the L1 misses went up slightly due to more complex tiled loop, I managed to reduce nearly 264M trips to the DRAM (LLd misses). This trade-off significantly improved performance because the cost of a DRAM access is orders of magnitude higher than an L1 cache miss.

Correctness and Memory

I ran a quick check against ONNX Runtime to make sure the math was right. The difference was less than 1e-3, which is standard floating-point noise. Also, the memory usage stayed around 55MB, proving that MLIR's buffer management is working as intended, it's operating on the data in place rather than making extra, unused copies of the tensors.

Lessons Learned

Throughout this project, I relied on foundational papers, open-source codebases, and technical articles. While these references are cited individually throughout the series, I want to highlight some that most directly shaped my understanding:

• MLIR paper: This provided the conceptual foundation for progressive lowering, demonstrating how MLIR bridges the gap between high-level expressions and machine code via isolated dialects.

  • Through this, I learned that specific optimizations are constrained to distinct levels of abstraction.

  • For instance, loop-nest transformations and cache-locality optimization are ideally handled in structured dialects like Affine or SCF, whereas memory management and bufferization are deferred to the MemRef dialect, and value-based optimizations occur on SSA-form representations

• Lei Zhang's MLIR Articles: My lowering pipeline architecture was largely derived from exploring his work. After pondering alternative pipelines documented in the core MLIR dialect pages, I found this specific progression to be the most effective structure for this project's workload.

• Algorithmica (optimizing Matrix Multiplication): Among the various resources on GeMM acceleration, I primarily used Algorithmica's guide to optimizing CPU matrix multiplication.

  • This resource provided the framework for analyzing cache hierarchy, SIMD vectorization, and minimizing data movement overhead on modern CPU architectures

• MLIR Toy Tutorial and Core Documentation: The infrastructure for this project was largely inspired by the architectural and directory layout from the official MLIR Toy tutorial

• Production MLIR Codebases (onnx-mlir, torch-mlir, IREE): I looked into their codebases to understand how production-grade systems work while working on this project.

• MLIR Code codebase and pass architecture: Many built-in compiler passes initially appeared to operate as black boxes. However, implementing custom transformation passes and resolving lower-level execution bugs forced me to study the underlying implementation details

References:

[1] ML Compiler definition - MLC.ai

[2] MLIR Paper

[3] MLIR Linalg Dialect Documentation

[4] MLIR Affine Dialect Design