diff --git a/Makefile b/Makefile index 17cbc513..81f00bf2 100644 --- a/Makefile +++ b/Makefile @@ -97,10 +97,10 @@ pages: $(MAKE) -C docs/site site agents: - scripts/llms/init_agents.py agents README.md "Links" "AI assistants" > AGENTS.md + scripts/llms/init_agents.py agents README.md "Links" "local LLVM" "AI assistants" > AGENTS.md claude: - scripts/llms/init_agents.py claude README.md "Links" "AI assistants" > CLAUDE.md + scripts/llms/init_agents.py claude README.md "Links" "local LLVM" "AI assistants" > CLAUDE.md run-tutorial: marimo run docs/tutorials/xtc_101.py diff --git a/README.md b/README.md index 8f49f91b..6936371a 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,38 @@ Available extensions are: - `[default]`: mlir + tvm - `[dev]`: default + test +### Using a local LLVM / MLIR build + +By default XTC uses the LLVM/MLIR toolchain shipped in its Python wheels. To +use a local LLVM checkout instead (e.g. to test compiler changes), point XTC +at your build directory (`$LLVM_BUILD` below). Two levels are possible. + +**Binaries only.** Override `opt`, `llc`, `mlir-opt` and `mlir-translate` +while keeping the wheel's Python bindings — the simplest option, e.g. when +your changes live in the LLVM middle-end / back-end: + +```bash +export XTC_LLVM_PREFIX=$LLVM_BUILD # providing bin/opt and bin/llc +export XTC_MLIR_PREFIX=$LLVM_BUILD # providing bin/mlir-opt and bin/mlir-translate +``` + +**Binaries and Python bindings.** The MLIR bindings are a native extension, +so the interpreter must match the Python version they were built for (check +the ABI tag under +`$LLVM_BUILD/tools/mlir/python_packages/mlir_core/mlir/_mlir_libs`). +Create a matching venv if, install XTC without the MLIR wheels, build the +runtime support libraries, then prepend the bindings to `PYTHONPATH`: + +```bash +uv venv -p 3.14 .venv-local && source .venv-local/bin/activate +uv pip install -e '.[tvm]' + +export PYTHONPATH=$LLVM_BUILD/tools/mlir/python_packages/mlir_core:$PYTHONPATH +export XTC_MLIR_PREFIX=$LLVM_BUILD +export XTC_LLVM_PREFIX=$LLVM_BUILD +export XTC_MLIR_TARGET=llvmir +``` + ### Code quality Code quality requirements: diff --git a/src/xtc/backends/mlir/MlirBindingsExtensions.py b/src/xtc/backends/mlir/MlirBindingsExtensions.py new file mode 100644 index 00000000..ca4fb6cc --- /dev/null +++ b/src/xtc/backends/mlir/MlirBindingsExtensions.py @@ -0,0 +1,43 @@ +# +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2024-2026 The XTC Project Authors +# +"""Optional XTC extensions to the MLIR Python bindings.""" + +from __future__ import annotations + +import importlib +import logging + +logger = logging.getLogger(__name__) + +# Optional XTC extensions to the MLIR Python bindings: the module to import +# mapped to the pass-pipeline entries it contributes. +_EXTENSIONS: dict[str, tuple[str, ...]] = { + "mlir.xtc_transform": ("func.func(reduce-extract-slices)",), +} + +# Reverse map from a contributed pass to its providing module, so a pass can be +# gated without the caller naming its extension. +_PASS_OWNER: dict[str, str] = {} +for _module, _module_passes in _EXTENSIONS.items(): + for _pass in _module_passes: + assert _pass not in _PASS_OWNER, f"pass {_pass!r} declared by two extensions" + _PASS_OWNER[_pass] = _module + +# Extension modules that imported successfully, resolved once at import time. +_loaded: set[str] = set() +for _module in _EXTENSIONS: + try: + importlib.import_module(_module) + _loaded.add(_module) + except ImportError as _exc: + logger.debug("MLIR binding extension %r unavailable: %s", _module, _exc) + + +def passes(pass_names: list[str]) -> list[str]: + """Return the subset of ``pass_names`` whose extension is available.""" + for pass_name in pass_names: + if pass_name not in _PASS_OWNER: + raise KeyError(f"unknown extension pass: {pass_name!r}") + return [p for p in pass_names if _PASS_OWNER[p] in _loaded] diff --git a/src/xtc/backends/mlir/MlirCompilerPasses.py b/src/xtc/backends/mlir/MlirCompilerPasses.py index 20cee233..ae4ed40f 100644 --- a/src/xtc/backends/mlir/MlirCompilerPasses.py +++ b/src/xtc/backends/mlir/MlirCompilerPasses.py @@ -33,7 +33,8 @@ ) from mlir.passmanager import PassManager from mlir.ir import Module -import mlir.xtc_transform + +import xtc.backends.mlir.MlirBindingsExtensions as binding_extensions # Import SDist if available try: @@ -779,7 +780,6 @@ def run(self, pass_names: list[str]) -> None: def apply_bufferization_passes(mlir_program: RawMlirProgram, mlir_install_dir: str): - assert mlir.xtc_transform bufferize_options = [ "bufferize-function-boundaries", "function-boundary-type-conversion=identity-layout-map", @@ -787,9 +787,8 @@ def apply_bufferization_passes(mlir_program: RawMlirProgram, mlir_install_dir: s ] MlirProgramApplyPasses(mlir_program).run( - [ - # xtc pass that folds extract slices to make smaller tensor.empty allocations - "func.func(reduce-extract-slices)", + binding_extensions.passes(["func.func(reduce-extract-slices)"]) + + [ "canonicalize", "cse", "eliminate-empty-tensors", # causes ops to write directly to out buffer diff --git a/src/xtc/backends/mlir/MlirTarget/MlirLLVMTarget.py b/src/xtc/backends/mlir/MlirTarget/MlirLLVMTarget.py index 4ecc57c3..dd9ca109 100644 --- a/src/xtc/backends/mlir/MlirTarget/MlirLLVMTarget.py +++ b/src/xtc/backends/mlir/MlirTarget/MlirLLVMTarget.py @@ -123,7 +123,7 @@ def _lowering_pipeline(self) -> list[str]: "sccp", # Data flow to LLVM "convert-math-to-llvm", - "convert-vector-to-llvm{enable-x86vector=true}", + "convert-vector-to-llvm", "convert-index-to-llvm", "convert-arith-to-llvm", "convert-ub-to-llvm", diff --git a/src/xtc/search/explore.py b/src/xtc/search/explore.py index e0cb9c57..58bedf25 100644 --- a/src/xtc/search/explore.py +++ b/src/xtc/search/explore.py @@ -118,9 +118,11 @@ def __post_init__(self): self.operator = None self.func_name = None - # Workaround to ensure that TVM backend is after MLIR backends, - # otherwise the import of tvm breaks the MLIR python bindings - self.backends = sorted(self.backends) + # Workaround to ensure that the TVM backend is imported before the MLIR + # backends: loading the MLIR wheel's LLVM first and tvm's LLVM second + # re-registers an LLVM command-line option and aborts. reverse=True puts + # "tvm" ahead of "mlir". + self.backends = sorted(self.backends, reverse=True) if self.operator: if not self.func_name: