Skip to content

Create package of OMPyCompile#3399

Open
chentong319 wants to merge 5 commits intoonnx:mainfrom
chentong319:OMPyCompile
Open

Create package of OMPyCompile#3399
chentong319 wants to merge 5 commits intoonnx:mainfrom
chentong319:OMPyCompile

Conversation

@chentong319
Copy link
Collaborator

This PR refactored the previous onnxmlirdocker.py to keep only the compilation functionality. It provide a function, compile, to return the compiled model.
Though the compilation is provided simply with one python script file, I created a package for it. The reason is that it would be easier for user to install and use.
Examples can be found in the tests directory.

import OMPyCompile
model = OMPyCompile.compile(
    "./test_add.onnx",
    compile_args="-O3",
    compiler_path="/Users/chentong/Projects/onnx-mlir/build/Debug/bin/onnx-mlir",
)
import OMPyInfer
sess = OMPyInfer.InferenceSession(model)
result = sess.run(inputs)

If container is used, change the compilation arguments to:

compiled_model = OMPyCompile.compile(
    "test_add.onnx",
    compile_args="-O3",
    container_engine="docker",
    compiler_image_name="ghcr.io/onnxmlir/onnx-mlir-dev",
    compiler_path="/workdir/onnx-mlir/build/Debug/bin/onnx-mlir",
)

Signed-off-by: chentong319 <chentong@us.ibm.com>

[tool.hatch.build]
include = [
"src/OMPyCompile/libs/*"
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Which libraries do we need?

@@ -0,0 +1,53 @@
<!--- SPDX-License-Identifier: Apache-2.0 -->
# OMPyInfer
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OMPyCompile? Look like this file was copied from PyInfer.

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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please remove this.

@@ -0,0 +1,3 @@
# IBM Confidential
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please remove this.


################# use_local_compiler.py #######################################
#
# Copyright 2021-2025 The IBM Research Authors.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should it be 2026?

class CompileSession:
def __init__(self, model_path, **kwargs):
self.debug = False
# self.output_dir = tempfile.TemporaryDirectory()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you initialize all members of this class here so that we know what are its members, e.g. self.compiler_image_name self.model_path, etc.?

# Compiled library
if self.compile_options != "":
command_str += " " + self.compile_options
print(self.compile_options)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we only print this in verbose/debug mode?

self.container_model_dirname, self.model_basename
)

print(command_str)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we only print this in verbose/debug mode?

def compile(model_path, **kwargs):
sess = CompileSession(model_path, **kwargs)
sess.Compile()
return sess.get_compiled_model_path()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about .constants.bin file or perhaps config file in the future? Do we need to return its path also?

# Alternative implementation: use env variable ONNX_MLIR_HOME?
# compile_args is the flags passed to onnx-mlir
r = OMPyCompile.compile(
"./test_add.onnx",
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we use test_add.mlir, we don't need to push test_add.onnx to the repository.

Signed-off-by: Tong Chen <chentong@us.ibm.com>
Signed-off-by: chentong319 <chentong@us.ibm.com>
Signed-off-by: chentong319 <chentong@us.ibm.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants