[TVM EP] Improved usability of TVM EP

setup.py

@@ -12,6 +12,7 @@
 import platform
 import subprocess
 import sys
+import textwrap
 import datetime
 
 from pathlib import Path
@@ -145,6 +146,33 @@ def _rewrite_ld_preload_tensorrt(self, to_preload):
                     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 @@ def run(self):
                 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]