Merge remote-tracking branch 'origin/master' into plugins

Author: peri044

.github/scripts/run_cpp_linter.py

@@ -26,4 +26,7 @@
 
 pr.create_review(commit, comment, approval)
 
-
+if output.returncode != 0:
+    exit(1)
+else:
+    exit(0)

.github/scripts/run_py_linter.py

@@ -25,3 +25,8 @@
     approval = 'REQUEST_CHANGES'
 
 pr.create_review(commit, comment, approval)
+
+if output.returncode != 0:
+    exit(1)
+else:
+    exit(0)

README.md

@@ -74,7 +74,7 @@ torch.jit.save(trt_ts_module, "trt_torchscript_module.ts")
 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 4.0.0
-- Libtorch 1.8.0 (built with CUDA 11.1)
+- Libtorch 1.8.1 (built with CUDA 11.1)
 - CUDA 11.1 (10.2 on Jetson)
 - cuDNN 8.1
 - TensorRT 7.2.3

WORKSPACE

@@ -3,23 +3,21 @@ workspace(name = "TRTorch")
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
 
-git_repository(
-    name = "rules_python",
-    remote = "https://github.com/bazelbuild/rules_python.git",
-    commit = "4fcc24fd8a850bdab2ef2e078b1de337eea751a6",
-    shallow_since = "1589292086 -0400"
-)
-
-load("@rules_python//python:repositories.bzl", "py_repositories")
-py_repositories()
+http_archive(
+        name = "rules_python",
+        url = "https://github.com/bazelbuild/rules_python/releases/download/0.2.0/rules_python-0.2.0.tar.gz",
+        sha256 = "778197e26c5fbeb07ac2a2c5ae405b30f6cb7ad1f5510ea6fdac03bded96cc6f",
+    )
 
-load("@rules_python//python:pip.bzl", "pip_repositories", "pip3_import")
-pip_repositories()
+load("@rules_python//python:pip.bzl", "pip_install")
 
 http_archive(
     name = "rules_pkg",
-    url = "https://github.com/bazelbuild/rules_pkg/releases/download/0.2.4/rules_pkg-0.2.4.tar.gz",
-    sha256 = "4ba8f4ab0ff85f2484287ab06c0d871dcb31cc54d439457d28fd4ae14b18450a",
+    urls = [
+    	"https://mirror.bazel.build/github.com/bazelbuild/rules_pkg/releases/download/0.4.0/rules_pkg-0.4.0.tar.gz",
+	"https://github.com/bazelbuild/rules_pkg/releases/download/0.4.0/rules_pkg-0.4.0.tar.gz",
+    ],
+    sha256 = "038f1caa773a7e35b3663865ffb003169c6a71dc995e39bf4815792f385d837d",
 )
 
 load("@rules_pkg//:deps.bzl", "rules_pkg_dependencies")
@@ -39,12 +37,6 @@ new_local_repository(
     build_file = "@//third_party/cuda:BUILD",
 )
 
-new_local_repository(
-    name = "cublas",
-    path = "/usr",
-    build_file = "@//third_party/cublas:BUILD",
-)
-
 #############################################################################################################
 # Tarballs and fetched dependencies (default - use in cases when building from precompiled bin and tarballs)
 #############################################################################################################
@@ -53,16 +45,16 @@ http_archive(
     name = "libtorch",
     build_file = "@//third_party/libtorch:BUILD",
     strip_prefix = "libtorch",
-    sha256 = "62a2c06761c32576b30f5884240cf675b937945d929e4b13cc776de8d9c2236c",
-    urls = ["https://download.pytorch.org/libtorch/cu111/libtorch-cxx11-abi-shared-with-deps-1.8.0%2Bcu111.zip"],
+    sha256 = "1f8aec376f9343538bd7c2fd3abb81ed3af11f575efe3aa72777c4d62044b832",
+    urls = ["https://download.pytorch.org/libtorch/cu111/libtorch-cxx11-abi-shared-with-deps-1.8.1%2Bcu111.zip"],
 )
 
 http_archive(
     name = "libtorch_pre_cxx11_abi",
     build_file = "@//third_party/libtorch:BUILD",
     strip_prefix = "libtorch",
-    sha256 = "1c8b0c0883dd17f5ce952d42ec5f7f0cc7ceb370307535cee26a66c10419f1f6",
-    urls = ["https://download.pytorch.org/libtorch/cu111/libtorch-shared-with-deps-1.8.0%2Bcu111.zip"],
+    sha256 = "3a6e0dc11859111e75caa640c8ce9bf904fbb6e9992b4345e444ed5410e4d77e",
+    urls = ["https://download.pytorch.org/libtorch/cu111/libtorch-shared-with-deps-1.8.1%2Bcu111.zip"],
 )
 
 # Download these tarballs manually from the NVIDIA website
@@ -71,15 +63,19 @@ http_archive(
 
 http_archive(
     name = "cudnn",
-    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",],
+    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 = "98a8784e92862f20018d20c281b30d4a0cd951f93694f6433ccf4ae9c502ba6a",
     strip_prefix = "cuda"
 )
 
 http_archive(
     name = "tensorrt",
-    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",],
+    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.3.4",
     sha256 = "d3a1f478e304b48878604fac70ce7920fece71f9cac62f925c9c59c197f5d087"
@@ -123,26 +119,17 @@ http_archive(
 #########################################################################
 # Testing Dependencies (optional - comment out on aarch64)
 #########################################################################
-pip3_import(
+pip_install(
     name = "trtorch_py_deps",
-    requirements = "//py:requirements.txt"
+    requirements = "//py:requirements.txt",
 )
 
-load("@trtorch_py_deps//:requirements.bzl", "pip_install")
-pip_install()
-
-pip3_import(
+pip_install(
     name = "py_test_deps",
-    requirements = "//tests/py:requirements.txt"
+    requirements = "//tests/py:requirements.txt",
 )
 
-load("@py_test_deps//:requirements.bzl", "pip_install")
-pip_install()
-
-pip3_import(
-   name = "pylinter_deps",
-   requirements = "//tools/linter:requirements.txt",
+pip_install(
+    name = "pylinter_deps",
+    requirements = "//tools/linter:requirements.txt",
 )
-
-load("@pylinter_deps//:requirements.bzl", "pip_install")
-pip_install()

core/compiler.cpp

@@ -46,7 +46,7 @@ c10::FunctionSchema GenerateGraphSchema(
 void AddEngineToGraph(
     torch::jit::script::Module mod,
     std::shared_ptr<torch::jit::Graph>& g,
-    std::string& serialized_engine) {
+    const std::string& serialized_engine) {
   auto engine_ptr = c10::make_intrusive<runtime::TRTEngine>(mod._ivalue()->name(), serialized_engine);
   // Get required metadata about the engine out
   auto num_io = engine_ptr->num_io;
@@ -173,6 +173,20 @@ torch::jit::script::Module CompileGraph(const torch::jit::script::Module& mod, C
   return new_mod;
 }
 
+torch::jit::script::Module EmbedEngineInNewModule(const std::string& engine) {
+  std::ostringstream engine_id;
+  engine_id << reinterpret_cast<const int*>(&engine);
+  torch::jit::script::Module new_mod("tensorrt_engine_mod_" + engine_id.str());
+  auto new_g = std::make_shared<torch::jit::Graph>();
+  AddEngineToGraph(new_mod, new_g, engine);
+  auto new_method = new_mod._ivalue()->compilation_unit()->create_function("forward", new_g);
+  auto schema = GenerateGraphSchema(new_mod, new_method->name(), new_g);
+  new_mod.type()->addMethod(new_method);
+  new_method->setSchema(schema);
+
+  return new_mod;
+}
+
 void set_device(const int gpu_id) {
   TRTORCH_ASSERT(cudaSetDevice(gpu_id) == cudaSuccess, "Unable to set CUDA device: " << gpu_id);
 }

core/compiler.h

@@ -19,6 +19,8 @@ std::string ConvertGraphToTRTEngine(const torch::jit::script::Module& mod, std::
 
 torch::jit::script::Module CompileGraph(const torch::jit::script::Module& module, CompileSpec cfg);
 
+torch::jit::script::Module EmbedEngineInNewModule(const std::string& engine);
+
 void set_device(const int gpu_id);
 
 } // namespace core

core/conversion/conversionctx/ConversionCtx.cpp

@@ -12,6 +12,7 @@ std::ostream& operator<<(std::ostream& os, const BuilderSettings& s) {
     os << "Settings requested for TensorRT engine:"                                        \
        << "\n    Operating Precision: " << s.op_precision                                  \
        << "\n    TF32 Floating Point Computation Enabled: " << !s.disable_tf32             \
+       << "\n    Truncate Long and Double: " << s.truncate_long_and_double                 \
        << "\n    Make Refittable Engine: " << s.refit                                      \
        << "\n    Debuggable Engine: " << s.debug                                           \
        << "\n    Strict Types: " << s.strict_types                                         \

core/conversion/converters/BUILD

@@ -45,6 +45,7 @@ cc_library(
         "impl/normalize.cpp",
         "impl/pooling.cpp",
         "impl/reduce.cpp",
+        "impl/replication_pad.cpp",
         "impl/shuffle.cpp",
         "impl/softmax.cpp",
         "impl/unary.cpp",

core/conversion/converters/impl/replication_pad.cpp

@@ -0,0 +1,117 @@
+#include <ATen/ATen.h>
+#include <vector>
+#include "NvInfer.h"
+#include "core/conversion/converters/converters.h"
+#include "core/util/prelude.h"
+#include "torch/torch.h"
+
+namespace trtorch {
+namespace core {
+namespace conversion {
+namespace converters {
+namespace impl {
+namespace {
+
+bool replication_padXd(ConversionCtx* ctx, const torch::jit::Node* n, args& args, int x_dim) {
+  auto in = args[0].ITensor();
+  auto inDims = in->getDimensions();
+  int64_t inRank = inDims.nbDims;
+  auto padding = args[1].unwrapToIntList().vec();
+  if (padding.size() == 1) {
+    for (int64_t i = 0; i < x_dim * 2 - 1; i++)
+      padding.push_back(padding[0]);
+  }
+  if (inRank == 3) {
+    TRTORCH_CHECK(padding.size() == 2, "3D tensors expect 2 values for padding");
+  } else if (inRank == 4) {
+    TRTORCH_CHECK(padding.size() == 4, "4D tensors expect 4 values for padding");
+  } else if (inRank == 5) {
+    TRTORCH_CHECK(padding.size() == 6, "5D tensors expect 6 values for padding");
+  } else {
+    TRTORCH_THROW_ERROR("Only 3D, 4D, 5D padding with non-constant padding are supported for now");
+  }
+
+  std::vector<nvinfer1::ITensor*> tensors_vec;
+  // input: (N, C, D_in, H_in, W_in).
+  // padding: (padding_left, padding_right, padding_top, padding_bottom, padding_front, padding_back)
+  // When axis is inRank - 1, making W_out = W_in + padding_left + padding_right.
+  // When axis is inRank - 2, making H_out = H_in + padding_top + padding_bottom.
+  // When axis is inRank - 1, making D_out = D_in + padding_front + padding_back.
+  for (int64_t i = 0; i < int(padding.size() / 2); i++) {
+    int64_t axis = inRank - (i + 1); // axis = {inRank - 1, inRank - 2, inRank - 3}
+    int64_t padding_index = i * 2;
+
+    if (padding[padding_index] > 0) { // left/top/front padding value
+      tensors_vec.clear();
+      at::Tensor left_indices = torch::tensor({0}, torch::kInt32);
+      auto indicesTensor = tensor_to_const(ctx, left_indices);
+      auto left_gather_layer = ctx->net->addGather(*in, *indicesTensor, axis);
+      auto left_gather_out = left_gather_layer->getOutput(0);
+      for (int i = 0; i < padding[padding_index]; i++) {
+        tensors_vec.push_back(left_gather_out);
+      }
+      tensors_vec.push_back(in);
+      auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
+      concat_layer->setAxis(axis);
+      in = concat_layer->getOutput(0);
+      inDims = in->getDimensions();
+    }
+
+    if (padding[padding_index + 1] > 0) { // right/bottom/back padding value
+      tensors_vec.clear();
+      tensors_vec.push_back(in);
+
+      nvinfer1::ITensor* indicesTensor = NULL;
+      if (inDims.d[axis] == -1) {
+        auto shapeTensor = ctx->net->addShape(*in)->getOutput(0);
+        at::Tensor dimValue = torch::tensor({axis}, torch::kInt32);
+        auto dimTensor = tensor_to_const(ctx, dimValue);
+        indicesTensor = ctx->net->addGather(*shapeTensor, *dimTensor, 0)->getOutput(0);
+      } else {
+        auto indices = torch::tensor({inDims.d[axis] - 1}, torch::kInt32);
+        indicesTensor = tensor_to_const(ctx, indices);
+      }
+      auto right_gather_layer = ctx->net->addGather(*in, *indicesTensor, axis);
+      auto right_gather_out = right_gather_layer->getOutput(0);
+
+      for (int i = 0; i < padding[padding_index + 1]; i++) {
+        tensors_vec.push_back(right_gather_out);
+      }
+
+      auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
+      concat_layer->setAxis(axis);
+      in = concat_layer->getOutput(0);
+      inDims = in->getDimensions();
+    }
+  }
+
+  auto out = ctx->AssociateValueAndTensor(n->outputs()[0], in);
+  LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+
+  return true;
+}
+
+auto replication_pad_registrations TRTORCH_UNUSED =
+    RegisterNodeConversionPatterns()
+        .pattern({"aten::replication_pad1d(Tensor self, int[2] padding) -> (Tensor)",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    replication_padXd(ctx, n, args, 1);
+                    return true;
+                  }})
+        .pattern({"aten::replication_pad2d(Tensor self, int[4] padding) -> (Tensor)",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    replication_padXd(ctx, n, args, 2);
+                    return true;
+                  }})
+        .pattern({"aten::replication_pad3d(Tensor self, int[6] padding) -> (Tensor)",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    replication_padXd(ctx, n, args, 3);
+                    return true;
+                  }});
+
+} // namespace
+} // namespace impl
+} // namespace converters
+} // namespace conversion
+} // namespace core
+} // namespace trtorch