Build light weight PyRuntime without llvm or onnx-mlir #3044

chentong319 · 2025-01-15T20:51:53Z

Motivation

Python driver is needed to run the compiled model. Currently, the driver is built with onnx-mlir and can only be run in the env where onnx-mlir is built, typically inside the onnx-mlir docker image. When the compilation can be done by calling the onnx-mlir docker image, we'd like to run the compiled .so with python driver in the local env, so that all the packages installed in the local env can be used, rather than installing them on top of docker.

In order to reach this goal, the PR tried to remove the unnecessary dependencies of pyruntime if the light-weight pyruntime is the target. Otherwise, there is no change in how onnx-mlir is built and used.
Details can be found in docs/build-pyruntime-lit.md.
I tried the build on a z16 machine: it takes less than 2 minutes.

Components in this PR

CMakefile and source code changes to cut the dependencies of pyruntime. An option ONNX_MLIR_ENABLE_PYRUNTIME_LIT is used to control Cmake, and consequently a compile definition ENABLE_PYRUNTIME_LIT issued to control the source code.
Wrap the built pyruntime driver into a python package
This python package use python docker package to call the compiler. This is equivalent to docker/onnx-mlir.py in functionality but with "import docker" interface.

Test
Run successfully with utils/BuildPyRuntimeLit.sh.

Future works:

Support of float16
Not all llvm utilities are replaced. Only the essential ones have been implemented.
Can third_party/onnx be removed?
Try to integrate the precompiled lib for different os-arch into the package. Enable user to use pip install the package from pip server.
Try to integrate the utils/build-pyruntime-lit.sh into python package and invoke the build when pip install is executed.
Add the test of light weight PyRuntime into the build.

Signed-off-by: Chen Tong <[email protected]>

AlexandreEichenberger

Still very early review; let me get the overall picture.

How would the new code base change if we were to adapt it for PyTorch as well? Namely, what is common to any "docker compilation + pyruntime_lit" vs what is specialized for "ORT" and what would be specialized for PyTorch?

Do you mind giving me this high level info, and then I will be able to better understand the overall approach. Thanks

AlexandreEichenberger · 2025-01-23T01:33:37Z

CMakeLists.txt

@@ -11,6 +11,7 @@ option(ONNX_MLIR_ENABLE_STABLEHLO "Enable StableHLO support." ON)
 option(ONNX_MLIR_ENABLE_WERROR "Enable warnings as errors." OFF)
 option(ONNX_MLIR_SUPPRESS_THIRD_PARTY_WARNINGS "Suppress warning in third_party code." ON)
 option(ONNX_MLIR_ENABLE_JAVA "Set to ON for building the Java runtime, tools, and tests" ON)
+option(ONNX_MLIR_ENABLE_PYRUNTIME_LIT "Set to ON for building Python driver of running the compiled model without llvm-project." OFF)


Maybe the name should be explained: if off, then no pyruntime is build? Or its build anyway, but when on, then it's only the pyruntime? Or when off, pyruntime is build one way, but when off, its build another way?

Maybe the name could be a bit more explicit, depending on what the answer is from the question above.

minor question: _LIT is it for "_LIGHT"?

When this option is off, the pyruntime is built with onnx-mlir and llvm-project, as it was previously.
When this option is on, only the pyruntime is built without llvm-project.

Yes, LIT for LIGHT.

AlexandreEichenberger · 2025-01-23T01:35:36Z

CMakeLists.txt

-add_subdirectory(src)
-add_subdirectory(docs)
-add_subdirectory(test)
+if (ONNX_MLIR_ENABLE_PYRUNTIME_LIT)


I see this (and above): there are some dir that added on both path. Is it that the order of them is important?

Only the src is added on both path. I do not think the order of add_subdirectory matters. Just to keep the original add_subdirectory together.

AlexandreEichenberger · 2025-01-23T01:37:57Z

MLIR.cmake

-    )
-endif()
+if (ONNX_MLIR_ENABLE_PYRUNTIME_LIT)
+  function(llvm_update_compile_flags name)


Can you add a one liner comment on why this function is defined here.

AlexandreEichenberger · 2025-01-23T01:54:29Z

src/Runtime/python/onnxmlirdocker.py

@@ -0,0 +1,151 @@
+import numpy as np


That file is totally new, is that right?

Yes, it is the file to use docker package.

chentong319 · 2025-01-27T15:19:12Z

Still very early review; let me get the overall picture.

How would the new code base change if we were to adapt it for PyTorch as well? Namely, what is common to any "docker compilation + pyruntime_lit" vs what is specialized for "ORT" and what would be specialized for PyTorch?

Do you mind giving me this high level info, and then I will be able to better understand the overall approach. Thanks

This docker compilation + pyruntime_lit provides the basic functionality for compilation and run. The interface for ORT or PyTorch will be in the python package built on top of them. This PR contains only the package, onnxmlir, with ORT interface. In future, another package for PyTorch interface will be provided.

Signed-off-by: Chen Tong <[email protected]>

AlexandreEichenberger

LGTM as a first pass, will wait for additional feedback from the team as we develop this further. Thanks for taking this great first step @chentong319

tungld

LGTM!

Hope that the environment variable ENABLE_PYRUNTIME_LIT will be removed completely in next PRs.

LIT in ONNX_MLIR_ENABLE_PYRUNTIME_LIT is confusing with LIT tests. It would be better to use other word, e.g. LIGHT or to remove it.

tungld

LGTM!

chentong319 · 2025-01-29T15:00:21Z

LGTM!

Hope that the environment variable ENABLE_PYRUNTIME_LIT will be removed completely in next PRs.

LIT in ONNX_MLIR_ENABLE_PYRUNTIME_LIT is confusing with LIT tests. It would be better to use other word, e.g. LIGHT or to remove it.

The definition, ENABLE_PYRUNTIME_LIT, will be away finally. Current version is the transition from the previous one and the new one.
Alex has the same concern about _LIT. I will change it to LIGHT in this PR.

Signed-off-by: Chen Tong <[email protected]>

jenkins-droid · 2025-01-29T18:38:32Z

Jenkins Linux s390x Build #16181 [push] Build light weight PyRun... started at 13:38

jenkins-droid · 2025-01-29T18:39:35Z

Jenkins Linux amd64 Build #16180 [push] Build light weight PyRun... started at 12:39

jenkins-droid · 2025-01-29T18:54:35Z

Jenkins Linux ppc64le Build #15209 [push] Build light weight PyRun... started at 13:58

jenkins-droid · 2025-01-29T20:10:16Z

Jenkins Linux amd64 Build #16180 [push] Build light weight PyRun... passed after 1 hr 30 min

jenkins-droid · 2025-01-29T20:25:48Z

Jenkins Linux s390x Build #16181 [push] Build light weight PyRun... passed after 1 hr 47 min

jenkins-droid · 2025-01-29T21:34:06Z

Jenkins Linux ppc64le Build #15209 [push] Build light weight PyRun... passed after 2 hr 54 min

* pass test Signed-off-by: Chen Tong <[email protected]> * package Signed-off-by: Chen Tong <[email protected]> * clean makefile Signed-off-by: Chen Tong <[email protected]> * document Signed-off-by: Chen Tong <[email protected]> * fix MLIR.cmake Signed-off-by: Chen Tong <[email protected]> * fix script Signed-off-by: Chen Tong <[email protected]> * fix Signed-off-by: Chen Tong <[email protected]> * add comments Signed-off-by: Chen Tong <[email protected]> * LIGHT Signed-off-by: Chen Tong <[email protected]> --------- Signed-off-by: Chen Tong <[email protected]>

chentong319 added 5 commits January 15, 2025 12:20

pass test

4ab6256

Signed-off-by: Chen Tong <[email protected]>

package

49034b5

Signed-off-by: Chen Tong <[email protected]>

clean makefile

45b1d9d

Signed-off-by: Chen Tong <[email protected]>

Merge remote-tracking branch 'upstream/main' into pyruntime-lit

0ee214f

document

beff2e3

Signed-off-by: Chen Tong <[email protected]>

chentong319 marked this pull request as draft January 15, 2025 21:00

chentong319 added 5 commits January 15, 2025 16:19

fix MLIR.cmake

28584dd

Signed-off-by: Chen Tong <[email protected]>

fix script

a92ccc2

Signed-off-by: Chen Tong <[email protected]>

Merge remote-tracking branch 'upstream/main' into pyruntime-lit

2037d8d

fix

8ecd7ff

Signed-off-by: Chen Tong <[email protected]>

Merge remote-tracking branch 'upstream/main' into pyruntime-lit

79ab54f

chentong319 marked this pull request as ready for review January 21, 2025 14:57

chentong319 requested review from gongsu832, AlexandreEichenberger and tungld January 21, 2025 14:59

AlexandreEichenberger reviewed Jan 23, 2025

View reviewed changes

Merge remote-tracking branch 'upstream/main' into pyruntime-lit

0bd61d2

chentong319 added 2 commits January 27, 2025 10:46

add comments

0694678

Signed-off-by: Chen Tong <[email protected]>

Merge remote-tracking branch 'upstream/main' into pyruntime-lit

0d049a5

AlexandreEichenberger approved these changes Jan 28, 2025

View reviewed changes

tungld reviewed Jan 29, 2025

View reviewed changes

tungld approved these changes Jan 29, 2025

View reviewed changes

chentong319 added 2 commits January 29, 2025 11:45

LIGHT

8c8a17f

Signed-off-by: Chen Tong <[email protected]>

Merge remote-tracking branch 'upstream/main' into pyruntime-lit

8ff28b1

chentong319 merged commit 2e4a46a into onnx:main Jan 29, 2025
6 of 7 checks passed

chentong319 deleted the pyruntime-lit branch January 29, 2025 18:38

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Build light weight PyRuntime without llvm or onnx-mlir #3044

Build light weight PyRuntime without llvm or onnx-mlir #3044

chentong319 commented Jan 15, 2025 •

edited

Loading

AlexandreEichenberger left a comment

AlexandreEichenberger Jan 23, 2025

chentong319 Jan 27, 2025

AlexandreEichenberger Jan 23, 2025

chentong319 Jan 27, 2025

AlexandreEichenberger Jan 23, 2025

AlexandreEichenberger Jan 23, 2025

chentong319 Jan 27, 2025

chentong319 commented Jan 27, 2025

AlexandreEichenberger left a comment

tungld left a comment

tungld left a comment

chentong319 commented Jan 29, 2025 •

edited

Loading

jenkins-droid commented Jan 29, 2025

jenkins-droid commented Jan 29, 2025

jenkins-droid commented Jan 29, 2025

jenkins-droid commented Jan 29, 2025

jenkins-droid commented Jan 29, 2025

jenkins-droid commented Jan 29, 2025

Build light weight PyRuntime without llvm or onnx-mlir #3044

Build light weight PyRuntime without llvm or onnx-mlir #3044

Conversation

chentong319 commented Jan 15, 2025 • edited Loading

AlexandreEichenberger left a comment

Choose a reason for hiding this comment

AlexandreEichenberger Jan 23, 2025

Choose a reason for hiding this comment

chentong319 Jan 27, 2025

Choose a reason for hiding this comment

AlexandreEichenberger Jan 23, 2025

Choose a reason for hiding this comment

chentong319 Jan 27, 2025

Choose a reason for hiding this comment

AlexandreEichenberger Jan 23, 2025

Choose a reason for hiding this comment

AlexandreEichenberger Jan 23, 2025

Choose a reason for hiding this comment

chentong319 Jan 27, 2025

Choose a reason for hiding this comment

chentong319 commented Jan 27, 2025

AlexandreEichenberger left a comment

Choose a reason for hiding this comment

tungld left a comment

Choose a reason for hiding this comment

tungld left a comment

Choose a reason for hiding this comment

chentong319 commented Jan 29, 2025 • edited Loading

jenkins-droid commented Jan 29, 2025

jenkins-droid commented Jan 29, 2025

jenkins-droid commented Jan 29, 2025

jenkins-droid commented Jan 29, 2025

jenkins-droid commented Jan 29, 2025

jenkins-droid commented Jan 29, 2025

chentong319 commented Jan 15, 2025 •

edited

Loading

chentong319 commented Jan 29, 2025 •

edited

Loading