torch-mlir: bridging PyTorch to MLIR for compiler engineers building hardware backends

torch-mlir replaces half a dozen redundant per-vendor PyTorch-to-MLIR lowering implementations

Before torch-mlir, hardware vendors who wanted to consume PyTorch models in MLIR-based compilers each wrote their own PyTorch frontend. The README documents this directly: "We have observed half a dozen custom lowerings from PyTorch to MLIR." torch-mlir replaces that fragmentation with a single shared lowering path.

Two known compiler projects use torch-mlir downstream: IREE (from iree-org) and Blade (from alibaba/BladeDISC). Both are listed as examples to refer to when starting a new integration.

An LLVM Incubator notice in the README is clear about project status: torch-mlir is not yet endorsed as an LLVM component and is not part of any official LLVM distribution. That status means stability guarantees and official support are not in scope.

ONNX and Fx are the 2 entry points that feed the Torch MLIR Dialect

Lowering a PyTorch model to MLIR through torch-mlir begins at 1 of 2 entry points. ONNX accepts models already exported to ONNX format. Fx uses PyTorch's Fx tracing to capture the computation graph and feeds it directly, skipping any ONNX conversion step.

Both paths converge on the Torch MLIR Dialect as a shared intermediate representation. From there, downstream compilers take the dialect and lower it further toward their target hardware. torch-mlir itself does not produce runnable code; that step belongs to the downstream.

For most engineers who want to run inference, that architecture places torch-mlir below the inference runtime layer. The README states this plainly: "Torch-MLIR is primarily a project that is integrated into compilers to bridge them to PyTorch and ONNX." Direct use without a compiler integration is possible but limited to examples like the FxImporter ResNet18 script.

Installing pre-built snapshot packages requires Python 3.11 or 3.10 and PyTorch nightly wheels

Pre-built snapshot packages are available for Python 3.11 and Python 3.10. Set up a virtual environment first.

python3.11 -m venv mlir_venv
source mlir_venv/bin/activate

Or using conda with Python 3.11.

conda create -n torch-mlir python=3.11
conda activate torch-mlir
python -m pip install --upgrade pip

With the environment active, install torch-mlir alongside PyTorch nightly wheels.

pip install --pre torch-mlir torchvision \
  --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
  -f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels

This install couples 2 separate nightly sources: torch-mlir's own snapshot packages and PyTorch's nightly server. Both can change between install attempts, and incompatibility between the 2 builds can cause the install to fail silently or produce an unusable combination.

FxImporter ResNet18 shows near-identical but not bit-identical output between PyTorch and torch-mlir

For an end user who wants to test torch-mlir without a full downstream integration, the repository provides an FxImporter example with ResNet18.

wget https://raw.githubusercontent.com/llvm/torch-mlir/main/projects/pt1/examples/fximporter_resnet18.py
python projects/pt1/examples/fximporter_resnet18.py

Running this script produces side-by-side predictions from PyTorch and torch-mlir. The README shows these specific output values for the top classification:

PyTorch: Labrador retriever 70.65674591064453 torch-mlir: Labrador retriever 70.6567153930664

The top-1 classification matches in both cases. But the confidence score differs: 70.65674591064453 from PyTorch versus 70.6567153930664 from torch-mlir. This is a small but nonzero difference that appears in the documented example output. Anyone building a pipeline where the downstream comparison depends on identical confidence values should confirm this behavior is acceptable before committing to the integration.

Source builds: cmake, ninja, nanobind 2.9 to 3.0 and a Windows-specific ninja cap

Building from source requires a C++ toolchain, CMake and Ninja. For wheel generation the project uses cibuildwheel. pyproject.toml lists the Python-side build dependencies with tight upper bounds: nanobind>=2.9, <3.0; numpy>=2.1.0, <=2.1.2; ml_dtypes>=0.5.0, <=0.6.0; and PyYAML>=5.4.0, <=6.0.1.

Ninja has a platform-specific cap: on Windows, pyproject.toml pins it as ninja<1.13.0. Non-Windows builds accept any ninja version. The reason for this Windows-only cap is not documented in pyproject.toml.

Environment variables from setup.py that control the build include LLVM_INSTALL_DIR for out-of-tree builds, CMAKE_BUILD_TYPE, TORCH_MLIR_CMAKE_BUILD_DIR and MAX_JOBS. setup.py recommends Ninja combined with ccache, setting CMAKE_GENERATOR=Ninja, CMAKE_C_COMPILER_LAUNCHER=ccache and CMAKE_CXX_COMPILER_LAUNCHER=ccache, which lets ccache reuse LLVM/MLIR object files across repeated builds.

IREE and Blade as downstreams; direct MLIR integration is a compiler engineering project, not a library swap

IREE and Blade illustrate what adopting torch-mlir actually entails. Both teams took the Torch MLIR Dialect output and built full lowering pipelines to target hardware. IREE covers CPU, Vulkan and Metal execution. Blade targets GPU execution at Alibaba scale.

A new team building a hardware compiler with torch-mlir inherits the full complexity of MLIR's dialect system: passes, conversions, lowering pipelines and type systems. The README frames the end goal as analogous to adding an LLVM target backend, a compiler engineering task, not adding a library to an existing application.

For projects that want PyTorch model portability without entering the MLIR stack, ONNX export combined with onnxruntime is a distinct approach. onnxruntime covers CPU, CUDA and some accelerators without any MLIR infrastructure. The trade-off is that onnxruntime's hardware target list is fixed, while MLIR-based compilers can define custom lowering passes for novel silicon.

pyproject.toml sets py38 as the Black target while wheel builds only cover cp310 through cp312

pyproject.toml contains a detail that stands out on close reading. Black's formatter configuration sets `target-version = ['py38']`, the line-length is 88, and that Python target has not been updated alongside the build system. Meanwhile, cibuildwheel is configured to build `cp310-* cp311-* cp312-*`, and the README's install section specifies Python 3.11 or 3.10 as the supported environment. The formatter target and the runtime targets disagree.

This gap is not documented anywhere in the README as intentional. It points to a project where configuration files evolve at different speeds. The nanobind upper bound (<3.0) and the numpy upper bound (<=2.1.2) are both tight, meaning a routine package upgrade could break the build without any change to torch-mlir itself.

The last push was on 2026-09-18, and the project is available under the Apache License v2.0 with LLVM Exceptions and also under a BSD-style licence as documented in setup.py.

Editorial conclusion

torch-mlir is the right integration point when building an MLIR-based compiler that consumes PyTorch or ONNX models, and IREE demonstrates it can work at production scale. Skip it when you need a stable versioned library: snapshot packages change daily without semantic versioning, and the project is not yet an official LLVM component. Before integrating from source, check the nanobind constraint (>=2.9, <3.0) against your build environment, confirm your Python version is in the cp310, cp311 or cp312 set from pyproject.toml, and on Windows verify that ninja is below version 1.13.0.

Frequently asked questions

Official sources

1. Issues
2. llvm/torch-mlir on GitHub
3. README
4. Releases