feat: Win/Linux Dual Compatible WORKSPACE + Upgrade CUDA + Upgrade PyT (#2907)

Author: gs-olive

.bazelrc

@@ -21,9 +21,9 @@
 # | Build Configurations                                       |
 # +------------------------------------------------------------+
 # Enable colorful output of GCC
-build --cxxopt="-fdiagnostics-color=always"
-build --cxxopt='-std=c++17'
-#build --linkopt="-Wl,--no-as-needed"
+build:default --cxxopt="-fdiagnostics-color=always"
+build:default --cxxopt='-std=c++17'
+#build:default --linkopt="-Wl,--no-as-needed"
 
 build:windows --cxxopt="/GS-" --cxxopt="/std:c++17" --cxxopt="/permissive-"
 build:windows --cxxopt="/wd4244" --cxxopt="/wd4267" --cxxopt="/wd4819"

.bazelversion

@@ -1 +1 @@
-6.2.1
+6.3.2

.github/workflows/docgen.yml

@@ -14,11 +14,11 @@ jobs:
     if: ${{ ! contains(github.actor, 'pytorchbot') }}
     environment: pytorchbot-env
     container:
-      image: docker.io/pytorch/manylinux-builder:cuda12.1
+      image: docker.io/pytorch/manylinux-builder:cuda12.4
       options: --gpus all
     env:
-      CUDA_HOME: /usr/local/cuda-12.1
-      VERSION_SUFFIX: cu121
+      CUDA_HOME: /usr/local/cuda-12.4
+      VERSION_SUFFIX: cu124
       CI_BUILD: 1
     steps:
       - uses: actions/checkout@v3
@@ -33,14 +33,14 @@ jobs:
       - name: Install base deps
         run: |
           python3 -m pip install pip --upgrade
-          python3 -m pip install pyyaml numpy torch --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu121
+          python3 -m pip install pyyaml numpy torch --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu124
           ./packaging/pre_build_script.sh
       - name: Get HEAD SHA
         id: vars
         run: echo "sha=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT
       - name: Build Python Package
         run: |
-          python3 -m pip install --pre . --extra-index-url https://download.pytorch.org/whl/nightly/cu121
+          python3 -m pip install --pre . --extra-index-url https://download.pytorch.org/whl/nightly/cu124
       - name: Generate New Docs
         run: |
           cd docsrc
@@ -61,4 +61,4 @@ jobs:
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref_name }}
-  cancel-in-progress: true
+  cancel-in-progress: true

README.md

@@ -15,7 +15,7 @@ Torch-TensorRT
 ---
 <div align="left">
 
-Torch-TensorRT brings the power of TensorRT to PyTorch. Accelerate inference latency by up to 5x compared to eager execution in just one line of code. 
+Torch-TensorRT brings the power of TensorRT to PyTorch. Accelerate inference latency by up to 5x compared to eager execution in just one line of code.
 </div></div>
 
 ## Installation
@@ -52,7 +52,7 @@ optimized_model(x) # this will be fast!
 ```
 
 ### Option 2: Export
-If you want to optimize your model ahead-of-time and/or deploy in a C++ environment, Torch-TensorRT provides an export-style workflow that serializes an optimized module. This module can be deployed in PyTorch or with libtorch (i.e. without a Python dependency). 
+If you want to optimize your model ahead-of-time and/or deploy in a C++ environment, Torch-TensorRT provides an export-style workflow that serializes an optimized module. This module can be deployed in PyTorch or with libtorch (i.e. without a Python dependency).
 
 #### Step 1: Optimize + serialize
 ```python
@@ -62,7 +62,7 @@ import torch_tensorrt
 model = MyModel().eval().cuda() # define your model here
 inputs = [torch.randn((1, 3, 224, 224)).cuda()] # define a list of representative inputs here
 
-trt_gm = torch_tensorrt.compile(model, ir="dynamo", inputs) 
+trt_gm = torch_tensorrt.compile(model, ir="dynamo", inputs)
 torch_tensorrt.save(trt_gm, "trt.ep", inputs=inputs) # PyTorch only supports Python runtime for an ExportedProgram. For C++ deployment, use a TorchScript file
 torch_tensorrt.save(trt_gm, "trt.ts", output_format="torchscript", inputs=inputs)
 ```
@@ -116,9 +116,9 @@ auto results = trt_mod.forward({input_tensor});
 
 These are the following dependencies used to verify the testcases. Torch-TensorRT can work with other versions, but the tests are not guaranteed to pass.
 
-- Bazel 5.2.0
-- Libtorch 2.4.0.dev (latest nightly) (built with CUDA 12.4)
-- CUDA 12.1
+- Bazel 6.3.2
+- Libtorch 2.5.0.dev (latest nightly) (built with CUDA 12.4)
+- CUDA 12.4
 - TensorRT 10.0.1.6
 
 ## Deprecation Policy

WORKSPACE

@@ -43,7 +43,13 @@ local_repository(
 new_local_repository(
     name = "cuda",
     build_file = "@//third_party/cuda:BUILD",
-    path = "/usr/local/cuda-12.1/",
+    path = "/usr/local/cuda-12.4/",
+)
+
+new_local_repository(
+    name = "cuda_win",
+    build_file = "@//third_party/cuda:BUILD",
+    path = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.4/",
 )
 
 #############################################################################################################
@@ -54,14 +60,21 @@ http_archive(
     name = "libtorch",
     build_file = "@//third_party/libtorch:BUILD",
     strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/nightly/cu121/libtorch-cxx11-abi-shared-with-deps-latest.zip"],
+    urls = ["https://download.pytorch.org/libtorch/nightly/cu124/libtorch-cxx11-abi-shared-with-deps-latest.zip"],
 )
 
 http_archive(
     name = "libtorch_pre_cxx11_abi",
     build_file = "@//third_party/libtorch:BUILD",
     strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/nightly/cu121/libtorch-shared-with-deps-latest.zip"],
+    urls = ["https://download.pytorch.org/libtorch/nightly/cu124/libtorch-shared-with-deps-latest.zip"],
+)
+
+http_archive(
+    name = "libtorch_win",
+    build_file = "@//third_party/libtorch:BUILD",
+    strip_prefix = "libtorch",
+    urls = ["https://download.pytorch.org/libtorch/nightly/cu124/libtorch-win-shared-with-deps-latest.zip"],
 )
 
 # Download these tarballs manually from the NVIDIA website
@@ -78,6 +91,18 @@ http_archive(
     ],
 )
 
+http_archive(
+    name = "tensorrt_win",
+    build_file = "@//third_party/tensorrt/archive:BUILD",
+    sha256 = "d667bd10b178e239b621a8929008ef3e27967d181bf07a39845a0f99edeec47a",
+    strip_prefix = "TensorRT-10.0.1.6",
+    urls = [
+        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.0.1/zip/TensorRT-10.0.1.6.Windows10.win10.cuda-12.4.zip",
+    ],
+)
+
+
+
 ####################################################################################
 # Locally installed dependencies (use in cases of custom dependencies or aarch64)
 ####################################################################################

core/BUILD

@@ -17,6 +17,13 @@ config_setting(
     },
 )
 
+config_setting(
+    name = "windows",
+    constraint_values = [
+        "@platforms//os:windows",
+    ],
+)
+
 cc_library(
     name = "core",
     srcs = [
@@ -31,10 +38,10 @@ cc_library(
         "//core/partitioning",
         "//core/runtime",
         "//core/util/logging",
-        "@tensorrt//:nvinfer",
     ] + select({
-        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default": ["@libtorch"],
+        ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
+        ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
     }),
     alwayslink = True,
 )

core/conversion/BUILD

@@ -10,6 +10,13 @@ config_setting(
     },
 )
 
+config_setting(
+    name = "windows",
+    constraint_values = [
+        "@platforms//os:windows",
+    ],
+)
+
 cc_library(
     name = "conversion",
     srcs = [
@@ -26,10 +33,10 @@ cc_library(
         "//core/conversion/var",
         "//core/ir",
         "//core/util:prelude",
-        "@tensorrt//:nvinfer",
     ] + select({
-        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default": ["@libtorch"],
+        ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
+        ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
     }),
     alwayslink = True,
 )

core/conversion/conversionctx/BUILD

@@ -10,6 +10,13 @@ config_setting(
     },
 )
 
+config_setting(
+    name = "windows",
+    constraint_values = [
+        "@platforms//os:windows",
+    ],
+)
+
 cc_library(
     name = "conversionctx",
     srcs = [
@@ -21,10 +28,10 @@ cc_library(
     deps = [
         "//core/ir",
         "//core/util:prelude",
-        "@tensorrt//:nvinfer",
     ] + select({
-        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default": ["@libtorch"],
+        ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
+        ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
     }),
     alwayslink = True,
 )

core/conversion/converters/BUILD

@@ -10,6 +10,13 @@ config_setting(
     },
 )
 
+config_setting(
+    name = "windows",
+    constraint_values = [
+        "@platforms//os:windows",
+    ],
+)
+
 cc_library(
     name = "weights",
     srcs = [
@@ -21,10 +28,10 @@ cc_library(
     deps = [
         "//core/conversion/conversionctx",
         "//core/util:prelude",
-        "@tensorrt//:nvinfer",
     ] + select({
-        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default": ["@libtorch"],
+        ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
+        ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
     }),
     alwayslink = True,
 )
@@ -41,10 +48,10 @@ cc_library(
         ":weights",
         "//core/conversion/conversionctx",
         "//core/util:prelude",
-        "@tensorrt//:nvinfer",
     ] + select({
-        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default": ["@libtorch"],
+        ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
+        ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
     }),
     alwayslink = True,
 )
@@ -98,10 +105,10 @@ cc_library(
         "//core/conversion/var",
         "//core/plugins:torch_tensorrt_plugins",
         "//core/util:prelude",
-        "@tensorrt//:nvinfer",
     ] + select({
-        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default": ["@libtorch"],
+        ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
+        ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
     }),
     alwayslink = True,
 )

core/conversion/evaluators/BUILD

@@ -10,6 +10,13 @@ config_setting(
     },
 )
 
+config_setting(
+    name = "windows",
+    constraint_values = [
+        "@platforms//os:windows",
+    ],
+)
+
 cc_library(
     name = "evaluators",
     srcs = [
@@ -28,6 +35,7 @@ cc_library(
         "//core/conversion/var",
         "//core/util:prelude",
     ] + select({
+        ":windows": ["@libtorch_win//:libtorch"],
         ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
         "//conditions:default": ["@libtorch"],
     }),

core/conversion/tensorcontainer/BUILD

@@ -10,6 +10,13 @@ config_setting(
     },
 )
 
+config_setting(
+    name = "windows",
+    constraint_values = [
+        "@platforms//os:windows",
+    ],
+)
+
 cc_library(
     name = "tensorcontainer",
     srcs = [
@@ -20,10 +27,10 @@ cc_library(
     ],
     deps = [
         "//core/util:prelude",
-        "@tensorrt//:nvinfer",
     ] + select({
-        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default": ["@libtorch"],
+        ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
+        ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
     }),
     alwayslink = True,
 )

core/conversion/var/BUILD

@@ -10,6 +10,13 @@ config_setting(
     },
 )
 
+config_setting(
+    name = "windows",
+    constraint_values = [
+        "@platforms//os:windows",
+    ],
+)
+
 cc_library(
     name = "var",
     srcs = [
@@ -23,10 +30,10 @@ cc_library(
         "//core/conversion/converters:converter_util",
         "//core/conversion/tensorcontainer",
         "//core/util:prelude",
-        "@tensorrt//:nvinfer",
     ] + select({
-        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default": ["@libtorch"],
+        ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
+        ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
     }),
     alwayslink = True,
 )

core/ir/BUILD

@@ -10,6 +10,13 @@ config_setting(
     },
 )
 
+config_setting(
+    name = "windows",
+    constraint_values = [
+        "@platforms//os:windows",
+    ],
+)
+
 cc_library(
     name = "ir",
     srcs = [
@@ -23,10 +30,10 @@ cc_library(
     ],
     deps = [
         "//core/util:prelude",
-        "@tensorrt//:nvinfer",
     ] + select({
-        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default": ["@libtorch"],
+        ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
+        ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
     }),
     alwayslink = True,
 )