Repository: alibaba/BladeDISC (language: python)

Issue:
TensorRT static linkage is not supported.



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/cpu_reusable.yml b/.github/workflows/cpu_reusable.yml
index 11f642e2fc8..a2aa0c56e40 100644
--- a/.github/workflows/cpu_reusable.yml
+++ b/.github/workflows/cpu_reusable.yml
@@ -46,6 +46,7 @@ jobs:
         docker build -t disc-dev-cpu \
           --build-arg BASEIMAGE=nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04 \
           --build-arg DISC_HOST_TF_VERSION="tensorflow==1.15.5" \
+          --build-arg DEVICE=cpu \
           -f docker/dev/Dockerfile .
     - name: Build And Test DISC
       run: |
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 789dbd4efb5..8d76039ab45 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -37,6 +37,7 @@ jobs:
         docker build -t disc-dev-cuda10.0 \
           --build-arg BASEIMAGE=nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04 \
           --build-arg DISC_HOST_TF_VERSION="tensorflow-gpu==1.15" \
+          --build-arg DEVICE=cu100 \
           -f docker/dev/Dockerfile .
     - name: Build And Test DISC
       run: |
@@ -82,6 +83,7 @@ jobs:
         docker build --cache-from bladedisc/bladedisc:latest-devel-cuda11.0 -t disc-dev-cuda11.0 \
           --build-arg BASEIMAGE=nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04 \
           --build-arg DISC_HOST_TF_VERSION="tensorflow-gpu==2.4" \
+          --build-arg DEVICE=cu110 \
           -f docker/dev/Dockerfile .
     - name: Build And Test DISC
       run: |
@@ -136,7 +138,9 @@ jobs:
         git submodule update --depth=1 --init --recursive
         docker pull bladedisc/bladedisc:latest-devel-cuda11.0
         docker build --cache-from bladedisc/bladedisc:latest-devel-cuda11.0 -t disc-dev-cuda11.0 \
-          --build-arg BASEIMAGE=nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04 -f docker/dev/Dockerfile .
+          --build-arg BASEIMAGE=nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04 \
+          --build-arg DEVICE=cu110 \
+          -f docker/dev/Dockerfile .
     - name: Build and Test DISC
       run: |
         set -e
@@ -164,6 +168,52 @@ jobs:
       run: |
         set -e
         bash ./scripts/ci/deploy_pytorch_blade.sh
+  CUDA10_2-TORCH181:
+    if: github.repository == 'alibaba/BladeDISC'
+    # The type of runner that the job will run on
+    runs-on: [self-hosted, gpu-t4]
+    # Steps represent a sequence of tasks that will be executed as part of the job
+    steps:
+    # Runs a single command using the runners shell
+    - name: Checkout
+      uses: actions/checkout@v2.4.0
+    - name: Build Dev Docker
+      run: |
+        set -e
+        git submodule sync
+        git submodule update --depth=1 --init --recursive
+        docker build -t disc-dev-cuda10.2 \
+          --build-arg BASEIMAGE=nvidia/cuda:10.2-cudnn8-devel-ubuntu18.04 \
+          --build-arg DEVICE=cu102 \
+          -f docker/dev/Dockerfile .
+    - name: Build and Test DISC
+      run: |
+        set -e
+        nvidia-docker run --rm -t --user $(id -u):$(id -g) \
+          -v $HOME/.cache:$HOME/.cache \
+          -v /etc/passwd:/etc/passwd:ro \
+          -v /etc/group:/etc/group:ro \
+          -v $PWD:/disc \
+          -e GITHUB_WORKFLOW=$GITHUB_WORKFLOW \
+          -e TORCH_BLADE_BUILD_TENSORRT_STATIC=ON \
+          -e TORCH_BLADE_CI_BUILD_TORCH_VERSION=1.8.1+cu102 \
+          -w /disc \
+          disc-dev-cuda10.2 bash ./scripts/ci/test_pytorch_blade.sh
+    - name: Deploy PyTorch Blade
+      if: github.event.ref == 'refs/heads/main'
+      env:
+        ALIYUN_DOCKER_USERNAME: ${{ secrets.ALIYUN_DOCKER_USERNAME }}
+        ALIYUN_DOCKER_PASSWORD: ${{ secrets.ALIYUN_DOCKER_PASSWORD }}
+        DOCKER_USERNAME: ${{ secrets.DOCKER_USERNAME }}
+        DOCKER_PASSWORD: ${{ secrets.DOCKER_PASSWORD }}
+        GITHUB_PULL_REQUEST: ${{ github.event.number }}
+        LOCAL_DEV_DOCKER: disc-dev-cuda10.2
+        REMOTE_DEV_DOCKER: bladedisc:latest-devel-cuda10.2
+        REMOTE_RUNTIME_DOCKER: bladedisc:latest-runtime-torch1.8.1+cu102
+        RUNTIME_BASEIMAGE: nvidia/cuda:10.2-cudnn8-devel-ubuntu18.04
+      run: |
+        set -e
+        bash ./scripts/ci/deploy_pytorch_blade.sh
   CPU-TF115:
     uses: ./.github/workflows/cpu_reusable.yml
     with:
diff --git a/docker/dev/Dockerfile b/docker/dev/Dockerfile
index 8d63b63befd..6b6e3dbc7b4 100644
--- a/docker/dev/Dockerfile
+++ b/docker/dev/Dockerfile
@@ -14,7 +14,9 @@ RUN bash /install/scripts/find-fastest-apt.sh && \
     bash /install/scripts/install-bazel.sh && \
     bash /install/scripts/install-python.sh
 
-RUN bash /install/scripts/install-tensorrt.sh
+ARG DEVICE=cu110
+RUN bash /install/scripts/install-tensorrt.sh "${DEVICE}"
+RUn bash /install/scripts/patch-cu102.sh "${DEVICE}"
 
 ENV PATH="/opt/cmake/bin:${PATH}"
-ENV LD_LIBRARY_PATH="/usr/local/TensorRT/lib/:${LD_LIBRARY_PATH}"
+ENV LD_LIBRARY_PATH="/usr/local/TensorRT/lib/:${LD_LIBRARY_PATH}"
\ No newline at end of file
diff --git a/docker/runtime/Dockerfile.pytorch b/docker/runtime/Dockerfile.pytorch
index aa0c8f6201c..3b2215daa61 100644
--- a/docker/runtime/Dockerfile.pytorch
+++ b/docker/runtime/Dockerfile.pytorch
@@ -1,9 +1,12 @@
 ARG BASEIMAGE=nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04
 FROM ${BASEIMAGE}
 
+ARG DEVICE=cu110
 ADD ./docker/scripts /install/scripts
-RUN bash /install/scripts/find-fastest-apt.sh
-RUN bash /install/scripts/install-tensorrt.sh
+RUN bash /install/scripts/find-fastest-apt.sh && \
+        apt-get install -y -qq wget curl unzip
+RUN bash /install/scripts/install-tensorrt.sh "${DEVICE}"
+RUn bash /install/scripts/patch-cu102.sh "${DEVICE}"
 
 ADD ./build/torch_blade*.whl  /install/python/
 
diff --git a/docker/scripts/install-cmake.sh b/docker/scripts/install-cmake.sh
index b3e64f2e840..dc9dc038add 100644
--- a/docker/scripts/install-cmake.sh
+++ b/docker/scripts/install-cmake.sh
@@ -12,8 +12,8 @@
 
 
 set -ex
-wget https://github.com/Kitware/CMake/releases/download/v3.20.0/cmake-3.20.0-linux-x86_64.sh \
-  -q -O /tmp/cmake-install.sh && \
+curl -sL https://github.com/Kitware/CMake/releases/download/v3.20.0/cmake-3.20.0-linux-x86_64.sh \
+    -o /tmp/cmake-install.sh && \
   chmod u+x /tmp/cmake-install.sh && \
   mkdir -p /opt/cmake && \
   /tmp/cmake-install.sh --skip-license --prefix=/opt/cmake && \
diff --git a/docker/scripts/install-tensorrt.sh b/docker/scripts/install-tensorrt.sh
index 3f7ab4f2052..84c4a899681 100644
--- a/docker/scripts/install-tensorrt.sh
+++ b/docker/scripts/install-tensorrt.sh
@@ -11,9 +11,20 @@
 # limitations under the License.
 set -ex
 
-apt-get update -y && apt-get install -y wget
-tensorrt_pkg=TensorRT-8.0.1.6.Linux.x86_64-gnu.cuda-11.3.cudnn8.2.tar.gz
-wget -q https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/tensorrt_versions/${tensorrt_pkg}
+[[ $# -eq 1 ]] || { echo "Usage: $0 <device>"; exit 1; }
+
+device=$1
+[[ -z "${device}" ]] && { echo "Device is required, but got '${device}'"; exit 1; }
+[[ ! "${device}" =~ ^cu[0-9]+ ]] && { echo "Skip installing TensorRT for non-CUDA image."; exit 0; }
+
+echo "Install TensorRT for: ${device}"
+
+if [[ ${device} == "cu102" ]]; then
+    tensorrt_pkg=TensorRT-8.0.1.6.Linux.x86_64-gnu.cuda-10.2.cudnn8.2.tar.gz
+else
+    tensorrt_pkg=TensorRT-8.0.1.6.Linux.x86_64-gnu.cuda-11.3.cudnn8.2.tar.gz
+fi
+curl -sL https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/tensorrt_versions/${tensorrt_pkg} -o ${tensorrt_pkg} 
 
 tar xvfz $tensorrt_pkg -C /usr/local/ 1>/dev/null 2>&1
 ln -s /usr/local/TensorRT-8.0.1.6 /usr/local/TensorRT
diff --git a/docker/scripts/patch-cu102.sh b/docker/scripts/patch-cu102.sh
new file mode 100644
index 00000000000..369fca0c1be
--- /dev/null
+++ b/docker/scripts/patch-cu102.sh
@@ -0,0 +1,38 @@
+#!/bin/bash
+# Copyright 2022 The BladeDISC Authors. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This script copies headers and libraries of cuDNN from /usr to /usr/local/cuda in
+# nvidia:cuda-10.2-xxx images. This is needed because tensorflow assumes that.
+
+[[ $1 != "cu102" ]] && { echo "Skip pathcing for non-cu102 images."; exit 0; }
+
+headers=$(ls /usr/include/*.h | grep -E 'cublas|cudnn|nvblas')
+for hdr in ${headers[@]}; do
+    echo "Copy ${hdr} to cuda home."
+	cp ${hdr} /usr/local/cuda/include/
+done
+
+libs=$(ls /usr/lib/x86_64-linux-gnu/ | grep -E 'cublas|cudnn|nvblas' | grep -E '\.so|\.a')
+for lib in ${libs[@]}; do
+    echo "Copy /usr/lib/x86_64-linux-gnu/${lib} to cuda home."
+	cp /usr/lib/x86_64-linux-gnu/${lib} /usr/local/cuda/lib64/
+done
+
+# patch 1
+curl -sL https://developer.download.nvidia.com/compute/cuda/10.2/Prod/patches/1/cuda_10.2.1_linux.run -o cuda_10.2.1_linux.run
+bash ./cuda_10.2.1_linux.run --silent --toolkit
+rm -f ./cuda_10.2.1_linux.run
+
+# patch 2
+curl -sL https://developer.download.nvidia.com/compute/cuda/10.2/Prod/patches/2/cuda_10.2.2_linux.run -o cuda_10.2.2_linux.run
+bash ./cuda_10.2.2_linux.run --silent --toolkit
+rm -f ./cuda_10.2.2_linux.run
\ No newline at end of file
diff --git a/pytorch_blade/.gitignore b/pytorch_blade/.gitignore
index 40ad6e9edfa..b6df1d53ed3 100644
--- a/pytorch_blade/.gitignore
+++ b/pytorch_blade/.gitignore
@@ -32,3 +32,4 @@ debug_bazel.sh
 py_test.out
 cpp_test.out
 torch_blade/*.so.*
+disc_compiler_main
diff --git a/pytorch_blade/WORKSPACE b/pytorch_blade/WORKSPACE
index 85776054ed3..9885536f64a 100644
--- a/pytorch_blade/WORKSPACE
+++ b/pytorch_blade/WORKSPACE
@@ -77,9 +77,11 @@ local_repository(
 )
 
 load("@org_third_party//bazel/tensorrt:repo.bzl", "tensorrt_configure")
+load("@org_third_party//bazel/cuda_supplement:cuda_supplement_configure.bzl", "cuda_supplement_configure")
 load("//bazel/torch:repo.bzl", "torch_configure", "torch_pybind11_configure")
 
 tensorrt_configure(name = "tensorrt")
+cuda_supplement_configure(name = "local_config_cuda_supplement")
 torch_configure(name = "local_org_torch")
 
 torch_pybind11_configure(
diff --git a/pytorch_blade/ci_build/build_pytorch_blade.sh b/pytorch_blade/ci_build/build_pytorch_blade.sh
index 338147cac60..9ee7e69834e 100755
--- a/pytorch_blade/ci_build/build_pytorch_blade.sh
+++ b/pytorch_blade/ci_build/build_pytorch_blade.sh
@@ -14,6 +14,7 @@ set -o pipefail
 set -e
 
 export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda/}
+export TF_CUDA_HOME=${CUDA_HOME} # for cuda_supplement_configure.bzl
 export CUDACXX=${CUDACXX:-"${CUDA_HOME}/bin/nvcc"}
 export PATH=${CUDA_HOME}/bin/:$PATH
 export TENSORRT_INSTALL_PATH=${TENSORRT_INSTALL_PATH:-/usr/local/TensorRT/}
diff --git a/pytorch_blade/requirements-dev-1.8.1+cu102.txt b/pytorch_blade/requirements-dev-1.8.1+cu102.txt
new file mode 100644
index 00000000000..5bcb15b6d93
--- /dev/null
+++ b/pytorch_blade/requirements-dev-1.8.1+cu102.txt
@@ -0,0 +1,6 @@
+pytest==5.0.0
+networkx
+onnx>=1.6.0
+torch==1.8.1+cu102
+torchvision==0.9.1+cu102
+
diff --git a/pytorch_blade/src/BUILD b/pytorch_blade/src/BUILD
index 7fcf99c0351..559d642a600 100644
--- a/pytorch_blade/src/BUILD
+++ b/pytorch_blade/src/BUILD
@@ -40,8 +40,14 @@ cc_library(
 
 cc_binary(
     name = "libtorch_blade.so",
+    linkopts = [
+        "-Wl,--version-script,$(location :torch_blade.lds)",
+    ],
     linkshared = 1,
-    deps = [":torch_blade"],
+    deps = [
+        ":torch_blade",
+        ":torch_blade.lds",
+    ],
 )
 
 cc_import(
diff --git a/pytorch_blade/src/compiler/mlir/converters/BUILD b/pytorch_blade/src/compiler/mlir/converters/BUILD
index 7e923de200c..8591b19bc3a 100644
--- a/pytorch_blade/src/compiler/mlir/converters/BUILD
+++ b/pytorch_blade/src/compiler/mlir/converters/BUILD
@@ -22,8 +22,8 @@ cc_library(
         "//visibility:public",
     ],
     deps = [
-	"//src:torch_blade_include",
-	"//src/common_utils:torch_blade_logging",
+        "//src:torch_blade_include",
+        "//src/common_utils:torch_blade_logging",
         "@local_org_torch//:libtorch",
         "@org_mhlo_builder//csrc:mhlo_library",
     ],
@@ -32,7 +32,10 @@ cc_library(
 
 cc_test(
     name = "mhlo_converter_register_test",
-    srcs = ["mhlo_converter_register_test.cpp"],
+    srcs = [
+        "mhlo_converter_register_test.cpp",
+        "//src:libtorch_blade.so",
+    ],
     linkopts = [
         "-lpthread",
         "-lm",
@@ -40,8 +43,10 @@ cc_test(
     ],
     linkstatic = True,
     deps = [
-        "//src:torch_blade",
+        "//src:libtorch_blade.so",
+        "@zlib",
         "@googltest//:gtest_main",
         "@local_org_torch//:libtorch",
-    ]
+        "@org_tensorflow//tensorflow/compiler/mlir/xla/ral:ral_base_context_lib",
+    ],
 )
diff --git a/pytorch_blade/src/compiler/mlir/runtime/BUILD b/pytorch_blade/src/compiler/mlir/runtime/BUILD
index c9790a6538a..3edce3ca1f5 100644
--- a/pytorch_blade/src/compiler/mlir/runtime/BUILD
+++ b/pytorch_blade/src/compiler/mlir/runtime/BUILD
@@ -38,7 +38,7 @@ cc_library(
     ],
     deps = [
         ":torch_blade_ral_context",
-	"//src/compiler/backends:torch_blade_backends",
+	    "//src/compiler/backends:torch_blade_backends",
         "@local_org_torch//:libtorch",
     ],
     copts = select({
diff --git a/pytorch_blade/src/compiler/mlir/runtime/ral_context.cpp b/pytorch_blade/src/compiler/mlir/runtime/ral_context.cpp
index 3fda979fac6..73f4ffad79b 100644
--- a/pytorch_blade/src/compiler/mlir/runtime/ral_context.cpp
+++ b/pytorch_blade/src/compiler/mlir/runtime/ral_context.cpp
@@ -14,7 +14,7 @@
 #include <dlfcn.h>
 
 #include <c10/core/CPUAllocator.h>
-#if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION >= 12
+#if PYTORCH_MAJOR_VERSION == 1 && PYTORCH_MINOR_VERSION >= 11
 #include <c10/core/impl/alloc_cpu.h>
 #endif
 
diff --git a/pytorch_blade/src/compiler/tensorrt/BUILD b/pytorch_blade/src/compiler/tensorrt/BUILD
index b5a78e0d05f..68a055ac0cb 100644
--- a/pytorch_blade/src/compiler/tensorrt/BUILD
+++ b/pytorch_blade/src/compiler/tensorrt/BUILD
@@ -1,4 +1,5 @@
 load("//bazel:build_defs.bzl", "if_tensorrt_static_linked")
+load("@local_config_cuda_supplement//:build_defs.bzl", "if_has_cublaslt", "if_has_cudnn_static")
 load("@tensorrt//:build_defs.bzl", "if_has_myelin")
 
 package(default_visibility = ["//visibility:public"])
@@ -26,14 +27,31 @@ cc_library(
             "@tensorrt//:nvinfer_static",
             "@tensorrt//:nvinferplugin_static",
             "@tensorrt//:nvonnxparser_static",
-        ],
+            "@local_config_cuda_supplement//:cublas_static",
+            "@local_config_cuda_supplement//:culibos_static",
+        ] + if_has_cudnn_static(
+            [
+                "@local_config_cuda_supplement//:cudnn_static",
+            ],
+            if_false = [
+                # cuda 10.2 has no libcudnn_static.a
+                "@local_config_cuda_supplement//:cudnn_adv_infer_static",
+                "@local_config_cuda_supplement//:cudnn_cnn_infer_static",
+                "@local_config_cuda_supplement//:cudnn_cnn_train_static",
+                "@local_config_cuda_supplement//:cudnn_ops_infer_static",
+                "@local_config_cuda_supplement//:cudnn_ops_train_static",
+            ],
+        ) + if_has_cublaslt([
+            "@local_config_cuda_supplement//:cublasLt_static",
+        ]) + if_has_myelin([
+            "@local_config_cuda_supplement//:nvrtc",
+            "@tensorrt//:myelin_static",
+        ]),
         if_false = [
             "@tensorrt//:nvinfer",
             "@tensorrt//:nvinferplugin",
             "@tensorrt//:nvonnxparser",
-        ] + if_has_myelin([
-            "@local_config_trt//:myelin_static",
-        ]),
+        ],
     ),
     alwayslink = 1,  # targets depending on it should carry all symbols in its children.
 )
diff --git a/pytorch_blade/src/compiler/tensorrt/bridge/tensorrt_onnx_parser.cpp b/pytorch_blade/src/compiler/tensorrt/bridge/tensorrt_onnx_parser.cpp
index e009ff07ef2..ad1513d0523 100644
--- a/pytorch_blade/src/compiler/tensorrt/bridge/tensorrt_onnx_parser.cpp
+++ b/pytorch_blade/src/compiler/tensorrt/bridge/tensorrt_onnx_parser.cpp
@@ -135,19 +135,21 @@ TrtUniquePtr<nvinfer1::ICudaEngine> TensorrtOnnxParser::BuildEngine(
     auto trt_engine = TrtUniquePtr<nvinfer1::ICudaEngine>(
         context.builder->buildEngineWithConfig(*context.network, *config));
 
-    // We had meet segfault when retrive error message after build engine
-    // failed. It's a issue of TensorRT.
-    // To make our process more robust we would like to enable log
-    // message only when TORCH_BLADE_DEBUG_LOG is set.
     bool debug_log_flag =
         env::ReadBoolFromEnvVar("TORCH_BLADE_DEBUG_LOG", false);
     if (trt_engine == nullptr && debug_log_flag) {
-      LOG(ERROR) << "Failed to build the engine, error message are:";
       auto error_recorder = context.builder->getErrorRecorder();
-      auto nb_error = error_recorder->getNbErrors();
-      for (int i = 0; i < nb_error; i++) {
-        auto error_msg = error_recorder->getErrorDesc(i);
-        LOG(ERROR) << error_msg;
+      if (error_recorder != nullptr) {
+        LOG(ERROR) << "Failed to build the engine, error message are:";
+        auto nb_error = error_recorder->getNbErrors();
+        for (int i = 0; i < nb_error; i++) {
+          auto error_msg = error_recorder->getErrorDesc(i);
+          LOG(ERROR) << error_msg;
+        }
+      } else {
+        // Keep quiet when `error_recorder` is nullptr, error's will be written
+        // to log stream said TensorRT doc:
+        //    https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_builder.html#aa83591ea175b212a9d05ad938ebc7e24
       }
     }
     return trt_engine;
diff --git a/pytorch_blade/src/torch_blade.lds b/pytorch_blade/src/torch_blade.lds
new file mode 100644
index 00000000000..b4ccc193cd3
--- /dev/null
+++ b/pytorch_blade/src/torch_blade.lds
@@ -0,0 +1,6 @@
+VERS_1.0 {
+    global:
+        *blade*;
+    local:
+        *;
+};
diff --git a/third_party/bazel/cuda_supplement/build_defs.bzl.tpl b/third_party/bazel/cuda_supplement/build_defs.bzl.tpl
index 693124f0088..1b3480a9842 100644
--- a/third_party/bazel/cuda_supplement/build_defs.bzl.tpl
+++ b/third_party/bazel/cuda_supplement/build_defs.bzl.tpl
@@ -2,3 +2,8 @@ def if_has_cublaslt(if_true, if_false = []):
     if %{IF_HAS_CUBLASLT}:
         return if_true
     return if_false
+
+def if_has_cudnn_static(if_true, if_false = []):
+    if %{IF_HAS_CUDNN_STATIC}:
+        return if_true
+    return if_false
diff --git a/third_party/bazel/cuda_supplement/cuda_supplement.BUILD.tpl b/third_party/bazel/cuda_supplement/cuda_supplement.BUILD.tpl
index f477d7f2835..3c19d350085 100644
--- a/third_party/bazel/cuda_supplement/cuda_supplement.BUILD.tpl
+++ b/third_party/bazel/cuda_supplement/cuda_supplement.BUILD.tpl
@@ -10,11 +10,14 @@ cc_import(
     static_library = "lib64/libcublas_static.a",
 )
 
-%{import_cublaslt}
+cc_import(
+    name = "cublasLt_static",
+    static_library = "lib64/libcublasLt_static.a",
+)
 
 cc_import(
     name = "culibos_static",
-    static_library = "lib64/libculibos_static.a",
+    static_library = "lib64/libculibos.a",
 )
 
 cc_import(
@@ -22,6 +25,31 @@ cc_import(
     static_library = "lib64/libcudnn_static.a",
 )
 
+cc_import(
+    name = "cudnn_adv_infer_static",
+    static_library = "lib64/libcudnn_adv_infer_static.a",
+)
+
+cc_import(
+    name = "cudnn_cnn_infer_static",
+    static_library = "lib64/libcudnn_cnn_infer_static.a",
+)
+
+cc_import(
+    name = "cudnn_cnn_train_static",
+    static_library = "lib64/libcudnn_cnn_train_static.a",
+)
+
+cc_import(
+    name = "cudnn_ops_infer_static",
+    static_library = "lib64/libcudnn_ops_infer_static.a",
+)
+
+cc_import(
+    name = "cudnn_ops_train_static",
+    static_library = "lib64/libcudnn_ops_train_static.a",
+)
+
 cc_import(
     name = "nvrtc",
     shared_library = "lib64/libnvrtc.so",
diff --git a/third_party/bazel/cuda_supplement/cuda_supplement_configure.bzl b/third_party/bazel/cuda_supplement/cuda_supplement_configure.bzl
index ae37f700494..63d1e703893 100644
--- a/third_party/bazel/cuda_supplement/cuda_supplement_configure.bzl
+++ b/third_party/bazel/cuda_supplement/cuda_supplement_configure.bzl
@@ -8,23 +8,19 @@ def _create_dummy_repository(repo_ctx):
         "%{IF_HAS_CUBLASLT}": "False",
     })
 
-def _cc_import_cublaslt():
-    return """
-cc_import(
-    name = "cublasLt_static",
-    static_library = "lib64/libcublasLt_static.a",
-)"""
 
 def _impl(repo_ctx):
     cuda_path = repo_ctx.os.environ.get(_TF_CUDA_HOME, None)
     if cuda_path != None:
-        if_has_cublaslt = files_exist(repo_ctx, ["lib64/libcublasLt_static.a"])[0]
-        repo_ctx.template("BUILD", Label("//bazel/cuda_supplement:cuda_supplement.BUILD.tpl"), {
-            "%{import_cublaslt}": _cc_import_cublaslt() if if_has_cublaslt else "",
-        })
+        if_has_cublaslt, if_has_cudnn_static = files_exist(repo_ctx, [
+            cuda_path + "/lib64/libcublasLt_static.a",
+            cuda_path + "/lib64/libcudnn_static.a",
+        ])
+        repo_ctx.template("BUILD", Label("//bazel/cuda_supplement:cuda_supplement.BUILD.tpl"), {})
         repo_ctx.symlink(cuda_path + "/lib64", "lib64")
         repo_ctx.template("build_defs.bzl", Label("//bazel/cuda_supplement:build_defs.bzl.tpl"), {
             "%{IF_HAS_CUBLASLT}": "True" if if_has_cublaslt else "False",
+            "%{IF_HAS_CUDNN_STATIC}": "True" if if_has_cudnn_static else "False",
         })
     else:
         _create_dummy_repository(repo_ctx)
diff --git a/third_party/bazel/tensorrt/trt.BUILD.tpl b/third_party/bazel/tensorrt/trt.BUILD.tpl
index 405ddf49f3d..2eabb18a61d 100644
--- a/third_party/bazel/tensorrt/trt.BUILD.tpl
+++ b/third_party/bazel/tensorrt/trt.BUILD.tpl
@@ -101,10 +101,7 @@ cc_library(
         "nvinfer_static_lib",
         "@local_config_cuda//cuda:cudart_static",
         #"@cudnn",
-    ] + select({
-        ":windows": ["@local_config_cuda//cuda:cublas"],
-        "//conditions:default": ["@local_config_cuda//cuda:cublas"],
-    }),
+    ] ,
 )
 
 %{myelin_static_rule}
@@ -400,10 +397,7 @@ cc_library(
         "nvinfer_static",
         "nvinferplugin_static_lib",
         "@local_config_cuda//cuda:cudart_static",
-    ] + select({
-        ":windows": ["@local_config_cuda//cuda:cublas"],
-        "//conditions:default": ["@local_config_cuda//cuda:cublas"],
-    }),
+    ],
     alwayslink = True,
     linkstatic = 1,
     copts = [


Test patch:
diff --git a/docker/scripts/find-fastest-apt.sh b/docker/scripts/find-fastest-apt.sh
index 8a4d10527a5..2a4483ef209 100755
--- a/docker/scripts/find-fastest-apt.sh
+++ b/docker/scripts/find-fastest-apt.sh
@@ -50,6 +50,7 @@ function find_fastest_apt_source() {
     # read command would return an non-zero exit code.
     read -r -d '\t' apt_sources <<EOM
 http://mirrors.cloud.aliyuncs.com
+http://mirrors.aliyun.com
 http://archive.ubuntu.com
 \t
 EOM
diff --git a/pytorch_blade/tests/tensorrt/test_support_info.py b/pytorch_blade/tests/tensorrt/test_support_info.py
index 8f570ec4f75..f706991277f 100644
--- a/pytorch_blade/tests/tensorrt/test_support_info.py
+++ b/pytorch_blade/tests/tensorrt/test_support_info.py
@@ -470,32 +470,56 @@ def test_aten_eq(self):
         self._make_check(graph, False)
 
     def test_const_fold_before_export(self):
-        graph = torch.parse_ir(
-            """
-            graph(%input0.2 : Float(1:165888, 512:324, 18:18, 18:1, requires_grad=0, device=cuda:0)):
-                %1 : None = prim::Constant() # :0:0
-                %2 : bool = prim::Constant[value=1]()
-                %3 : float[] = prim::Constant[value=[2., 2.]]()
-                %x1.3 : Float(1:663552, 512:1296, 36:36, 36:1, requires_grad=0, device=cuda:0) = aten::upsample_bilinear2d(%input0.2, %1, %2, %3)
-                return (%x1.3)
-            """
-        )
+        if utils.torch_version_number() >= utils.parse_version("1.8.1"):
+            graph = torch.parse_ir(
+                """
+                graph(%input0.2 : Float(1, 512, 18, 18, requires_grad=0, device=cuda:0)):
+                    %1 : None = prim::Constant() # :0:0
+                    %2 : bool = prim::Constant[value=1]()
+                    %3 : float[] = prim::Constant[value=[2., 2.]]()
+                    %x1.3 : Float(1, 512, 36, 36, requires_grad=0, device=cuda:0) = aten::upsample_bilinear2d(%input0.2, %1, %2, %3)
+                    return (%x1.3)
+                """
+            )
+        else:
+            graph = torch.parse_ir(
+                """
+                graph(%input0.2 : Float(1:165888, 512:324, 18:18, 18:1, requires_grad=0, device=cuda:0)):
+                    %1 : None = prim::Constant() # :0:0
+                    %2 : bool = prim::Constant[value=1]()
+                    %3 : float[] = prim::Constant[value=[2., 2.]]()
+                    %x1.3 : Float(1:663552, 512:1296, 36:36, 36:1, requires_grad=0, device=cuda:0) = aten::upsample_bilinear2d(%input0.2, %1, %2, %3)
+                    return (%x1.3)
+                """
+            )
         cfg = Config.get_current_context_or_new().clone()
         cfg.customize_onnx_opset_version = 11
         with cfg:
             self._make_check(graph, True)
 
     def test_scalar_input_on_graph(self):
-        graph = torch.parse_ir(
-            """
-            graph(%x.3 : Float(1:64, 64:1, 1:1, 1:1, requires_grad=0, device=cuda:0),
-                    %1 : int):
-                %2 : int = prim::Constant[value=-1]()
-                %3 : int[] = prim::ListConstruct(%1, %2)
-                %input.14 : Float(1:64, 64:1, requires_grad=0, device=cuda:0) = aten::view(%x.3, %3)
-                return (%input.14)
-            """
-        )
+        if utils.torch_version_number() >= utils.parse_version("1.8.1"):
+            graph = torch.parse_ir(
+                """
+                graph(%x.3 : Float(1, 64, 1, 1, requires_grad=0, device=cuda:0),
+                        %1 : int):
+                    %2 : int = prim::Constant[value=-1]()
+                    %3 : int[] = prim::ListConstruct(%1, %2)
+                    %input.14 : Float(1, 64, requires_grad=0, device=cuda:0) = aten::view(%x.3, %3)
+                    return (%input.14)
+                """
+            )
+        else:
+            graph = torch.parse_ir(
+                """
+                graph(%x.3 : Float(1:64, 64:1, 1:1, 1:1, requires_grad=0, device=cuda:0),
+                        %1 : int):
+                    %2 : int = prim::Constant[value=-1]()
+                    %3 : int[] = prim::ListConstruct(%1, %2)
+                    %input.14 : Float(1:64, 64:1, requires_grad=0, device=cuda:0) = aten::view(%x.3, %3)
+                    return (%input.14)
+                """
+            )
         self._make_check(graph, True)
 
 
diff --git a/scripts/ci/test_pytorch_blade.sh b/scripts/ci/test_pytorch_blade.sh
index 1bb06b14426..c5bfbc383c9 100644
--- a/scripts/ci/test_pytorch_blade.sh
+++ b/scripts/ci/test_pytorch_blade.sh
@@ -17,11 +17,13 @@ SCRIPT_DIR="$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 export CXXFLAGS=${CXXFLAGS:-"-Wno-deprecated-declarations"}
 export CFLAGS=${CFLAGS:-"-Wno-deprecated-declarations"}
 export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
+export TF_CUDA_HOME=${CUDA_HOME}
 export CUDACXX=${CUDACXX:-"${CUDA_HOME}/bin/nvcc"}
 export PATH=${CUDA_HOME}/bin/:$PATH
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 export LIBRARY_PATH=${CUDA_HOME}/lib64:$LIBRARY_PATH
 export TF_REMOTE_CACHE=${TF_REMOTE_CACHE}
+export TORCH_BLADE_BUILD_TENSORRT_STATIC=${TORCH_BLADE_BUILD_TENSORRT_STATIC:-OFF}
 
 if [[ -f ~/.cache/proxy_config ]]; then
   source ~/.cache/proxy_config