Experimental MLIR (IREE llvm-cpu) backend feasibility spike

pyproject.toml

@@ -167,6 +167,9 @@ known-first-party = ["pytensor", "tests"]
 "tests/compile/debug/test_monitormode.py" = ["T201"]
 "scripts/run_mypy.py" = ["T201"]
 "scripts/bump_numba_upper_bound.py" = ["T201"]
+"scripts/mlir_demo.py" = ["T201"]
+"scripts/mlir_toolchain_hello.py" = ["T201"]
+"scripts/benchmark_mlx_vs_mlir_metal.py" = ["T201"]
 # Test modules of optional backends that use `pytest.importorskip`, skip "E402"
 "tests/link/jax/**/test_*.py" = ["E402"]
 "tests/link/numba/**/test_*.py" = ["E402"]

pytensor/compile/mode.py

@@ -28,6 +28,7 @@
 from pytensor.link.basic import Linker, PerformLinker
 from pytensor.link.c.basic import CLinker, OpWiseCLinker
 from pytensor.link.jax.linker import JAXLinker
+from pytensor.link.mlir.linker import MLIRLinker
 from pytensor.link.mlx.linker import MLXLinker
 from pytensor.link.numba.linker import NumbaLinker
 from pytensor.link.pytorch.linker import PytorchLinker
@@ -53,6 +54,7 @@
     "pytorch": PytorchLinker(),
     "numba": NumbaLinker(),
     "mlx": MLXLinker(),
+    "mlir": MLIRLinker(),
 }
 
 
@@ -534,6 +536,16 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
     RewriteDatabaseQuery(include=["fast_run", "mlx"], exclude=["fusion"]),
 )
 
+MLIR = Mode(
+    MLIRLinker(),
+    RewriteDatabaseQuery(include=["fast_compile"], exclude=["fusion"]),
+)
+
+MLIR_METAL = Mode(
+    MLIRLinker(target_backend="metal-spirv"),
+    RewriteDatabaseQuery(include=["fast_compile"], exclude=["fusion"]),
+)
+
 FAST_COMPILE = Mode(
     VMLinker(use_cloop=False, c_thunks=False),
     RewriteDatabaseQuery(include=["fast_compile", "py_only"]),
@@ -553,6 +565,8 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
     "NUMBA": NUMBA,
     "PYTORCH": PYTORCH,
     "MLX": MLX,
+    "MLIR": MLIR,
+    "MLIR_METAL": MLIR_METAL,
 }
 
 _CACHED_RUNTIME_MODES: dict[Any, Mode] = {}

pytensor/configdefaults.py

@@ -49,6 +49,8 @@ def _filter_mode(val):
         "NUMBA",
         "PYTORCH",
         "MLX",
+        "MLIR",
+        "MLIR_METAL",
     ]
     if val in str_options:
         return val

pytensor/link/mlir/__init__.py

@@ -0,0 +1 @@
+from pytensor.link.mlir.linker import MLIRLinker

pytensor/link/mlir/dispatch/__init__.py

@@ -0,0 +1 @@
+from pytensor.link.mlir.dispatch.basic import mlir_funcify, mlir_typify