Skip to content

Commit

Permalink
improve readability of inserted code
Browse files Browse the repository at this point in the history
  • Loading branch information
KJlaccHoeUM9l committed Jan 21, 2022
1 parent 313e26d commit 908a5b4
Showing 1 changed file with 25 additions and 19 deletions.
44 changes: 25 additions & 19 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import platform
import subprocess
import sys
import textwrap
import datetime

from pathlib import Path
Expand Down Expand Up @@ -147,25 +148,30 @@ def _rewrite_ld_preload_tensorrt(self, to_preload):

def _rewrite_ld_preload_tvm(self):
with open('onnxruntime/capi/_ld_preload.py', 'a') as f:
f.write('import warnings\n\n')
f.write('try:\n')
f.write(' # This import is necessary in order to delegate the loading of libtvm.so to TVM.\n')
f.write(' import tvm\n')
f.write('except ImportError as e:\n')
f.write(' warnings.warn(\n')
f.write(' f"WARNING: Failed to import TVM, libtvm.so was not loaded. More details: {e}"\n')
f.write(' )\n\n')
f.write('try:\n')
f.write(' # Working between the C++ and Python parts in TVM EP is done using\n')
f.write(' # the PackedFunc and Registry classes. In order to use a Python function in C++ code,\n')
f.write(' # it must be registered in the global table of functions. Registration is carried out\n')
f.write(' # through JIT interface, so it is necessary to call special functions for registration.\n')
f.write(' # To do this, we need to make the following import.\n')
f.write(' import onnxruntime.providers.stvm\n')
f.write('except ImportError as e:\n')
f.write(' warnings.warn(\n')
f.write(' f"WARNING: Failed to register python functions for TVM EP. More details: {e}"\n')
f.write(' )\n')
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:
Expand Down

0 comments on commit 908a5b4

Please sign in to comment.