From de1a5b302718a1f7910a10715d05679ad0cb89bd Mon Sep 17 00:00:00 2001 From: Adam Siemieniuk Date: Wed, 17 Jun 2026 12:21:03 +0200 Subject: [PATCH 1/3] [pipeline] Support relative python imports Improves python transforms importer to handle relative imports present in python files defining MLIR transforms. --- lighthouse/pipeline/helper.py | 51 +++++++++++++++++++++++++++++++++++ lighthouse/pipeline/stage.py | 12 ++------- 2 files changed, 53 insertions(+), 10 deletions(-) diff --git a/lighthouse/pipeline/helper.py b/lighthouse/pipeline/helper.py index 9c4820d5..75fb8618 100644 --- a/lighthouse/pipeline/helper.py +++ b/lighthouse/pipeline/helper.py @@ -1,4 +1,10 @@ +import importlib +import importlib.util import os +import sys +from functools import lru_cache +from pathlib import Path +from types import ModuleType from mlir import ir from mlir.dialects import transform @@ -15,6 +21,51 @@ def import_mlir_module(path: str, context: ir.Context) -> ir.Module: return ir.Module.parse(f.read(), context=context) +@lru_cache(maxsize=None) +def _resolve_package(directory: Path) -> tuple[str, str]: + """ + Resolve the enclosing package for a directory. + + Results are cached per directory so the walk runs at most once + for each directory. + """ + package_parts: list[str] = [] + parent = directory + while (parent / "__init__.py").exists(): + package_parts.insert(0, parent.name) + parent = parent.parent + return str(parent), ".".join(package_parts) + + +def import_python_module(path: str) -> ModuleType: + """Import a Python module from a file.""" + if path is None: + raise ValueError("Path to the module must be provided.") + if not os.path.exists(path): + raise ValueError(f"Path to the module does not exist: {path}") + + filepath = Path(path).resolve() + + # Resolve the enclosing package to enable relative imports. + package_root, dotted_prefix = _resolve_package(filepath.parent) + + module_name = filepath.stem + if dotted_prefix: + # The file is part of a package: build its dotted name and make sure + # the package root is importable so relative imports resolve. + qualified_name = f"{dotted_prefix}.{module_name}" + if package_root not in sys.path: + sys.path.insert(0, package_root) + return importlib.import_module(qualified_name) + + # Standalone file: load it directly from its location. + spec = importlib.util.spec_from_file_location(module_name, str(filepath)) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + def apply_registered_pass(*args, **kwargs): """Utility function to add a bundle of passes to a Transform Schedule""" return transform.apply_registered_pass(transform.AnyOpType.get(), *args, **kwargs) diff --git a/lighthouse/pipeline/stage.py b/lighthouse/pipeline/stage.py index 00a5c660..a54eddef 100644 --- a/lighthouse/pipeline/stage.py +++ b/lighthouse/pipeline/stage.py @@ -1,13 +1,10 @@ from abc import abstractmethod -import importlib from enum import Enum -from pathlib import Path -import os from mlir import ir from mlir.passmanager import PassManager from mlir.dialects import transform -from lighthouse.pipeline.helper import import_mlir_module +from lighthouse.pipeline.helper import import_mlir_module, import_python_module from lighthouse.pipeline.descriptor import Descriptor, PipelineDescriptor @@ -170,12 +167,7 @@ def __init__(self, transform: Transform | ir.Module, context: ir.Context): elif transform.type == Transform.Type.Python: # For Python transforms, we expect the file to define a function # that takes an MLIR module and returns a transformed MLIR module. - module_name = Path(os.path.basename(transform.filename)).stem - spec = importlib.util.spec_from_file_location( - module_name, transform.filename - ) - transform_module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(transform_module) + transform_module = import_python_module(transform.filename) if not hasattr(transform_module, transform.generator): raise ValueError( f"Transform module '{transform.filename}' does not define a '{transform.generator}' generator function." From 0b746c5e88c59cf5bc66a714020d26a62660fc3e Mon Sep 17 00:00:00 2001 From: Adam Siemieniuk Date: Wed, 17 Jun 2026 14:18:39 +0200 Subject: [PATCH 2/3] Move import funcs to utils --- lighthouse/ingress/torch/importer.py | 8 +--- lighthouse/pipeline/driver.py | 2 +- lighthouse/pipeline/helper.py | 63 --------------------------- lighthouse/pipeline/stage.py | 2 +- lighthouse/utils/importer.py | 64 ++++++++++++++++++++++++++++ 5 files changed, 68 insertions(+), 71 deletions(-) create mode 100644 lighthouse/utils/importer.py diff --git a/lighthouse/ingress/torch/importer.py b/lighthouse/ingress/torch/importer.py index 7e419bb7..3e7f6817 100644 --- a/lighthouse/ingress/torch/importer.py +++ b/lighthouse/ingress/torch/importer.py @@ -1,5 +1,3 @@ -import importlib -import importlib.util from pathlib import Path from typing import Iterable, Mapping @@ -7,6 +5,7 @@ load_and_run_callable, maybe_load_and_run_callable, ) +from lighthouse.utils.importer import import_python_module try: import torch @@ -124,11 +123,8 @@ def get_inputs(): """ if isinstance(filepath, str): filepath = Path(filepath) - module_name = filepath.stem - spec = importlib.util.spec_from_file_location(module_name, filepath) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) + module = import_python_module(filepath) model = getattr(module, model_class_name, None) if model is None: diff --git a/lighthouse/pipeline/driver.py b/lighthouse/pipeline/driver.py index 1eb32f72..0fc84ad0 100644 --- a/lighthouse/pipeline/driver.py +++ b/lighthouse/pipeline/driver.py @@ -2,8 +2,8 @@ from mlir import ir import lighthouse.pipeline.stage as lhs -from lighthouse.pipeline.helper import import_mlir_module from lighthouse.pipeline.descriptor import PipelineDescriptor, Descriptor +from lighthouse.utils.importer import import_mlir_module import lighthouse.dialects as lh_dialects diff --git a/lighthouse/pipeline/helper.py b/lighthouse/pipeline/helper.py index 75fb8618..e74945e6 100644 --- a/lighthouse/pipeline/helper.py +++ b/lighthouse/pipeline/helper.py @@ -1,71 +1,8 @@ -import importlib -import importlib.util -import os -import sys -from functools import lru_cache -from pathlib import Path -from types import ModuleType - from mlir import ir from mlir.dialects import transform from mlir.dialects.transform import structured -def import_mlir_module(path: str, context: ir.Context) -> ir.Module: - """Import an MLIR text file into an MLIR module""" - if path is None: - raise ValueError("Path to the module must be provided.") - if not os.path.exists(path): - raise ValueError(f"Path to the module does not exist: {path}") - with open(path, "r") as f: - return ir.Module.parse(f.read(), context=context) - - -@lru_cache(maxsize=None) -def _resolve_package(directory: Path) -> tuple[str, str]: - """ - Resolve the enclosing package for a directory. - - Results are cached per directory so the walk runs at most once - for each directory. - """ - package_parts: list[str] = [] - parent = directory - while (parent / "__init__.py").exists(): - package_parts.insert(0, parent.name) - parent = parent.parent - return str(parent), ".".join(package_parts) - - -def import_python_module(path: str) -> ModuleType: - """Import a Python module from a file.""" - if path is None: - raise ValueError("Path to the module must be provided.") - if not os.path.exists(path): - raise ValueError(f"Path to the module does not exist: {path}") - - filepath = Path(path).resolve() - - # Resolve the enclosing package to enable relative imports. - package_root, dotted_prefix = _resolve_package(filepath.parent) - - module_name = filepath.stem - if dotted_prefix: - # The file is part of a package: build its dotted name and make sure - # the package root is importable so relative imports resolve. - qualified_name = f"{dotted_prefix}.{module_name}" - if package_root not in sys.path: - sys.path.insert(0, package_root) - return importlib.import_module(qualified_name) - - # Standalone file: load it directly from its location. - spec = importlib.util.spec_from_file_location(module_name, str(filepath)) - module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - spec.loader.exec_module(module) - return module - - def apply_registered_pass(*args, **kwargs): """Utility function to add a bundle of passes to a Transform Schedule""" return transform.apply_registered_pass(transform.AnyOpType.get(), *args, **kwargs) diff --git a/lighthouse/pipeline/stage.py b/lighthouse/pipeline/stage.py index a54eddef..7d85fd9d 100644 --- a/lighthouse/pipeline/stage.py +++ b/lighthouse/pipeline/stage.py @@ -4,8 +4,8 @@ from mlir import ir from mlir.passmanager import PassManager from mlir.dialects import transform -from lighthouse.pipeline.helper import import_mlir_module, import_python_module from lighthouse.pipeline.descriptor import Descriptor, PipelineDescriptor +from lighthouse.utils.importer import import_mlir_module, import_python_module class Pass: diff --git a/lighthouse/utils/importer.py b/lighthouse/utils/importer.py new file mode 100644 index 00000000..4b60b6d3 --- /dev/null +++ b/lighthouse/utils/importer.py @@ -0,0 +1,64 @@ +import importlib +import importlib.util +import os +import sys +from functools import lru_cache +from pathlib import Path +from types import ModuleType + +from mlir import ir + + +def import_mlir_module(path: str, context: ir.Context) -> ir.Module: + """Import an MLIR text file into an MLIR module""" + if path is None: + raise ValueError("Path to the module must be provided.") + if not os.path.exists(path): + raise ValueError(f"Path to the module does not exist: {path}") + with open(path, "r") as f: + return ir.Module.parse(f.read(), context=context) + + +@lru_cache(maxsize=None) +def _resolve_package(directory: Path) -> tuple[str, str]: + """ + Resolve the enclosing package for a directory. + + Results are cached per directory so the walk runs at most once + for each directory. + """ + package_parts: list[str] = [] + parent = directory + while (parent / "__init__.py").exists(): + package_parts.insert(0, parent.name) + parent = parent.parent + return str(parent), ".".join(package_parts) + + +def import_python_module(path: str) -> ModuleType: + """Import a Python module from a file.""" + if path is None: + raise ValueError("Path to the module must be provided.") + if not os.path.exists(path): + raise ValueError(f"Path to the module does not exist: {path}") + + filepath = Path(path).resolve() + + # Resolve the enclosing package to enable relative imports. + package_root, dotted_prefix = _resolve_package(filepath.parent) + + module_name = filepath.stem + if dotted_prefix: + # The file is part of a package: build its dotted name and make sure + # the package root is importable so relative imports resolve. + qualified_name = f"{dotted_prefix}.{module_name}" + if package_root not in sys.path: + sys.path.insert(0, package_root) + return importlib.import_module(qualified_name) + + # Standalone file: load it directly from its location. + spec = importlib.util.spec_from_file_location(module_name, str(filepath)) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module From 1f24031015be112d941417b96f3bbdd7e1c582cf Mon Sep 17 00:00:00 2001 From: Adam Siemieniuk Date: Wed, 17 Jun 2026 14:26:22 +0200 Subject: [PATCH 3/3] Improve docs --- lighthouse/utils/importer.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/lighthouse/utils/importer.py b/lighthouse/utils/importer.py index 4b60b6d3..c8552d0c 100644 --- a/lighthouse/utils/importer.py +++ b/lighthouse/utils/importer.py @@ -10,7 +10,15 @@ def import_mlir_module(path: str, context: ir.Context) -> ir.Module: - """Import an MLIR text file into an MLIR module""" + """ + Import an MLIR text file into an MLIR module + + Args: + path: Path to the MLIR file + context: MLIR context + Returns: + MLIR module + """ if path is None: raise ValueError("Path to the module must be provided.") if not os.path.exists(path): @@ -36,7 +44,14 @@ def _resolve_package(directory: Path) -> tuple[str, str]: def import_python_module(path: str) -> ModuleType: - """Import a Python module from a file.""" + """ + Import a Python module from a file. + + Args: + path: Path to the Python file + Returns: + Imported Python module + """ if path is None: raise ValueError("Path to the module must be provided.") if not os.path.exists(path):