chore: Move shape mode to enum, fix CI tests by storing input_ivalues map in a PartioningCtx object

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

core/compiler.cpp

@@ -127,6 +127,8 @@ bool CheckMethodOperatorSupport(const torch::jit::script::Module& mod, std::stri
   return conversion::VerifyConverterSupportForBlock(g->block());
 }
 
+
+
 partitioning::GraphAndMapping BuildHybridGraph(
     torch::jit::script::Module& new_mod,
     torch::jit::Block* block,
@@ -138,6 +140,11 @@ partitioning::GraphAndMapping BuildHybridGraph(
 
   auto partitioning_ctx = partitioning::PartitioningCtx(block, partitioning_info);
   partitioning_ctx.input_types_map = first_use_types;
+
+  // Generate a dictionary of input torch::jit::Value's to their min, opt, max tensors and store in ctx
+  // TODO: Combine this within partition call
+  partitioning::populateInputIValues(&partitioning_ctx);
+
   partitioning::partition(&partitioning_ctx);
 
   for (auto& partitioned_block : partitioning_ctx.partitioned_blocks) {

core/ir/ir.h

@@ -11,6 +11,12 @@ namespace torch_tensorrt {
 namespace core {
 namespace ir {
 
+enum class ShapeMode {
+    kMIN,
+    kOPT,
+    kMAX,
+};
+
 struct Input : torch::CustomClassHolder {
   Input(){};
   Input(

core/partitioning/partitioning.cpp

@@ -450,6 +450,20 @@ bool isInputDynamic(PartitioningCtx* ctx) {
   return input_is_dynamic;
 }
 
+void populateInputIValues(PartitioningCtx* ctx){
+  if (isInputDynamic(ctx)) {
+    ctx->min_input_ivalues_map =
+        partitioning::generateRandomInputs(ctx->settings.collection_input_spec_map, ctx->input_types_map, ir::ShapeMode::kMIN);
+    ctx->opt_input_ivalues_map =
+        partitioning::generateRandomInputs(ctx->settings.collection_input_spec_map, ctx->input_types_map, ir::ShapeMode::kOPT);
+    ctx->max_input_ivalues_map =
+        partitioning::generateRandomInputs(ctx->settings.collection_input_spec_map, ctx->input_types_map, ir::ShapeMode::kMAX);
+  } else {
+    ctx->opt_input_ivalues_map =
+        partitioning::generateRandomInputs(ctx->settings.collection_input_spec_map, ctx->input_types_map, ir::ShapeMode::kOPT);
+  }
+}
+
 void partition(PartitioningCtx* ctx) {
   LOG_DEBUG(ctx->settings);
 
@@ -471,21 +485,12 @@ void partition(PartitioningCtx* ctx) {
     // output shapes for each block accordingly
     if (isInputDynamic(ctx)) {
       LOG_DEBUG("Performing shape analysis for segmented blocks using min/opt/max shapes for inputs");
-      auto min_input_ivalues_map =
-          partitioning::generateRandomInputs(ctx->settings.collection_input_spec_map, ctx->input_types_map, "min");
-      auto opt_input_ivalues_map =
-          partitioning::generateRandomInputs(ctx->settings.collection_input_spec_map, ctx->input_types_map, "opt");
-      auto max_input_ivalues_map =
-          partitioning::generateRandomInputs(ctx->settings.collection_input_spec_map, ctx->input_types_map, "max");
-
-      runShapeAnalysis(ctx, block, min_input_ivalues_map, "min");
-      runShapeAnalysis(ctx, block, opt_input_ivalues_map, "opt");
-      runShapeAnalysis(ctx, block, max_input_ivalues_map, "max");
+      runShapeAnalysis(ctx, block, ctx->min_input_ivalues_map, ir::ShapeMode::kMIN);
+      runShapeAnalysis(ctx, block, ctx->opt_input_ivalues_map, ir::ShapeMode::kOPT);
+      runShapeAnalysis(ctx, block, ctx->max_input_ivalues_map, ir::ShapeMode::kMAX);
     } else {
       LOG_DEBUG("Performing shape analysis for segmented blocks using static shapes for inputs");
-      auto opt_input_ivalues_map =
-          partitioning::generateRandomInputs(ctx->settings.collection_input_spec_map, ctx->input_types_map, "opt");
-      runShapeAnalysis(ctx, block, opt_input_ivalues_map, "opt");
+      runShapeAnalysis(ctx, block, ctx->opt_input_ivalues_map, ir::ShapeMode::kOPT);
     }
   }
 }

core/partitioning/partitioning.h

@@ -21,13 +21,15 @@ typedef std::pair<std::shared_ptr<torch::jit::Graph>, std::unordered_map<torch::
 ExampleIValues generateRandomInputs(
     ir::CollectionInputSpecMap& input_ranges,
     ir::CollectionTypeMap& input_types,
-    const std::string& shape_mode = std::string("opt"));
+    const ir::ShapeMode& shape_mode = ir::ShapeMode::kOPT);
+
+void populateInputIValues(PartitioningCtx* ctx);
 
 void runShapeAnalysis(
     PartitioningCtx* ctx,
     torch::jit::Block* block,
     ExampleIValues& ivalues_maps,
-    const std::string& shape_mode);
+    const ir::ShapeMode& shape_mode);
 
 void segmentGraph(PartitioningCtx* ctx, torch::jit::Block* block);
 

core/partitioning/partitioningctx/PartitioningCtx.h

@@ -47,6 +47,9 @@ struct UsageInfo {
 struct PartitioningCtx {
   // TODO: Make the set a part of settings not stand alone
   PartitioningInfo settings;
+  std::unordered_map<const torch::jit::Value*, torch::jit::IValue> min_input_ivalues_map;
+  std::unordered_map<const torch::jit::Value*, torch::jit::IValue> opt_input_ivalues_map;
+  std::unordered_map<const torch::jit::Value*, torch::jit::IValue> max_input_ivalues_map;
   // records all the original blocks topologically in the module
   std::vector<torch::jit::Block*> original_blocks;
   // mapping: node=> execution status

core/partitioning/segmentedblock/SegmentedBlock.h

@@ -73,10 +73,10 @@ struct SegmentedBlock {
   bool contain_raw_value(torch::jit::Value* input) const {
     return old_to_new_.count(input);
   }
-  void register_inshapes(std::vector<std::vector<int64_t>>& in_shapes, const std::string& shape_mode) {
-    if (shape_mode.compare("min") == 0) {
+  void register_inshapes(std::vector<std::vector<int64_t>>& in_shapes, const ir::ShapeMode& shape_mode) {
+    if (shape_mode == ir::ShapeMode::kMIN) {
       min_shapes_ = in_shapes;
-    } else if (shape_mode.compare("opt") == 0) {
+    } else if (shape_mode == ir::ShapeMode::kOPT) {
       opt_shapes_ = in_shapes;
     } else {
       max_shapes_ = in_shapes;

core/partitioning/shape_analysis.cpp

@@ -12,12 +12,12 @@ namespace partitioning {
 at::Tensor generateSingleInput(
     ir::Input& input,
     c10::optional<at::ScalarType>& type_opt,
-    const std::string& shape_mode) {
+    const ir::ShapeMode& shape_mode) {
   nvinfer1::Dims input_shape = input.input_shape;
   if (input.input_is_dynamic) {
-    if (shape_mode.compare("min") == 0) {
+    if (shape_mode == ir::ShapeMode::kMIN) {
       input_shape = input.min;
-    } else if (shape_mode.compare("opt") == 0) {
+    } else if (shape_mode == ir::ShapeMode::kOPT) {
       input_shape = input.opt;
     } else {
       input_shape = input.max;
@@ -38,7 +38,7 @@ at::Tensor generateSingleInput(
 std::unordered_map<const torch::jit::Value*, torch::jit::IValue> generateRandomInputs(
     std::unordered_map<const torch::jit::Value*, std::vector<ir::Input>>& inputs,
     std::unordered_map<const torch::jit::Value*, std::vector<c10::optional<at::ScalarType>>>& types,
-    const std::string& shape_mode) {
+    const ir::ShapeMode& shape_mode) {
   // generate random inputs for running pytorch segments
   std::unordered_map<const torch::jit::Value*, torch::jit::IValue> ivalue_map;
 
@@ -72,7 +72,7 @@ void getSegmentsOutputByRunning(
     SegmentedBlock& seg_block,
     std::unordered_map<const torch::jit::Value*, torch::jit::IValue>& ivalues_maps,
     const PartitioningInfo& partitioning_info,
-    const std::string& shape_mode) {
+    const ir::ShapeMode& shape_mode) {
   // create a module to run the graph
   auto g = seg_block.g();
   auto copy_g = g->copy();
@@ -195,7 +195,7 @@ void runShapeAnalysis(
     PartitioningCtx* ctx,
     torch::jit::Block* block,
     ExampleIValues& example_tensor_map,
-    const std::string& shape_mode) {
+    const ir::ShapeMode& shape_mode) {
   // register every segment's input shape, and it's running output IValues
   for (auto& seg_block : ctx->partitioned_blocks[block]) {
     torch::jit::ConstantPooling(seg_block.g());

tests/core/partitioning/test_conditionals.cpp

@@ -43,33 +43,33 @@ TEST(Partitioning, FallbackOnConditionalsCorrectly) {
   ASSERT_TRUE(conditional_engines_count == 2);
 }
 
-// TEST(Partitioning, FallbackInplaceOPInConditionalsCorrectly) {
-//   torch::jit::script::Module mod;
-//   try {
-//     mod = torch::jit::load("tests/modules/inplace_op_if_scripted.jit.pt");
-//   } catch (const c10::Error& e) {
-//     std::cerr << "error loading the model\n";
-//     return;
-//   }
-//
-//   const std::vector<std::vector<int64_t>> input_shapes = {{4, 4}, {4, 4}};
-//   std::vector<torch::jit::IValue> jit_inputs_ivalues;
-//   std::vector<torch::jit::IValue> trt_inputs_ivalues;
-//   for (auto in_shape : input_shapes) {
-//     auto in = at::randint(5, in_shape, {at::kCUDA});
-//     jit_inputs_ivalues.push_back(in.clone());
-//     trt_inputs_ivalues.push_back(in.clone());
-//   }
-//
-//   std::vector<torch_tensorrt::core::ir::Input> inputs{
-//       torch_tensorrt::core::ir::Input({4, 4}), torch_tensorrt::core::ir::Input({4, 4})};
-//   auto g = mod.get_method("forward").graph();
-//   torch_tensorrt::core::CompileSpec cfg(inputs);
-//   cfg.partitioning_info.enabled = true;
-//   cfg.partitioning_info.forced_fallback_operators.push_back("prim::ListConstruct");
-//
-//   auto jit_results = mod.forward(jit_inputs_ivalues).toTensor();
-//   auto trt_mod = torch_tensorrt::core::CompileGraph(mod, cfg);
-//   auto trt_results = trt_mod.forward(trt_inputs_ivalues).toTensor();
-//   ASSERT_TRUE(torch_tensorrt::tests::util::cosineSimEqual(jit_results, trt_results));
-// }
+TEST(Partitioning, FallbackInplaceOPInConditionalsCorrectly) {
+  torch::jit::script::Module mod;
+  try {
+    mod = torch::jit::load("tests/modules/inplace_op_if_scripted.jit.pt");
+  } catch (const c10::Error& e) {
+    std::cerr << "error loading the model\n";
+    return;
+  }
+
+  const std::vector<std::vector<int64_t>> input_shapes = {{4, 4}, {4, 4}};
+  std::vector<torch::jit::IValue> jit_inputs_ivalues;
+  std::vector<torch::jit::IValue> trt_inputs_ivalues;
+  for (auto in_shape : input_shapes) {
+    auto in = at::randint(5, in_shape, {at::kCUDA});
+    jit_inputs_ivalues.push_back(in.clone());
+    trt_inputs_ivalues.push_back(in.clone());
+  }
+
+  std::vector<torch_tensorrt::core::ir::Input> inputs{
+      torch_tensorrt::core::ir::Input({4, 4}), torch_tensorrt::core::ir::Input({4, 4})};
+  auto g = mod.get_method("forward").graph();
+  torch_tensorrt::core::CompileSpec cfg(inputs);
+  cfg.partitioning_info.enabled = true;
+  cfg.partitioning_info.forced_fallback_operators.push_back("prim::ListConstruct");
+
+  auto jit_results = mod.forward(jit_inputs_ivalues).toTensor();
+  auto trt_mod = torch_tensorrt::core::CompileGraph(mod, cfg);
+  auto trt_results = trt_mod.forward(trt_inputs_ivalues).toTensor();
+  ASSERT_TRUE(torch_tensorrt::tests::util::cosineSimEqual(jit_results, trt_results));
+}

tests/core/partitioning/test_resolve_nontensor_inputs.cpp

@@ -123,8 +123,11 @@ TEST(Partitioning, ResolveNonTensorInputsCorrectly) {
     input_types.insert({g->inputs()[i], {{at::kFloat}}});
   }
 
+  partitioning_info.collection_input_spec_map = inputs_map;
   torch_tensorrt::core::partitioning::PartitioningCtx ctx(g->block(), partitioning_info);
   ctx.input_types_map = input_types;
+
+  torch_tensorrt::core::partitioning::populateInputIValues(&ctx);
   torch_tensorrt::core::partitioning::partition(&ctx);
   std::vector<torch_tensorrt::core::partitioning::SegmentedBlock> segmented_blocks =
       ctx.partitioned_blocks.begin()->second;
@@ -184,8 +187,10 @@ TEST(Partitioning, ResolveTensorListInputsInTrtCorrectly) {
     input_types.insert({g->inputs()[i], {{at::kFloat}}});
   }
 
+  partitioning_info.collection_input_spec_map = inputs_map;
   torch_tensorrt::core::partitioning::PartitioningCtx ctx(g->block(), partitioning_info);
   ctx.input_types_map = input_types;
+  torch_tensorrt::core::partitioning::populateInputIValues(&ctx);
   torch_tensorrt::core::partitioning::partition(&ctx);
   std::vector<torch_tensorrt::core::partitioning::SegmentedBlock> segmented_blocks =
       ctx.partitioned_blocks.begin()->second;
@@ -263,7 +268,7 @@ TEST(Partitioning, ConvertForTensorListInputsInFallbackCorrectly) {
   int count = count_trt_engines(fallback_g);
   ASSERT_TRUE(count == 1);
 }
-
+//
 TEST(Partitioning, ResolveOnlyNeccessaryNonTensorInputs) {
   /* parseIR does not support "= aten::_set_item" so we will build this graph manually
     const auto graph = R"IR(
@@ -377,9 +382,10 @@ TEST(Partitioning, ResolveOnlyNeccessaryNonTensorInputs) {
     inputs_map.insert({g->inputs()[i], {inputs[i]}});
     input_types.insert({g->inputs()[i], {{at::kFloat}}});
   }
-  // auto input_ivalues_map = torch_tensorrt::core::partitioning::generateRandomInputs(inputs_map, input_types);
+  partitioning_info.collection_input_spec_map = inputs_map;
   torch_tensorrt::core::partitioning::PartitioningCtx ctx(g->block(), partitioning_info);
   ctx.input_types_map = input_types;
+  torch_tensorrt::core::partitioning::populateInputIValues(&ctx);
   torch_tensorrt::core::partitioning::partition(&ctx);
   auto segmented_blocks = ctx.partitioned_blocks.begin()->second;
 

tests/core/partitioning/test_shape_analysis.cpp

@@ -66,10 +66,12 @@ TEST(Partitioning, InferSequentialModelSegmentedBlockShapeCorrectly) {
     inputs_map.insert({g->inputs()[i], {inputs[i]}});
     input_types.insert({g->inputs()[i], {{at::kFloat}}});
   }
-
+  // Store a map of torch::jit::Value to ir::Input
+  partitioning_info.collection_input_spec_map = inputs_map;
   torch_tensorrt::core::partitioning::PartitioningCtx ctx(g->block(), partitioning_info);
   ctx.input_types_map = input_types;
-  ctx.settings.collection_input_spec_map = inputs_map;
+
+  torch_tensorrt::core::partitioning::populateInputIValues(&ctx);
   torch_tensorrt::core::partitioning::partition(&ctx);
   auto segmented_blocks = ctx.partitioned_blocks.begin()->second;
 
@@ -120,9 +122,12 @@ TEST(Partitioning, InferBranchModelSegmentedBlockShapeCorrectly) {
     input_types.insert({g->inputs()[i], {{at::kFloat}}});
   }
 
+  // Store a map of torch::jit::Value to ir::Input
+  partitioning_info.collection_input_spec_map = inputs_map;
   torch_tensorrt::core::partitioning::PartitioningCtx ctx(g->block(), partitioning_info);
   ctx.input_types_map = input_types;
-  ctx.settings.collection_input_spec_map = inputs_map;
+
+  torch_tensorrt::core::partitioning::populateInputIValues(&ctx);
   torch_tensorrt::core::partitioning::partition(&ctx);
   auto segmented_blocks = ctx.partitioned_blocks.begin()->second;