fix: Resolve issue in isInputDynamic with mixed static/dynamic shapes (#1883)

Author: mfeliz-cruise

core/partitioning/partitioning.cpp

@@ -527,16 +527,15 @@ void segmentGraph(PartitioningCtx* ctx, torch::jit::Block* block) {
 
 bool isInputDynamic(PartitioningCtx* ctx) {
   // Check if inputs have dynamic shapes
-  bool input_is_dynamic = true;
   auto inputs_map = ctx->settings.collection_input_spec_map;
   for (auto inputs : inputs_map) {
     for (auto input : inputs.second) {
-      if (!input.input_is_dynamic) {
-        input_is_dynamic = false;
+      if (input.input_is_dynamic) {
+        return true;
       }
     }
   }
-  return input_is_dynamic;
+  return false;
 }
 
 void populateInputIValues(PartitioningCtx* ctx) {

tests/core/partitioning/test_shape_analysis.cpp

@@ -137,3 +137,37 @@ TEST(Partitioning, InferBranchModelSegmentedBlockShapeCorrectly) {
        {{3, 32, 16, 16}},
        {{3, 32, 16, 16}, {16, 32, 3, 3}, {16}, {3, 16, 16, 16}}}));
 }
+
+TEST(Partitioning, PopulateInputIValuesDynamic) {
+  const auto graph = R"IR(
+          graph(%0 : Tensor, %1 : Tensor):
+            %2 : float = prim::Constant[value=1]()
+            %30 : Tensor = aten::add(%0, %1, %2)
+            return (%30))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get(), true);
+
+  torch_tensorrt::core::partitioning::PartitioningInfo partitioning_info;
+  partitioning_info.enabled = true;
+  partitioning_info.truncate_long_and_double = true;
+  std::vector<torch_tensorrt::core::ir::Input> inputs;
+
+  inputs.push_back(torch_tensorrt::core::ir::Input({1}, {2}, {3}));
+  inputs.push_back(torch_tensorrt::core::ir::Input({1}));
+
+  std::unordered_map<const torch::jit::Value*, std::vector<torch_tensorrt::core::ir::Input>> inputs_map;
+  std::unordered_map<const torch::jit::Value*, std::vector<c10::optional<at::ScalarType>>> input_types;
+  inputs_map.insert({g->inputs()[0], {inputs[0]}});
+  inputs_map.insert({g->inputs()[1], {inputs[1]}});
+  input_types.insert({g->inputs()[0], {{at::kFloat}}});
+  input_types.insert({g->inputs()[1], {{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);
+  ASSERT_EQ(ctx.min_input_ivalues_map.size(), 2UL);
+  ASSERT_EQ(ctx.max_input_ivalues_map.size(), 2UL);
+}