refactor(//core/partitioning): Reorganizing partitioning deps

Reorganizing the partitioning dependencies so that there is a clearer
relationship between major compiler modules

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

core/compiler.cpp

@@ -214,7 +214,7 @@ torch::jit::script::Module CompileGraphWithFallback(const torch::jit::script::Mo
 
       // segment the graph and convert segmented TensorRT block
       auto segmented_blocks =
-          partitioning::Partition(g, convert_cfg.input_ranges, convert_cfg.engine_settings.torch_fallback);
+          partitioning::Partition(g, convert_cfg.input_ranges, cfg.partition_info);
       if (segmented_blocks.size() == 1 && segmented_blocks[0].target() == partitioning::SegmentedBlock::kTorch) {
         return mod;
       }
@@ -223,9 +223,9 @@ torch::jit::script::Module CompileGraphWithFallback(const torch::jit::script::Mo
       std::unordered_map<torch::jit::Value*, torch::jit::Value*> old_to_new_g;
       for (auto& seg_block : segmented_blocks) {
         if (seg_block.target() == partitioning::SegmentedBlock::kTensorRT) {
-          std::vector<conversion::InputRange> input_ranges;
+          std::vector<ir::InputRange> input_ranges;
           for (auto& shape : seg_block.in_shape()) {
-            input_ranges.push_back(conversion::InputRange(util::toVec(shape)));
+            input_ranges.push_back(ir::InputRange(util::toVec(shape)));
           }
           // update the input ranges for each segments
           convert_cfg.input_ranges = input_ranges;
@@ -258,7 +258,7 @@ torch::jit::script::Module CompileGraphWithFallback(const torch::jit::script::Mo
 
 torch::jit::script::Module CompileGraph(const torch::jit::script::Module& mod, CompileSpec cfg) {
   // TODO: not sure how to deal with duplicated code here, so just cut out a branch temporally
-  if (cfg.convert_info.engine_settings.torch_fallback.enabled) {
+  if (cfg.partition_info.enabled) {
     return CompileGraphWithFallback(mod, cfg);
   }
   // TODO: Should be doing a functional transform but need PR #31978

core/compiler.h

@@ -3,14 +3,17 @@
 #include <cuda_runtime.h>
 #include <vector>
 #include "core/conversion/conversion.h"
+#include "core/ir/ir.h"
+#include "core/partitioning/partitioning.h"
 #include "torch/csrc/jit/api/module.h"
 
 namespace trtorch {
 namespace core {
 
 struct CompileSpec {
-  CompileSpec(std::vector<conversion::InputRange> input_ranges) : convert_info(std::move(input_ranges)) {}
+  CompileSpec(std::vector<ir::InputRange> input_ranges) : convert_info(std::move(input_ranges)) {}
   conversion::ConversionInfo convert_info;
+  partitioning::PartitionInfo partition_info;
 };
 
 bool CheckMethodOperatorSupport(const torch::jit::script::Module& mod, std::string method_name);

core/conversion/BUILD

@@ -23,7 +23,8 @@ cc_library(
         "//core/conversion/conversionctx",
         "//core/conversion/converters",
         "//core/conversion/evaluators",
-        "//core/util:prelude"
+        "//core/util:prelude",
+        "//core/ir",
     ] + select({
         ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
         "//conditions:default":  ["@libtorch//:libtorch"],

core/conversion/InterfaceTypes.cpp

@@ -23,55 +23,6 @@ GraphParams get_named_params(c10::ArrayRef<torch::jit::Value*> inputs, std::vect
   return std::move(named_params);
 }
 
-InputRange::InputRange(std::vector<int64_t> d) {
-  if (d.size() > 5) {
-    LOG_WARNING("Verify that this dim size is accepted");
-  }
-
-  opt = util::toDims(d);
-  min = util::toDims(d);
-  max = util::toDims(d);
-  input_shape = util::toDims(d);
-  input_is_dynamic = false;
-}
-
-InputRange::InputRange(std::vector<int64_t> min_shape, std::vector<int64_t> opt_shape, std::vector<int64_t> max_shape) {
-  if (min_shape.size() > 5 || opt_shape.size() > 5 || max_shape.size() > 5) {
-    LOG_WARNING("Verify that this dim size is accepted");
-  }
-
-  std::set<size_t> sizes;
-  sizes.insert(min_shape.size());
-  sizes.insert(opt_shape.size());
-  sizes.insert(max_shape.size());
-
-  if (sizes.size() != 1) {
-    LOG_ERROR(
-        "Expected all input sizes have the same dimensions, but found dimensions: min("
-        << min_shape.size() << "), opt(" << opt_shape.size() << "), max(" << max_shape.size() << ")");
-  }
-
-  min = util::toDims(min_shape);
-  opt = util::toDims(opt_shape);
-  max = util::toDims(max_shape);
-
-  std::vector<int64_t> dyn_shape;
-  for (size_t i = 0; i < opt_shape.size(); i++) {
-    std::set<uint64_t> dim;
-    dim.insert(min_shape[i]);
-    dim.insert(opt_shape[i]);
-    dim.insert(max_shape[i]);
-    if (dim.size() != 1) {
-      dyn_shape.push_back(-1);
-      input_is_dynamic = true;
-    } else {
-      dyn_shape.push_back(opt_shape[i]);
-    }
-  }
-
-  input_shape = util::toDims(dyn_shape);
-}
-
 } // namespace conversion
 } // namespace core
 } // namespace trtorch

core/conversion/conversion.cpp

@@ -118,7 +118,7 @@ void AddLayer(ConversionCtx* ctx, const torch::jit::Node* n) {
                        << "please report this error to https://www.github.com/NVIDIA/TRTorch/issues");
 }
 
-void AddInputs(ConversionCtx* ctx, at::ArrayRef<const torch::jit::Value*> inputs, std::vector<InputRange>& input_dims) {
+void AddInputs(ConversionCtx* ctx, at::ArrayRef<const torch::jit::Value*> inputs, std::vector<ir::InputRange>& input_dims) {
   std::vector<const torch::jit::Value*> input_tensors;
   for (auto in : inputs) {
     // Disregarding inputs that are not tensors

core/conversion/conversion.h

@@ -4,27 +4,17 @@
 
 #include "NvInfer.h"
 #include "core/conversion/conversionctx/ConversionCtx.h"
+#include "core/ir/ir.h"
 #include "torch/csrc/jit/ir/ir.h"
 
 namespace trtorch {
 namespace core {
 namespace conversion {
 
-struct InputRange {
-  nvinfer1::Dims min;
-  nvinfer1::Dims max;
-  nvinfer1::Dims opt;
-  nvinfer1::Dims input_shape;
-  bool input_is_dynamic = false;
-  // Should we restrict to unsigned?
-  InputRange(std::vector<int64_t> d);
-  InputRange(std::vector<int64_t> min_shape, std::vector<int64_t> opt_shape, std::vector<int64_t> max_shape);
-};
-
 struct ConversionInfo {
-  std::vector<InputRange> input_ranges;
+  std::vector<ir::InputRange> input_ranges;
   BuilderSettings engine_settings;
-  ConversionInfo(std::vector<InputRange> input_ranges)
+  ConversionInfo(std::vector<ir::InputRange> input_ranges)
       : input_ranges(std::move(input_ranges)), engine_settings(BuilderSettings()) {}
 };
 

core/conversion/conversionctx/ConversionCtx.cpp

@@ -36,17 +36,6 @@ std::ostream& operator<<(std::ostream& os, const BuilderSettings& s) {
     }
     os << "\n    Engine Capability: " << s.capability                                      \
        << "\n    Calibrator Created: " << (s.calibrator != nullptr);
-
-    os << "\n    Torch Fallback: " << s.torch_fallback.enabled;
-    if (s.torch_fallback.enabled) {
-      os << "\n    Fallback Min Block Size: " << s.torch_fallback.min_block_size;
-      if (!s.torch_fallback.forced_fallback_operators.empty()) {
-        os << "\n    Forced Fallback Operators:";
-        for (auto it = s.torch_fallback.forced_fallback_operators.begin(); it != s.torch_fallback.forced_fallback_operators.end(); ++it) {
-          os << " " << *it;
-        }
-      }
-    }
     return os;
 }
 // clang-format on

core/conversion/conversionctx/ConversionCtx.h

@@ -22,12 +22,6 @@ struct Device {
   Device() : device_type(nvinfer1::DeviceType::kGPU), gpu_id(0), dla_core(0), allow_gpu_fallback(false) {}
 };
 
-struct TorchFallback {
-  bool enabled = false;
-  uint64_t min_block_size = 1;
-  std::vector<std::string> forced_fallback_operators;
-};
-
 struct BuilderSettings {
   nvinfer1::DataType op_precision = nvinfer1::DataType::kFLOAT;
   bool disable_tf32 = false;
@@ -36,7 +30,6 @@ struct BuilderSettings {
   bool strict_types = false;
   bool truncate_long_and_double = false;
   Device device;
-  TorchFallback torch_fallback;
   nvinfer1::EngineCapability capability = nvinfer1::EngineCapability::kDEFAULT;
   nvinfer1::IInt8Calibrator* calibrator = nullptr;
   uint64_t num_min_timing_iters = 2;

core/ir/BUILD

@@ -0,0 +1,35 @@
+package(default_visibility = ["//visibility:public"])
+
+config_setting(
+    name = "use_pre_cxx11_abi",
+    values = {
+        "define": "abi=pre_cxx11_abi",
+    }
+)
+
+cc_library(
+    name = "ir",
+    hdrs = [
+        "ir.h"
+    ],
+    srcs = [
+        "InputRange.cpp",
+    ],
+    deps = [
+        "@tensorrt//:nvinfer",
+        "//core/util:prelude",
+    ] + select({
+        ":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/ir/",
+    srcs = [
+        "ir.h",
+    ],
+)

core/ir/InputRange.cpp

@@ -0,0 +1,59 @@
+#include "core/ir/ir.h"
+#include "core/util/prelude.h"
+
+namespace trtorch {
+namespace core {
+namespace ir {
+
+InputRange::InputRange(std::vector<int64_t> d) {
+  if (d.size() > 5) {
+    LOG_WARNING("Verify that this dim size is accepted");
+  }
+
+  opt = util::toDims(d);
+  min = util::toDims(d);
+  max = util::toDims(d);
+  input_shape = util::toDims(d);
+  input_is_dynamic = false;
+}
+
+InputRange::InputRange(std::vector<int64_t> min_shape, std::vector<int64_t> opt_shape, std::vector<int64_t> max_shape) {
+  if (min_shape.size() > 5 || opt_shape.size() > 5 || max_shape.size() > 5) {
+    LOG_WARNING("Verify that this dim size is accepted");
+  }
+
+  std::set<size_t> sizes;
+  sizes.insert(min_shape.size());
+  sizes.insert(opt_shape.size());
+  sizes.insert(max_shape.size());
+
+  if (sizes.size() != 1) {
+    LOG_ERROR(
+        "Expected all input sizes have the same dimensions, but found dimensions: min("
+        << min_shape.size() << "), opt(" << opt_shape.size() << "), max(" << max_shape.size() << ")");
+  }
+
+  min = util::toDims(min_shape);
+  opt = util::toDims(opt_shape);
+  max = util::toDims(max_shape);
+
+  std::vector<int64_t> dyn_shape;
+  for (size_t i = 0; i < opt_shape.size(); i++) {
+    std::set<uint64_t> dim;
+    dim.insert(min_shape[i]);
+    dim.insert(opt_shape[i]);
+    dim.insert(max_shape[i]);
+    if (dim.size() != 1) {
+      dyn_shape.push_back(-1);
+      input_is_dynamic = true;
+    } else {
+      dyn_shape.push_back(opt_shape[i]);
+    }
+  }
+
+  input_shape = util::toDims(dyn_shape);
+}
+
+} // namespace ir
+} // namespace core
+} // namespace trtorch

core/ir/ir.h

@@ -0,0 +1,23 @@
+#pragma once
+
+#include <vector>
+#include "NvInfer.h"
+
+namespace trtorch {
+namespace core {
+namespace ir {
+
+struct InputRange {
+  nvinfer1::Dims min;
+  nvinfer1::Dims max;
+  nvinfer1::Dims opt;
+  nvinfer1::Dims input_shape;
+  bool input_is_dynamic = false;
+  // Should we restrict to unsigned?
+  InputRange(std::vector<int64_t> d);
+  InputRange(std::vector<int64_t> min_shape, std::vector<int64_t> opt_shape, std::vector<int64_t> max_shape);
+};
+
+} // namespace ir
+} // namespace core
+} // namespace trtorch

core/partitioning/BUILD

@@ -12,6 +12,7 @@ cc_library(
     hdrs = [
         "SegmentedBlock.h",
         "shape_analysis.h",
+        "PartitionInfo.h",
         "partitioning.h",
     ],
     srcs = [
@@ -20,8 +21,9 @@ cc_library(
         "partitioning.cpp",
     ],
     deps = [
-        "//core/conversion",
         "//core/util:prelude",
+        "//core/ir",
+        "//core/conversion",
         "//core/lowering"
     ] + select({
         ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
@@ -35,6 +37,11 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar")
 pkg_tar(
     name = "include",
     package_dir = "core/partitioning/",
-    srcs = ["partitioning.h"],
+    srcs = [
+        "SegmentedBlock.h",
+        "shape_analysis.h",
+        "PartitionInfo.h",
+        "partitioning.h",
+    ],
 )
 

core/partitioning/PartitionInfo.h

@@ -0,0 +1,19 @@
+#pragma once
+
+#include <cstdint>
+#include <vector>
+#include <string>
+
+namespace trtorch {
+namespace core {
+namespace partitioning {
+
+struct PartitionInfo {
+    bool enabled = false;
+    uint64_t min_block_size = 1;
+    std::vector<std::string> forced_fallback_operators;
+};
+
+} // namespace partitioning
+} // namespace core
+} // namespace trtorch