chore: reorg in prep for state centralizing

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

Author: narendasan

core/compiler.cpp

@@ -227,11 +227,11 @@ GraphAndMapping ConstructFallbackGraph(
     ir::StaticParams static_params,
     std::unordered_map<torch::jit::Node*, int>& fallback_nodes) {
   auto convert_cfg = cfg.convert_info;
-  auto partition_info = cfg.partition_info;
+  auto partitioning_info = cfg.partitioning_info;
 
   auto new_g = std::make_shared<torch::jit::Graph>();
 
-  auto segmented_blocks = partitioning::Partition(block, example_tensor_map, partition_info, fallback_nodes);
+  auto segmented_blocks = partitioning::Partition(block, example_tensor_map, partitioning_info, fallback_nodes);
 
   // the mapping from lowering graph => fallback global graph
   std::unordered_map<torch::jit::Value*, torch::jit::Value*> old_to_new_g;
@@ -339,7 +339,7 @@ void MapInputsAndDetermineDTypes(
             "Cannot infer input type from calcuations in graph for input "
             << in->debugName() << ". Assuming it is Float32. If not, specify input type explicity");
         spec[i].dtype = nvinfer1::DataType::kFLOAT;
-      } else if (spec[i].dtype_is_user_defined && cfg.partition_info.enabled) {
+      } else if (spec[i].dtype_is_user_defined && cfg.partitioning_info.enabled) {
         if (!est_type_opt[i]) {
           LOG_INFO("Cannot infer input tensor dtype in graph, compiler is going to use the user setting");
           std::stringstream ss;
@@ -424,15 +424,15 @@ torch::jit::Module CompileGraph(const torch::jit::Module& mod, CompileSpec cfg)
       MapInputsAndDetermineDTypes(cfg, g, static_params, first_use_types);
       auto isBlockConvertible = conversion::VerifyConverterSupportForBlock(g->block(), true);
       auto outputIsCollection = conversion::OutputIsCollection(g->block());
-      if (cfg.partition_info.enabled &&
+      if (cfg.partitioning_info.enabled &&
           (cfg.lower_info.forced_fallback_modules.size() == 0 &&
-           cfg.partition_info.forced_fallback_operators.size() == 0 && isBlockConvertible)) {
+           cfg.partitioning_info.forced_fallback_operators.size() == 0 && isBlockConvertible)) {
         LOG_INFO("Skipping partitioning since model is fully supported");
       }
 
-      if (cfg.partition_info.enabled &&
+      if (cfg.partitioning_info.enabled &&
           (!(cfg.lower_info.forced_fallback_modules.size() == 0 &&
-             cfg.partition_info.forced_fallback_operators.size() == 0 && isBlockConvertible) ||
+             cfg.partitioning_info.forced_fallback_operators.size() == 0 && isBlockConvertible) ||
            outputIsCollection)) {
         std::unordered_map<torch::jit::Node*, int> fallback_nodes;
         auto collection_input_ivalues_map =

core/compiler.h

@@ -19,7 +19,7 @@ struct CompileSpec {
   ir::GraphInputs graph_inputs;
   conversion::ConversionInfo convert_info;
   lowering::LowerInfo lower_info;
-  partitioning::PartitionInfo partition_info;
+  partitioning::PartitioningInfo partitioning_info;
 };
 
 bool CheckMethodOperatorSupport(const torch::jit::script::Module& mod, std::string method_name);

core/lowering/passes/CMakeLists.txt

@@ -17,6 +17,7 @@ target_sources(${lib_name}
             "${CMAKE_CURRENT_SOURCE_DIR}/silu_to_sigmoid_multiplication.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/unpack_addmm.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/unpack_batch_norm.cpp"
+            "${CMAKE_CURRENT_SOURCE_DIR}/unpack_hardsigmoid.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/unpack_hardswish.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/unpack_log_softmax.cpp"
             "${CMAKE_CURRENT_SOURCE_DIR}/unpack_std.cpp"

core/partitioning/BUILD

@@ -13,22 +13,19 @@ config_setting(
 cc_library(
     name = "partitioning",
     srcs = [
-        "PartitionInfo.cpp",
-        "SegmentedBlock.cpp",
         "partitioning.cpp",
         "shape_analysis.cpp",
     ],
     hdrs = [
-        "PartitionInfo.h",
-        "SegmentedBlock.h",
         "partitioning.h",
-        "shape_analysis.h",
     ],
     deps = [
         "//core/util:prelude",
         "//core/ir",
         "//core/conversion",
         "//core/lowering",
+        "//core/partitioning/partitioninginfo",
+        "//core/partitioning/segmentedblock",
     ] + select({
         ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
         "//conditions:default": ["@libtorch//:libtorch"],
@@ -39,10 +36,7 @@ cc_library(
 pkg_tar(
     name = "include",
     srcs = [
-        "PartitionInfo.h",
-        "SegmentedBlock.h",
         "partitioning.h",
-        "shape_analysis.h",
     ],
     package_dir = "core/partitioning/",
 )

core/partitioning/CMakeLists.txt

@@ -1,33 +1,38 @@
 set(lib_name "core_partitioning")
 add_library(${lib_name} OBJECT)
 
-target_sources(${lib_name}
-    PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/SegmentedBlock.cpp"
-            "${CMAKE_CURRENT_SOURCE_DIR}/shape_analysis.cpp"
-            "${CMAKE_CURRENT_SOURCE_DIR}/partitioning.cpp"
-            "${CMAKE_CURRENT_SOURCE_DIR}/PartitionInfo.cpp"
-            $<TARGET_OBJECTS:core_conversion>
-    PUBLIC  $<TARGET_OBJECTS:core_ir>
-            $<TARGET_OBJECTS:core_util>
+set(CXX_SRCS
+    "${CMAKE_CURRENT_SOURCE_DIR}/partitioning.cpp"
+    "${CMAKE_CURRENT_SOURCE_DIR}/shape_analysis.cpp"
 )
 
 set(HEADER_FILES
-    "${CMAKE_CURRENT_SOURCE_DIR}/SegmentedBlock.h"
-    "${CMAKE_CURRENT_SOURCE_DIR}/shape_analysis.h"
-    "${CMAKE_CURRENT_SOURCE_DIR}/PartitionInfo.h"
     "${CMAKE_CURRENT_SOURCE_DIR}/partitioning.h"
 )
 
-target_include_directories(${lib_name} PUBLIC "$<BUILD_INTERFACE:${CMAKE_SOURCE_DIR}>")
+target_sources(${lib_name}
+    PRIVATE
+        ${CXX_SRCS}
+    PUBLIC
+        $<TARGET_OBJECTS:core_conversion>
+        $<TARGET_OBJECTS:core_ir>
+        $<TARGET_OBJECTS:core_util>
+)
+
 target_link_libraries(${lib_name}
     PUBLIC
-        torch
         TensorRT::nvinfer
+        torch
         core_ir
         core_util
-    PRIVATE
         core_conversion
 )
 
-# Install headers
-install(FILES ${HEADER_FILES} DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/torch_tensorrt/core/partitioning/")
+target_include_directories(${lib_name}
+    PUBLIC "$<BUILD_INTERFACE:${CMAKE_SOURCE_DIR}>"
+)
+
+add_subdirectory(partitioninginfo)
+add_subdirectory(segmentedblock)
+
+install(FILES ${HEADER_FILES} DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/torch_tensorrt/core/partitioning")

core/partitioning/partitioning.cpp

@@ -1,12 +1,12 @@
-#include "partitioning.h"
-
 #include <queue>
-#include "core/conversion/conversion.h"
-#include "core/conversion/evaluators/evaluators.h"
-#include "core/partitioning/shape_analysis.h"
+
 #include "torch/csrc/jit/passes/constant_pooling.h"
 #include "torch/csrc/jit/passes/dead_code_elimination.h"
 
+#include "core/conversion/conversion.h"
+#include "core/conversion/evaluators/evaluators.h"
+#include "core/partitioning/partitioning.h"
+
 namespace torch_tensorrt {
 namespace core {
 namespace partitioning {
@@ -357,11 +357,11 @@ void find_min_block_size_fallback_nodes(
 
 PartitionedGraph segment_graph(
     torch::jit::Block* block,
-    const PartitionInfo& partition_info,
+    const PartitioningInfo& partitioning_info,
     std::unordered_map<torch::jit::Node*, int>& global_fallback_nodes) {
-  auto min_block_size = partition_info.min_block_size;
+  auto min_block_size = partitioning_info.min_block_size;
   std::unordered_set<std::string> forced_fallback_ops(
-      partition_info.forced_fallback_operators.begin(), partition_info.forced_fallback_operators.end());
+      partitioning_info.forced_fallback_operators.begin(), partitioning_info.forced_fallback_operators.end());
 
   // get the initial fallback nodes (nodes that are unsupported or forced fallback)
   get_fallback_nodes(block, forced_fallback_ops, global_fallback_nodes);
@@ -450,16 +450,16 @@ PartitionedGraph segment_graph(
 PartitionedGraph Partition(
     torch::jit::Block* block,
     std::unordered_map<const torch::jit::Value*, torch::jit::IValue>& example_tensor_map,
-    const PartitionInfo& partition_info,
+    const PartitioningInfo& partitioning_info,
     std::unordered_map<torch::jit::Node*, int>& global_fallback_nodes) {
-  LOG_DEBUG(partition_info);
+  LOG_DEBUG(partitioning_info);
   // if there is nonTensor input/output for the entire graph, fallback the node that consumes/produces this nonTensor
   // output
   fallback_graph_nontensor_in_out(block, global_fallback_nodes);
 
   // segment lowering global graph into blocks
   LOG_DEBUG("Parititioning source module into PyTorch and TensorRT sub blocks");
-  PartitionedGraph segmented_blocks = segment_graph(block, partition_info, global_fallback_nodes);
+  PartitionedGraph segmented_blocks = segment_graph(block, partitioning_info, global_fallback_nodes);
 
   // It's possible that some TensorRT blocks have nonTensor inputs/output because they are interleaved by Torch blocks
 
@@ -471,7 +471,7 @@ PartitionedGraph Partition(
   registerSegmentsOutputs(segmented_blocks, block);
 
   // run shape analysis on each segmented block
-  runShapeAnalysis(segmented_blocks, example_tensor_map, partition_info);
+  runShapeAnalysis(segmented_blocks, example_tensor_map, partitioning_info);
 
   for (uint64_t i = 0; i < segmented_blocks.size(); i++) {
     segmented_blocks[i].update_id(i);

core/partitioning/partitioning.h

@@ -3,12 +3,12 @@
 #include <iostream>
 #include <vector>
 
+#include "torch/csrc/jit/ir/ir.h"
+
 #include "core/ir/ir.h"
-#include "core/partitioning/PartitionInfo.h"
-#include "core/partitioning/SegmentedBlock.h"
-#include "core/partitioning/shape_analysis.h"
+#include "core/partitioning/partitioninginfo/PartitioningInfo.h"
+#include "core/partitioning/segmentedblock/SegmentedBlock.h"
 #include "core/util/prelude.h"
-#include "torch/csrc/jit/ir/ir.h"
 
 namespace torch_tensorrt {
 namespace core {
@@ -30,15 +30,24 @@ enum FallbackNodeType {
   kNON_TENSOR,
 };
 
+std::unordered_map<const torch::jit::Value*, torch::jit::IValue> generateRandomInputs(
+    std::unordered_map<const torch::jit::Value*, std::vector<ir::Input>>& input_ranges,
+    std::unordered_map<const torch::jit::Value*, std::vector<c10::optional<at::ScalarType>>>& input_types);
+
+void runShapeAnalysis(
+    std::vector<SegmentedBlock>& segmented_blocks,
+    std::unordered_map<const torch::jit::Value*, torch::jit::IValue>& ivalues_maps,
+    const PartitioningInfo& partitioning_info);
+
 PartitionedGraph segment_graph(
     torch::jit::Block* block,
-    const PartitionInfo& partition_info,
+    const PartitioningInfo& partitioning_info,
     std::unordered_map<torch::jit::Node*, int>& fallback_nodes);
 
 PartitionedGraph Partition(
     torch::jit::Block* block,
     std::unordered_map<const torch::jit::Value*, torch::jit::IValue>& example_tensor_map,
-    const PartitionInfo& partition_info,
+    const PartitioningInfo& partitioning_info,
     std::unordered_map<torch::jit::Node*, int>& fallback_nodes);
 
 std::ostream& operator<<(std::ostream& os, const PartitionedGraph& g);

core/partitioning/partitioninginfo/BUILD

@@ -0,0 +1,39 @@
+load("@rules_cc//cc:defs.bzl", "cc_library")
+load("@rules_pkg//:pkg.bzl", "pkg_tar")
+
+package(default_visibility = ["//visibility:public"])
+
+config_setting(
+    name = "use_pre_cxx11_abi",
+    values = {
+        "define": "abi=pre_cxx11_abi",
+    },
+)
+
+cc_library(
+    name = "partitioninginfo",
+    srcs = [
+        "PartitioningInfo.cpp",
+    ],
+    hdrs = [
+        "PartitioningInfo.h",
+    ],
+    deps = [
+        "//core/util:prelude",
+        "//core/ir",
+        "//core/conversion",
+        "//core/lowering",
+    ] + select({
+        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@libtorch//:libtorch"],
+    }),
+    alwayslink = True,
+)
+
+pkg_tar(
+    name = "include",
+    srcs = [
+        "PartitioningInfo.h",
+    ],
+    package_dir = "core/partitioning/partitioninginfo",
+)

core/partitioning/partitioninginfo/CMakeLists.txt

@@ -0,0 +1,12 @@
+set(sub_lib_name "partitioninginfo")
+
+target_sources(${lib_name}
+    PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/PartitioningInfo.cpp"
+)
+
+set(HEADER_FILES
+    "${CMAKE_CURRENT_SOURCE_DIR}/PartitioningInfo.h"
+)
+
+# Install headers
+install(FILES ${HEADER_FILES} DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/torch_tensorrt/core/partitioning/${sub_lib_name}")

core/partitioning/PartitionInfo.cpp renamed to core/partitioning/partitioninginfo/PartitioningInfo.cpp

@@ -2,13 +2,13 @@
 #include <sstream>
 #include <utility>
 
-#include "core/partitioning/PartitionInfo.h"
+#include "core/partitioning/partitioninginfo/PartitioningInfo.h"
 
 namespace torch_tensorrt {
 namespace core {
 namespace partitioning {
 // clang-format off
-std::ostream& operator<<(std::ostream& os, const PartitionInfo& s) {
+std::ostream& operator<<(std::ostream& os, const PartitioningInfo& s) {
   os << "Settings requested for Torch Fallback:" \
      << "\n    \"enabled\": ";
   if (s.enabled) {