Update release branch

.github/actions/assigner/package.json

@@ -19,7 +19,7 @@
   "license": "BSD-3-Clause",
   "repository": "https://www.github.com/pytorch/TensorRT",
   "dependencies": {
-    "@actions/core": "^1.8.2",
+    "@actions/core": "^1.9.1",
     "@actions/github": "^5.0.3",
     "fs": "^0.0.1-security",
     "js-yaml": "^4.1.0"

.github/workflows/docgen.yml

@@ -10,9 +10,9 @@ on:
 
 jobs:
   build-docs:
-    runs-on: ubuntu-18.04
+    runs-on: ubuntu-20.04
     container:
-      image: docker.pkg.github.com/pytorch/tensorrt/docgen:latest
+      image: ghcr.io/pytorch/tensorrt/docgen:latest
       credentials:
         username: ${{ github.actor }}
         password: ${{ secrets.GITHUB_TOKEN }}
@@ -22,23 +22,66 @@ jobs:
           rm -rf /usr/share/dotnet
           rm -rf /opt/ghc
           rm -rf "/usr/local/share/boost"
+          rm -rf /usr/local/cuda/cuda-*
+      - name: Install base deps
+        run: |
+          apt update
+          DEBIAN_FRONTEND=noninteractive apt install -y software-properties-common gcc git curl wget make zlib1g-dev bzip2 libbz2-dev lzma lzma-dev libreadline-dev libsqlite3-dev libssl-dev libffi-dev doxygen pandoc
+          git config --global --add safe.directory '*'
+      - name: Set up Python 3.9.4
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.9.4 
       - uses: actions/checkout@v2
         with:
           ref: ${{github.head_ref}}
       - name: Get HEAD SHA
         id: vars
         run: echo "::set-output name=sha::$(git rev-parse --short HEAD)"
+      - name: Get Bazel version
+        id: bazel_info
+        run: echo "::set-output name=version::$(cat .bazelversion)"
+      - name: Install Bazel
+        run: |
+          wget -q https://github.com/bazelbuild/bazel/releases/download/${{ steps.bazel_info.outputs.version }}/bazel-${{ steps.bazel_info.outputs.version }}-linux-x86_64 -O /usr/bin/bazel
+          chmod a+x /usr/bin/bazel
+      - name: Install cudnn + tensorrt
+        run: |
+          apt-get update
+          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
+          mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
+          apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub
+          apt-key adv --keyserver keyserver.ubuntu.com --recv-keys 536F8F1DE80F6A35
+          apt-key adv --keyserver keyserver.ubuntu.com --recv-keys A4B469963BF863CC
+          add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"
+          apt-get update
+          apt-get install -y libcudnn8 libcudnn8-dev
+
+          apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub
+          add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"
+          apt-get update
+
+          apt-get install -y libnvinfer8 libnvinfer-plugin8 libnvinfer-dev libnvinfer-plugin-dev
+      - name: Install Torch
+        run: |
+          python3 -m pip install -r py/requirements.txt
       - name: Build Python Package
         run: |
-          cp docker/WORKSPACE.docker WORKSPACE
+          cp toolchains/ci_workspaces/WORKSPACE.x86_64 WORKSPACE
           cd py
+          echo $(which python3)
+          echo $(python3 -c 'import site; print(site.getsitepackages()[0])')
+          mkdir -p /opt/circleci/.pyenv/versions/3.9.4/lib/python3.9/
+          ln -s $(python3 -c 'import site; print(site.getsitepackages()[0])') /opt/circleci/.pyenv/versions/3.9.4/lib/python3.9/site-packages
           python3 setup.py install
+          cd ..
       - name: Generate New Docs
         run: |
           cd docsrc
-          pip3 install -r requirements.txt
+          python3 -m pip install -r requirements.txt
           python3 -c "import torch_tensorrt; print(torch_tensorrt.__version__)"
           make html
+          cd ..
       - uses: stefanzweifel/git-auto-commit-action@v4
         with:
         # Required

.github/workflows/linter.yml

@@ -25,25 +25,25 @@ jobs:
       - uses: actions/checkout@v2
         with:
           ref: ${{ github.event.pull_request.head.sha }}
-      - name: Docker login
-        uses: docker/login-action@v1
+      - name: Set up Python 3.9.4
+        uses: actions/setup-python@v2
         with:
-          registry: docker.pkg.github.com
-          username: ${{ github.actor }}
-          password: ${{ secrets.GITHUB_TOKEN }}
-      - name: Run image
+          python-version: 3.9.4
+      - name: Setup env
         run: |
-          docker pull docker.pkg.github.com/pytorch/tensorrt/docgen:latest
-          docker run -it -d --name cpplinter -e GITHUB_TOKEN=$GITHUB_TOKEN -v $GITHUB_WORKSPACE:/workspace -v $GITHUB_EVENT_PATH:/GITHUB_EVENT.json -w /workspace  docker.pkg.github.com/pytorch/tensorrt/docgen:latest
-          docker exec cpplinter bash -c "cp /workspace/docker/WORKSPACE.docker /workspace/WORKSPACE"
+          sudo ln -s $GITHUB_EVENT_PATH /GITHUB_EVENT.json
         env:
           GITHUB_TOKEN: ${{secrets.GITHUB_TOKEN}}
       - name: Install dependencies
         run: |
-          docker exec cpplinter bash -c "pip3 install -r /workspace/.github/scripts/requirements.txt"
-          docker exec cpplinter bash -c "pip3 install -r /workspace/requirements-dev.txt"
+          pip3 install -r $GITHUB_WORKSPACE/.github/scripts/requirements.txt
+          pip3 install -r $GITHUB_WORKSPACE/requirements-dev.txt
       - name: Lint C++
-        run: docker exec cpplinter bash -c "cd /workspace && python3 /workspace/.github/scripts/run_cpp_linter.py"
+        run: | 
+          cd $GITHUB_WORKSPACE
+          python3 $GITHUB_WORKSPACE/.github/scripts/run_cpp_linter.py
+        env:
+          GITHUB_TOKEN: ${{secrets.GITHUB_TOKEN}}
 
   py-linting:
     name: Python Linting
@@ -65,22 +65,28 @@ jobs:
       - uses: actions/checkout@v2
         with:
           ref: ${{ github.event.pull_request.head.sha }}
+      - name: Set up Python 3.9.4
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.9.4
       - name: Docker login
         uses: docker/login-action@v1
         with:
           registry: docker.pkg.github.com
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
-      - name: Run image
+      - name: Setup env
         run: |
-          docker pull docker.pkg.github.com/pytorch/tensorrt/docgen:latest
-          docker run -it -d --name pylinter -e GITHUB_TOKEN=$GITHUB_TOKEN -v $GITHUB_WORKSPACE:/workspace -v $GITHUB_EVENT_PATH:/GITHUB_EVENT.json -w /workspace  docker.pkg.github.com/pytorch/tensorrt/docgen:latest
-          docker exec pylinter bash -c "cp /workspace/docker/WORKSPACE.docker /workspace/WORKSPACE"
+          sudo ln -s $GITHUB_EVENT_PATH /GITHUB_EVENT.json
         env:
           GITHUB_TOKEN: ${{secrets.GITHUB_TOKEN}}
       - name: Install dependencies
         run: |
-          docker exec pylinter bash -c "pip3 install -r /workspace/.github/scripts/requirements.txt"
-          docker exec pylinter bash -c "pip3 install -r /workspace/requirements-dev.txt"
+          pip3 install -r $GITHUB_WORKSPACE/.github/scripts/requirements.txt
+          pip3 install -r $GITHUB_WORKSPACE/requirements-dev.txt
       - name: Lint Python
-        run: docker exec pylinter bash -c "cd /workspace && python3 /workspace/.github/scripts/run_py_linter.py"
+        run: |
+          cd $GITHUB_WORKSPACE
+          python3 $GITHUB_WORKSPACE/.github/scripts/run_py_linter.py
+        env:
+          GITHUB_TOKEN: ${{secrets.GITHUB_TOKEN}}

core/compiler.cpp

@@ -426,7 +426,8 @@ torch::jit::Module CompileGraph(const torch::jit::Module& mod, CompileSpec cfg)
       auto outputIsCollection = conversion::OutputIsCollection(g->block());
       if (cfg.partition_info.enabled &&
           (cfg.lower_info.forced_fallback_modules.size() == 0 &&
-           cfg.partition_info.forced_fallback_operators.size() == 0 && isBlockConvertible)) {
+           cfg.partition_info.forced_fallback_operators.size() == 0 && isBlockConvertible) &&
+          !outputIsCollection) {
         LOG_INFO("Skipping partitioning since model is fully supported");
       }
 

core/conversion/converters/impl/constant_pad.cpp

@@ -21,7 +21,8 @@ auto constant_pad_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns
        auto padding = args[1].unwrapToIntList().vec();
        int64_t padSize = padding.size();
        auto value = args[2].unwrapToScalar().to<float>();
-
+       at::Tensor value_tensor = torch::tensor(value, util::TRTDataTypeToScalarType(in->getType()));
+       auto valueTensor = tensor_to_const(ctx, value_tensor);
        TORCHTRT_CHECK(padSize % 2 == 0, "Length of pad must be even but instead it equals " << padSize);
 
        int64_t l_pad = padSize / 2;
@@ -55,10 +56,8 @@ auto constant_pad_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns
              auto fill_layer = ctx->net->addFill(nvinfer1::Dims{1, {1}}, nvinfer1::FillOperation::kLINSPACE);
              auto shape_gather_out = ctx->net->addShape(*left_gather_out)->getOutput(0);
              fill_layer->setInput(0, *shape_gather_out);
-             at::Tensor value_tensor = torch::tensor(value, torch::kFloat32);
-             auto valueTensor = tensor_to_const(ctx, value_tensor);
              fill_layer->setInput(1, *valueTensor);
-             at::Tensor delta_tensor = torch::zeros(inRank);
+             at::Tensor delta_tensor = torch::zeros(inRank, util::TRTDataTypeToScalarType(in->getType()));
              auto deltaTensor = tensor_to_const(ctx, delta_tensor);
              fill_layer->setInput(2, *deltaTensor);
              auto padTensor = fill_layer->getOutput(0);
@@ -69,10 +68,8 @@ auto constant_pad_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns
            } else {
              inDims.d[axis] = padding[padding_index];
              auto fill_layer = ctx->net->addFill(inDims, nvinfer1::FillOperation::kLINSPACE);
-             at::Tensor value_tensor = torch::tensor(value, torch::kFloat32);
-             auto valueTensor = tensor_to_const(ctx, value_tensor);
              fill_layer->setInput(1, *valueTensor);
-             at::Tensor delta_tensor = torch::zeros(inRank);
+             at::Tensor delta_tensor = torch::zeros(inRank, util::TRTDataTypeToScalarType(in->getType()));
              auto deltaTensor = tensor_to_const(ctx, delta_tensor);
              fill_layer->setInput(2, *deltaTensor);
              auto padTensor = fill_layer->getOutput(0);
@@ -112,10 +109,8 @@ auto constant_pad_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns
              auto fill_layer = ctx->net->addFill(nvinfer1::Dims{1, {1}}, nvinfer1::FillOperation::kLINSPACE);
              auto shape_gather_out = ctx->net->addShape(*right_gather_out)->getOutput(0);
              fill_layer->setInput(0, *shape_gather_out);
-             at::Tensor value_tensor = torch::tensor(value, torch::kFloat32);
-             auto valueTensor = tensor_to_const(ctx, value_tensor);
              fill_layer->setInput(1, *valueTensor);
-             at::Tensor delta_tensor = torch::zeros(inRank);
+             at::Tensor delta_tensor = torch::zeros(inRank, util::TRTDataTypeToScalarType(in->getType()));
              auto deltaTensor = tensor_to_const(ctx, delta_tensor);
              fill_layer->setInput(2, *deltaTensor);
              auto padTensor = fill_layer->getOutput(0);
@@ -126,10 +121,8 @@ auto constant_pad_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns
            } else {
              inDims.d[axis] = padding[padding_index + 1];
              auto fill_layer = ctx->net->addFill(inDims, nvinfer1::FillOperation::kLINSPACE);
-             at::Tensor value_tensor = torch::tensor(value, torch::kFloat32);
-             auto valueTensor = tensor_to_const(ctx, value_tensor);
              fill_layer->setInput(1, *valueTensor);
-             at::Tensor delta_tensor = torch::zeros(inRank);
+             at::Tensor delta_tensor = torch::zeros(inRank, util::TRTDataTypeToScalarType(in->getType()));
              auto deltaTensor = tensor_to_const(ctx, delta_tensor);
              fill_layer->setInput(2, *deltaTensor);
              auto padTensor = fill_layer->getOutput(0);

core/conversion/converters/impl/element_wise.cpp

@@ -390,6 +390,18 @@ auto element_wise_registrations TORCHTRT_UNUSED =
                LOG_DEBUG("Output tensor shape: " << out->getDimensions());
                return true;
              }})
+        .pattern(
+            {"aten::square(Tensor self) -> Tensor",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               auto self = args[0].ITensorOrFreeze(ctx);
+               auto mul = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, self, util::node_info(n));
+               TORCHTRT_CHECK(mul, "Unable to create mul layer from node: " << *n);
+
+               mul->setName(util::node_info(n).c_str());
+               auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0));
+               LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+               return true;
+             }})
         .pattern(
             {"aten::mul.Tensor(Tensor self, Tensor other) -> Tensor",
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {

core/partitioning/partitioning.cpp

@@ -178,7 +178,8 @@ void resolveTRTNonTensorInputs(PartitionedGraph& segmented_blocks) {
 
 void registerSegmentsOutputs(PartitionedGraph& segmented_blocks, torch::jit::Block* block) {
   // find the corresponding raw values in original global graph for this segmented block's inputs/outputs
-  std::set<torch::jit::Value*> input_values;
+  auto cmp = [](torch::jit::Value* a, torch::jit::Value* b) { return a->unique() < b->unique(); };
+  std::set<torch::jit::Value*, decltype(cmp)> input_values(cmp);
   for (auto& seg_block : segmented_blocks) {
     for (auto& input : seg_block.raw_inputs()) {
       input_values.insert(input);