Convolution cast

.circleci/config.yml

@@ -508,6 +508,22 @@ commands:
       - store_artifacts:
           path: /tmp/testlogs
 
+  test-fx_core:
+    description: "Test the fx core"
+    steps:
+      - run:
+          name: Run FX core tests
+          command: |
+            cd py/torch_tensorrt/fx/test
+            pushd core/
+            pytest --junitxml=/tmp/artifacts/test_results/fx/core/test_results.xml
+            popd
+
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
   test-fx_converters:
     description: "Test the fx converters"
     steps:
@@ -586,6 +602,22 @@ commands:
       - store_artifacts:
           path: /tmp/testlogs
 
+  test-fx_quant:
+    description: "Test the fx quant"
+    steps:
+      - run:
+          name: Run FX quant tests
+          command: |
+            cd py/torch_tensorrt/fx/test
+            pushd quant/
+            pytest --junitxml=/tmp/artifacts/test_results/fx/quant/test_results.xml
+            popd
+
+      - store_test_results:
+          path: /tmp/artifacts
+      - store_artifacts:
+          path: /tmp/testlogs
+
   test-fx:
     description: "Test the fx backend"
     steps:
@@ -598,6 +630,8 @@ commands:
       - test-fx_tools
       - test-fx_trt_lower
       - test-fx_tracer
+      - test-fx_core
+      - test-fx_quant
       - store_test_results:
           path: /tmp/artifacts
       - store_artifacts:

core/conversion/converters/impl/conv_deconv.cpp

@@ -13,7 +13,12 @@ namespace {
 bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args) {
   // Input to conv/deconv
   auto in = args[0].ITensor();
-
+  if (in->getType() == nvinfer1::DataType::kINT32) {
+    LOG_WARNING(
+        "Found type  " << in->getType() << "in aten::convolution, casting to" << nvinfer1::DataType::kFLOAT
+                       << " for compatibility.");
+    in = castITensor(ctx, in, nvinfer1::DataType::kFLOAT);
+  }
   // Conv /deconv parameters
   auto stride = util::toDims(args[3].unwrapToIntList());
   auto padding = util::toDims(args[4].unwrapToIntList());

core/lowering/passes/remove_dropout.cpp

@@ -1,4 +1,4 @@
-#include <torch/csrc/jit/passes/subgraph_rewrite.h>
+#include "torch/csrc/jit/passes/dead_code_elimination.h"
 
 #include "core/util/prelude.h"
 
@@ -7,86 +7,52 @@ namespace core {
 namespace lowering {
 namespace passes {
 
-void RemoveDropout(std::shared_ptr<torch::jit::Graph>& graph) {
-  std::string dropout_pattern = R"IR(
-        graph(%input, %4, %5):
-            %6 = aten::dropout(%input, %4, %5)
-            return (%6))IR";
-  std::string no_dropout_pattern = R"IR(
-        graph(%input, %4, %5):
-            return (%input))IR";
-
-  torch::jit::SubgraphRewriter remove_dropout;
-  remove_dropout.RegisterRewritePattern(dropout_pattern, no_dropout_pattern);
-  remove_dropout.runOnGraph(graph);
-
-  std::string dropout_inplace_pattern = R"IR(
-        graph(%input, %4, %5):
-            %6 = aten::dropout_(%input, %4, %5)
-            return (%6))IR";
-  std::string no_dropout_inplace_pattern = R"IR(
-        graph(%input, %4, %5):
-            return (%input))IR";
-
-  torch::jit::SubgraphRewriter remove_dropout_inplace_pattern;
-  remove_dropout_inplace_pattern.RegisterRewritePattern(dropout_inplace_pattern, no_dropout_inplace_pattern);
-  remove_dropout_inplace_pattern.runOnGraph(graph);
-
-  // remove feature_dropout
-  std::string feature_dropout_pattern = R"IR(
-        graph(%input, %4, %5):
-            %6 = aten::feature_dropout(%input, %4, %5)
-            return (%6))IR";
-  std::string no_feature_dropout_pattern = R"IR(
-        graph(%input, %4, %5):
-            return (%input))IR";
-
-  torch::jit::SubgraphRewriter remove_feature_dropout_pattern;
-  remove_feature_dropout_pattern.RegisterRewritePattern(feature_dropout_pattern, no_feature_dropout_pattern);
-  remove_feature_dropout_pattern.runOnGraph(graph);
-
-  // remove feature_dropout inplace
-  std::string feature_dropout_inplace_pattern = R"IR(
-        graph(%input, %4, %5):
-            %6 = aten::feature_dropout_(%input, %4, %5)
-            return (%6))IR";
-  std::string no_feature_dropout_inplace_pattern = R"IR(
-        graph(%input, %4, %5):
-            return (%input))IR";
-
-  torch::jit::SubgraphRewriter remove_feature_dropout_inplace_pattern;
-  remove_feature_dropout_inplace_pattern.RegisterRewritePattern(
-      feature_dropout_inplace_pattern, no_feature_dropout_inplace_pattern);
-  remove_feature_dropout_inplace_pattern.runOnGraph(graph);
-
-  // remove feature_alpha_dropout
-  std::string feature_alpha_dropout_pattern = R"IR(
-        graph(%input, %4, %5):
-            %6 = aten::feature_alpha_dropout(%input, %4, %5)
-            return (%6))IR";
-  std::string no_feature_alpha_dropout_pattern = R"IR(
-        graph(%input, %4, %5):
-            return (%input))IR";
-
-  torch::jit::SubgraphRewriter remove_feature_alpha_dropout_pattern;
-  remove_feature_alpha_dropout_pattern.RegisterRewritePattern(
-      feature_alpha_dropout_pattern, no_feature_alpha_dropout_pattern);
-  remove_feature_alpha_dropout_pattern.runOnGraph(graph);
-
-  // remove feature_alpha_dropout inplace
-  std::string feature_alpha_dropout_inplace_pattern = R"IR(
-        graph(%input, %4, %5):
-            %6 = aten::feature_alpha_dropout_(%input, %4, %5)
-            return (%6))IR";
-  std::string no_feature_alpha_dropout_inplace_pattern = R"IR(
-        graph(%input, %4, %5):
-            return (%input))IR";
-
-  torch::jit::SubgraphRewriter remove_feature_alpha_dropout_inplace_pattern;
-  remove_feature_alpha_dropout_inplace_pattern.RegisterRewritePattern(
-      feature_alpha_dropout_inplace_pattern, no_feature_alpha_dropout_inplace_pattern);
-  remove_feature_alpha_dropout_inplace_pattern.runOnGraph(graph);
+// Schemas for dropout variants
+const std::unordered_set<c10::Symbol> DropoutNodeKinds = {
+    c10::Symbol::fromQualString("aten::dropout"),
+    c10::Symbol::fromQualString("aten::dropout_"),
+    c10::Symbol::fromQualString("aten::feature_dropout"),
+    c10::Symbol::fromQualString("aten::feature_dropout_"),
+    c10::Symbol::fromQualString("aten::feature_alpha_dropout"),
+    c10::Symbol::fromQualString("aten::feature_alpha_dropout_"),
+};
+
+void removeDropoutInBlock(torch::jit::Block* block) {
+  /*
+  Function adapted from:
+  torch/csrc/jit/passes/remove_dropout.cpp
+
+  Modified for conciseness, documentation, and allowing new variants of dropout operators to be quickly added
+  */
+  std::vector<torch::jit::Node*> dropout_nodes_to_remove;
+
+  for (auto node : block->nodes()) {
+    // Remove dropout for each member block within a node
+    for (auto block : node->blocks()) {
+      removeDropoutInBlock(block);
+    }
+
+    // For each node having a dropout-variant Schema, remove the node
+    if (DropoutNodeKinds.find(node->kind()) != DropoutNodeKinds.end()) {
+      // Extract input and output tensors of dropout operator
+      auto input_value = node->inputs()[0];
+      auto output_value = node->outputs()[0];
+
+      output_value->replaceAllUsesWith(input_value);
+      dropout_nodes_to_remove.push_back(node);
+    }
+  }
+
+  // Delete dropout nodes
+  for (auto del_node : dropout_nodes_to_remove) {
+    del_node->destroy();
+  }
+}
 
+void RemoveDropout(std::shared_ptr<torch::jit::Graph>& graph) {
+  // Remove all instances of dropout variants from graph
+  removeDropoutInBlock(graph->block());
+  torch::jit::EliminateDeadCode(graph);
   LOG_GRAPH("Post remove dropout: " << *graph);
 }
 

docker/Dockerfile

@@ -1,62 +1,79 @@
-ARG BASE=22.04
-ARG BASE_IMG=nvcr.io/nvidia/tensorrt:${BASE}-py3
+# Base image starts with CUDA
+ARG BASE_IMG=nvidia/cuda:11.7.1-devel-ubuntu18.04
 FROM ${BASE_IMG} as base
 
-FROM base as torch-tensorrt-builder-base
+# Install basic dependencies
+RUN apt-get update
+RUN apt install -y build-essential manpages-dev wget zlib1g software-properties-common git
+RUN add-apt-repository ppa:deadsnakes/ppa
+RUN apt install -y python3.8 python3.8-distutils python3.8-dev
+RUN wget https://bootstrap.pypa.io/get-pip.py
+RUN ln -s /usr/bin/python3.8 /usr/bin/python
+RUN python get-pip.py
+RUN pip3 install wheel
+
+# Install Pytorch
+RUN pip3 install torch==2.0.0.dev20230103+cu117 torchvision==0.15.0.dev20230103+cu117 --extra-index-url https://download.pytorch.org/whl/nightly/cu117
+
+# Install CUDNN + TensorRT
+RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-ubuntu1804.pin
+RUN mv cuda-ubuntu1804.pin /etc/apt/preferences.d/cuda-repository-pin-600
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
+RUN apt-key adv --keyserver keyserver.ubuntu.com --recv-keys 536F8F1DE80F6A35
+RUN apt-key adv --keyserver keyserver.ubuntu.com --recv-keys A4B469963BF863CC
+RUN add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/ /"
+RUN apt-get update
+RUN apt-get install -y libcudnn8=8.5.0* libcudnn8-dev=8.5.0*
+
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub
+RUN add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/ /"
+RUN apt-get update
+
+RUN apt-get install -y libnvinfer8=8.5.1* libnvinfer-plugin8=8.5.1* libnvinfer-dev=8.5.1* libnvinfer-plugin-dev=8.5.1* libnvonnxparsers8=8.5.1-1* libnvonnxparsers-dev=8.5.1-1* libnvparsers8=8.5.1-1*  libnvparsers-dev=8.5.1-1*
+
+# Setup Bazel
+ARG BAZEL_VERSION=5.2.0
+RUN wget -q https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION}/bazel-${BAZEL_VERSION}-linux-x86_64 -O /usr/bin/bazel \
+ && chmod a+x /usr/bin/bazel
+
+# Remove cuda symlink to avoid bazel circle symlink errors
+RUN rm /usr/local/cuda-11.7/cuda-11.7
 
-# Removing any bazel or torch-tensorrt pre-installed from the base image
-RUN rm -rf /opt/pytorch/torch_tensorrt /usr/bin/bazel
+# Build Torch-TensorRT in an auxillary container
+FROM base as torch-tensorrt-builder-base
 
 ARG ARCH="x86_64"
 ARG TARGETARCH="amd64"
-ARG BAZEL_VERSION=5.2.0
-
-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 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 install -y python3-setuptools
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub
+RUN  apt-get update
 
 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
 
 FROM torch-tensorrt-builder-base as torch-tensorrt-builder
 
-# Removing any bazel or torch-tensorrt pre-installed from the base image
-RUN rm -rf /opt/pytorch/torch_tensorrt
-
 COPY . /workspace/torch_tensorrt/src
 WORKDIR /workspace/torch_tensorrt/src
 RUN cp ./docker/WORKSPACE.docker WORKSPACE
 
 # This script builds both libtorchtrt bin/lib/include tarball and the Python wheel, in dist/
 RUN ./docker/dist-build.sh
 
+# Copy and install Torch-TRT into the main container
 FROM base as torch-tensorrt
 
-# Removing any bazel or torch-tensorrt pre-installed from the base image
-RUN rm -rf /opt/pytorch/torch_tensorrt
-
-# copy source repo
-COPY . /workspace/torch_tensorrt
+COPY . /opt/torch_tensorrt
 COPY --from=torch-tensorrt-builder  /workspace/torch_tensorrt/src/py/dist/ .
 
-RUN pip3 install ipywidgets --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org
-RUN jupyter nbextension enable --py widgetsnbextension
-
+RUN cp /opt/torch_tensorrt/docker/WORKSPACE.docker /opt/torch_tensorrt/WORKSPACE
 RUN pip3 install *.whl && rm -fr /workspace/torch_tensorrt/py/dist/* *.whl
 
-ENV LD_LIBRARY_PATH /usr/local/lib/python3.8/dist-packages/torch/lib:/usr/local/lib/python3.8/dist-packages/torch_tensorrt/lib:${LD_LIBRARY_PATH}
+# Install native tensorrt python package required by torch_tensorrt whl file
+RUN pip install tensorrt==8.5.1.7
+
+WORKDIR /opt/torch_tensorrt
+ENV LD_LIBRARY_PATH /usr/local/lib/python3.8/dist-packages/torch/lib:/usr/local/lib/python3.8/dist-packages/torch_tensorrt/lib:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH}
 ENV PATH /usr/local/lib/python3.8/dist-packages/torch_tensorrt/bin:${PATH}
-#
-WORKDIR /workspace
-RUN mv /workspace/torch_tensorrt /opt/torch_tensorrt
-RUN cp /opt/torch_tensorrt/docker/WORKSPACE.docker /opt/torch_tensorrt/WORKSPACE
-RUN mkdir torch_tensorrt
-RUN ln -s /opt/torch_tensorrt/notebooks /workspace/torch_tensorrt/notebooks
 
-CMD /bin/bash
+CMD /bin/bash