IREE: from PyTorch model to running executable

These are my notes from learning IREE's architecture. The diagram above is the source, pan and zoom into it directly. Below that is a longer write-up of how a model actually moves from PyTorch down to a compiled executable running on real hardware.

What is IREE?

IREE is an MLIR-based compiler and runtime system. The compiler turns machine learning models into optimized executable modules. The runtime loads those modules, feeds them input tensors, executes them on the target device, and returns the results.

Architecture overview

Three components carry the entire system: the compiler, the module, and the runtime. Everything else, flow, stream, HAL, the codegen backends, the drivers, lives inside one of those three. A model goes in one end as a framework-specific program and comes out the other end as compiled code running on real hardware. Every stage in between exists to strip away framework-specific and hardware-specific detail, then add hardware-specific detail back. That happens deliberately, and only once, at the very end.

Rendering diagram…

Model import & MLIR

Input: a model from PyTorch, ONNX, JAX, TensorFlow, or a custom plugin importer. Output: MLIR in the linalg and arith dialects.

Each framework gets its own importer, and each one converts that framework's operation set into the same common target: generic structured tensor operations (linalg) and generic arithmetic (arith). No trace of the source framework's own naming is left in it.

Flow layer

Input: the imported MLIR. Output: the same computation, split into dispatch regions, with the data dependencies between them made explicit as a graph.

A dispatch region is the chunk of work IREE will eventually hand to a device as one unit, not unlike a single GPU kernel launch or a single CPU task. Flow's job is deciding where those boundaries go: which operations fuse into one region, and which split apart. It answers exactly one question: what produces which data, and what consumes it.

Stream layer

Input: the dispatch regions and dependency graph from flow. Output: an explicitly scheduled, asynchronous program.

Where flow knows what depends on what, stream decides how execution actually happens: memory gets allocated for each region's inputs and outputs, dispatches get ordered and assigned to run concurrently wherever the dependency graph allows it, and explicit synchronization points get inserted wherever one region genuinely has to wait on another. This is also where a computation gets assigned to a specific device, if more than one is available.

HAL layer

Input: the scheduled stream program. Output: command buffers and device executables.

HAL, the hardware abstraction layer, is the one representation every backend plugs into, so a CPU target and a GPU target don't require the compiler to fork into separate paths. It produces two things together: command buffers, host-side sequences describing what to execute and in what order, and device executables, the compute itself, described uniformly regardless of which device will eventually run it. Buffers, devices, and synchronization are all modeled here in backend-neutral terms, hal.buffer, hal.device, hal.executable, hal.fence, and stay neutral right up until code generation gives them a real target.

This isn't a strictly one-directional pipeline. Target configuration and profiling feedback loop back into flow, stream, and HAL. How work gets partitioned, scheduled, and buffered can be shaped by what a specific target actually supports and by how earlier runs performed, not only by the input program considered in isolation.

Code generation

Input: command buffers and device executables from HAL. Output: compiled code, generated along two separate, parallel paths.

Host code generation compiles the command-buffer side into code that runs on the host and drives execution, through LLVM, direct C emission, or IREE's own VM bytecode format, targeting a CPU or WebAssembly. This path barely changes based on which device the compute ultimately runs on; its job is orchestration, not computation.

Device code generation compiles the executable side into the actual compute, and this path is entirely target-specific: LLVM for native CPU instruction sets (ARM, RISC-V, x86) and for GPU instruction sets reached through LLVM backends (PTX for CUDA, ROCDL for AMD/HIP), SPIR-V for Vulkan, or VMVX, IREE's own portable, software-only execution format for targets without native code generation support at all.

Modules & runtime

Input: the generated host code and device code. Output: a single compiled artifact, in one of four formats: VM bytecode, C source, a static library, or a shared library.

VM bytecode, packaged as a .vmfb file, is the common case. The others exist for deployment situations where dynamic module loading isn't available, such as small embedded targets. Whichever format is used, the module is tied to the backend it was generated for. The same source model compiled for CPU and for CUDA produces two different, non-interchangeable files.

At runtime, the module is loaded by two internal layers working together. The VM interprets its bytecode: resolving and dispatching function calls, managing loaded modules, exposing reflection information about them. The HAL, now doing real work instead of just providing a compile-time abstraction, hands actual compute off to whichever driver is active for the current backend. Plugins sit at both the top level of the runtime and inside the HAL specifically, as an extension point for adding capability without touching the runtime's core. The whole thing is deliberately small: IREE's own documentation puts a typical build at roughly 25 to 150 KB, the direct payoff of how much work the compiler already finished ahead of time.

Runtime backends & deployment

Input: a loaded module and a chosen backend. Output: execution on real target hardware.

A driver exists per backend (CPU, CUDA, Vulkan, Metal, HIP, WebGPU), and each one implements the same HAL interface against genuinely different hardware underneath: a CPU driver manages a thread pool and runs natively compiled code directly; a CUDA driver manages GPU device memory and launches compiled PTX kernels through NVIDIA's driver API; Vulkan and Metal drivers do the equivalent through their own graphics/compute APIs. Swapping which backend a program targets is, in principle, a matter of swapping which driver gets asked for. Nothing about the VM or the module's host code needs to change.

Two things sit beside the runtime rather than inside it. Tools (iree-run-module, iree-check-module, iree-benchmark-module) are command-line utilities for running, testing, and benchmarking a module without writing any host code at all. Bindings (C, Python, Rust, TFLite, or a custom binding) are the language-level interface an application actually links against.

A worked example: one model, two backends

Take the architecture above and walk it once as a single example: a PyTorch model compiled for two backends at once, so the host/bytecode side and a real GPU device side can both be followed to the end.

Up to HAL, there's only one path. The model gets imported, its signature rewritten to hal.buffer_view, split into dispatch regions by flow, scheduled by stream, and turned into command buffers and device executables by HAL, identically, regardless of what hardware it will eventually run on. Nothing backend-specific has happened yet.

At code generation, the path splits into two, and stays split all the way to the end:

Rendering diagram…

The bytecode path doesn't change no matter what device the compute runs on: host code generation takes the command-buffer side and compiles it into VM bytecode, the instructions that decide which dispatch runs when and marshal data in and out. Its whole job is orchestration.

The device path, concretely for an NVIDIA GPU target, takes the executable side and, because the target backend was set to CUDA, routes it through LLVM's PTX backend, producing real compiled GPU kernel code, not a hardware-agnostic description anymore. This PTX code and the host bytecode from the other path both get packaged into the same IREE module.

At runtime, the two paths reunite. The VM loads and interprets the host bytecode, and when it reaches a dispatch instruction, it doesn't execute anything itself. It hands that instruction to the HAL, which for a CUDA target means the CUDA driver: allocating GPU memory, copying inputs across, launching the embedded PTX kernel through NVIDIA's driver API, and copying the result back once the GPU signals it's done. Swap the target backend to CPU instead, and the bytecode path is untouched. Only the device path and which driver the HAL hands off to would change.

Reference Material : https://iree.dev/