feat(//py): Use TensorRT to fill in .so libraries automatically if

possible

Signed-off-by: Naren Dasan <naren@narendasan.com>
Signed-off-by: Naren Dasan <narens@nvidia.com>

py/setup.py

@@ -23,7 +23,9 @@
 JETPACK_VERSION = None
 
 __version__ = '1.2.0a0'
-
+__cuda_version__ = '11.3'
+__cudnn_version__ = '8.2'
+__tensorrt_version__ = '8.2'
 
 def get_git_revision_short_hash() -> str:
     return subprocess.check_output(['git', 'rev-parse', '--short', 'HEAD']).decode('ascii').strip()
@@ -114,7 +116,10 @@ def gen_version_file():
 
     with open(dir_path + '/torch_tensorrt/_version.py', 'w') as f:
         print("creating version file")
-        f.write("__version__ = \"" + __version__ + '\"')
+        f.write("__version__ = \"" + __version__ + '\"\n')
+        f.write("__cuda_version__ = \"" + __cuda_version__ + '\"\n')
+        f.write("__cudnn_version__ = \"" + __cudnn_version__ + '\"\n')
+        f.write("__tensorrt_version__ = \"" + __tensorrt_version__ + '\"\n')
 
 
 def copy_libtorchtrt(multilinux=False):

py/torch_tensorrt/__init__.py

@@ -1,13 +1,78 @@
+import ctypes
+import glob
 import os
 import sys
+import warnings
+from torch_tensorrt._version import __version__, __cuda_version__, __cudnn_version__, __tensorrt_version__
+
 
 if sys.version_info < (3,):
     raise Exception("Python 2 has reached end-of-life and is not supported by Torch-TensorRT")
 
-import ctypes
+def _parse_semver(version):
+    split = version.split(".")
+    if len(split) < 3:
+        split.append("")
+
+    return {
+        "major": split[0],
+        "minor": split[1],
+        "patch": split[2]
+    }
+
+def _find_lib(name, paths):
+    for path in paths:
+        libpath = os.path.join(path, name)
+        if os.path.isfile(libpath):
+            return libpath
+
+    raise FileNotFoundError(
+        f"Could not find {name}\n  Search paths: {paths}"
+    )
+
+try:
+    import tensorrt
+except:
+    cuda_version = _parse_semver(__cuda_version__)
+    cudnn_version = _parse_semver(__cudnn_version__)
+    tensorrt_version = _parse_semver(__tensorrt_version__)
+
+    CUDA_MAJOR = cuda_version["major"]
+    CUDNN_MAJOR = cudnn_version["major"]
+    TENSORRT_MAJOR = tensorrt_version["major"]
+
+    if sys.platform.startswith("win"):
+        WIN_LIBS = [
+            f"cublas64_{CUDA_MAJOR}.dll",
+            f"cublasLt64_{CUDA_MAJOR}.dll",
+            f"cudnn64_{CUDNN_MAJOR}.dll",
+            "nvinfer.dll",
+            "nvinfer_plugin.dll",
+        ]
+
+        WIN_PATHS = os.environ["PATH"].split(os.path.pathsep)
+
+
+        for lib in WIN_LIBS:
+            ctypes.CDLL(_find_lib(lib, WIN_PATHS))
+
+    elif sys.platform.startswith("linux"):
+        LINUX_PATHS = [
+            "/usr/lib/x86_64-linux-gnu",
+            "/usr/local/cuda/lib64",
+        ] + os.environ["LD_LIBRARY_PATH"].split(os.path.pathsep)
+
+        LINUX_LIBS = [
+            f"libcudnn.so.{CUDNN_MAJOR}",
+            f"libnvinfer.so.{TENSORRT_MAJOR}",
+            f"libnvinfer_plugin.so.{TENSORRT_MAJOR}",
+        ]
+
+        for lib in LINUX_LIBS:
+            ctypes.CDLL(_find_lib(lib, LINUX_PATHS))
+
 import torch
 
-from torch_tensorrt._version import __version__
 from torch_tensorrt._compile import *
 from torch_tensorrt._util import *
 from torch_tensorrt import ts