Experimental MLIR (IREE llvm-cpu) backend feasibility spike

[cetagostini]

Summary

Feasibility spike: can PyTensor grow an MLIR backend? This PR adds an experimental pytensor/link/mlir/ that lowers a FunctionGraph to textual MLIR (linalg + arith on dynamic tensor<?x...>), compiles it with IREE's llvm-cpu target, and executes it through the IREE runtime — end to end, with correct results on simple CPU graphs.

This is not a backend proposal. It is evidence for a discussion: the working slice, the exact blockers found, and where they live.

The simple graph

import numpy as np, pytensor, pytensor.tensor as pt

x = pt.vector("x", dtype="float64")
y = pt.vector("y", dtype="float64")
f = pytensor.function([x, y], x + y * 2, mode="MLIR")
f(np.array([1.0, 2.0, 3.0]), np.array([4.0, 5.0, 6.0]))
# array([ 9., 12., 15.])

scripts/mlir_demo.py runs the full matrix (every case checked with np.testing.assert_allclose against FAST_COMPILE):

vector_float64    [ 9. 12. 15.]
multiply_float64  [ 4. 10. 18.]
vector_float32    [ 9. 12. 15.]
broadcast_float64 [[11. 22. 33.] [14. 25. 36.]]
transpose_float64 [[1. 4.] [2. 5.] [3. 6.]]
dot_float64       [[ 58.  64.] [139. 154.]]
dot_float32       [[ 58.  64.] [139. 154.]]

Because the emitted tensor types use ? dims, one compiled function runs at multiple input lengths without retracing (exercised at 3 and 5).

Design

Follows the house backend pattern (mirrors link/mlx / link/jax):

• MLIRLinker(JITLinker) converts the whole FunctionGraph via a singledispatch mlir_funcify, emits one MLIR module, calls iree.compiler.compile_str(target_backends=["llvm-cpu"]), and invokes the loaded VM module. No Python-op fallback — unsupported Ops raise NotImplementedError loudly.
• Dispatch covers what the demo needs: Elemwise Add/Mul/Second → linalg.generic + arith, DimShuffle/DeepCopyOp → linalg.generic affine maps, Dot/Dot22 → linalg.fill + linalg.matmul. float32/float64, rank ≤ 2.
• Mode("MLIR") registered in compile/mode.py; config.mode accepts MLIR.
• float64 preserved by passing --iree-input-demote-f64-to-f32=false.
• A runtime validator enforces PyTensor's static-size-1 broadcasting rule and Dot inner-dim equality before IREE (IREE does not implement our runtime-broadcast policy itself).

Toolchain (macOS arm64)

• Works: pip iree-base-compiler / iree-base-runtime 3.11.0 — clean install, compile_str → VM flatbuffer → runtime invocation. Deps are only installed in the experiment env, intentionally not declared in pyproject.
• Doesn't: conda-forge mlir-python-bindings 22.1.8 — imports fine but any mlir.ir.Context() aborts the process (LLVM ERROR: Attempting to attach an interface to an unregistered operation builtin.unrealized_conversion_cast), which is why the IR is string-templated rather than built through the bindings.

Blockers found (why this stays a spike)

1. Subnormal flush-to-zero: IREE llvm-cpu silently flushes subnormal results (float32 tiny*0.5 → 0.0 instead of 5.877e-39; same for float64). No exposed IEEE-denormal control in iree-compile --help-hidden. Principal semantic divergence from the C/Numba/JAX backends.
2. Host-mapped buffer leak: returning NumPy outputs requires IREE host mapping, which leaks MappedMemory/HalBufferView instances per call (~9.8 MB RSS growth over a 100-call probe; reproduced in raw IREE without PyTensor). Rules out iterative workloads until fixed upstream.
3. Textual-IR emission is workable but fragile long-term; proper bindings (or IREE's Python IR APIs) would be the real path.

Test plan

• scripts/mlir_demo.py — all cases green vs FAST_COMPILE (exit 0).
• scripts/mlir_toolchain_hello.py — raw IREE hello-world (dynamic tensor<?xf64>, linalg.generic/arith.addf) prints [5. 7. 9.] and asserts float64 preservation.
• Edge paths exercised during review: rank-0, empty vector, transpose/broadcast, non-contiguous input, duplicate/multi-output functions, scientific-notation literals, x*0 rewrites.

🤖 Explored and implemented with AI assistance (Claude Code), including adversarial/correctness/edge-case review passes and an independent fact-check of the claims above.

Metal/MPS results

Verdict: yes, with a float32-only contract. The same simple graphs now run on the M3 Max GPU through IREE's Metal HAL:

textual linalg/arith MLIR
  -> IREE metal-spirv
  -> SPIR-V -> MSL -> .metallib
  -> IREE runtime driver="metal"
  -> Apple M3 Max GPU

MLIR_METAL selects metal-spirv and driver="metal"; it is deliberately distinct from the CPU MLIR mode. scripts/mlir_demo.py enumerates and prints the actual IREE device, rejects float64 graphs before compilation, and checks all GPU outputs against FAST_COMPILE with rtol=atol=1e-6.

Observed run (IREE 3.11.0):

metal_devices [{'path': '0000000100000635', 'name': 'Apple M3 Max', ...}]
metal_float64_rejected MLIR Metal only supports float32 graphs
metal_float32_tolerances rtol=1e-6 atol=1e-6
metal_vector_float32 [ 9. 12. 15.]
metal_multiply_float32 [ 4. 10. 18.]
metal_broadcast_float32 [[11. 22. 33.]
 [14. 25. 36.]]
metal_transpose_float32 [[1. 4.]
 [2. 5.]
 [3. 6.]]
metal_dot_float32 [[ 58.  64.]
 [139. 154.]]

The raw IREE hello-world was also compiled outside PyTensor from the existing linalg.generic source after changing its element type to f32; target_backends=["metal-spirv"] plus load_vm_flatbuffer(..., driver="metal") printed [5. 7. 9.]. The original f64 source fails at the intended compiler link with:

failed to convert SPIR-V type: 'memref<?xf64, #spirv.storage_class<StorageBuffer>>'
failed to legalize operation 'hal.interface.binding.subspan'

The Metal mode rejects every live float64 graph value and passes --iree-input-demote-f64-to-f32=false as a compiler backstop, so it does not silently change PyTensor's dtype contract. The existing MLIR CPU mode still preserves float64.

What “MPS dialect / Metal backend” actually means

This working path is not MPSGraph. It is IREE's actively maintained Metal HAL target, which lowers through SPIR-V and cross-compiles to Metal Shading Language; see IREE's Metal HAL design and target implementation.

There was a real Apple proposal for an MLIR MPS dialect targeting MetalPerformanceShadersGraph, but it was an RFC rather than a merged upstream MLIR dialect. The corresponding generic public linalg/arith -> MPSGraph Python route is not available. On this host, /usr/bin/mpsgraphtool --help exposes convert -coremlpackage, not a generic MLIR/bytecode input. Apple's jax-metal documentation describes a separate JAX/OpenXLA StableHLO -> MPSGraph/PJRT route; it is not an ingress for this spike and it likewise lists np.float64 as unsupported.

Caveats

• This is a local macOS/Apple-silicon experiment, not a portable backend proposal or a performance result. No GPU timing comparison was made.
• Metal execution needs the IREE metal-spirv compiler target, metal runtime driver, and Xcode's Metal toolchain. The local wheel has all of them; iree-run-module --dump_devices=metal identifies the M3 Max as Metal 3/unified-memory.
• The linker materializes each returned value as a NumPy copy. The demo also emits IREE/nanobind reference-leak warnings at interpreter shutdown; this needs upstream investigation before iterative workloads.
• Input/output float32 is supported in this slice. Float64 is intentionally rejected; complex and broader Op/dtype coverage remain outside the spike.