Merge branch 'anuragd/dockerfile_platform_agnostics' into 'main'

(//docker): Dockerfile platform agnostic

See merge request nvidia/trtorch/Torch-TensorRT-Preview!7

Author: narendasan

docker/Dockerfile

@@ -8,15 +8,20 @@ FROM base as torch-tensorrt-builder-base
 
 RUN rm -rf /opt/torch-tensorrt /usr/bin/bazel
 
-RUN apt-get update && apt-get install --no-install-recommends -y curl gnupg
-RUN curl -fsSL https://bazel.build/bazel-release.pub.gpg | gpg --dearmor > /etc/apt/trusted.gpg.d/bazel.gpg
-RUN echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list
+ARG ARCH="x86_64"
+ARG TARGETARCH="amd64"
+ARG BAZEL_VERSION=4.0.0
 
-RUN apt-get update && apt-get install -y --no-install-recommends bazel-4.0.0
-RUN ln -s /usr/bin/bazel-4.0.0 /usr/bin/bazel
+RUN [[ "$TARGETARCH" == "amd64" ]] && ARCH="x86_64" || ARCH="${TARGETARCH}" \
+ && wget -q https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION}/bazel-${BAZEL_VERSION}-linux-${ARCH} -O /usr/bin/bazel \
+ && chmod a+x /usr/bin/bazel
 
 # Workaround for bazel expecting both static and shared versions, we only use shared libraries inside container
-RUN cp /usr/lib/x86_64-linux-gnu/libnvinfer.so /usr/lib/x86_64-linux-gnu/libnvinfer_static.a
+RUN touch /usr/lib/$HOSTTYPE-linux-gnu/libnvinfer_static.a
+
+RUN rm -rf /usr/local/cuda/lib* /usr/local/cuda/include \
+  && ln -sf /usr/local/cuda/targets/$HOSTTYPE-linux/lib /usr/local/cuda/lib64 \
+  && ln -sf /usr/local/cuda/targets/$HOSTTYPE-linux/include /usr/local/cuda/include
 
 RUN apt-get update && apt-get install -y --no-install-recommends locales ninja-build && rm -rf /var/lib/apt/lists/* && locale-gen en_US.UTF-8
 
@@ -47,7 +52,5 @@ ENV PATH /opt/conda/lib/python3.8/site-packages/torch_tensorrt/bin:${PATH}
 
 WORKDIR /workspace
 RUN mv /workspace/torch_tensorrt /opt/pytorch/torch_tensorrt
-RUN mkdir torch_tensorrt
-RUN ln -s /opt/pytorch/torch_tensorrt/notebooks /workspace/torch_tensorrt/notebooks
 
 CMD /bin/bash

third_party/cuda/BUILD

@@ -14,7 +14,7 @@ cc_library(
             "lib/x64/cudart.lib",
         ],
         "//conditions:default": [
-            "targets/x86_64-linux/lib/libcudart.so",
+            "lib64/libcudart.so",
         ],
     }),
     hdrs = glob([
@@ -32,7 +32,7 @@ cc_library(
             "bin/nvToolsExt64_1.dll",
         ],
         "//conditions:default": [
-            "targets/x86_64-linux/lib/libnvToolsExt.so.1",
+            "lib64/libnvToolsExt.so.1",
         ],
     }),
 )
@@ -44,7 +44,7 @@ cc_library(
             "bin/*.dll",
         ],
         "//conditions:default": glob([
-            "targets/x86_64-linux/lib/**/lib*.so",
+            "lib64/**/lib*.so",
         ]),
     }),
     hdrs = glob([
@@ -63,7 +63,7 @@ cc_library(
             "lib/x64/cublas.lib",
         ]),
         "//conditions:default": glob([
-            "targets/x86_64-linux/lib/**/*libcublas.so",
+            "lib64/**/*libcublas.so",
         ]),
     }),
     hdrs = glob([