cross compile guard for windows

Author: apbose

py/torch_tensorrt/_features.py

@@ -3,7 +3,10 @@
 from collections import namedtuple
 from typing import Any, Callable, Dict, List, Optional, Type, TypeVar
 
-from torch_tensorrt._utils import sanitized_torch_version
+from torch_tensorrt._utils import (
+    check_cross_compile_trt_win_lib,
+    sanitized_torch_version,
+)
 
 from packaging import version
 
@@ -15,6 +18,7 @@
         "dynamo_frontend",
         "fx_frontend",
         "refit",
+        "windows_cross_compile",
     ],
 )
 
@@ -38,9 +42,15 @@
 _DYNAMO_FE_AVAIL = version.parse(sanitized_torch_version()) >= version.parse("2.1.dev")
 _FX_FE_AVAIL = True
 _REFIT_AVAIL = version.parse(sys.version.split()[0]) < version.parse("3.13")
+_WINDOWS_CROSS_COMPILE = check_cross_compile_trt_win_lib
 
 ENABLED_FEATURES = FeatureSet(
-    _TS_FE_AVAIL, _TORCHTRT_RT_AVAIL, _DYNAMO_FE_AVAIL, _FX_FE_AVAIL, _REFIT_AVAIL
+    _TS_FE_AVAIL,
+    _TORCHTRT_RT_AVAIL,
+    _DYNAMO_FE_AVAIL,
+    _FX_FE_AVAIL,
+    _REFIT_AVAIL,
+    _WINDOWS_CROSS_COMPILE,
 )
 
 
@@ -80,6 +90,22 @@ def not_implemented(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
     return wrapper
 
 
+def needs_cross_compile(f: Callable[..., Any]) -> Callable[..., Any]:
+    def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
+        if ENABLED_FEATURES.windows_cross_compile:
+            return f(*args, **kwargs)
+        else:
+
+            def not_implemented(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
+                raise NotImplementedError(
+                    "Windows cross compilation feature is not available"
+                )
+
+            return not_implemented(*args, **kwargs)
+
+    return wrapper
+
+
 T = TypeVar("T")
 
 

py/torch_tensorrt/_utils.py

@@ -1,6 +1,10 @@
+import ctypes
+import platform
+import sys
 from typing import Any
 
 import torch
+from torch_tensorrt import _find_lib
 
 
 def sanitized_torch_version() -> Any:
@@ -9,3 +13,27 @@ def sanitized_torch_version() -> Any:
         if ".nv" not in torch.__version__
         else torch.__version__.split(".nv")[0]
     )
+
+
+def check_cross_compile_trt_win_lib(lib_name):
+    if sys.platform.startswith("linux"):
+        LINUX_PATHS = ["/usr/local/cuda-12.8/lib64", "/usr/lib", "/usr/lib64"]
+
+    if platform.uname().processor == "x86_64":
+        LINUX_PATHS += [
+            "/usr/lib/x86_64-linux-gnu",
+        ]
+    elif platform.uname().processor == "aarch64":
+        LINUX_PATHS += ["/usr/lib/aarch64-linux-gnu"]
+
+    LINUX_LIBS = [
+        f"libnvinfer_builder_resource_win.so.*",
+    ]
+
+    for lib in LINUX_LIBS:
+        try:
+            ctypes.CDLL(_find_lib(lib, LINUX_PATHS))
+            return True
+        except:
+            continue
+    return False

tests/py/dynamo/runtime/test_003_cross_compile_for_windows.py

@@ -7,6 +7,7 @@
 import torch
 import torch_tensorrt
 from torch.testing._internal.common_utils import TestCase
+from torch_tensorrt._utils import check_cross_compile_trt_win_lib
 
 from ..testing_utilities import DECIMALS_OF_AGREEMENT
 
@@ -16,6 +17,9 @@ class TestCrossCompileSaveForWindows(TestCase):
         platform.system() != "Linux" or platform.architecture()[0] != "64bit",
         "Cross compile for windows can only be enabled on linux x86-64 platform",
     )
+    @unittest.skipIf(
+        not (check_cross_compile_trt_win_lib()),
+    )
     @pytest.mark.unit
     def test_cross_compile_for_windows(self):
         class Add(torch.nn.Module):
@@ -40,6 +44,9 @@ def forward(self, a, b):
         platform.system() != "Linux" or platform.architecture()[0] != "64bit",
         "Cross compile for windows can only be enabled on linux x86-64 platform",
     )
+    @unittest.skipIf(
+        not (check_cross_compile_trt_win_lib()),
+    )
     @pytest.mark.unit
     def test_dynamo_cross_compile_for_windows(self):
         class Add(torch.nn.Module):