From 1a4b755b795ee6039b2e8d060bda26c3b8c0846a Mon Sep 17 00:00:00 2001 From: dchigarev Date: Mon, 20 Oct 2025 09:07:24 +0000 Subject: [PATCH 01/16] [ingress][torch-mlir] Add utility functions to import models using torch-mlir Signed-off-by: dchigarev --- .../ingress/torch/01-dummy-mlp-from-file.py | 58 +++++ .../ingress/torch/02-dummy-mlp-from-model.py | 50 ++++ .../examples/ingress/torch/DummyMLP/model.py | 33 +++ python/lighthouse/__init__.py | 1 + python/lighthouse/ingress/__init__.py | 0 python/lighthouse/ingress/torch/__init__.py | 1 + .../lighthouse/ingress/torch/torch_import.py | 215 ++++++++++++++++++ 7 files changed, 358 insertions(+) create mode 100644 python/examples/ingress/torch/01-dummy-mlp-from-file.py create mode 100644 python/examples/ingress/torch/02-dummy-mlp-from-model.py create mode 100644 python/examples/ingress/torch/DummyMLP/model.py create mode 100644 python/lighthouse/__init__.py create mode 100644 python/lighthouse/ingress/__init__.py create mode 100644 python/lighthouse/ingress/torch/__init__.py create mode 100644 python/lighthouse/ingress/torch/torch_import.py diff --git a/python/examples/ingress/torch/01-dummy-mlp-from-file.py b/python/examples/ingress/torch/01-dummy-mlp-from-file.py new file mode 100644 index 0000000..421b712 --- /dev/null +++ b/python/examples/ingress/torch/01-dummy-mlp-from-file.py @@ -0,0 +1,58 @@ +""" +Example demonstrating how to load a PyTorch model to MLIR using Lighthouse, +without instantiating the model on our side. + +The script uses 'lighthouse.ingress.torch.import_from_file' function that +takes a filepath to a Python file containing the model definition, along with +the names of functions to get model init arguments and sample inputs. The function +imports the model class on its own, instantiates it, and passes it torch_mlir +to get a MLIR module in the specified dialect. + +The script uses model from 'DummyMLP/model.py' as an example. +""" + +import os +from pyparsing import Path + +# MLIR infrastructure imports (only needed if you want to manipulate the MLIR module) +import mlir.dialects.func as func +from mlir import ir, passmanager + +# Lighthouse imports +from lighthouse.ingress.torch import import_from_file + +# Step 1: Setup paths to locate the model definition file +script_dir = Path(os.path.dirname(os.path.abspath(__file__))) +model_path = script_dir / "DummyMLP" / "model.py" + +ir_context = ir.Context() + +# Step 2: Convert PyTorch model to MLIR +# Conversion step where Lighthouse: +# - Loads the DummyMLP class and instantiates with arguments obtained from 'get_init_inputs()' +# - Calls get_sample_inputs() to get sample input tensors for shape inference +# - Converts PyTorch model to linalg-on-tensors dialect operations using torch_mlir +mlir_module_ir : ir.Module = import_from_file( + model_path, # Path to the Python file containing the model + model_class_name="DummyMLP", # Name of the PyTorch nn.Module class to convert + init_args_fn_name="get_init_inputs", # Function that returns args for model.__init__() + inputs_args_fn_name="get_sample_inputs", # Function that returns sample inputs for tracing + dialect="linalg-on-tensors", # Target MLIR dialect (linalg ops on tensor types) + ir_context=ir_context # MLIR context for the conversion +) + +# Step 3: Extract the main function operation from the MLIR module and print its metadata +func_op : func.FuncOp = mlir_module_ir.operation.regions[0].blocks[0].operations[0] +print(f"entry-point name: {func_op.name}") +print(f"entry-point type: {func_op.type}") + +# Step 4: Apply some MLIR passes using a PassManager +pm = passmanager.PassManager(context=ir_context) + +pm.add("one-shot-bufferize") +pm.add("linalg-specialize-generic-ops") +pm.run(mlir_module_ir.operation) + +# Step 5: Output the final MLIR +print("\n\nModule dump after running pm.Pipeline:") +mlir_module_ir.dump() diff --git a/python/examples/ingress/torch/02-dummy-mlp-from-model.py b/python/examples/ingress/torch/02-dummy-mlp-from-model.py new file mode 100644 index 0000000..fb4de89 --- /dev/null +++ b/python/examples/ingress/torch/02-dummy-mlp-from-model.py @@ -0,0 +1,50 @@ +""" +Example demonstrating how to load an already instantiated PyTorch model +to MLIR using Lighthouse. + +The script uses 'lighthouse.ingress.torch.import_from_model' function that +takes an already instantiated PyTorch model, along with its sample inputs. +The function passes the model to torch_mlir to get a MLIR module in the +specified dialect. + +The script uses model from 'DummyMLP/model.py' as an example. +""" + +import torch + +# MLIR infrastructure imports (only needed if you want to manipulate the MLIR module) +import mlir.dialects.func as func +from mlir import ir, passmanager + +# Lighthouse imports +from lighthouse.ingress.torch import import_from_model + +# Import a sample model definition +from .DummyMLP.model import DummyMLP + +# Step 1: Instantiate a model and prepare sample input +model = DummyMLP() +sample_input = torch.randn(1, 10) + +ir_context = ir.Context() +# Step 2: Convert PyTorch model to MLIR +mlir_module_ir : ir.Module = import_from_model( + model, + sample_args=(sample_input,), + ir_context=ir_context +) + +# Step 3: Extract the main function operation from the MLIR module and print its metadata +func_op : func.FuncOp = mlir_module_ir.operation.regions[0].blocks[0].operations[0] +print(f"entry-point name: {func_op.name}") +print(f"entry-point type: {func_op.type}") + +# Step 4: Apply some MLIR passes using a PassManager (optional) +pm = passmanager.PassManager(context=ir_context) +pm.add("one-shot-bufferize") +pm.add("linalg-specialize-generic-ops") +pm.run(mlir_module_ir.operation) + +# Step 5: Output the final MLIR +print("\n\nModule dump after running pm.Pipeline:") +mlir_module_ir.dump() diff --git a/python/examples/ingress/torch/DummyMLP/model.py b/python/examples/ingress/torch/DummyMLP/model.py new file mode 100644 index 0000000..905ecc8 --- /dev/null +++ b/python/examples/ingress/torch/DummyMLP/model.py @@ -0,0 +1,33 @@ +"""Defines a simple PyTorch model to be used in lighthouse's ingress examples.""" + +import torch +import torch.nn as nn + +import os + +class DummyMLP(nn.Module): + def __init__(self): + super().__init__() + self.net = nn.Sequential( + nn.Linear(10, 32), + nn.ReLU(), + nn.Linear(32, 2) + ) + + def forward(self, x): + return self.net(x) + + +def get_init_inputs(): + """Function to return args to pass to DummyMLP.__init__()""" + return () + + +def get_sample_inputs(): + """Arguments to pass to DummyMLP.forward()""" + return (torch.randn(1, 10),) + + +if __name__ == "__main__": + script_dir = os.path.dirname(os.path.abspath(__file__)) + torch.save(DummyMLP().state_dict(), os.path.join(script_dir, "dummy_mlp.pth")) diff --git a/python/lighthouse/__init__.py b/python/lighthouse/__init__.py new file mode 100644 index 0000000..1ac008e --- /dev/null +++ b/python/lighthouse/__init__.py @@ -0,0 +1 @@ +__version__ = "0.1.0a1" diff --git a/python/lighthouse/ingress/__init__.py b/python/lighthouse/ingress/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/lighthouse/ingress/torch/__init__.py b/python/lighthouse/ingress/torch/__init__.py new file mode 100644 index 0000000..c3a6e72 --- /dev/null +++ b/python/lighthouse/ingress/torch/__init__.py @@ -0,0 +1 @@ +from .torch_import import import_from_file, import_from_model diff --git a/python/lighthouse/ingress/torch/torch_import.py b/python/lighthouse/ingress/torch/torch_import.py new file mode 100644 index 0000000..f64b9ff --- /dev/null +++ b/python/lighthouse/ingress/torch/torch_import.py @@ -0,0 +1,215 @@ +import importlib +import importlib.util +from pathlib import Path +from typing import Iterable, Mapping + +try: + import torch + import torch.nn as nn +except ImportError as e: + raise ImportError( + "PyTorch is required to use the torch import functionality. " + "Please run 'uv pip install .[torch-mlir]'" + ) from e + +try: + from torch_mlir import fx + from torch_mlir.fx import OutputType +except ImportError as e: + raise ImportError( + "torch-mlir is required to use the torch import functionality. " + "Please run 'uv pip install .[torch-mlir]'" + ) from e + +from mlir import ir + +def import_from_model( + model: nn.Module, + sample_args : Iterable, + sample_kwargs : Mapping = None, + dialect : OutputType | str = OutputType.LINALG_ON_TENSORS, + ir_context : ir.Context | None = None, + **kwargs, +) -> str | ir.Module: + """ + Import a PyTorch nn.Module into MLIR. + + The function uses torch-mlir's FX importer to convert the given PyTorch model + into an MLIR module in the specified dialect. The user has to provide sample + input arguments (e.g. a torch.Tensor with the correct shape). + + Parameters + ---------- + model : nn.Module + The PyTorch model to import. + sample_args : Iterable + Sample input arguments to the model. + sample_kwargs : Mapping, optional + Sample keyword arguments to the model. + dialect : torch_mlir.fx.OutputType | {"linalg-on-tensors", "torch", "tosa"}, default: OutputType.LINALG_ON_TENSORS + The target dialect for the imported MLIR module. + ir_context : ir.Context, optional + An optional MLIR context to use for parsing the module. + If not provided, the module is returned as a string. + **kwargs + Additional keyword arguments passed to the fx.export_and_import function. + + Returns + ------- + str | ir.Module + The imported MLIR module as a string or an ir.Module if `ir_context` is provided. + + Examples + -------- + >>> import torch + >>> import torch.nn as nn + >>> from lighthouse.ingress.torch_import import import_from_model + >>> class SimpleModel(nn.Module): + ... def __init__(self): + ... super().__init__() + ... self.fc = nn.Linear(10, 5) + ... def forward(self, x): + ... return self.fc(x) + >>> model = SimpleModel() + >>> sample_input = (torch.randn(1, 10),) + >>> # + >>> # option 1: get MLIR module as a string + >>> mlir_module : str = import_from_model(model, sample_input, dialect="linalg-on-tensors") + >>> print(mlir_module) # prints the MLIR module in linalg-on-tensors dialect + >>> # option 2: get MLIR module as an ir.Module + >>> ir_context = ir.Context() + >>> mlir_module_ir : ir.Module = import_from_model(model, sample_input, dialect="tosa", ir_context=ir_context) + >>> # ... run pm.Pipeline on the ir.Module ... + """ + if dialect == "linalg": + raise ValueError( + "Dialect 'linalg' is not supported. Did you mean 'linalg-on-tensors'?" + ) + + if sample_kwargs is None: + sample_kwargs = {} + + model.eval() + module = fx.export_and_import( + model, *sample_args, output_type=dialect, **sample_kwargs, **kwargs + ) + + text_module = str(module) + if ir_context is None: + return text_module + return ir.Module.parse(text_module, context=ir_context) + + +def import_from_file( + filepath: str | Path, + model_class_name: str = "Model", + init_args_fn_name: str | None = "get_init_inputs", + inputs_args_fn_name: str = "get_inputs", + state_path : str | Path | None = None, + dialect : OutputType | str = OutputType.LINALG_ON_TENSORS, + ir_context : ir.Context | None = None, + **kwargs, +) -> str | ir.Module: + """ + Load a PyTorch nn.Module from a file and import it into MLIR. + + The function takes a `filepath` to a Python file containing the model definition, + along with the names of functions to get model init arguments and sample inputs. + The function imports the model class on its own, instantiates it, and passes + it ``torch_mlir`` to get a MLIR module in the specified `dialect`. + + Parameters + ---------- + filepath : str | Path + Path to the Python file containing the model definition. + model_class_name : str, default: "Model" + The name of the model class in the file. + init_args_fn_name : str | None, default: "get_init_inputs" + The name of the function in the file that returns the arguments for + initializing the model. If None, the model is initialized without arguments. + inputs_args_fn_name : str, default: "get_inputs" + The name of the function in the file that returns the sample input arguments + for the model. + state_path : str | Path | None, default: None + Optional path to a file containing the model's ``state_dict``. + dialect: torch_mlir.fx.OutputType | {"linalg-on-tensors", "torch", "tosa"}, default: OutputType.LINALG_ON_TENSORS + The target dialect for the imported MLIR module. + ir_context : ir.Context, optional + An optional MLIR context to use for parsing the module. + If not provided, the module is returned as a string. + **kwargs + Additional keyword arguments passed to the fx.export_and_import function. + + Returns + ------- + str | ir.Module + The imported MLIR module as a string or an ir.Module if `ir_context` is provided. + + Examples + -------- + Given a file `path/to/model_file.py` with the following content: + ```python + import torch + import torch.nn as nn + + class MyModel(nn.Module): + def __init__(self): + super().__init__() + self.fc = nn.Linear(10, 5) + def forward(self, x): + return self.fc(x) + + def get_inputs(): + return (torch.randn(1, 10),) + ``` + + The import script would look line: + >>> from lighthouse.ingress.torch_import import import_from_file + >>> # option 1: get MLIR module as a string + >>> mlir_module : str = import_from_file( + ... "path/to/model_file.py", + ... model_class_name="MyModel", + ... init_args_fn_name=None, + ... dialect="linalg-on-tensors" + ... ) + >>> print(mlir_module) # prints the MLIR module in linalg-on-tensors dialect + >>> # option 2: get MLIR module as an ir.Module + >>> ir_context = ir.Context() + >>> mlir_module_ir : ir.Module = import_from_file( + ... "path/to/model_file.py", + ... model_class_name="MyModel", + ... init_args_fn_name=None, + ... dialect="linalg-on-tensors", + ... ir_context=ir_context + ... ) + >>> # ... run pm.Pipeline on the ir.Module ... + """ + 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) + + model = getattr(module, model_class_name, None) + if model is None: + raise ValueError(f"Model class '{model_class_name}' not found in {filepath}") + + if init_args_fn_name is None: + init_args_fn = lambda *args, **kwargs: () + else: + init_args_fn = getattr(module, init_args_fn_name, None) + if init_args_fn is None: + raise ValueError(f"Init args function '{init_args_fn_name}' not found in {filepath}") + + inputs_args_fn = getattr(module, inputs_args_fn_name, None) + if inputs_args_fn is None: + raise ValueError(f"Inputs args function '{inputs_args_fn_name}' not found in {filepath}") + + nn_model : nn.Module = model(*init_args_fn()) + if state_path is not None: + state_dict = torch.load(state_path) + nn_model.load_state_dict(state_dict) + + return import_from_model(nn_model, *inputs_args_fn(), dialect=dialect, ir_context=ir_context, **kwargs) From 1c6df474874ffae686727f75fe79d34c3dc1bf75 Mon Sep 17 00:00:00 2001 From: dchigarev Date: Mon, 20 Oct 2025 09:10:29 +0000 Subject: [PATCH 02/16] add .gitignore for python files Signed-off-by: dchigarev --- .gitignore | 9 +++++++++ 1 file changed, 9 insertions(+) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..6ada2b0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +# Python caches +__pycache__/ +*.py[cod] +.pytest_cache + +# Lighthouse Python module builds +dist/ +python/lighthouse*.egg-info/ +*.whl From 370e3c035882c0ec05393da3f53da97dfefed229 Mon Sep 17 00:00:00 2001 From: dchigarev Date: Mon, 20 Oct 2025 09:10:43 +0000 Subject: [PATCH 03/16] delete old torch ingress Signed-off-by: dchigarev --- ingress/Torch-MLIR/generate-mlir.py | 72 ------------------------ ingress/Torch-MLIR/generate-mlir.sh | 42 -------------- ingress/Torch-MLIR/install-virtualenv.sh | 58 ------------------- 3 files changed, 172 deletions(-) delete mode 100644 ingress/Torch-MLIR/generate-mlir.py delete mode 100755 ingress/Torch-MLIR/generate-mlir.sh delete mode 100755 ingress/Torch-MLIR/install-virtualenv.sh diff --git a/ingress/Torch-MLIR/generate-mlir.py b/ingress/Torch-MLIR/generate-mlir.py deleted file mode 100644 index 888e6dd..0000000 --- a/ingress/Torch-MLIR/generate-mlir.py +++ /dev/null @@ -1,72 +0,0 @@ -#!/usr/bin/env python3 - -import argparse -import os -import torch -import torch.nn as nn -from torch_mlir import fx -from torch_mlir.fx import OutputType - -# Parse arguments for selecting which model to load and which MLIR dialect to generate -def parse_args(): - parser = argparse.ArgumentParser(description="Generate MLIR for Torch-MLIR models.") - parser.add_argument( - "--model", - type=str, - required=True, - help="Path to the Torch model file.", - ) - parser.add_argument( - "--dialect", - type=str, - choices=["torch", "linalg", "stablehlo", "tosa"], - default="linalg", - help="MLIR dialect to generate.", - ) - return parser.parse_args() - -# Functin to load the Torch model -def load_torch_model(model_path): - - if not os.path.exists(model_path): - raise FileNotFoundError(f"Model file {model_path} does not exist.") - - model = torch.load(model_path) - return model - -# Function to generate MLIR from the Torch model -# See: https://github.com/MrSidims/PytorchExplorer/blob/main/backend/server.py#L237 -def generate_mlir(model, dialect): - - # Convert the Torch model to MLIR - output_type = None - if dialect == "torch": - output_type = OutputType.TORCH - elif dialect == "linalg": - output_type = OutputType.LINALG - elif dialect == "stablehlo": - output_type = OutputType.STABLEHLO - elif dialect == "tosa": - output_type = OutputType.TOSA - else: - raise ValueError(f"Unsupported dialect: {dialect}") - - module = fx.export_and_import(model, "", output_type=output_type) - return module - -# Main function to execute the script -def main(): - args = parse_args() - - # Load the Torch model - model = load_torch_model(args.model) - - # Generate MLIR from the model - mlir_module = generate_mlir(model, args.dialect) - - # Print or save the MLIR module - print(mlir_module) - -# Entry point for the script -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/ingress/Torch-MLIR/generate-mlir.sh b/ingress/Torch-MLIR/generate-mlir.sh deleted file mode 100755 index 0a079c6..0000000 --- a/ingress/Torch-MLIR/generate-mlir.sh +++ /dev/null @@ -1,42 +0,0 @@ -#!/usr/bin/env bash - -# Command line argument for model to load and MLIR dialect to generate -while getopts "m:d:" opt; do - case $opt in - m) - MODEL=$OPTARG - ;; - d) - DIALECT=$OPTARG - ;; - *) - echo "Usage: $0 [-m model] [-d dialect]" - exit 1 - ;; - esac -done -if [ -z "$MODEL" ]; then - echo "Model not specified. Please provide a model using -m option." - exit 1 -fi -if [ -z "$DIALECT" ]; then - DIALECT="linalg" -fi - -# Enable local virtualenv created by install-virtualenv.sh -if [ ! -d "torch-mlir-venv" ]; then - echo "Virtual environment not found. Please run install-virtualenv.sh first." - exit 1 -fi -source torch-mlir-venv/bin/activate - -# Find script directory -SCRIPT_DIR=$(dirname "$(readlink -f "$0")") - -# Use the Python script to generate MLIR -echo "Generating MLIR for model '$MODEL' with dialect '$DIALECT'..." -python $SCRIPT_DIR/generate-mlir.py --model "$MODEL" --dialect "$DIALECT" -if [ $? -ne 0 ]; then - echo "Failed to generate MLIR for model '$MODEL'." - exit 1 -fi diff --git a/ingress/Torch-MLIR/install-virtualenv.sh b/ingress/Torch-MLIR/install-virtualenv.sh deleted file mode 100755 index f48019c..0000000 --- a/ingress/Torch-MLIR/install-virtualenv.sh +++ /dev/null @@ -1,58 +0,0 @@ -#!/usr/bin/env bash - -# Command line argument for device type -# Options are: "AMD", "NVIDIA", "Intel", "CPU" -# If no argument is provided, it will detect the device type using lspci -if [ $# -eq 1 ]; then - DEVICE_TYPE=$1 -else - DEVICE_TYPE=$(lspci | grep VGA) -fi - - -# Install torch-mlir inside a virtual environment -echo "First ensure uv is installed" - -python -m pip install uv --upgrade - -echo "Preparing the virtual environment" -python -m uv venv torch-mlir-venv --python 3.12 - -#echo "Preparing the virtual environment" -#python3 -m venv torch-mlir-venv -source torch-mlir-venv/bin/activate -uv pip install --upgrade pip wheel - -# GPU support ("AMD", "NVIDIA", "Intel") -EXTRA_INDEX_URL="" -if [[ $DEVICE_TYPE == *"NVIDIA"* ]]; then - # This is the default pytorch - echo "Installing PyTorch for GPU: NVIDIA" -elif [[ $DEVICE_TYPE == *"AMD"* ]]; then - # https://rocm.docs.amd.com/projects/radeon/en/latest/docs/install/native_linux/install-pytorch.html - echo "Installing PyTorch for GPU: AMD" - EXTRA_INDEX_URL="--extra-index-url https://download.pytorch.org/whl/rocm6.4" -elif [[ $DEVICE_TYPE == *"Intel"* ]]; then - echo "Installing PyTorch for GPU: Intel" - EXTRA_INDEX_URL="--index-url https://download.pytorch.org/whl/xpu" -else - echo "No GPU support detected, installing CPU version." - EXTRA_INDEX_URL="--extra-index-url https://download.pytorch.org/whl/cpu" -fi - -echo "Installing torch and models" -uv pip install --pre torch torchvision torchrec transformers $EXTRA_INDEX_URL -if [ $? != 0 ]; then - exit 1 -fi - - -echo "Installing torch-mlir" -# This only seems to work on Ubuntu -uv pip install --pre torch-mlir \ - --extra-index-url https://download.pytorch.org/whl/nightly/cpu \ - -f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels -if [ $? != 0 ]; then - exit 1 -fi - From b984314ad777ba272f5b3de7bc695b8560c2ec8f Mon Sep 17 00:00:00 2001 From: dchigarev Date: Mon, 20 Oct 2025 10:04:13 +0000 Subject: [PATCH 04/16] add readme Signed-off-by: dchigarev --- python/lighthouse/ingress/README.md | 10 ++++++++++ python/lighthouse/ingress/torch/torch_import.py | 1 + 2 files changed, 11 insertions(+) create mode 100644 python/lighthouse/ingress/README.md diff --git a/python/lighthouse/ingress/README.md b/python/lighthouse/ingress/README.md new file mode 100644 index 0000000..98b7605 --- /dev/null +++ b/python/lighthouse/ingress/README.md @@ -0,0 +1,10 @@ +# Lighthouse Ingress + +The `lighthouse.ingress` module converts various input formats to MLIR modules. + +## Supported Formats + +#### Torch +Convert PyTorch models to MLIR using `lighthouse.ingress.torch`. + +**Examples:** [torch examples](https://github.com/llvm/lighthouse/tree/main/python/examples/ingress/torch) diff --git a/python/lighthouse/ingress/torch/torch_import.py b/python/lighthouse/ingress/torch/torch_import.py index f64b9ff..9d15aad 100644 --- a/python/lighthouse/ingress/torch/torch_import.py +++ b/python/lighthouse/ingress/torch/torch_import.py @@ -97,6 +97,7 @@ def import_from_model( text_module = str(module) if ir_context is None: return text_module + # Cross boundary from torch-mlir's mlir to environment's mlir return ir.Module.parse(text_module, context=ir_context) From 2897e0cf50383560f9d1d243eee31a8120585c3c Mon Sep 17 00:00:00 2001 From: dchigarev Date: Mon, 20 Oct 2025 10:53:40 +0000 Subject: [PATCH 05/16] fix 02-* tutorial Signed-off-by: dchigarev --- python/examples/ingress/torch/02-dummy-mlp-from-model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/examples/ingress/torch/02-dummy-mlp-from-model.py b/python/examples/ingress/torch/02-dummy-mlp-from-model.py index fb4de89..f46f07e 100644 --- a/python/examples/ingress/torch/02-dummy-mlp-from-model.py +++ b/python/examples/ingress/torch/02-dummy-mlp-from-model.py @@ -20,7 +20,7 @@ from lighthouse.ingress.torch import import_from_model # Import a sample model definition -from .DummyMLP.model import DummyMLP +from DummyMLP.model import DummyMLP # Step 1: Instantiate a model and prepare sample input model = DummyMLP() From fa7d1de3a57a0a0c583b7ce0829038e879c94d72 Mon Sep 17 00:00:00 2001 From: dchigarev Date: Mon, 20 Oct 2025 11:31:53 +0000 Subject: [PATCH 06/16] fix grammar & typos Signed-off-by: dchigarev --- .../ingress/torch/01-dummy-mlp-from-file.py | 21 ++++++++++++------- .../ingress/torch/02-dummy-mlp-from-model.py | 14 +++++++++---- python/lighthouse/ingress/README.md | 2 +- .../lighthouse/ingress/torch/torch_import.py | 8 +++---- 4 files changed, 28 insertions(+), 17 deletions(-) diff --git a/python/examples/ingress/torch/01-dummy-mlp-from-file.py b/python/examples/ingress/torch/01-dummy-mlp-from-file.py index 421b712..cdafa60 100644 --- a/python/examples/ingress/torch/01-dummy-mlp-from-file.py +++ b/python/examples/ingress/torch/01-dummy-mlp-from-file.py @@ -1,14 +1,14 @@ """ -Example demonstrating how to load a PyTorch model to MLIR using Lighthouse, -without instantiating the model on our side. +Example demonstrating how to load a PyTorch model to MLIR using Lighthouse +without instantiating the model on the user's side. The script uses 'lighthouse.ingress.torch.import_from_file' function that -takes a filepath to a Python file containing the model definition, along with +takes a path to a Python file containing the model definition, along with the names of functions to get model init arguments and sample inputs. The function -imports the model class on its own, instantiates it, and passes it torch_mlir +imports the model class on its own, instantiates it, and passes it to torch_mlir to get a MLIR module in the specified dialect. -The script uses model from 'DummyMLP/model.py' as an example. +The script uses the model from 'DummyMLP/model.py' as an example. """ import os @@ -21,7 +21,7 @@ # Lighthouse imports from lighthouse.ingress.torch import import_from_file -# Step 1: Setup paths to locate the model definition file +# Step 1: Set up paths to locate the model definition file script_dir = Path(os.path.dirname(os.path.abspath(__file__))) model_path = script_dir / "DummyMLP" / "model.py" @@ -29,18 +29,23 @@ # Step 2: Convert PyTorch model to MLIR # Conversion step where Lighthouse: -# - Loads the DummyMLP class and instantiates with arguments obtained from 'get_init_inputs()' +# - Loads the DummyMLP class and instantiates it with arguments obtained from 'get_init_inputs()' # - Calls get_sample_inputs() to get sample input tensors for shape inference # - Converts PyTorch model to linalg-on-tensors dialect operations using torch_mlir mlir_module_ir : ir.Module = import_from_file( model_path, # Path to the Python file containing the model model_class_name="DummyMLP", # Name of the PyTorch nn.Module class to convert init_args_fn_name="get_init_inputs", # Function that returns args for model.__init__() - inputs_args_fn_name="get_sample_inputs", # Function that returns sample inputs for tracing + inputs_args_fn_name="get_sample_inputs", # Function that returns sample inputs to pass to 'model(...)' dialect="linalg-on-tensors", # Target MLIR dialect (linalg ops on tensor types) ir_context=ir_context # MLIR context for the conversion ) +# The model is converted at this point. You can now convert the MLIR module +# to a text form (e.g. 'str(mlir_module_ir)') and save it to a file. +# Futher are optional MLIR processing steps to give you an idea of what can +# also be done with the MLIR module. + # Step 3: Extract the main function operation from the MLIR module and print its metadata func_op : func.FuncOp = mlir_module_ir.operation.regions[0].blocks[0].operations[0] print(f"entry-point name: {func_op.name}") diff --git a/python/examples/ingress/torch/02-dummy-mlp-from-model.py b/python/examples/ingress/torch/02-dummy-mlp-from-model.py index f46f07e..88290c6 100644 --- a/python/examples/ingress/torch/02-dummy-mlp-from-model.py +++ b/python/examples/ingress/torch/02-dummy-mlp-from-model.py @@ -2,12 +2,12 @@ Example demonstrating how to load an already instantiated PyTorch model to MLIR using Lighthouse. -The script uses 'lighthouse.ingress.torch.import_from_model' function that -takes an already instantiated PyTorch model, along with its sample inputs. +The script uses the 'lighthouse.ingress.torch.import_from_model' function that +takes a PyTorch model that has already been instantiated, along with its sample inputs. The function passes the model to torch_mlir to get a MLIR module in the specified dialect. -The script uses model from 'DummyMLP/model.py' as an example. +The script uses a model from 'DummyMLP/model.py' as an example. """ import torch @@ -27,13 +27,19 @@ sample_input = torch.randn(1, 10) ir_context = ir.Context() -# Step 2: Convert PyTorch model to MLIR +# Step 2: Convert the PyTorch model to MLIR mlir_module_ir : ir.Module = import_from_model( model, sample_args=(sample_input,), ir_context=ir_context ) +# The PyTorch model is now converted to MLIR at this point. You can now convert +# the MLIR module to a text form (e.g. 'str(mlir_module_ir)') and save it to a file. +# +# Further are optional MLIR processing steps to give you an idea of what can +# also be done with the MLIR module. + # Step 3: Extract the main function operation from the MLIR module and print its metadata func_op : func.FuncOp = mlir_module_ir.operation.regions[0].blocks[0].operations[0] print(f"entry-point name: {func_op.name}") diff --git a/python/lighthouse/ingress/README.md b/python/lighthouse/ingress/README.md index 98b7605..61236ef 100644 --- a/python/lighthouse/ingress/README.md +++ b/python/lighthouse/ingress/README.md @@ -5,6 +5,6 @@ The `lighthouse.ingress` module converts various input formats to MLIR modules. ## Supported Formats #### Torch -Convert PyTorch models to MLIR using `lighthouse.ingress.torch`. +Converts PyTorch models to MLIR using `lighthouse.ingress.torch`. **Examples:** [torch examples](https://github.com/llvm/lighthouse/tree/main/python/examples/ingress/torch) diff --git a/python/lighthouse/ingress/torch/torch_import.py b/python/lighthouse/ingress/torch/torch_import.py index 9d15aad..b168901 100644 --- a/python/lighthouse/ingress/torch/torch_import.py +++ b/python/lighthouse/ingress/torch/torch_import.py @@ -164,7 +164,7 @@ def get_inputs(): return (torch.randn(1, 10),) ``` - The import script would look line: + The import script would look like: >>> from lighthouse.ingress.torch_import import import_from_file >>> # option 1: get MLIR module as a string >>> mlir_module : str = import_from_file( @@ -200,9 +200,9 @@ def get_inputs(): if init_args_fn_name is None: init_args_fn = lambda *args, **kwargs: () else: - init_args_fn = getattr(module, init_args_fn_name, None) - if init_args_fn is None: - raise ValueError(f"Init args function '{init_args_fn_name}' not found in {filepath}") + init_args_fn = getattr(module, init_args_fn_name, None) + if init_args_fn is None: + raise ValueError(f"Init args function '{init_args_fn_name}' not found in {filepath}") inputs_args_fn = getattr(module, inputs_args_fn_name, None) if inputs_args_fn is None: From eaf8c9e8bbaf65d7d08ac9ffd60efa8ada8711f3 Mon Sep 17 00:00:00 2001 From: dchigarev Date: Tue, 21 Oct 2025 12:35:31 +0000 Subject: [PATCH 07/16] change docstrings from numpy to google format Signed-off-by: dchigarev --- .../lighthouse/ingress/torch/torch_import.py | 197 ++++++++---------- 1 file changed, 90 insertions(+), 107 deletions(-) diff --git a/python/lighthouse/ingress/torch/torch_import.py b/python/lighthouse/ingress/torch/torch_import.py index b168901..8b5bb9a 100644 --- a/python/lighthouse/ingress/torch/torch_import.py +++ b/python/lighthouse/ingress/torch/torch_import.py @@ -31,55 +31,46 @@ def import_from_model( ir_context : ir.Context | None = None, **kwargs, ) -> str | ir.Module: - """ - Import a PyTorch nn.Module into MLIR. + """Import a PyTorch nn.Module into MLIR. The function uses torch-mlir's FX importer to convert the given PyTorch model into an MLIR module in the specified dialect. The user has to provide sample input arguments (e.g. a torch.Tensor with the correct shape). - Parameters - ---------- - model : nn.Module - The PyTorch model to import. - sample_args : Iterable - Sample input arguments to the model. - sample_kwargs : Mapping, optional - Sample keyword arguments to the model. - dialect : torch_mlir.fx.OutputType | {"linalg-on-tensors", "torch", "tosa"}, default: OutputType.LINALG_ON_TENSORS - The target dialect for the imported MLIR module. - ir_context : ir.Context, optional - An optional MLIR context to use for parsing the module. - If not provided, the module is returned as a string. - **kwargs - Additional keyword arguments passed to the fx.export_and_import function. + Args: + model (nn.Module): The PyTorch model to import. + sample_args (Iterable): Sample input arguments to the model. + sample_kwargs (Mapping, optional): Sample keyword arguments to the model. + dialect (torch_mlir.fx.OutputType | {"linalg-on-tensors", "torch", "tosa"}): + The target dialect for the imported MLIR module. Defaults to + ``OutputType.LINALG_ON_TENSORS``. + ir_context (ir.Context, optional): An optional MLIR context to use for parsing + the module. If not provided, the module is returned as a string. + **kwargs: Additional keyword arguments passed to the ``torch_mlir.fx.export_and_import`` function. - Returns - ------- - str | ir.Module - The imported MLIR module as a string or an ir.Module if `ir_context` is provided. + Returns: + str | ir.Module: The imported MLIR module as a string or an ir.Module if `ir_context` is provided. - Examples - -------- - >>> import torch - >>> import torch.nn as nn - >>> from lighthouse.ingress.torch_import import import_from_model - >>> class SimpleModel(nn.Module): - ... def __init__(self): - ... super().__init__() - ... self.fc = nn.Linear(10, 5) - ... def forward(self, x): - ... return self.fc(x) - >>> model = SimpleModel() - >>> sample_input = (torch.randn(1, 10),) - >>> # - >>> # option 1: get MLIR module as a string - >>> mlir_module : str = import_from_model(model, sample_input, dialect="linalg-on-tensors") - >>> print(mlir_module) # prints the MLIR module in linalg-on-tensors dialect - >>> # option 2: get MLIR module as an ir.Module - >>> ir_context = ir.Context() - >>> mlir_module_ir : ir.Module = import_from_model(model, sample_input, dialect="tosa", ir_context=ir_context) - >>> # ... run pm.Pipeline on the ir.Module ... + Examples: + >>> import torch + >>> import torch.nn as nn + >>> from lighthouse.ingress.torch_import import import_from_model + >>> class SimpleModel(nn.Module): + ... def __init__(self): + ... super().__init__() + ... self.fc = nn.Linear(10, 5) + ... def forward(self, x): + ... return self.fc(x) + >>> model = SimpleModel() + >>> sample_input = (torch.randn(1, 10),) + >>> # + >>> # option 1: get MLIR module as a string + >>> mlir_module : str = import_from_model(model, sample_input, dialect="linalg-on-tensors") + >>> print(mlir_module) # prints the MLIR module in linalg-on-tensors dialect + >>> # option 2: get MLIR module as an ir.Module + >>> ir_context = ir.Context() + >>> mlir_module_ir : ir.Module = import_from_model(model, sample_input, dialect="tosa", ir_context=ir_context) + >>> # ... run pm.Pipeline on the ir.Module ... """ if dialect == "linalg": raise ValueError( @@ -111,79 +102,71 @@ def import_from_file( ir_context : ir.Context | None = None, **kwargs, ) -> str | ir.Module: - """ - Load a PyTorch nn.Module from a file and import it into MLIR. + """Load a PyTorch nn.Module from a file and import it into MLIR. The function takes a `filepath` to a Python file containing the model definition, along with the names of functions to get model init arguments and sample inputs. The function imports the model class on its own, instantiates it, and passes it ``torch_mlir`` to get a MLIR module in the specified `dialect`. - Parameters - ---------- - filepath : str | Path - Path to the Python file containing the model definition. - model_class_name : str, default: "Model" - The name of the model class in the file. - init_args_fn_name : str | None, default: "get_init_inputs" - The name of the function in the file that returns the arguments for - initializing the model. If None, the model is initialized without arguments. - inputs_args_fn_name : str, default: "get_inputs" - The name of the function in the file that returns the sample input arguments - for the model. - state_path : str | Path | None, default: None - Optional path to a file containing the model's ``state_dict``. - dialect: torch_mlir.fx.OutputType | {"linalg-on-tensors", "torch", "tosa"}, default: OutputType.LINALG_ON_TENSORS - The target dialect for the imported MLIR module. - ir_context : ir.Context, optional - An optional MLIR context to use for parsing the module. - If not provided, the module is returned as a string. - **kwargs - Additional keyword arguments passed to the fx.export_and_import function. - - Returns - ------- - str | ir.Module - The imported MLIR module as a string or an ir.Module if `ir_context` is provided. + Args: + filepath (str | Path): Path to the Python file containing the model definition. + model_class_name (str, optional): The name of the model class in the file. + Defaults to "Model". + init_args_fn_name (str | None, optional): The name of the function in the file + that returns the arguments for initializing the model. If None, the model + is initialized without arguments. Defaults to "get_init_inputs". + inputs_args_fn_name (str, optional): The name of the function in the file that + returns the sample input arguments for the model. Defaults to "get_inputs". + state_path (str | Path | None, optional): Optional path to a file containing + the model's ``state_dict``. Defaults to None. + dialect (torch_mlir.fx.OutputType | {"linalg-on-tensors", "torch", "tosa"}, optional): + The target dialect for the imported MLIR module. Defaults to + ``OutputType.LINALG_ON_TENSORS``. + ir_context (ir.Context, optional): An optional MLIR context to use for parsing + the module. If not provided, the module is returned as a string. + **kwargs: Additional keyword arguments passed to the ``torch_mlir.fx.export_and_import`` function. - Examples - -------- - Given a file `path/to/model_file.py` with the following content: - ```python - import torch - import torch.nn as nn - - class MyModel(nn.Module): - def __init__(self): - super().__init__() - self.fc = nn.Linear(10, 5) - def forward(self, x): - return self.fc(x) + Returns: + str | ir.Module: The imported MLIR module as a string or an ir.Module if `ir_context` is provided. - def get_inputs(): - return (torch.randn(1, 10),) - ``` - - The import script would look like: - >>> from lighthouse.ingress.torch_import import import_from_file - >>> # option 1: get MLIR module as a string - >>> mlir_module : str = import_from_file( - ... "path/to/model_file.py", - ... model_class_name="MyModel", - ... init_args_fn_name=None, - ... dialect="linalg-on-tensors" - ... ) - >>> print(mlir_module) # prints the MLIR module in linalg-on-tensors dialect - >>> # option 2: get MLIR module as an ir.Module - >>> ir_context = ir.Context() - >>> mlir_module_ir : ir.Module = import_from_file( - ... "path/to/model_file.py", - ... model_class_name="MyModel", - ... init_args_fn_name=None, - ... dialect="linalg-on-tensors", - ... ir_context=ir_context - ... ) - >>> # ... run pm.Pipeline on the ir.Module ... + Examples: + Given a file `path/to/model_file.py` with the following content: + ```python + import torch + import torch.nn as nn + + class MyModel(nn.Module): + def __init__(self): + super().__init__() + self.fc = nn.Linear(10, 5) + def forward(self, x): + return self.fc(x) + + def get_inputs(): + return (torch.randn(1, 10),) + ``` + + The import script would look like: + >>> from lighthouse.ingress.torch_import import import_from_file + >>> # option 1: get MLIR module as a string + >>> mlir_module : str = import_from_file( + ... "path/to/model_file.py", + ... model_class_name="MyModel", + ... init_args_fn_name=None, + ... dialect="linalg-on-tensors" + ... ) + >>> print(mlir_module) # prints the MLIR module in linalg-on-tensors dialect + >>> # option 2: get MLIR module as an ir.Module + >>> ir_context = ir.Context() + >>> mlir_module_ir : ir.Module = import_from_file( + ... "path/to/model_file.py", + ... model_class_name="MyModel", + ... init_args_fn_name=None, + ... dialect="linalg-on-tensors", + ... ir_context=ir_context + ... ) + >>> # ... run pm.Pipeline on the ir.Module ... """ if isinstance(filepath, str): filepath = Path(filepath) From 889314df0dbf03dacc20728b430edbd266c0bf72 Mon Sep 17 00:00:00 2001 From: dchigarev Date: Tue, 21 Oct 2025 13:02:17 +0000 Subject: [PATCH 08/16] add kwarg functions to 'import_from_file' fn Signed-off-by: dchigarev --- .../lighthouse/ingress/torch/torch_import.py | 60 ++++++++++++++----- python/lighthouse/ingress/torch/utils.py | 47 +++++++++++++++ 2 files changed, 93 insertions(+), 14 deletions(-) create mode 100644 python/lighthouse/ingress/torch/utils.py diff --git a/python/lighthouse/ingress/torch/torch_import.py b/python/lighthouse/ingress/torch/torch_import.py index 8b5bb9a..92f2f9e 100644 --- a/python/lighthouse/ingress/torch/torch_import.py +++ b/python/lighthouse/ingress/torch/torch_import.py @@ -3,6 +3,8 @@ from pathlib import Path from typing import Iterable, Mapping +from lighthouse.ingress.torch.utils import load_and_run_callable + try: import torch import torch.nn as nn @@ -96,7 +98,9 @@ def import_from_file( filepath: str | Path, model_class_name: str = "Model", init_args_fn_name: str | None = "get_init_inputs", - inputs_args_fn_name: str = "get_inputs", + init_kwargs_fn_name: str | None = None, + sample_args_fn_name: str = "get_inputs", + sample_kwargs_fn_name: str | None = None, state_path : str | Path | None = None, dialect : OutputType | str = OutputType.LINALG_ON_TENSORS, ir_context : ir.Context | None = None, @@ -116,8 +120,13 @@ def import_from_file( init_args_fn_name (str | None, optional): The name of the function in the file that returns the arguments for initializing the model. If None, the model is initialized without arguments. Defaults to "get_init_inputs". - inputs_args_fn_name (str, optional): The name of the function in the file that + init_kwargs_fn_name (str | None, optional): The name of the function in the file + that returns the keyword arguments for initializing the model. If None, the model + is initialized without keyword arguments. + sample_args_fn_name (str, optional): The name of the function in the file that returns the sample input arguments for the model. Defaults to "get_inputs". + sample_kwargs_fn_name (str, optional): The name of the function in the file that + returns the sample keyword input arguments for the model. Defaults to None. state_path (str | Path | None, optional): Optional path to a file containing the model's ``state_dict``. Defaults to None. dialect (torch_mlir.fx.OutputType | {"linalg-on-tensors", "torch", "tosa"}, optional): @@ -180,20 +189,43 @@ def get_inputs(): if model is None: raise ValueError(f"Model class '{model_class_name}' not found in {filepath}") - if init_args_fn_name is None: - init_args_fn = lambda *args, **kwargs: () - else: - init_args_fn = getattr(module, init_args_fn_name, None) - if init_args_fn is None: - raise ValueError(f"Init args function '{init_args_fn_name}' not found in {filepath}") - - inputs_args_fn = getattr(module, inputs_args_fn_name, None) - if inputs_args_fn is None: - raise ValueError(f"Inputs args function '{inputs_args_fn_name}' not found in {filepath}") + model_init_args = load_and_run_callable( + module, + init_args_fn_name, + raise_on_missing=False, + default=tuple(), + error_msg=f"Init args function '{init_args_fn_name}' not found in {filepath}" + ) + model_init_kwargs = load_and_run_callable( + module, + init_kwargs_fn_name, + raise_on_missing=False, + default=tuple(), + error_msg=f"Init kwargs function '{init_kwargs_fn_name}' not found in {filepath}" + ) + sample_args = load_and_run_callable( + module, + sample_args_fn_name, + f"Sample args function '{sample_args_fn_name}' not found in {filepath}" + ) + sample_kwargs = load_and_run_callable( + module, + sample_kwargs_fn_name, + raise_on_missing=False, + default={}, + error_msg=f"Sample kwargs function '{sample_kwargs_fn_name}' not found in {filepath}" + ) - nn_model : nn.Module = model(*init_args_fn()) + nn_model : nn.Module = model(*model_init_args, **model_init_kwargs) if state_path is not None: state_dict = torch.load(state_path) nn_model.load_state_dict(state_dict) - return import_from_model(nn_model, *inputs_args_fn(), dialect=dialect, ir_context=ir_context, **kwargs) + return import_from_model( + nn_model, + sample_args=sample_args, + sample_kwargs=sample_kwargs, + dialect=dialect, + ir_context=ir_context, + **kwargs, + ) diff --git a/python/lighthouse/ingress/torch/utils.py b/python/lighthouse/ingress/torch/utils.py new file mode 100644 index 0000000..cd957fc --- /dev/null +++ b/python/lighthouse/ingress/torch/utils.py @@ -0,0 +1,47 @@ +from types import ModuleType +from typing import Any + + +def load_and_run_callable( + module: ModuleType, + symbol_name: str | None, + raise_on_missing: bool = True, + default: Any | None = None, + error_msg: str | None = None, +): + """Helper to load and run a callable from a module by its symbol name. + + Args: + module (ModuleType): The python module to load the callable from. + symbol_name (str | None): The name of the callable symbol to load. + If not specified (None), the function will either return the `default` value + or raise an error based on `raise_on_missing`. + raise_on_missing (bool): Whether to raise an error if the symbol is missing + or not specified. Default is True. + default (Any | None): The default value to return if the symbol is missing + and `raise_on_missing` is False. + error_msg (str | None): Custom error message to use when raising an error + for missing symbol. If not provided, a default message will be used. + + Returns: + Any: The result of calling the loaded callable or `default` if the symbol is + missing and `raise_on_missing` is False. + """ + if symbol_name is None: + if raise_on_missing: + if error_msg: + raise ValueError(error_msg) + raise ValueError("No symbol name provided") + return default + func = getattr(module, symbol_name, None) + if func is None: + if raise_on_missing: + if error_msg: + raise ValueError(error_msg) + raise ValueError( + f"Symbol '{symbol_name}' not found in module '{module.__name__}'" + ) + return default + if not callable(func): + raise ValueError(f"Symbol '{symbol_name}' is not callable") + return func() From 0a188e564e19b7d2138210a11b061f480b24e1ce Mon Sep 17 00:00:00 2001 From: dchigarev Date: Tue, 21 Oct 2025 13:10:27 +0000 Subject: [PATCH 09/16] fix incode-comments and doc-strings wording Signed-off-by: dchigarev --- .../ingress/torch/01-dummy-mlp-from-file.py | 14 +++++++------- .../ingress/torch/02-dummy-mlp-from-model.py | 8 ++++---- python/lighthouse/ingress/torch/torch_import.py | 4 +--- 3 files changed, 12 insertions(+), 14 deletions(-) diff --git a/python/examples/ingress/torch/01-dummy-mlp-from-file.py b/python/examples/ingress/torch/01-dummy-mlp-from-file.py index cdafa60..9b3c827 100644 --- a/python/examples/ingress/torch/01-dummy-mlp-from-file.py +++ b/python/examples/ingress/torch/01-dummy-mlp-from-file.py @@ -12,7 +12,7 @@ """ import os -from pyparsing import Path +from pathlib import Path # MLIR infrastructure imports (only needed if you want to manipulate the MLIR module) import mlir.dialects.func as func @@ -41,9 +41,10 @@ ir_context=ir_context # MLIR context for the conversion ) -# The model is converted at this point. You can now convert the MLIR module -# to a text form (e.g. 'str(mlir_module_ir)') and save it to a file. -# Futher are optional MLIR processing steps to give you an idea of what can +# The PyTorch model is now converted to MLIR at this point. You can now convert +# the MLIR module to a text form (e.g. 'str(mlir_module_ir)') and save it to a file. +# +# The following optional MLIR-processing steps are to give you an idea of what can # also be done with the MLIR module. # Step 3: Extract the main function operation from the MLIR module and print its metadata @@ -53,11 +54,10 @@ # Step 4: Apply some MLIR passes using a PassManager pm = passmanager.PassManager(context=ir_context) - -pm.add("one-shot-bufferize") pm.add("linalg-specialize-generic-ops") +pm.add("one-shot-bufferize") pm.run(mlir_module_ir.operation) # Step 5: Output the final MLIR -print("\n\nModule dump after running pm.Pipeline:") +print("\n\nModule dump after running the pipeline:") mlir_module_ir.dump() diff --git a/python/examples/ingress/torch/02-dummy-mlp-from-model.py b/python/examples/ingress/torch/02-dummy-mlp-from-model.py index 88290c6..fc62150 100644 --- a/python/examples/ingress/torch/02-dummy-mlp-from-model.py +++ b/python/examples/ingress/torch/02-dummy-mlp-from-model.py @@ -37,7 +37,7 @@ # The PyTorch model is now converted to MLIR at this point. You can now convert # the MLIR module to a text form (e.g. 'str(mlir_module_ir)') and save it to a file. # -# Further are optional MLIR processing steps to give you an idea of what can +# The following optional MLIR-processing steps are to give you an idea of what can # also be done with the MLIR module. # Step 3: Extract the main function operation from the MLIR module and print its metadata @@ -45,12 +45,12 @@ print(f"entry-point name: {func_op.name}") print(f"entry-point type: {func_op.type}") -# Step 4: Apply some MLIR passes using a PassManager (optional) +# Step 4: Apply some MLIR passes using a PassManager pm = passmanager.PassManager(context=ir_context) -pm.add("one-shot-bufferize") pm.add("linalg-specialize-generic-ops") +pm.add("one-shot-bufferize") pm.run(mlir_module_ir.operation) # Step 5: Output the final MLIR -print("\n\nModule dump after running pm.Pipeline:") +print("\n\nModule dump after running the pipeline:") mlir_module_ir.dump() diff --git a/python/lighthouse/ingress/torch/torch_import.py b/python/lighthouse/ingress/torch/torch_import.py index 92f2f9e..64d0236 100644 --- a/python/lighthouse/ingress/torch/torch_import.py +++ b/python/lighthouse/ingress/torch/torch_import.py @@ -72,7 +72,6 @@ def import_from_model( >>> # option 2: get MLIR module as an ir.Module >>> ir_context = ir.Context() >>> mlir_module_ir : ir.Module = import_from_model(model, sample_input, dialect="tosa", ir_context=ir_context) - >>> # ... run pm.Pipeline on the ir.Module ... """ if dialect == "linalg": raise ValueError( @@ -111,7 +110,7 @@ def import_from_file( The function takes a `filepath` to a Python file containing the model definition, along with the names of functions to get model init arguments and sample inputs. The function imports the model class on its own, instantiates it, and passes - it ``torch_mlir`` to get a MLIR module in the specified `dialect`. + it to ``torch_mlir`` to get a MLIR module in the specified `dialect`. Args: filepath (str | Path): Path to the Python file containing the model definition. @@ -175,7 +174,6 @@ def get_inputs(): ... dialect="linalg-on-tensors", ... ir_context=ir_context ... ) - >>> # ... run pm.Pipeline on the ir.Module ... """ if isinstance(filepath, str): filepath = Path(filepath) From 1068bf93a39ea85353652b7d1d4db69809752172 Mon Sep 17 00:00:00 2001 From: dchigarev Date: Tue, 21 Oct 2025 14:20:22 +0000 Subject: [PATCH 10/16] fix tutorials Signed-off-by: dchigarev --- .../ingress/torch/01-dummy-mlp-from-file.py | 2 +- .../lighthouse/ingress/torch/torch_import.py | 13 ++--- python/lighthouse/ingress/torch/utils.py | 53 ++++++++++--------- 3 files changed, 34 insertions(+), 34 deletions(-) diff --git a/python/examples/ingress/torch/01-dummy-mlp-from-file.py b/python/examples/ingress/torch/01-dummy-mlp-from-file.py index 9b3c827..efa2983 100644 --- a/python/examples/ingress/torch/01-dummy-mlp-from-file.py +++ b/python/examples/ingress/torch/01-dummy-mlp-from-file.py @@ -36,7 +36,7 @@ model_path, # Path to the Python file containing the model model_class_name="DummyMLP", # Name of the PyTorch nn.Module class to convert init_args_fn_name="get_init_inputs", # Function that returns args for model.__init__() - inputs_args_fn_name="get_sample_inputs", # Function that returns sample inputs to pass to 'model(...)' + sample_args_fn_name="get_sample_inputs", # Function that returns sample inputs to pass to 'model(...)' dialect="linalg-on-tensors", # Target MLIR dialect (linalg ops on tensor types) ir_context=ir_context # MLIR context for the conversion ) diff --git a/python/lighthouse/ingress/torch/torch_import.py b/python/lighthouse/ingress/torch/torch_import.py index 64d0236..4ed367f 100644 --- a/python/lighthouse/ingress/torch/torch_import.py +++ b/python/lighthouse/ingress/torch/torch_import.py @@ -3,7 +3,7 @@ from pathlib import Path from typing import Iterable, Mapping -from lighthouse.ingress.torch.utils import load_and_run_callable +from lighthouse.ingress.torch.utils import load_and_run_callable, maybe_load_and_run_callable try: import torch @@ -187,18 +187,16 @@ def get_inputs(): if model is None: raise ValueError(f"Model class '{model_class_name}' not found in {filepath}") - model_init_args = load_and_run_callable( + model_init_args = maybe_load_and_run_callable( module, init_args_fn_name, - raise_on_missing=False, default=tuple(), error_msg=f"Init args function '{init_args_fn_name}' not found in {filepath}" ) - model_init_kwargs = load_and_run_callable( + model_init_kwargs = maybe_load_and_run_callable( module, init_kwargs_fn_name, - raise_on_missing=False, - default=tuple(), + default={}, error_msg=f"Init kwargs function '{init_kwargs_fn_name}' not found in {filepath}" ) sample_args = load_and_run_callable( @@ -206,10 +204,9 @@ def get_inputs(): sample_args_fn_name, f"Sample args function '{sample_args_fn_name}' not found in {filepath}" ) - sample_kwargs = load_and_run_callable( + sample_kwargs = maybe_load_and_run_callable( module, sample_kwargs_fn_name, - raise_on_missing=False, default={}, error_msg=f"Sample kwargs function '{sample_kwargs_fn_name}' not found in {filepath}" ) diff --git a/python/lighthouse/ingress/torch/utils.py b/python/lighthouse/ingress/torch/utils.py index cd957fc..9a464b3 100644 --- a/python/lighthouse/ingress/torch/utils.py +++ b/python/lighthouse/ingress/torch/utils.py @@ -4,44 +4,47 @@ def load_and_run_callable( module: ModuleType, - symbol_name: str | None, - raise_on_missing: bool = True, - default: Any | None = None, + symbol_name: str, error_msg: str | None = None, ): """Helper to load and run a callable from a module by its symbol name. Args: module (ModuleType): The python module to load the callable from. - symbol_name (str | None): The name of the callable symbol to load. - If not specified (None), the function will either return the `default` value - or raise an error based on `raise_on_missing`. - raise_on_missing (bool): Whether to raise an error if the symbol is missing - or not specified. Default is True. - default (Any | None): The default value to return if the symbol is missing - and `raise_on_missing` is False. + symbol_name (str): The name of the callable symbol to load. error_msg (str | None): Custom error message to use when raising an error for missing symbol. If not provided, a default message will be used. Returns: - Any: The result of calling the loaded callable or `default` if the symbol is - missing and `raise_on_missing` is False. + Any: The result of calling the loaded callable. """ - if symbol_name is None: - if raise_on_missing: - if error_msg: - raise ValueError(error_msg) - raise ValueError("No symbol name provided") - return default func = getattr(module, symbol_name, None) if func is None: - if raise_on_missing: - if error_msg: - raise ValueError(error_msg) - raise ValueError( - f"Symbol '{symbol_name}' not found in module '{module.__name__}'" - ) - return default + if error_msg: + raise ValueError(error_msg) + raise ValueError( + f"Symbol '{symbol_name}' not found in module '{module.__name__}'" + ) if not callable(func): raise ValueError(f"Symbol '{symbol_name}' is not callable") return func() + + +def maybe_load_and_run_callable( + module: ModuleType, + symbol_name: str | None, + default: Any, + error_msg: str | None = None, +): + """Helper to conditionally load and run a callable from a module by its symbol name. + + If `symbol_name` is None, the function returns the provided default value. Otherwise + it calls ``load_and_run_callable`` with the provided arguments. + """ + if symbol_name is None: + return default + return load_and_run_callable( + module, + symbol_name, + error_msg=error_msg + ) From 0912de15b458a78a88f111045fe0e54618ae83a1 Mon Sep 17 00:00:00 2001 From: dchigarev Date: Tue, 21 Oct 2025 14:27:58 +0000 Subject: [PATCH 11/16] fix whitespaces for type annotations according to pep Signed-off-by: dchigarev --- .../ingress/torch/01-dummy-mlp-from-file.py | 4 ++-- .../ingress/torch/02-dummy-mlp-from-model.py | 4 ++-- python/lighthouse/ingress/torch/torch_import.py | 16 ++++++++-------- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/python/examples/ingress/torch/01-dummy-mlp-from-file.py b/python/examples/ingress/torch/01-dummy-mlp-from-file.py index efa2983..bf7c15c 100644 --- a/python/examples/ingress/torch/01-dummy-mlp-from-file.py +++ b/python/examples/ingress/torch/01-dummy-mlp-from-file.py @@ -32,7 +32,7 @@ # - Loads the DummyMLP class and instantiates it with arguments obtained from 'get_init_inputs()' # - Calls get_sample_inputs() to get sample input tensors for shape inference # - Converts PyTorch model to linalg-on-tensors dialect operations using torch_mlir -mlir_module_ir : ir.Module = import_from_file( +mlir_module_ir: ir.Module = import_from_file( model_path, # Path to the Python file containing the model model_class_name="DummyMLP", # Name of the PyTorch nn.Module class to convert init_args_fn_name="get_init_inputs", # Function that returns args for model.__init__() @@ -48,7 +48,7 @@ # also be done with the MLIR module. # Step 3: Extract the main function operation from the MLIR module and print its metadata -func_op : func.FuncOp = mlir_module_ir.operation.regions[0].blocks[0].operations[0] +func_op: func.FuncOp = mlir_module_ir.operation.regions[0].blocks[0].operations[0] print(f"entry-point name: {func_op.name}") print(f"entry-point type: {func_op.type}") diff --git a/python/examples/ingress/torch/02-dummy-mlp-from-model.py b/python/examples/ingress/torch/02-dummy-mlp-from-model.py index fc62150..246deda 100644 --- a/python/examples/ingress/torch/02-dummy-mlp-from-model.py +++ b/python/examples/ingress/torch/02-dummy-mlp-from-model.py @@ -28,7 +28,7 @@ ir_context = ir.Context() # Step 2: Convert the PyTorch model to MLIR -mlir_module_ir : ir.Module = import_from_model( +mlir_module_ir: ir.Module = import_from_model( model, sample_args=(sample_input,), ir_context=ir_context @@ -41,7 +41,7 @@ # also be done with the MLIR module. # Step 3: Extract the main function operation from the MLIR module and print its metadata -func_op : func.FuncOp = mlir_module_ir.operation.regions[0].blocks[0].operations[0] +func_op: func.FuncOp = mlir_module_ir.operation.regions[0].blocks[0].operations[0] print(f"entry-point name: {func_op.name}") print(f"entry-point type: {func_op.type}") diff --git a/python/lighthouse/ingress/torch/torch_import.py b/python/lighthouse/ingress/torch/torch_import.py index 4ed367f..dfaccd0 100644 --- a/python/lighthouse/ingress/torch/torch_import.py +++ b/python/lighthouse/ingress/torch/torch_import.py @@ -27,10 +27,10 @@ def import_from_model( model: nn.Module, - sample_args : Iterable, - sample_kwargs : Mapping = None, - dialect : OutputType | str = OutputType.LINALG_ON_TENSORS, - ir_context : ir.Context | None = None, + sample_args: Iterable, + sample_kwargs: Mapping = None, + dialect: OutputType | str = OutputType.LINALG_ON_TENSORS, + ir_context: ir.Context | None = None, **kwargs, ) -> str | ir.Module: """Import a PyTorch nn.Module into MLIR. @@ -100,9 +100,9 @@ def import_from_file( init_kwargs_fn_name: str | None = None, sample_args_fn_name: str = "get_inputs", sample_kwargs_fn_name: str | None = None, - state_path : str | Path | None = None, - dialect : OutputType | str = OutputType.LINALG_ON_TENSORS, - ir_context : ir.Context | None = None, + state_path: str | Path | None = None, + dialect: OutputType | str = OutputType.LINALG_ON_TENSORS, + ir_context: ir.Context | None = None, **kwargs, ) -> str | ir.Module: """Load a PyTorch nn.Module from a file and import it into MLIR. @@ -211,7 +211,7 @@ def get_inputs(): error_msg=f"Sample kwargs function '{sample_kwargs_fn_name}' not found in {filepath}" ) - nn_model : nn.Module = model(*model_init_args, **model_init_kwargs) + nn_model: nn.Module = model(*model_init_args, **model_init_kwargs) if state_path is not None: state_dict = torch.load(state_path) nn_model.load_state_dict(state_dict) From d5c710c69b454e3c7046f07e93476c5ffbd0453d Mon Sep 17 00:00:00 2001 From: dchigarev Date: Tue, 21 Oct 2025 21:29:28 +0000 Subject: [PATCH 12/16] fix ImportError message Signed-off-by: dchigarev --- python/lighthouse/ingress/torch/torch_import.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/lighthouse/ingress/torch/torch_import.py b/python/lighthouse/ingress/torch/torch_import.py index dfaccd0..87c7655 100644 --- a/python/lighthouse/ingress/torch/torch_import.py +++ b/python/lighthouse/ingress/torch/torch_import.py @@ -11,7 +11,7 @@ except ImportError as e: raise ImportError( "PyTorch is required to use the torch import functionality. " - "Please run 'uv pip install .[torch-mlir]'" + "Make sure to install ingress-torch dependencies e.g. 'uv sync --extra ingress-torch-cpu'" ) from e try: @@ -20,7 +20,7 @@ except ImportError as e: raise ImportError( "torch-mlir is required to use the torch import functionality. " - "Please run 'uv pip install .[torch-mlir]'" + "Make sure to install ingress-torch dependencies e.g. 'uv sync --extra ingress-torch-cpu'" ) from e from mlir import ir From 029b8ad1c3f27c79aba45785e8eee842c93f2ec0 Mon Sep 17 00:00:00 2001 From: dchigarev Date: Wed, 22 Oct 2025 08:58:30 +0000 Subject: [PATCH 13/16] apply review suggestions Signed-off-by: dchigarev --- .../torch/{DummyMLP => MLPModel}/model.py | 8 ++--- ...ummy-mlp-from-file.py => mlp_from_file.py} | 34 +++++++++++-------- ...my-mlp-from-model.py => mlp_from_model.py} | 32 +++++++++-------- python/lighthouse/ingress/__init__.py | 1 + python/lighthouse/ingress/torch/__init__.py | 4 ++- .../torch/{torch_import.py => importer.py} | 0 6 files changed, 45 insertions(+), 34 deletions(-) rename python/examples/ingress/torch/{DummyMLP => MLPModel}/model.py (73%) rename python/examples/ingress/torch/{01-dummy-mlp-from-file.py => mlp_from_file.py} (73%) rename python/examples/ingress/torch/{02-dummy-mlp-from-model.py => mlp_from_model.py} (65%) rename python/lighthouse/ingress/torch/{torch_import.py => importer.py} (100%) diff --git a/python/examples/ingress/torch/DummyMLP/model.py b/python/examples/ingress/torch/MLPModel/model.py similarity index 73% rename from python/examples/ingress/torch/DummyMLP/model.py rename to python/examples/ingress/torch/MLPModel/model.py index 905ecc8..af1bde1 100644 --- a/python/examples/ingress/torch/DummyMLP/model.py +++ b/python/examples/ingress/torch/MLPModel/model.py @@ -5,7 +5,7 @@ import os -class DummyMLP(nn.Module): +class MLPModel(nn.Module): def __init__(self): super().__init__() self.net = nn.Sequential( @@ -19,15 +19,15 @@ def forward(self, x): def get_init_inputs(): - """Function to return args to pass to DummyMLP.__init__()""" + """Function to return args to pass to MLPModel.__init__()""" return () def get_sample_inputs(): - """Arguments to pass to DummyMLP.forward()""" + """Arguments to pass to MLPModel.forward()""" return (torch.randn(1, 10),) if __name__ == "__main__": script_dir = os.path.dirname(os.path.abspath(__file__)) - torch.save(DummyMLP().state_dict(), os.path.join(script_dir, "dummy_mlp.pth")) + torch.save(MLPModel().state_dict(), os.path.join(script_dir, "dummy_mlp.pth")) diff --git a/python/examples/ingress/torch/01-dummy-mlp-from-file.py b/python/examples/ingress/torch/mlp_from_file.py similarity index 73% rename from python/examples/ingress/torch/01-dummy-mlp-from-file.py rename to python/examples/ingress/torch/mlp_from_file.py index bf7c15c..9c4c848 100644 --- a/python/examples/ingress/torch/01-dummy-mlp-from-file.py +++ b/python/examples/ingress/torch/mlp_from_file.py @@ -1,14 +1,14 @@ """ Example demonstrating how to load a PyTorch model to MLIR using Lighthouse -without instantiating the model on the user's side. +without initializing the model on the user's side. The script uses 'lighthouse.ingress.torch.import_from_file' function that takes a path to a Python file containing the model definition, along with the names of functions to get model init arguments and sample inputs. The function -imports the model class on its own, instantiates it, and passes it to torch_mlir +imports the model class on its own, initializes it, and passes it to torch_mlir to get a MLIR module in the specified dialect. -The script uses the model from 'DummyMLP/model.py' as an example. +The script uses the model from 'MLPModel/model.py' as an example. """ import os @@ -16,25 +16,25 @@ # MLIR infrastructure imports (only needed if you want to manipulate the MLIR module) import mlir.dialects.func as func -from mlir import ir, passmanager +from mlir import ir # Lighthouse imports from lighthouse.ingress.torch import import_from_file # Step 1: Set up paths to locate the model definition file script_dir = Path(os.path.dirname(os.path.abspath(__file__))) -model_path = script_dir / "DummyMLP" / "model.py" +model_path = script_dir / "MLPModel" / "model.py" ir_context = ir.Context() # Step 2: Convert PyTorch model to MLIR # Conversion step where Lighthouse: -# - Loads the DummyMLP class and instantiates it with arguments obtained from 'get_init_inputs()' +# - Loads the MLPModel class and instantiates it with arguments obtained from 'get_init_inputs()' # - Calls get_sample_inputs() to get sample input tensors for shape inference # - Converts PyTorch model to linalg-on-tensors dialect operations using torch_mlir mlir_module_ir: ir.Module = import_from_file( model_path, # Path to the Python file containing the model - model_class_name="DummyMLP", # Name of the PyTorch nn.Module class to convert + model_class_name="MLPModel", # Name of the PyTorch nn.Module class to convert init_args_fn_name="get_init_inputs", # Function that returns args for model.__init__() sample_args_fn_name="get_sample_inputs", # Function that returns sample inputs to pass to 'model(...)' dialect="linalg-on-tensors", # Target MLIR dialect (linalg ops on tensor types) @@ -52,12 +52,16 @@ print(f"entry-point name: {func_op.name}") print(f"entry-point type: {func_op.type}") -# Step 4: Apply some MLIR passes using a PassManager -pm = passmanager.PassManager(context=ir_context) -pm.add("linalg-specialize-generic-ops") -pm.add("one-shot-bufferize") -pm.run(mlir_module_ir.operation) - -# Step 5: Output the final MLIR -print("\n\nModule dump after running the pipeline:") +# Step 4: output the imported MLIR module +print("\n\nModule dump:") mlir_module_ir.dump() + +# You can alternatively write the MLIR module to a file: +# with open("output.mlir", "w") as f: +# f.write(str(mlir_module_ir)) +# +# Or apply some MLIR passes using a PassManager: +# pm = passmanager.PassManager(context=ir_context) +# pm.add("linalg-specialize-generic-ops") +# pm.add(...) +# pm.run(mlir_module_ir.operation) diff --git a/python/examples/ingress/torch/02-dummy-mlp-from-model.py b/python/examples/ingress/torch/mlp_from_model.py similarity index 65% rename from python/examples/ingress/torch/02-dummy-mlp-from-model.py rename to python/examples/ingress/torch/mlp_from_model.py index 246deda..c847bdd 100644 --- a/python/examples/ingress/torch/02-dummy-mlp-from-model.py +++ b/python/examples/ingress/torch/mlp_from_model.py @@ -1,29 +1,29 @@ """ -Example demonstrating how to load an already instantiated PyTorch model +Example demonstrating how to load an already initialized PyTorch model to MLIR using Lighthouse. The script uses the 'lighthouse.ingress.torch.import_from_model' function that -takes a PyTorch model that has already been instantiated, along with its sample inputs. +takes a PyTorch model that has already been initialized, along with its sample inputs. The function passes the model to torch_mlir to get a MLIR module in the specified dialect. -The script uses a model from 'DummyMLP/model.py' as an example. +The script uses a model from 'MLPModel/model.py' as an example. """ import torch # MLIR infrastructure imports (only needed if you want to manipulate the MLIR module) import mlir.dialects.func as func -from mlir import ir, passmanager +from mlir import ir # Lighthouse imports from lighthouse.ingress.torch import import_from_model # Import a sample model definition -from DummyMLP.model import DummyMLP +from MLPModel.model import MLPModel # Step 1: Instantiate a model and prepare sample input -model = DummyMLP() +model = MLPModel() sample_input = torch.randn(1, 10) ir_context = ir.Context() @@ -45,12 +45,16 @@ print(f"entry-point name: {func_op.name}") print(f"entry-point type: {func_op.type}") -# Step 4: Apply some MLIR passes using a PassManager -pm = passmanager.PassManager(context=ir_context) -pm.add("linalg-specialize-generic-ops") -pm.add("one-shot-bufferize") -pm.run(mlir_module_ir.operation) - -# Step 5: Output the final MLIR -print("\n\nModule dump after running the pipeline:") +# Step 4: output the imported MLIR module +print("\n\nModule dump:") mlir_module_ir.dump() + +# You can alternatively write the MLIR module to a file: +# with open("output.mlir", "w") as f: +# f.write(str(mlir_module_ir)) +# +# Or apply some MLIR passes using a PassManager: +# pm = passmanager.PassManager(context=ir_context) +# pm.add("linalg-specialize-generic-ops") +# pm.add(...) +# pm.run(mlir_module_ir.operation) diff --git a/python/lighthouse/ingress/__init__.py b/python/lighthouse/ingress/__init__.py index e69de29..2aa3859 100644 --- a/python/lighthouse/ingress/__init__.py +++ b/python/lighthouse/ingress/__init__.py @@ -0,0 +1 @@ +"""Provides functions to convert source objects (code, models, designs) into MLIR files that the MLIR project can consume""" diff --git a/python/lighthouse/ingress/torch/__init__.py b/python/lighthouse/ingress/torch/__init__.py index c3a6e72..d73f426 100644 --- a/python/lighthouse/ingress/torch/__init__.py +++ b/python/lighthouse/ingress/torch/__init__.py @@ -1 +1,3 @@ -from .torch_import import import_from_file, import_from_model +"""Provides functions to convert PyTorch models to MLIR.""" + +from .importer import import_from_file, import_from_model diff --git a/python/lighthouse/ingress/torch/torch_import.py b/python/lighthouse/ingress/torch/importer.py similarity index 100% rename from python/lighthouse/ingress/torch/torch_import.py rename to python/lighthouse/ingress/torch/importer.py From 2215980b60a8f13a7aea742385bd315117ee6920 Mon Sep 17 00:00:00 2001 From: dchigarev Date: Wed, 22 Oct 2025 09:03:33 +0000 Subject: [PATCH 14/16] remove passManager from ingress examples Signed-off-by: dchigarev --- python/examples/ingress/torch/mlp_from_file.py | 6 ------ python/examples/ingress/torch/mlp_from_model.py | 6 ------ 2 files changed, 12 deletions(-) diff --git a/python/examples/ingress/torch/mlp_from_file.py b/python/examples/ingress/torch/mlp_from_file.py index 9c4c848..764d750 100644 --- a/python/examples/ingress/torch/mlp_from_file.py +++ b/python/examples/ingress/torch/mlp_from_file.py @@ -59,9 +59,3 @@ # You can alternatively write the MLIR module to a file: # with open("output.mlir", "w") as f: # f.write(str(mlir_module_ir)) -# -# Or apply some MLIR passes using a PassManager: -# pm = passmanager.PassManager(context=ir_context) -# pm.add("linalg-specialize-generic-ops") -# pm.add(...) -# pm.run(mlir_module_ir.operation) diff --git a/python/examples/ingress/torch/mlp_from_model.py b/python/examples/ingress/torch/mlp_from_model.py index c847bdd..f86f598 100644 --- a/python/examples/ingress/torch/mlp_from_model.py +++ b/python/examples/ingress/torch/mlp_from_model.py @@ -52,9 +52,3 @@ # You can alternatively write the MLIR module to a file: # with open("output.mlir", "w") as f: # f.write(str(mlir_module_ir)) -# -# Or apply some MLIR passes using a PassManager: -# pm = passmanager.PassManager(context=ir_context) -# pm.add("linalg-specialize-generic-ops") -# pm.add(...) -# pm.run(mlir_module_ir.operation) From 8bec558971dc09f75af84a0c198472055f6e6997 Mon Sep 17 00:00:00 2001 From: dchigarev Date: Wed, 22 Oct 2025 09:54:45 +0000 Subject: [PATCH 15/16] do not create .pth for model-mlp Signed-off-by: dchigarev --- python/examples/ingress/torch/MLPModel/model.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/python/examples/ingress/torch/MLPModel/model.py b/python/examples/ingress/torch/MLPModel/model.py index af1bde1..5bbdbbe 100644 --- a/python/examples/ingress/torch/MLPModel/model.py +++ b/python/examples/ingress/torch/MLPModel/model.py @@ -26,8 +26,3 @@ def get_init_inputs(): def get_sample_inputs(): """Arguments to pass to MLPModel.forward()""" return (torch.randn(1, 10),) - - -if __name__ == "__main__": - script_dir = os.path.dirname(os.path.abspath(__file__)) - torch.save(MLPModel().state_dict(), os.path.join(script_dir, "dummy_mlp.pth")) From a1a1dddce60f7b67054b4ca7b4466c8b4da4aff5 Mon Sep 17 00:00:00 2001 From: dchigarev Date: Wed, 22 Oct 2025 10:42:01 +0000 Subject: [PATCH 16/16] make example's docstring clearer on what 'pytorch mode' means Signed-off-by: dchigarev --- python/examples/ingress/torch/mlp_from_file.py | 6 +++--- python/examples/ingress/torch/mlp_from_model.py | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/python/examples/ingress/torch/mlp_from_file.py b/python/examples/ingress/torch/mlp_from_file.py index 764d750..01edeec 100644 --- a/python/examples/ingress/torch/mlp_from_file.py +++ b/python/examples/ingress/torch/mlp_from_file.py @@ -1,10 +1,10 @@ """ Example demonstrating how to load a PyTorch model to MLIR using Lighthouse -without initializing the model on the user's side. +without initializing the model class on the user's side. The script uses 'lighthouse.ingress.torch.import_from_file' function that -takes a path to a Python file containing the model definition, along with -the names of functions to get model init arguments and sample inputs. The function +takes a path to a Python file containing the model definition (a Python class derived from 'nn.Module'), +along with the names of functions to get model init arguments and sample inputs. The function imports the model class on its own, initializes it, and passes it to torch_mlir to get a MLIR module in the specified dialect. diff --git a/python/examples/ingress/torch/mlp_from_model.py b/python/examples/ingress/torch/mlp_from_model.py index f86f598..b51af80 100644 --- a/python/examples/ingress/torch/mlp_from_model.py +++ b/python/examples/ingress/torch/mlp_from_model.py @@ -3,9 +3,9 @@ to MLIR using Lighthouse. The script uses the 'lighthouse.ingress.torch.import_from_model' function that -takes a PyTorch model that has already been initialized, along with its sample inputs. -The function passes the model to torch_mlir to get a MLIR module in the -specified dialect. +takes an initialized PyTorch model (an instance of a Python class derived from 'nn.Module'), +along with its sample inputs. The function passes the model to torch_mlir +to get a MLIR module in the specified dialect. The script uses a model from 'MLPModel/model.py' as an example. """ @@ -22,7 +22,7 @@ # Import a sample model definition from MLPModel.model import MLPModel -# Step 1: Instantiate a model and prepare sample input +# Step 1: Instantiate a model class and prepare sample input model = MLPModel() sample_input = torch.randn(1, 10)