saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[TVM EP] Improved usability of TVM EP (#10241)

Browse source 
* improved usability of TVM EP
* moved technical import under a condition related to TVM EP only
* Revert "moved technical import under a condition related to TVM EP only"
* add conditional _ld_preload.py file extension for TVM EP
* improve readability of inserted code
This commit is contained in:
Alexey Gladyshev 2022-01-25 20:48:08 +03:00 committed by GitHub
parent 6e95c0316d
commit a0fe4a7c1c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 30 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
setup.py
View file

@ -12,6 +12,7 @@ from shutil import copyfile
import platform
import subprocess
import sys
import textwrap
import datetime
from pathlib import Path
@ -145,6 +146,33 @@ try:
f.write(' import os\n')
f.write(' os.environ["ORT_TENSORRT_UNAVAILABLE"] = "1"\n')
def _rewrite_ld_preload_tvm(self):
with open('onnxruntime/capi/_ld_preload.py', 'a') as f:
f.write(textwrap.dedent(
"""
import warnings
try:
# This import is necessary in order to delegate the loading of libtvm.so to TVM.
import tvm
except ImportError as e:
warnings.warn(
f"WARNING: Failed to import TVM, libtvm.so was not loaded. More details: {e}"
)
try:
# Working between the C++ and Python parts in TVM EP is done using the PackedFunc and
# Registry classes. In order to use a Python function in C++ code, it must be registered in
# the global table of functions. Registration is carried out through the JIT interface,
# so it is necessary to call special functions for registration.
# To do this, we need to make the following import.
import onnxruntime.providers.stvm
except ImportError as e:
warnings.warn(
f"WARNING: Failed to register python functions to work with TVM EP. More details: {e}"
)
"""
))
def run(self):
if is_manylinux:
source = 'onnxruntime/capi/onnxruntime_pybind11_state.so'
@ -207,6 +235,8 @@ try:
self._rewrite_ld_preload(to_preload)
self._rewrite_ld_preload_cuda(to_preload_cuda)
self._rewrite_ld_preload_tensorrt(to_preload_tensorrt)
if package_name == 'onnxruntime-tvm':
self._rewrite_ld_preload_tvm()
_bdist_wheel.run(self)
if is_manylinux and not disable_auditwheel_repair:
file = glob(path.join(self.dist_dir, '*linux*.whl'))[0]