[ingress][torch-mlir] Initial version of wrapper around torch-mlir's fx-importer (#4)

The PR adds two utility functions as part of lighthouse's python package to convert torch models to a mlir module using `torch_mlir`.

A user can import one of the functions (`import_from_model` or `import_from_file`) and get a `mlir.ir.Module` that they can use to run passes on or simply write its content into a file.

Some use cases:

1. Import from an instance of a model:

```python
from lighthouse.ingress.torch import import_from_model
from mlir import ir

ctx = ir.Context()

module : ir.Module = import_from_model(torch_model_instance, sample_args=(torch.rand(1, 10),), ir_context=ctx)

# can now run some passes on the module
```

2. Import from a file where a torch model is defined:

Imagine we want to import a [model from `KernelBench`](https://github.com/ScalingIntelligence/KernelBench/blob/main/KernelBench/level1/10_3D_tensor_matrix_multiplication.py). They ship models as python files where models and their arguments are uniformly defined.

```python
from lighthouse.ingress.torch import import_from_file
from mlir import ir

ctx = ir.Context()
kernel_bench_root = Path(...)

module : ir.Module = import_from_model(
    filepath=kernel_bench_root / "level1" / "10_3D_tensor_matrix_multiplication.py",
    ir_context=ctx
)

# can now run some passes on the module
```

Author: dchigarev

ingress/Torch-MLIR/generate-mlir.py

@@ -0,0 +1,28 @@
+"""Defines a simple PyTorch model to be used in lighthouse's ingress examples."""
+
+import torch
+import torch.nn as nn
+
+import os
+
+class MLPModel(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 MLPModel.__init__()"""
+    return ()
+
+
+def get_sample_inputs():
+    """Arguments to pass to MLPModel.forward()"""
+    return (torch.randn(1, 10),)

ingress/Torch-MLIR/generate-mlir.sh

@@ -0,0 +1,61 @@
+"""
+Example demonstrating how to load a PyTorch model to MLIR using Lighthouse
+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 (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.
+
+The script uses the model from 'MLPModel/model.py' as an example.
+"""
+
+import os
+from pathlib 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
+
+# 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 / "MLPModel" / "model.py"
+
+ir_context = ir.Context()
+
+# Step 2: Convert PyTorch model to MLIR
+# Conversion step where Lighthouse:
+# - 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="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)
+    ir_context=ir_context                    # MLIR context for the conversion
+)
+
+# 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
+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: 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))

python/examples/ingress/torch/MLPModel/model.py

@@ -0,0 +1,54 @@
+"""
+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 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.
+"""
+
+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
+
+# Lighthouse imports
+from lighthouse.ingress.torch import import_from_model
+
+# Import a sample model definition
+from MLPModel.model import MLPModel
+
+# Step 1: Instantiate a model class and prepare sample input
+model = MLPModel()
+sample_input = torch.randn(1, 10)
+
+ir_context = ir.Context()
+# 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.
+#
+# 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
+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: 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))

python/examples/ingress/torch/mlp_from_file.py

@@ -0,0 +1,10 @@
+# Lighthouse Ingress
+
+The `lighthouse.ingress` module converts various input formats to MLIR modules.
+
+## Supported Formats
+
+#### Torch
+Converts PyTorch models to MLIR using `lighthouse.ingress.torch`.
+
+**Examples:** [torch examples](https://github.com/llvm/lighthouse/tree/main/python/examples/ingress/torch)

python/examples/ingress/torch/mlp_from_model.py

@@ -0,0 +1 @@
+"""Provides functions to convert source objects (code, models, designs) into MLIR files that the MLIR project can consume"""

python/lighthouse/ingress/README.md

@@ -0,0 +1,3 @@
+"""Provides functions to convert PyTorch models to MLIR."""
+
+from .importer import import_from_file, import_from_model