Merge branch 'master' into double_long_ival

Author: narendasan

.github/workflows/docgen.yml

@@ -10,26 +10,32 @@ on:
 
 jobs:
   build-docs:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-18.04
+    container:
+      image: docker.pkg.github.com/nvidia/trtorch/docgen:latest
+      credentials:
+        username: $GITHUB_ACTOR
+        password: ${{secrets.GITHUB_TOKEN}}
     steps:
       - uses: actions/checkout@v2
         with:
           ref: ${{github.head_ref}}
-      - name: Docker login
-        run: docker login docker.pkg.github.com -u $GITHUB_ACTOR -p $GITHUB_TOKEN
-        env:
-          GITHUB_TOKEN: ${{secrets.GITHUB_TOKEN}}
-      - name: Run image
-        run: docker run -it -d --name builder -v $GITHUB_WORKSPACE:/workspace -w /workspace docker.pkg.github.com/nvidia/trtorch/docgen:latest
-      - name: Build python package
-        run: docker exec builder bash -c "cp docker/WORKSPACE.docker WORKSPACE && cd py && python3 setup.py install"
-      - name: Generate new docs
-        run: docker exec builder bash -c "cd docsrc && make html"
+      - name: Get HEAD SHA
+        id: vars
+        run: echo "::set-output name=sha::$(git rev-parse --short HEAD)"
+      - name: Build Python Package
+        run: |
+          cp docker/WORKSPACE.docs WORKSPACE
+          cd py
+          python3 setup.py install
+      - name: Generate New Docs
+        run: |
+          cd docsrc
+          make html
       - uses: stefanzweifel/git-auto-commit-action@v4
         with:
         # Required
-          commit_message: >
-            docs: [Automated] Regenerating documenation from ${{github.head_ref}}
+          commit_message: "docs: [Automated] Regenerating documenation for ${{ steps.vars.outputs.sha }}"
           commit_options: "--no-verify --signoff"
           file_pattern: docs/
           commit_user_name: TRTorch Github Bot

.github/workflows/linter.yml

@@ -19,7 +19,7 @@ jobs:
       - name: Run image
         run: |
           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/nvidia/trtorch/docgen:latest
-          docker exec cpplinter bash -c "cp /workspace/docker/WORKSPACE.docker /workspace/WORKSPACE"
+          docker exec cpplinter bash -c "cp /workspace/docker/WORKSPACE.docs /workspace/WORKSPACE"
         env:
           GITHUB_TOKEN: ${{secrets.GITHUB_TOKEN}}
       - name: Install dependencies
@@ -41,7 +41,7 @@ jobs:
       - name: Run image
         run: |
           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/nvidia/trtorch/docgen:latest
-          docker exec pylinter bash -c "cp /workspace/docker/WORKSPACE.docker /workspace/WORKSPACE"
+          docker exec pylinter bash -c "cp /workspace/docker/WORKSPACE.docs /workspace/WORKSPACE"
         env:
           GITHUB_TOKEN: ${{secrets.GITHUB_TOKEN}}
       - name: Install dependencies

CONTRIBUTING.md

@@ -16,7 +16,7 @@ We use the PyTorch Slack for communication about core development, integration w
 
     - Use the built in linting tools to ensure that your code matches the style guidelines
       ```sh
-      # C++ Linting (After installing clang-format)
+      # C++ Linting (After installing clang-format [Version 9.0.0])
       # Print non-conforming sections of code
       bazel run //tools/linter:cpplint_diff -- //...
       # Modify code to conform with style guidelines

README.md

@@ -73,11 +73,11 @@ torch.jit.save(trt_ts_module, "trt_torchscript_module.ts")
 ### Dependencies
 These are the following dependencies used to verify the testcases. TRTorch can work with other versions, but the tests are not guaranteed to pass.
 
-- Bazel 3.7.0
-- Libtorch 1.7.x (built with CUDA 11.0)
-- CUDA 11.0 (10.2 on Jetson)
-- cuDNN 8
-- TensorRT 7.2
+- Bazel 4.0.0
+- Libtorch 1.8.0 (built with CUDA 11.1)
+- CUDA 11.1 (10.2 on Jetson)
+- cuDNN 8.1
+- TensorRT 7.2.3
 
 ## Prebuilt Binaries and Wheel files
 

WORKSPACE

@@ -35,7 +35,7 @@ git_repository(
 # CUDA should be installed on the system locally
 new_local_repository(
     name = "cuda",
-    path = "/usr/local/cuda-11.0/",
+    path = "/usr/local/cuda-11.1/",
     build_file = "@//third_party/cuda:BUILD",
 )
 
@@ -53,16 +53,16 @@ http_archive(
     name = "libtorch",
     build_file = "@//third_party/libtorch:BUILD",
     strip_prefix = "libtorch",
-    sha256 = "117f6dd65b7267839197397edd0b10fd2900b0f291e3e54b0b800caefc31bcb6",
-    urls = ["https://download.pytorch.org/libtorch/cu110/libtorch-cxx11-abi-shared-with-deps-1.7.1%2Bcu110.zip"],
+    sha256 = "62a2c06761c32576b30f5884240cf675b937945d929e4b13cc776de8d9c2236c",
+    urls = ["https://download.pytorch.org/libtorch/cu111/libtorch-cxx11-abi-shared-with-deps-1.8.0%2Bcu111.zip"],
 )
 
 http_archive(
     name = "libtorch_pre_cxx11_abi",
     build_file = "@//third_party/libtorch:BUILD",
     strip_prefix = "libtorch",
-    sha256 = "c77f926afd55d7e860ec9c7abc992c25be77c89771c3ec6fcc13ea42f07d46df",
-    urls = ["https://download.pytorch.org/libtorch/cu110/libtorch-shared-with-deps-1.7.1%2Bcu110.zip"],
+    sha256 = "1c8b0c0883dd17f5ce952d42ec5f7f0cc7ceb370307535cee26a66c10419f1f6",
+    urls = ["https://download.pytorch.org/libtorch/cu111/libtorch-shared-with-deps-1.8.0%2Bcu111.zip"],
 )
 
 # Download these tarballs manually from the NVIDIA website
@@ -71,18 +71,18 @@ http_archive(
 
 http_archive(
     name = "cudnn",
-    urls = ["https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.0.5/11.0_20201106/cudnn-11.0-linux-x64-v8.0.5.39.tgz",],
+    urls = ["https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.1.1.33/11.2_20210301/cudnn-11.2-linux-x64-v8.1.1.33.tgz",],
     build_file = "@//third_party/cudnn/archive:BUILD",
-    sha256 = "4e16ee7895deb4a8b1c194b812ba49586ef7d26902051401d3717511898a9b73",
+    sha256 = "98a8784e92862f20018d20c281b30d4a0cd951f93694f6433ccf4ae9c502ba6a",
     strip_prefix = "cuda"
 )
 
 http_archive(
     name = "tensorrt",
-    urls = ["https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/7.2.2/tars/TensorRT-7.2.2.3.Ubuntu-18.04.x86_64-gnu.cuda-11.0.cudnn8.0.tar.gz",],
+    urls = ["https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/7.2.3/tars/TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-11.1.cudnn8.1.tar.gz",],
     build_file = "@//third_party/tensorrt/archive:BUILD",
-    strip_prefix = "TensorRT-7.2.2.3",
-    sha256 = "b5c325e38e1d92ce1ce92ca8b54ede9c224bf128c9a53eb0b9022f1ee4313ee0"
+    strip_prefix = "TensorRT-7.2.3.4",
+    sha256 = "d3a1f478e304b48878604fac70ce7920fece71f9cac62f925c9c59c197f5d087"
 )
 
 ####################################################################################

core/conversion/converters/impl/batch_norm.cpp

@@ -43,7 +43,7 @@ auto batch_norm_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().
       auto should_unpack = util::toVec(orig_shape).size() < 4;
       if (should_unpack) {
         // expand spatial dims from 1D to 2D
-        auto new_shape = util::toDimsPad(util::toVec(orig_shape), 4);
+        auto new_shape = util::toDimsTailPad(util::toVec(orig_shape), 4);
         LOG_DEBUG(
             "Input shape is less than 4D got: "
             << orig_shape << ", inserting shuffle layer to reshape to 4D tensor shape: " << new_shape);

core/conversion/converters/impl/plugins/interpolate_plugin.cpp

@@ -254,7 +254,8 @@ int InterpolatePlugin::enqueue(
     cudaStream_t stream) {
 #if NV_TENSORRT_MAJOR < 7 || (NV_TENSORRT_MAJOR == 7 && NV_TENSORRT_MINOR < 1)
   at::Tensor input = at::from_blob((void*)inputs[0], util::toVec(inputDesc->dims), [](void*) {}, tensor_options_);
-  at::Tensor output = at::from_blob(outputs[0], util::volume(outputDesc->dims), [](void*) {}, tensor_options_);
+  at::Tensor output = at::from_blob(
+      outputs[0], util::volume(outputDesc->dims), [](void*) {}, tensor_options_);
 
   at::cuda::CUDAStream torch_stream = at::cuda::getStreamFromPool();
   at::cuda::CUDAStreamGuard torch_guard(torch_stream);

core/conversion/converters/impl/select.cpp

@@ -18,7 +18,7 @@ bool add_split(ConversionCtx* ctx, const torch::jit::Node* n, args& args, bool s
   auto in = args[0].ITensor();
   auto axis = args[2].unwrapToInt();
   auto inDimSize = in->getDimensions().d[axis];
-  auto numOutputs = 1;
+  auto numOutputs = 1, numRemainder = 0;
   std::vector<int64_t> sizes;
 
   if (split_list) {
@@ -27,10 +27,13 @@ bool add_split(ConversionCtx* ctx, const torch::jit::Node* n, args& args, bool s
   } else {
     auto split_size = args[1].unwrapToInt();
     numOutputs = inDimSize / split_size;
-    if (numOutputs == 1) {
+    numRemainder = inDimSize % split_size;
+    for (int64_t i = 0; i < numOutputs; i++) {
       sizes.push_back(split_size);
-    } else {
-      sizes = std::vector<int64_t>(numOutputs, 1);
+    }
+    if (numRemainder) {
+      numOutputs += 1;
+      sizes.push_back(numRemainder);
     }
   }
 
@@ -42,7 +45,7 @@ bool add_split(ConversionCtx* ctx, const torch::jit::Node* n, args& args, bool s
   list.reserve(numOutputs);
 
   int start_idx = 0;
-  for (int i = 0; i < numOutputs; i++) {
+  for (int64_t i = 0; i < numOutputs; i++) {
     at::Tensor indices = torch::arange(start_idx, start_idx + sizes[i], 1).to(torch::kI32);
     auto indicesTensor = tensor_to_const(ctx, indices);
 
@@ -68,7 +71,9 @@ auto select_registrations TRTORCH_UNUSED =
         .pattern({"aten::select.int(Tensor(a) self, int dim, int index) -> (Tensor(a))",
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                     auto in = args[0].ITensor();
+                    auto maxDim = static_cast<int64_t>(in->getDimensions().nbDims);
                     auto axis = args[1].unwrapToInt();
+                    axis = axis < 0 ? axis + maxDim : axis;
                     auto ind = (int32_t)args[2].unwrapToInt();
 
                     // index to access needs to be an at::Tensor
@@ -89,7 +94,7 @@ auto select_registrations TRTORCH_UNUSED =
                     // IShuffleLayer removes redundant dimensions
                     auto shuffle_layer = ctx->net->addShuffle(*gather_out);
                     TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
-                    shuffle_layer->setReshapeDimensions(util::unpadDims(gather_out->getDimensions()));
+                    shuffle_layer->setReshapeDimensions(util::squeezeDims(gather_out->getDimensions(), axis));
                     shuffle_layer->setName(util::node_info(n).c_str());
                     auto shuffle_out = shuffle_layer->getOutput(0);
 
@@ -174,9 +179,11 @@ auto select_registrations TRTORCH_UNUSED =
             {"aten::embedding(Tensor weight, Tensor indices, int padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> (Tensor)",
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                auto embeddingTensor = args[0].ITensorOrFreeze(ctx);
-               auto indicesTensor = args[1].ITensor();
+               auto indicesTensor = args[1].ITensorOrFreeze(ctx);
                // Set datatype for indices tensor to INT32
-               indicesTensor->setType(nvinfer1::DataType::kINT32);
+               auto identity = ctx->net->addIdentity(*indicesTensor);
+               identity->setOutputType(0, nvinfer1::DataType::kINT32);
+               indicesTensor = identity->getOutput(0);
 
                // IGatherLayer takes in input tensor, the indices, and the axis of input tensor to take indices from
                auto gather_layer = ctx->net->addGather(*embeddingTensor, *indicesTensor, 0);

core/conversion/converters/impl/unsqueeze.cpp

@@ -21,8 +21,13 @@ auto unsqueeze_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().p
        auto dim = args[1].unwrapToInt();
 
        auto selfDim = util::toVec(self->getDimensions());
+       int64_t nbDims = selfDim.size();
+       TRTORCH_CHECK(
+           dim <= nbDims && dim >= -(nbDims + 1),
+           "Dimension out of range (expected to be in range of [" << -(nbDims + 1) << ", " << nbDims << "], but got "
+                                                                  << dim << ")");
        if (dim < 0) {
-         dim = selfDim.size() + dim;
+         dim = nbDims + 1 + dim;
        }
 
        auto shuffle_layer = ctx->net->addShuffle(*self);

core/conversion/evaluators/aten.cpp

@@ -396,6 +396,20 @@ auto aten_registrations TRTORCH_UNUSED =
                     EvalOptions().validSchemas({
                         "aten::numel(Tensor self) -> int",
                     })})
+        .evaluator({c10::Symbol::fromQualString("aten::t"),
+                    [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+                      auto tensor_var = args.at(n->input(0));
+                      if (tensor_var.IValue()->isTensor()) {
+                        auto tensor = tensor_var.unwrapToTensor();
+                        return tensor.t();
+                      } else {
+                        TRTORCH_THROW_ERROR("Unimplemented data type for aten::t evaluator: ITensor");
+                        return {};
+                      }
+                    },
+                    EvalOptions().validSchemas({
+                        "aten::t(Tensor self) -> Tensor",
+                    })})
         .evaluator({c10::Symbol::fromQualString("aten::dim"),
                     [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
                       auto tensor_var = args.at(n->input(0));
@@ -467,7 +481,61 @@ auto aten_registrations TRTORCH_UNUSED =
                       LOG_WARNING("Warning from TorchScript: " << *warning);
                       return {};
                     },
-                    EvalOptions()});
+                    EvalOptions()})
+        .evaluator({c10::Symbol::fromQualString("aten::arange"),
+                    [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+                      int input_size = n->inputs().size();
+                      int scalar_count = 0;
+                      for (int i = 0; i < input_size; i++) {
+                        if (args.at(n->input(i)).IValue()->isScalar()) {
+                          scalar_count += 1;
+                        }
+                      }
+                      if (scalar_count == 1) {
+                        if (args.at(n->input(0)).IValue()->isInt()) {
+                          int end_scalar = args.at(n->input(0)).unwrapToInt();
+                          return torch::arange(end_scalar);
+                        } else if (args.at(n->input(0)).IValue()->isDouble()) {
+                          float end_scalar = args.at(n->input(0)).unwrapToScalar().to<float>();
+                          return torch::arange(end_scalar);
+                        }
+                      } else if (scalar_count == 2) {
+                        if (args.at(n->input(0)).IValue()->isDouble() || args.at(n->input(1)).IValue()->isDouble()) {
+                          float start_scalar = args.at(n->input(0)).unwrapToScalar().to<float>();
+                          float end_scalar = args.at(n->input(1)).unwrapToScalar().to<float>();
+                          return torch::arange(start_scalar, end_scalar);
+                        } else {
+                          int start_scalar = args.at(n->input(0)).unwrapToInt();
+                          int end_scalar = args.at(n->input(1)).unwrapToInt();
+                          return torch::arange(start_scalar, end_scalar);
+                        }
+                      } else if (scalar_count == 3) {
+                        if (args.at(n->input(0)).IValue()->isDouble() || args.at(n->input(1)).IValue()->isDouble() ||
+                            args.at(n->input(2)).IValue()->isDouble()) {
+                          float start_scalar = args.at(n->input(0)).unwrapToScalar().to<float>();
+                          float end_scalar = args.at(n->input(1)).unwrapToScalar().to<float>();
+                          float step_scalar = args.at(n->input(2)).unwrapToScalar().to<float>();
+                          return torch::arange(start_scalar, end_scalar, step_scalar);
+                        } else {
+                          int start_scalar = args.at(n->input(0)).unwrapToInt();
+                          int end_scalar = args.at(n->input(1)).unwrapToInt();
+                          int step_scalar = args.at(n->input(2)).unwrapToInt();
+                          return torch::arange(start_scalar, end_scalar, step_scalar);
+                        }
+                      } else {
+                        TRTORCH_THROW_ERROR(
+                            "Invalid input argument size for aten::arange, input argument size: " << input_size);
+                      }
+                      return {};
+                    },
+                    EvalOptions().validSchemas({
+                        R"SIG(aten::arange(Scalar end, *, int? dtype=None, int? layout=None,
+                            Device? device=None, bool? pin_memory=None) -> (Tensor))SIG",
+                        R"SIG(aten::arange.start(Scalar start, Scalar end, *, ScalarType? dtype=None,
+                            Layout? layout=None, Device? device=None, bool? pin_memory=None) -> (Tensor))SIG",
+                        R"SIG(aten::arange.start_step(Scalar start, Scalar end, Scalar step, *, ScalarType? dtype=None,
+                        Layout? layout=None, Device? device=None, bool? pin_memory=None) -> (Tensor))SIG",
+                    })});
 } // namespace
 } // namespace evaluators
 } // namespace conversion