PYTHON LAB — driving MLIR from Python (how Triton backends are wired)

TRACK: python | DIFFICULTY: 2 | NEEDS: MLIR Python bindings (MLIR_ENABLE_BINDINGS_PYTHON=ON) and a venv on the Python version they were built for | FILE: mlir_python.py

Why this matters

A Triton backend is a Python class whose add_stages() calls into C++ MLIR through Python bindings: a Triton CPU backend’s compiler.py calls a C++ translate_linalg_to_obj(mlir_mod, options_map) entry point (lab 02_lowering/03 showed its C++ side), and Triton's own make_ttir builds a PassManager in Python (passes.common.add_canonicalizer(pm), ...). The upstream MLIR bindings expose the same objects: ir.Module, PassManager, the transform interpreter, and a JIT (ExecutionEngine).

Run it

source ../tools/env.sh
PYTHONPATH=$MLIR_PYTHONPATH ../tools/pin.sh -q $(cd .. && pwd)/.venv-mlirpy/bin/python mlir_python.py

Real output (tail):

== 2+3. run a pipeline from Python, JIT the kernel, call it on numpy arrays
   y = [ 1.  3.  5.  7.  9. 11. 13. 15. 17. 19.] | matches numpy: True
== 4. time matmul builds from Python (512^3 f32), numpy BLAS as reference
   naive                      209.63 ms     1.3 GFLOP/s  OK
   tile+vec+hoist 4x16          8.19 ms    32.8 GFLOP/s  OK
   numpy (OpenBLAS/MKL)         4.89 ms    54.9 GFLOP/s  OK

What each part shows

1. Walk IR like a pass (op.name, op.operands, op.attributes["iterator_types"], op.parent up to the func.func). Handy for quick analyses over hundreds of dumped kernels ("which kernels have a reduction?").
2. Pass pipelines from Python: PassManager.parse(LOWER_STD).run(module.operation), the same strings as the shell labs.
3. JIT + numpy. ExecutionEngine(module, opt_level=3, shared_libs=[...]) and invoke("saxpy", ...). Arguments follow the C-interface ABI from lab 02_lowering/04: each memref is a pointer to a pointer to a descriptor (get_ranked_memref_descriptor(np_array)), scalars by pointer.
4. The transform interpreter from Python: interpreter.apply_named_sequence(payload, entry, schedule_module), keeping the schedule as a separate module. (The test-transform-dialect-erase-schedule pass the shell labs use is a test pass. Test passes aren't registered in the Python bindings, and trying it fails with does not refer to a registered pass.) Then a third reference point for lab 04/01: numpy's BLAS, single core because the process is pinned. That's roughly what a hand-tuned library reaches; IREE with data tiling (lab 08) and our schedule are at ~60–80% of it.

Gotcha: the bindings are tied to one Python version

_mlir.cpython-314-x86_64-linux-gnu.so

build-full built the bindings for Python 3.14 (the system Python). From the 3.13 venv, from mlir import ir fails with ImportError: cannot import name 'ir' from 'mlir._mlir_libs._mlir'. Hence the separate .venv-mlirpy (3.14 + numpy). Triton has the same property: its wheels bundle their own MLIR bindings built for each Python version. Never mix a compiled extension with a different interpreter.

Try this yourself

1. In part 1, also print each generic's indexing_maps, and flag any non-identity input map (a transpose or broadcast).
2. Add the SVE schedule (lab 05/02) as a variant in part 4. What happens when you JIT it on x86? Why?
3. Build a tiny function with the IR builder API (func.FuncOp, arith.AddFOp inside with ir.InsertionPoint(...)) instead of parsing text.