merge: merge and resolve the conficts from fallback branch

Signed-off-by: Bo Wang <[email protected]>

Author: bowang007

BUILD

@@ -15,29 +15,28 @@ pkg_tar(
         "//core/conversion:include",
         "//core/conversion/conversionctx:include",
         "//core/conversion/converters:include",
-        "//core/conversion/var:include",
-        "//core/conversion/tensorcontainer:include",
-        "//core/conversion/evaluators:include",
         "//core/conversion/converters/impl/plugins:include",
-        "//core/runtime:include",
+        "//core/conversion/evaluators:include",
+        "//core/conversion/tensorcontainer:include",
+        "//core/conversion/var:include",
         "//core/lowering:include",
         "//core/lowering/passes:include",
+        "//core/runtime:include",
         "//core/util:include",
-        "//core/util/logging:include"
+        "//core/util/logging:include",
     ],
 )
 
 pkg_tar(
     name = "include",
-    package_dir = "include/trtorch/",
     srcs = [
         "//cpp/api:api_headers",
     ],
+    package_dir = "include/trtorch/",
 )
 
 pkg_tar(
     name = "lib",
-    package_dir = "lib/",
     srcs = select({
         ":windows": ["//cpp/api/lib:trtorch.dll"],
         "//conditions:default": [
@@ -46,26 +45,25 @@ pkg_tar(
         ],
     }),
     mode = "0755",
+    package_dir = "lib/",
 )
 
-
 pkg_tar(
     name = "bin",
-    package_dir = "bin/",
     srcs = [
-        "//cpp/trtorchc:trtorchc",
+        "//cpp/trtorchc",
     ],
     mode = "0755",
+    package_dir = "bin/",
 )
 
-
 pkg_tar(
     name = "libtrtorch",
-    extension = "tar.gz",
-    package_dir = "trtorch",
     srcs = [
-        "//:LICENSE"
+        "//:LICENSE",
     ],
+    extension = "tar.gz",
+    package_dir = "trtorch",
     deps = [
         ":lib",
         ":include",

WORKSPACE

@@ -4,37 +4,38 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
 
 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",
-    )
+    name = "rules_python",
+    sha256 = "778197e26c5fbeb07ac2a2c5ae405b30f6cb7ad1f5510ea6fdac03bded96cc6f",
+    url = "https://github.com/bazelbuild/rules_python/releases/download/0.2.0/rules_python-0.2.0.tar.gz",
+)
 
 load("@rules_python//python:pip.bzl", "pip_install")
 
 http_archive(
     name = "rules_pkg",
+    sha256 = "038f1caa773a7e35b3663865ffb003169c6a71dc995e39bf4815792f385d837d",
     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",
+        "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")
+
 rules_pkg_dependencies()
 
 git_repository(
     name = "googletest",
-    remote = "https://github.com/google/googletest",
     commit = "703bd9caab50b139428cea1aaff9974ebee5742e",
-    shallow_since = "1570114335 -0400"
+    remote = "https://github.com/google/googletest",
+    shallow_since = "1570114335 -0400",
 )
 
 # CUDA should be installed on the system locally
 new_local_repository(
     name = "cuda",
-    path = "/usr/local/cuda-11.1/",
     build_file = "@//third_party/cuda:BUILD",
+    path = "/usr/local/cuda-11.1/",
 )
 
 #############################################################################################################
@@ -44,16 +45,16 @@ new_local_repository(
 http_archive(
     name = "libtorch",
     build_file = "@//third_party/libtorch:BUILD",
-    strip_prefix = "libtorch",
     sha256 = "1f8aec376f9343538bd7c2fd3abb81ed3af11f575efe3aa72777c4d62044b832",
+    strip_prefix = "libtorch",
     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 = "3a6e0dc11859111e75caa640c8ce9bf904fbb6e9992b4345e444ed5410e4d77e",
+    strip_prefix = "libtorch",
     urls = ["https://download.pytorch.org/libtorch/cu111/libtorch-shared-with-deps-1.8.1%2Bcu111.zip"],
 )
 
@@ -63,22 +64,22 @@ http_archive(
 
 http_archive(
     name = "cudnn",
+    build_file = "@//third_party/cudnn/archive:BUILD",
+    sha256 = "98a8784e92862f20018d20c281b30d4a0cd951f93694f6433ccf4ae9c502ba6a",
+    strip_prefix = "cuda",
     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",
+    build_file = "@//third_party/tensorrt/archive:BUILD",
+    sha256 = "d3a1f478e304b48878604fac70ce7920fece71f9cac62f925c9c59c197f5d087",
+    strip_prefix = "TensorRT-7.2.3.4",
     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"
 )
 
 ####################################################################################

core/BUILD

@@ -4,44 +4,44 @@ config_setting(
     name = "use_pre_cxx11_abi",
     values = {
         "define": "abi=pre_cxx11_abi",
-    }
+    },
 )
 
 config_setting(
     name = "python_core",
     values = {
-        "define": "target_lang=python"
-    }
+        "define": "target_lang=python",
+    },
 )
 
 cc_library(
     name = "core",
-    hdrs = [
-        "compiler.h",
-    ],
     srcs = [
         "compiler.cpp",
     ],
+    hdrs = [
+        "compiler.h",
+    ],
     deps = [
         "//core/conversion",
         "//core/runtime",
         "//core/lowering",
         "//core/partitioning",
         "//core/util/logging",
-        "@tensorrt//:nvinfer"
+        "@tensorrt//:nvinfer",
     ] + select({
-        ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default":  ["@libtorch//:libtorch"],
+        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@libtorch//:libtorch"],
     }),
-    alwayslink=True,
+    alwayslink = True,
 )
 
 load("@rules_pkg//:pkg.bzl", "pkg_tar")
 
 pkg_tar(
     name = "include",
-    package_dir = "core/",
     srcs = [
         "compiler.h",
     ],
+    package_dir = "core/",
 )

core/compiler.cpp

@@ -13,6 +13,7 @@
 #include "torch/csrc/jit/ir/ir.h"
 #include "torch/csrc/jit/ir/ir_views.h"
 #include "torch/csrc/jit/passes/graph_fuser.h"
+#include "torch/csrc/jit/passes/loop_unrolling.h"
 #include "torch/csrc/jit/passes/lower_graph.h"
 #include "torch/csrc/jit/passes/pass_manager.h"
 #include "torch/custom_class.h"
@@ -31,10 +32,9 @@ void AddEngineToGraph(
     torch::jit::script::Module mod,
     std::shared_ptr<torch::jit::Graph>& g,
     const std::string& serialized_engine,
-    int engine_id = 0,
+    std::string engine_id = "",
     bool fallback = false) {
-  auto engine_ptr =
-      c10::make_intrusive<runtime::TRTEngine>(mod._ivalue()->name() + std::to_string(engine_id), serialized_engine);
+  auto engine_ptr = c10::make_intrusive<runtime::TRTEngine>(mod._ivalue()->name() + engine_id, serialized_engine);
   // Get required metadata about the engine out
   auto num_io = engine_ptr->num_io;
   auto name = engine_ptr->name;
@@ -237,7 +237,6 @@ GraphAndMapping ConstructFallbackGraph(
     torch::jit::Block* block,
     std::unordered_map<torch::jit::Value*, torch::jit::IValue> input_ivalues_map,
     CompileSpec cfg,
-    int& trt_engine_id,
     conversion::GraphParams named_params) {
   auto convert_cfg = cfg.convert_info;
   auto partition_info = cfg.partition_info;
@@ -254,6 +253,8 @@ GraphAndMapping ConstructFallbackGraph(
 
   for (auto& seg_block : segmented_blocks) {
     LOG_INFO(*seg_block.g() << "(GraphInSegmentedBlock)\n");
+    std::ostringstream trt_engine_id;
+    trt_engine_id << reinterpret_cast<const int*>(&seg_block);
 
     if (seg_block.target() == partitioning::SegmentedBlock::kTensorRT) {
       std::vector<ir::InputRange> input_ranges;
@@ -264,7 +265,7 @@ GraphAndMapping ConstructFallbackGraph(
       convert_cfg.input_ranges = input_ranges;
       auto engine = conversion::ConvertBlockToEngine(seg_block.block(), convert_cfg, named_params);
       auto temp_g = std::make_shared<torch::jit::Graph>();
-      AddEngineToGraph(new_mod, temp_g, engine, trt_engine_id++, true);
+      AddEngineToGraph(new_mod, temp_g, engine, trt_engine_id.str(), true);
 
       seg_block.update_graph(temp_g);
       AddSegmentedBlockToGraph(new_g, seg_block, old_to_new_g);
@@ -276,7 +277,7 @@ GraphAndMapping ConstructFallbackGraph(
         std::vector<GraphAndMapping> graph_and_mappings;
         for (auto cur_block : if_node->blocks()) {
           graph_and_mappings.push_back(
-              ConstructFallbackGraph(new_mod, cur_block, input_ivalues_map, cfg, trt_engine_id, named_params));
+              ConstructFallbackGraph(new_mod, cur_block, input_ivalues_map, cfg, named_params));
         }
         AddIfBlockToGraph(new_g, if_node, graph_and_mappings, old_to_new_g);
 
@@ -311,14 +312,13 @@ torch::jit::script::Module CompileGraphWithFallback(const torch::jit::script::Mo
       auto named_params = conversion::get_named_params(g->inputs(), params);
       LOG_INFO(*g << "(LoweringGraph)\n");
 
-      int trt_engine_id = 0;
       std::unordered_map<torch::jit::Value*, ir::InputRange> input_ranges;
       for (size_t i = 0; i < g->inputs().size(); ++i) {
         input_ranges.insert({g->inputs()[i], cfg.convert_info.input_ranges[i]});
       }
       auto input_ivalues_map = partitioning::generateRandomInputs(input_ranges);
       auto graph_and_mapping =
-          ConstructFallbackGraph(new_mod, g->block(), input_ivalues_map, cfg, trt_engine_id, named_params);
+          ConstructFallbackGraph(new_mod, g->block(), input_ivalues_map, cfg, named_params);
       new_g = graph_and_mapping.first;
       LOG_INFO(*new_g << "(FallbackGraph)\n");
 

core/conversion/BUILD

@@ -4,37 +4,37 @@ config_setting(
     name = "use_pre_cxx11_abi",
     values = {
         "define": "abi=pre_cxx11_abi",
-    }
+    },
 )
 
 cc_library(
     name = "conversion",
-    hdrs = [
-        "conversion.h",
-    ],
     srcs = [
+        "InterfaceTypes.cpp",
         "conversion.cpp",
         "conversion_ignorelist.cpp",
-        "InterfaceTypes.cpp"
+    ],
+    hdrs = [
+        "conversion.h",
     ],
     deps = [
         "@tensorrt//:nvinfer",
         "//core/conversion/var",
         "//core/conversion/conversionctx",
         "//core/conversion/converters",
         "//core/conversion/evaluators",
-        "//core/util:prelude",
         "//core/ir",
+        "//core/util:prelude",
     ] + select({
-        ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default":  ["@libtorch//:libtorch"],
+        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@libtorch//:libtorch"],
     }),
 )
 
 load("@rules_pkg//:pkg.bzl", "pkg_tar")
 
 pkg_tar(
     name = "include",
-    package_dir = "core/conversion/",
     srcs = ["conversion.h"],
+    package_dir = "core/conversion/",
 )

core/conversion/conversionctx/BUILD

@@ -4,30 +4,30 @@ config_setting(
     name = "use_pre_cxx11_abi",
     values = {
         "define": "abi=pre_cxx11_abi",
-    }
+    },
 )
 
 cc_library(
     name = "conversionctx",
-    hdrs = [
-        "ConversionCtx.h",
-    ],
     srcs = [
         "ConversionCtx.cpp",
     ],
+    hdrs = [
+        "ConversionCtx.h",
+    ],
     deps = [
         "@tensorrt//:nvinfer",
         "//core/util:prelude",
     ] + select({
-        ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default":  ["@libtorch//:libtorch"],
+        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@libtorch//:libtorch"],
     }),
 )
 
 load("@rules_pkg//:pkg.bzl", "pkg_tar")
 
 pkg_tar(
     name = "include",
-    package_dir = "core/conversion/conversionctx/",
     srcs = ["ConversionCtx.h"],
+    package_dir = "core/conversion/conversionctx/",
 )