diff --git a/src/Runtime/python/OMPyCompile/.gitignore b/src/Runtime/python/OMPyCompile/.gitignore new file mode 100644 index 00000000000..ff5cacba37d --- /dev/null +++ b/src/Runtime/python/OMPyCompile/.gitignore @@ -0,0 +1,3 @@ +# Ignore directories for build or install with -e +dist/ +src/OMPyCompile/__pycache__/ diff --git a/src/Runtime/python/OMPyCompile/README.md b/src/Runtime/python/OMPyCompile/README.md new file mode 100644 index 00000000000..07091b637e5 --- /dev/null +++ b/src/Runtime/python/OMPyCompile/README.md @@ -0,0 +1,52 @@ + +# OMPyCompile +This package provides a python driver to compile a ONNX model with onnx-mlir. +The compilation can be done with onnx-mlir compiler at local directory or in +a container. +To use +``` +import numpy as np +import OMPyInfer + +# Initialize the inference session +# The onnx model simply performs tensor add on two 3x4x5xf32 tensors +# It is compiled into test_add.so with zDLC +sess = OMPyInfer.InferenceSession("./test_add.so") + +# Prepare the inputs +a = np.arange(3 * 4 * 5, dtype=np.float32).reshape((3, 4, 5)) +b = a + 4 + +# Run inference +r = sess.run([a, b]) + +# Print output +print(r) +``` + +## Compile onnx model to shared library +TBD + + +## Pre-requisites for OMPyInfer +These pre-requisities are currently provided as part of the OMPyInfer package for Python versions 3.9 until 3.13. +Prebuilt libraries for Linux on Z is provided. +Follow these instructions (TBD) to build the libraries for your own system. + +## Install +Currently, only local installation is supported. +Suppose you have onnx-mlir cloned on your machine. Install OMPyInfer with the following command: +``` +python onnx-mlir/src/Runtime/python/OMPyInfer +``` + + +## Verify +``` +cd OMPyInfer/tests +python helloworld.py +``` + +## VERSIONS +Version 1.0.0 supports the model copied with onnx-mlir before 29bde823f, 2026-02-04. + diff --git a/src/Runtime/python/OMPyCompile/VERSION_NUMBER b/src/Runtime/python/OMPyCompile/VERSION_NUMBER new file mode 100644 index 00000000000..f15d4cd8753 --- /dev/null +++ b/src/Runtime/python/OMPyCompile/VERSION_NUMBER @@ -0,0 +1,3 @@ + + +1.0.0 diff --git a/src/Runtime/python/OMPyCompile/pyproject.toml b/src/Runtime/python/OMPyCompile/pyproject.toml new file mode 100644 index 00000000000..c1d7264f244 --- /dev/null +++ b/src/Runtime/python/OMPyCompile/pyproject.toml @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" +[project] +name = "OMPyCompile" +version = "1.0.0" +authors = [ + { name="Tong Chen", email="chentong@us.ibm.com" }, + { name="Sunny Anand", email="sunny.anand79@ibm.com" }, +] +description = "Python driver to run inference on onnx model compiled with zdlc" +readme = "README.md" +requires-python = ">=3.9" +license = "Apache-2.0" +classifiers = [ + "Programming Language :: Python :: 3", + "Operating System :: POSIX :: Linux", +] + +dependencies = [ + "docker" +] + +[project.urls] +Homepage = "https://github.com/onnx/onnx-mlir" +Issues = "https://github.com/onnx/onnx-mlir/issues" + +[tool.hatch.build.targets.wheel] +packages = ["src/OMPyCompile"] + +[tool.hatch.build] +include = [ + "src/OMPyCompile/libs/*" +] + + diff --git a/src/Runtime/python/OMPyCompile/src/OMPyCompile/__init__.py b/src/Runtime/python/OMPyCompile/src/OMPyCompile/__init__.py new file mode 100644 index 00000000000..c2fb3f45a0e --- /dev/null +++ b/src/Runtime/python/OMPyCompile/src/OMPyCompile/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 + +from .onnxmlirdockercompile import compile + +__all__ = ["compile"] diff --git a/src/Runtime/python/OMPyCompile/src/OMPyCompile/__main__.py b/src/Runtime/python/OMPyCompile/src/OMPyCompile/__main__.py new file mode 100644 index 00000000000..bfdcd0c1158 --- /dev/null +++ b/src/Runtime/python/OMPyCompile/src/OMPyCompile/__main__.py @@ -0,0 +1,4 @@ +from .cli import main + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/Runtime/python/OMPyCompile/src/OMPyCompile/cli.py b/src/Runtime/python/OMPyCompile/src/OMPyCompile/cli.py new file mode 100644 index 00000000000..c18c162141f --- /dev/null +++ b/src/Runtime/python/OMPyCompile/src/OMPyCompile/cli.py @@ -0,0 +1,70 @@ +import argparse +from .onnxmlirdockercompile import compile + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="python -m OMPyCompile", + description="Command-line interface for package OMPyCompile.", + ) + parser.add_argument( + "model", + type=str, + help="model file", + ) + parser.add_argument( + "--debug", + action="store_true", + default=False, + help="Enable debug output (default: False)", + ) + parser.add_argument( + "--compile_tag", + type=str, + default="NONE", + help="Tag to pass to the compiler via --tag= (default: NONE)", + ) + parser.add_argument( + "--compile_options", + type=str, + default="", + help="Additional options to pass to the onnx-mlir compiler (default: '')", + ) + parser.add_argument( + "--compiler_image_name", + type=str, + default="ghcr.io/onnxmlir/onnx-mlir-dev", + help="Container image name to use for compilation", + ) + parser.add_argument( + "--container_engine", + type=str, + default=None, + choices=["docker", "podman", None], + help=( + "Container engine to use: 'docker', 'podman', or None to auto-detect " + "(default: None)" + ), + ) + parser.add_argument( + "--compiler_path", + type=str, + default=None, + help=( + "Path to the onnx-mlir compiler binary. Required when using a custom " + "compiler_image_name that is not in the built-in image dictionary." + ), + ) + return parser + + +def main(argv=None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + + kwargs = vars(args) + model = kwargs.pop("model") + compiled_model = compile(model, **kwargs) + print(f"Compiled to {compiled_model}") + + return 0 diff --git a/src/Runtime/python/OMPyCompile/src/OMPyCompile/onnxmlirdockercompile.py b/src/Runtime/python/OMPyCompile/src/OMPyCompile/onnxmlirdockercompile.py new file mode 100644 index 00000000000..3cfce7fc930 --- /dev/null +++ b/src/Runtime/python/OMPyCompile/src/OMPyCompile/onnxmlirdockercompile.py @@ -0,0 +1,294 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 + +##################### onnxmlirdockercompile.py ################################# +# +# Copyright 2019-2026 The IBM Research Authors. +# +################################################################################ +# +# This script define a class to compile an onnx model with local compiler, +# or with a container image. +################################################################################ + +import numpy as np +import os +import sys +import tempfile +import json +import subprocess + + +class config: + image_path_dictionary = { + "ghcr.io/onnxmlir/onnx-mlir": "/usr/local/bin/bin/onnx-mlir", + "ghcr.io/onnxmlir/onnx-mlir-dev": "/workdir/onnx-mlir/build/Debug/bin/onnx-mlir", + "onnxmlir/onnx-mlir-dev": "/workdir/onnx-mlir/build/Debug/bin/onnx-mlir", + # The entry point of zDLC image is the compiler + "icr.io/ibmz/zdlc:5.0.0": "", + } + + default_compiler_image_name = "ghcr.io/onnxmlir/onnx-mlir-dev" + + +def get_names_in_signature(signature): + names = [] + # Load the input signature. + signature_dict = json.loads(signature) + for sig in signature_dict: + names.append(sig["name"]) + return names + + +# Return the compiler path of an image based on the image_path_dictionary +def find_compiler_path(image_name): + dict = config.image_path_dictionary + if image_name in dict: + return dict[image_name] + else: + return None + + +# Class to compile ONNX or MLIR models to shared libraries (.so files) +# using either a local compiler or a container image. +class CompileSession: + def __init__(self, model_path, **kwargs): + self.debug = False + # self.output_dir = tempfile.TemporaryDirectory() + self.handleParameters(model_path, **kwargs) + self.checkCompiler() + + def handleParameters(self, model_path, **kwargs): + if "debug" in kwargs.keys(): + self.debug = kwargs["debug"] + if "compile_tag" in kwargs.keys(): + self.compile_tag = kwargs["compile_tag"] + else: + # Conform with the default of onnx-mlir + self.compile_tag = "" + + self.model_path = model_path + if model_path.endswith(".mlir"): + self.model_suffix = ".mlir" + elif model_path.endswith(".onnx"): + self.model_suffix = ".onnx" + else: + print("Invalid input model path. Must end with .onnx or .mlir") + exit(1) + + absolute_path = os.path.abspath(self.model_path) + self.model_basename = os.path.basename(absolute_path) + self.model_dirname = os.path.dirname(absolute_path) + + # By default, the compiled model is in the directory of the model + # Perhaps, the current directory? + self.output_dirname = self.model_dirname + + if "compile_options" in kwargs.keys(): + self.compile_options = kwargs["compile_options"] + else: + self.compile_options = "" + + # Parse the compile_options to handle the -o option. + # If user does not specify the output, the compilation will be done in + # the temporary directory. + options_list = self.compile_options.split() + + # Logically, container_output_dirname should be used for -o + if "-o" in options_list: + # Convert the output to absolute path so that the compilation + # can be done with compiler image. + self.compiled_model = os.path.abspath( + options_list[options_list.index("-o") + 1] + ) + options_list[options_list.index("-o") + 1] = self.compiled_model + + self.compile_options = " ".join(options_list) + if not self.compiled_model.endswith(".so"): + self.compiled_model += ".so" + self.output_dirname = os.path.dirname(self.compiled_model) + else: + # self.output_dirname = self.output_dir.name + self.compiled_model = os.path.join( + self.output_dirname, self.model_basename.removesuffix(self.model_suffix) + ) + self.compile_options += f" -o {self.compiled_model}" + # Compile will automatically append .so to the -o target + # Session need the suffix .so + self.compiled_model += ".so" + + if "compiler_image_name" in kwargs.keys(): + self.compiler_image_name = kwargs["compiler_image_name"] + if ( + self.compiler_image_name == "local" + or self.compiler_image_name == "None" + ): + self.compiler_image_name = None + self.compiler_path = find_compiler_path(self.compiler_image_name) + if self.compiler_path is None and "compiler_path" not in kwargs.keys(): + print( + "Please specify the path to your compiler when you are not using the default image" + ) + exit(1) + else: + # Default image + self.compiler_image_name = config.default_compiler_image_name + self.compiler_path = find_compiler_path(self.compiler_image_name) + + if "container_engine" in kwargs.keys(): + self.container_engine = kwargs["container_engine"] + if ( + self.container_engine != "docker" + and self.container_engine != "podman" + and self.container_engine != None + ): + print( + "Container_engine has to be either 'docker' or 'podman', or None to let system choose the availabe one." + ) + exit(1) + else: + self.container_engine = None + + if "compiler_path" in kwargs.keys(): + self.compiler_path = kwargs["compiler_path"] + + def checkCompiler(self): + if self.compiler_image_name == None: + if not os.path.exists(self.compiler_path): + print("the compiler path does not exist: ", self.compiler_path) + exit(-1) + else: + # Import container tool, either docker or podman package + if self.container_engine is None: + try: + import docker as ce + except ImportError: + try: + import podman as ce + except ImportError: + raise ImportError( + "Failure to load docker or podman package. To install docker, you can either do 'pip install docker', or install onnxmlir with `pip install -e path/onnxmlir[docker]`. Similar for podman" + ) + elif self.container_engine == "docker": + try: + import docker as ce + except ImportError: + raise ImportError( + "Failure to load docker package. you can either do 'pip install docker', or install onnxmlir with `pip install -e path/onnxmlir[docker]`" + ) + else: + try: + import podman as ce + except ImportError: + raise ImportError( + "Failure to load podman package. you can either do 'pip install podman', or install onnxmlir with `pip install -e path/onnxmlir[podman]`" + ) + # The docker and podman package has the same interface + # Get container client using env setting. + self.container_client = ce.from_env() + + # Pull the image if not already available + try: + image = self.container_client.images.get(self.compiler_image_name) + except ce.errors.ImageNotFound: + image = self.container_client.images.pull(self.compiler_image_name) + + # Chek whether the specified compiler exists or not + if os.path.exists(self.compiler_path): + self.mount_compiler = True + else: + self.mount_compiler = False + try: + msg = self.container_client.containers.run( + self.compiler_image_name, self.compiler_path + " --version" + ) + # Could check the valid version of compiler + except Exception as e: + print( + "the compiler path does not exist in container: ", + self.compiler_path, + ) + exit(-1) + + def Compile(self): + # Logically use different variable for path in current env and + # container env. They could be different. + # However, if the caller is already in a container, the path + # on the host has to be used. Therefore, it is easier to always use + # the same path to mount on container. + self.container_model_dirname = self.model_dirname + self.container_output_dirname = self.output_dirname + + # Construct compilation command + command_str = self.compiler_path + + # Compiled library + if self.compile_options != "": + command_str += " " + self.compile_options + + # When tag is used, the name of entry function changed. + command_str += " --tag=" + self.compile_tag + + command_str += " " + os.path.join( + self.container_model_dirname, self.model_basename + ) + + if self.debug: + print(command_str) + + # Logically, the model directory could be mounted as read only. + # But wrong time error occurred with "r" mode + if self.compiler_image_name is None: + subprocess.run(command_str.split(" ")) + self.container = None + else: + # ToFix: try detach=True? + try: + msg = self.container_client.containers.run( + self.compiler_image_name, + command_str, + volumes={ + self.model_dirname: { + "bind": self.container_model_dirname, + "mode": "rw", + }, + self.output_dirname: { + "bind": self.container_output_dirname, + "mode": "rw", + }, + # Mount the compiler if needed. + # ToFix: compiler built with container may fail, possibly because + # the directory is mounted differently when the compiler is built + # with container image. + # A good practice is to always mount with the original path + **( + { + os.path.dirname(self.compiler_path): { + "bind": os.path.dirname(self.compiler_path), + "mode": "rw", + } + } + if self.mount_compiler + else {} + ), + }, + ) + except Exception as e: + print( + f"Failed in compilation with docker image: {self.compiler_image_name}" + ) + print(f"Here is the compiling command: {command_str}") + print( + "Check whether the flags are correct, the input file exists, and the output file is writable" + ) + exit(-1) + + def get_compiled_model_path(self): + """Return the path to the compiled .so file""" + return self.compiled_model + + +def compile(model_path, **kwargs): + sess = CompileSession(model_path, **kwargs) + sess.Compile() + return sess.get_compiled_model_path() diff --git a/src/Runtime/python/OMPyCompile/tests/compile.sh b/src/Runtime/python/OMPyCompile/tests/compile.sh new file mode 100644 index 00000000000..c163d6bb8c3 --- /dev/null +++ b/src/Runtime/python/OMPyCompile/tests/compile.sh @@ -0,0 +1 @@ +python -m OMPyCompile test_add.mlir --container_engine=docker --compiler_image_name="ghcr.io/onnxmlir/onnx-mlir-dev:latest" --compiler_path="/workdir/onnx-mlir/build/Debug/bin/onnx-mlir" diff --git a/src/Runtime/python/OMPyCompile/tests/test_add.mlir b/src/Runtime/python/OMPyCompile/tests/test_add.mlir new file mode 100644 index 00000000000..bd989a687b4 --- /dev/null +++ b/src/Runtime/python/OMPyCompile/tests/test_add.mlir @@ -0,0 +1,7 @@ +module { + func.func @main_graph(%arg0: tensor<3x4x5xf32> {onnx.name = "x"}, %arg1: tensor<3x4x5xf32> {onnx.name = "y"}) -> (tensor<3x4x5xf32> {onnx.name = "sum"}) { + %0 = "onnx.Add"(%arg0, %arg1) : (tensor<3x4x5xf32>, tensor<3x4x5xf32>) -> tensor<3x4x5xf32> + onnx.Return %0 : tensor<3x4x5xf32> + } + "onnx.EntryPoint"() {func = @main_graph} : () -> () +} diff --git a/src/Runtime/python/OMPyCompile/tests/test_add.onnx b/src/Runtime/python/OMPyCompile/tests/test_add.onnx new file mode 100644 index 00000000000..8aa4ab2bc38 Binary files /dev/null and b/src/Runtime/python/OMPyCompile/tests/test_add.onnx differ diff --git a/src/Runtime/python/OMPyCompile/tests/use_compiler_container.py b/src/Runtime/python/OMPyCompile/tests/use_compiler_container.py new file mode 100644 index 00000000000..2ebb0be14f8 --- /dev/null +++ b/src/Runtime/python/OMPyCompile/tests/use_compiler_container.py @@ -0,0 +1,44 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 + +############# use_compiler_container.py ####################################### +# +# Copyright 2021-2025 The IBM Research Authors. +# +################################################################################ +# Test case to compile a model with compiler container +################################################################################ + +# Local model file +from pathlib import Path + +script_dir = Path(__file__).resolve().parent +model_file = str(script_dir / "test_add.mlir") + +# To use compiler container, the image name and compiler path in the image +# need to be provided. +# compile_args is the flags passed to onnx-mlir +import OMPyCompile + +compiled_model = OMPyCompile.compile( + model_file, + compile_options="-O3", + container_engine="docker", + compiler_image_name="ghcr.io/onnxmlir/onnx-mlir-dev:latest", + compiler_path="/workdir/onnx-mlir/build/Debug/bin/onnx-mlir", +) +print(compiled_model) + +# Prepare input data +import numpy as np + +a = np.arange(3 * 4 * 5, dtype=np.float32).reshape((3, 4, 5)) +b = a + 4 + +# Run inference +import OMPyInfer + +sess = OMPyInfer.InferenceSession(compiled_model) + +r = sess.run([a, b]) +print(r) diff --git a/src/Runtime/python/OMPyCompile/tests/use_local_compiler.py b/src/Runtime/python/OMPyCompile/tests/use_local_compiler.py new file mode 100644 index 00000000000..cd9fd085765 --- /dev/null +++ b/src/Runtime/python/OMPyCompile/tests/use_local_compiler.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 + +################# use_local_compiler.py ####################################### +# +# Copyright 2021-2025 The IBM Research Authors. +# +################################################################################ +# Test case to use the local compiler to compile the model +################################################################################ + +# Local model file +from pathlib import Path + +script_dir = Path(__file__).resolve().parent +model_file = str(script_dir / "test_add.mlir") + +# When compiler_image_name is None, local compiler will be used. +# The compiler_path is used to locate the compiler. +# compile_args is the flags passed to onnx-mlir +import OMPyCompile + +compiled_model = OMPyCompile.compile( + "./test_add.mlir", + compile_options="-O3", + compiler_image_name=None, + compiler_path="/Users/chentong/Projects/onnx-mlir/build/Debug/bin/onnx-mlir", +) +print(compiled_model) + +# Prepare input data +import numpy as np + +a = np.arange(3 * 4 * 5, dtype=np.float32).reshape((3, 4, 5)) +b = a + 4 + +# Run inference +import OMPyInfer + +sess = OMPyInfer.InferenceSession(compiled_model) +r = sess.run([a, b]) +print(r) diff --git a/src/Runtime/python/OMPyInfer/README.md b/src/Runtime/python/OMPyInfer/README.md index eb6fd296013..1f98f0a85fa 100644 --- a/src/Runtime/python/OMPyInfer/README.md +++ b/src/Runtime/python/OMPyInfer/README.md @@ -3,9 +3,6 @@ This package provides a python driver to run inference on ONNX model compiled onnx-mlir. There is a helloworld example in the tests folder with the package: ``` -# IBM Confidential -# © Copyright IBM Corp. 2025 - import numpy as np import OMPyInfer diff --git a/src/Runtime/python/OMPyInfer/VERSION_NUMBER b/src/Runtime/python/OMPyInfer/VERSION_NUMBER index 5e77dfbd302..f15d4cd8753 100644 --- a/src/Runtime/python/OMPyInfer/VERSION_NUMBER +++ b/src/Runtime/python/OMPyInfer/VERSION_NUMBER @@ -1,3 +1,3 @@ -# IBM Confidential -# © Copyright IBM Corp. 2025 + + 1.0.0