From 996200cf19cf2532e230cca52fb5e608f289c61a Mon Sep 17 00:00:00 2001 From: xintin Date: Wed, 13 Nov 2024 01:15:05 -0500 Subject: [PATCH 1/3] Added readme to aot_mlp Signed-off-by: Gaurav Verma Signed-off-by: xintin --- examples/aot_mlp/README.md | 23 +++++++++++ examples/aot_mlp/mlp_export_simple.py | 56 +++++++++++++++++++++------ examples/aot_mlp/requirements.txt | 4 ++ 3 files changed, 72 insertions(+), 11 deletions(-) create mode 100644 examples/aot_mlp/README.md create mode 100644 examples/aot_mlp/requirements.txt diff --git a/examples/aot_mlp/README.md b/examples/aot_mlp/README.md new file mode 100644 index 000000000..e8057b70e --- /dev/null +++ b/examples/aot_mlp/README.md @@ -0,0 +1,23 @@ +# AOT MLP Example + +This example demonstrates export, compilation, and inference of +a simple Multi-Layer Perceptron (MLP) model. +The model is a four-layer neural network. + +To run this example, you should clone the repository to your local device and +install the requirements in a virtual environment: + +```bash +git clone https://github.com/iree-org/iree-turbine.git +cd iree-turbine/examples/aot_mlp +python -m venv mlp.venv +source ./mlp.venv/bin/activate +pip install -r requirements.txt +``` + +Once the requirements are installed, you should be able to run the example. + +```bash +python mlp_export_simple.py +``` + diff --git a/examples/aot_mlp/mlp_export_simple.py b/examples/aot_mlp/mlp_export_simple.py index 30d7ae952..b9ae79c9f 100644 --- a/examples/aot_mlp/mlp_export_simple.py +++ b/examples/aot_mlp/mlp_export_simple.py @@ -4,23 +4,44 @@ # See https://llvm.org/LICENSE.txt for license information. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +import iree.runtime as rt import logging -import unittest +import numpy as np import torch import torch.nn as nn +import unittest import iree.turbine.aot as aot +logging.basicConfig(level=logging.DEBUG) +logger = logging.getLogger(__name__) + + class MLP(nn.Module): - def __init__(self): + """ + Multi-Layer Perceptron (MLP) model class. + Defines a neural network with four linear layers and sigmoid activations. + """ + + def __init__(self) -> None: super().__init__() + # Define model layers self.layer0 = nn.Linear(8, 8, bias=True) self.layer1 = nn.Linear(8, 4, bias=True) self.layer2 = nn.Linear(4, 2, bias=True) self.layer3 = nn.Linear(2, 2, bias=True) - def forward(self, x: torch.Tensor): + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass of the MLP model. + + Args: + x (torch.Tensor): Input tensor. + + Returns: + torch.Tensor: Output tensor after forward pass. + """ x = self.layer0(x) x = torch.sigmoid(x) x = self.layer1(x) @@ -38,25 +59,38 @@ def forward(self, x: torch.Tensor): compiled_binary = exported.compile(save_to=None) -def infer(): - import numpy as np - import iree.runtime as rt +def run_inference() -> np.ndarray: + """ + Runs inference on the compiled model. + Returns: + np.ndarray: The result of inference as a NumPy array. + """ config = rt.Config("local-task") vmm = rt.load_vm_module( - rt.VmModule.wrap_buffer(config.vm_instance, compiled_binary.map_memory()), + rt.VmModule.wrap_buffer( + config.vm_instance, compiled_binary.map_memory() + ), config, ) x = np.random.rand(97, 8).astype(np.float32) y = vmm.main(x) - print(y.to_host()) + logger.debug(f"Inference result: {y.to_host()}") + return y.to_host() class ModelTest(unittest.TestCase): - def testMLPExportSimple(selfs): - infer() + def test_mlp_export_simple(self) -> None: + """Tests if the model export and inference work as expected.""" + output = run_inference() + self.assertIsNotNone(output, "Inference output should not be None") + self.assertEqual( + output.shape, (97, 2), + "Output shape doesn't match the expected (97, 2)" + ) if __name__ == "__main__": - logging.basicConfig(level=logging.DEBUG) + # Run unit tests unittest.main() + diff --git a/examples/aot_mlp/requirements.txt b/examples/aot_mlp/requirements.txt new file mode 100644 index 000000000..aa30906ea --- /dev/null +++ b/examples/aot_mlp/requirements.txt @@ -0,0 +1,4 @@ +numpy +torch +iree-turbine + From 55c2e92d3c48cb9fefea2e62db6cf47c2b8021e1 Mon Sep 17 00:00:00 2001 From: xintin Date: Fri, 14 Feb 2025 16:29:34 -0500 Subject: [PATCH 2/3] updated readme Signed-off-by: xintin --- examples/aot_mlp/README.md | 11 ++++++++--- examples/aot_mlp/mlp_export_simple.py | 17 ++++++++--------- examples/aot_mlp/requirements.txt | 2 -- 3 files changed, 16 insertions(+), 14 deletions(-) diff --git a/examples/aot_mlp/README.md b/examples/aot_mlp/README.md index e8057b70e..e74142ecf 100644 --- a/examples/aot_mlp/README.md +++ b/examples/aot_mlp/README.md @@ -1,7 +1,7 @@ # AOT MLP Example -This example demonstrates export, compilation, and inference of -a simple Multi-Layer Perceptron (MLP) model. +This example demonstrates export, compilation, and inference of +a simple Multi-Layer Perceptron (MLP) model. The model is a four-layer neural network. To run this example, you should clone the repository to your local device and @@ -12,6 +12,12 @@ git clone https://github.com/iree-org/iree-turbine.git cd iree-turbine/examples/aot_mlp python -m venv mlp.venv source ./mlp.venv/bin/activate +``` + +To install `torch`, follow the instructions available [here](https://github.com/iree-org/iree-turbine/tree/main?tab=readme-ov-file#install-pytorch-for-your-system). + +Other requirements can be installed via the following command: +``` pip install -r requirements.txt ``` @@ -20,4 +26,3 @@ Once the requirements are installed, you should be able to run the example. ```bash python mlp_export_simple.py ``` - diff --git a/examples/aot_mlp/mlp_export_simple.py b/examples/aot_mlp/mlp_export_simple.py index b9ae79c9f..5260715ee 100644 --- a/examples/aot_mlp/mlp_export_simple.py +++ b/examples/aot_mlp/mlp_export_simple.py @@ -24,7 +24,7 @@ class MLP(nn.Module): Defines a neural network with four linear layers and sigmoid activations. """ - def __init__(self) -> None: + def __init__(self) -> object: super().__init__() # Define model layers self.layer0 = nn.Linear(8, 8, bias=True) @@ -68,9 +68,7 @@ def run_inference() -> np.ndarray: """ config = rt.Config("local-task") vmm = rt.load_vm_module( - rt.VmModule.wrap_buffer( - config.vm_instance, compiled_binary.map_memory() - ), + rt.VmModule.wrap_buffer(config.vm_instance, compiled_binary.map_memory()), config, ) x = np.random.rand(97, 8).astype(np.float32) @@ -78,19 +76,20 @@ def run_inference() -> np.ndarray: logger.debug(f"Inference result: {y.to_host()}") return y.to_host() + return y.to_host() + class ModelTest(unittest.TestCase): def test_mlp_export_simple(self) -> None: - """Tests if the model export and inference work as expected.""" output = run_inference() - self.assertIsNotNone(output, "Inference output should not be None") + + self.assertIsNotNone(output, "inference output should not be None") self.assertEqual( - output.shape, (97, 2), - "Output shape doesn't match the expected (97, 2)" + output.shape, (97, 2), "output shape doesn't match the expected (97, 2)" ) if __name__ == "__main__": + logging.basicConfig(level=logging.DEBUG) # Run unit tests unittest.main() - diff --git a/examples/aot_mlp/requirements.txt b/examples/aot_mlp/requirements.txt index aa30906ea..a161f0d16 100644 --- a/examples/aot_mlp/requirements.txt +++ b/examples/aot_mlp/requirements.txt @@ -1,4 +1,2 @@ numpy -torch iree-turbine - From b4d4ddff5cc83964d4af1346bb12caf3f2183ad9 Mon Sep 17 00:00:00 2001 From: Scott Todd Date: Fri, 14 Feb 2025 15:33:22 -0800 Subject: [PATCH 3/3] Remove extra return. --- examples/aot_mlp/mlp_export_simple.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/examples/aot_mlp/mlp_export_simple.py b/examples/aot_mlp/mlp_export_simple.py index 5260715ee..8f118298c 100644 --- a/examples/aot_mlp/mlp_export_simple.py +++ b/examples/aot_mlp/mlp_export_simple.py @@ -76,8 +76,6 @@ def run_inference() -> np.ndarray: logger.debug(f"Inference result: {y.to_host()}") return y.to_host() - return y.to_host() - class ModelTest(unittest.TestCase): def test_mlp_export_simple(self) -> None: