tests: add test for inplace op in if block

Signed-off-by: Bo Wang <[email protected]>
Christina-Young-NVIDIA · Jun 16, 2022 · da9b13c · da9b13c
1 parent 8d9a8d4
commit da9b13c
Show file tree

Hide file tree

Showing 8 changed files with 72 additions and 13 deletions.
diff --git a/core/partitioning/partitioning.h b/core/partitioning/partitioning.h
@@ -16,7 +16,7 @@ namespace partitioning {
 
 typedef std::vector<SegmentedBlock> PartitionedGraph;
 
-PartitionedGraph segment_graph(torch::jit::Block* block, const PartitionInfo& partition_info);
+PartitionedGraph segment_graph(torch::jit::Block* block, const PartitionInfo& partition_info, std::unordered_map<torch::jit::Node*, int> &fallback_nodes);
 
 PartitionedGraph Partition(
     torch::jit::Block* block,

diff --git a/tests/core/partitioning/BUILD b/tests/core/partitioning/BUILD
@@ -13,7 +13,8 @@ filegroup(
             "//tests/modules:mobilenet_v2_traced.jit.pt",
             "//tests/modules:conditional_scripted.jit.pt",
             "//tests/modules:loop_fallback_eval_scripted.jit.pt",
-            "//tests/modules:loop_fallback_no_eval_scripted.jit.pt"]
+            "//tests/modules:loop_fallback_no_eval_scripted.jit.pt",
+            "//tests/modules:inplace_op_if_scripted.jit.pt"]
 )
 
 partitioning_test(

diff --git a/tests/core/partitioning/test_conditionals.cpp b/tests/core/partitioning/test_conditionals.cpp
@@ -42,3 +42,34 @@ 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.partition_info.enabled = true;
+  cfg.partition_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::almostEqual(jit_results, trt_results, 2e-6));
+}
diff --git a/tests/core/partitioning/test_resolve_nontensor_inputs.cpp b/tests/core/partitioning/test_resolve_nontensor_inputs.cpp
@@ -123,8 +123,9 @@ TEST(Partitioning, ResolveNonTensorInputsCorrectly) {
     input_types.insert({g->inputs()[i], {at::kFloat}});
   }
   auto input_ivalues_map = torch_tensorrt::core::partitioning::generateRandomInputs(inputs_map, input_types);
+  std::unordered_map<torch::jit::Node*, int> fallback_nodes;
   std::vector<torch_tensorrt::core::partitioning::SegmentedBlock> segmented_blocks =
-      torch_tensorrt::core::partitioning::Partition(g->block(), input_ivalues_map, partition_info);
+      torch_tensorrt::core::partitioning::Partition(g->block(), input_ivalues_map, partition_info, fallback_nodes);
 
   int torch_block_cnt = 0, trt_block_cnt = 0;
   for (const auto& segmented_block : segmented_blocks) {
@@ -181,8 +182,9 @@ TEST(Partitioning, ResolveTensorListInputsInTrtCorrectly) {
     input_types.insert({g->inputs()[i], {at::kFloat}});
   }
   auto input_ivalues_map = torch_tensorrt::core::partitioning::generateRandomInputs(inputs_map, input_types);
+  std::unordered_map<torch::jit::Node*, int> fallback_nodes;
   std::vector<torch_tensorrt::core::partitioning::SegmentedBlock> segmented_blocks =
-      torch_tensorrt::core::partitioning::Partition(g->block(), input_ivalues_map, partition_info);
+      torch_tensorrt::core::partitioning::Partition(g->block(), input_ivalues_map, partition_info, fallback_nodes);
 
   int torch_block_cnt = 0, trt_block_cnt = 0;
   for (const auto& segmented_block : segmented_blocks) {
@@ -372,7 +374,8 @@ TEST(Partitioning, ResolveOnlyNeccessaryNonTensorInputs) {
     input_types.insert({g->inputs()[i], {at::kFloat}});
   }
   auto input_ivalues_map = torch_tensorrt::core::partitioning::generateRandomInputs(inputs_map, input_types);
-  auto segmented_blocks = torch_tensorrt::core::partitioning::Partition(g->block(), input_ivalues_map, partition_info);
+  std::unordered_map<torch::jit::Node*, int> fallback_nodes;
+  auto segmented_blocks = torch_tensorrt::core::partitioning::Partition(g->block(), input_ivalues_map, partition_info, fallback_nodes);
 
   int torch_block_cnt = 0, trt_block_cnt = 0;
   for (const auto& segmented_block : segmented_blocks) {

diff --git a/tests/core/partitioning/test_segmentation.cpp b/tests/core/partitioning/test_segmentation.cpp
@@ -74,8 +74,9 @@ TEST(Partitioning, SegmentSequentialModelCorrectly) {
 
   torch_tensorrt::core::partitioning::PartitionInfo partition_info;
   partition_info.enabled = true;
+  std::unordered_map<torch::jit::Node*, int> fallback_nodes;
   std::vector<torch_tensorrt::core::partitioning::SegmentedBlock> segmented_blocks =
-      torch_tensorrt::core::partitioning::segment_graph(g->block(), partition_info);
+      torch_tensorrt::core::partitioning::segment_graph(g->block(), partition_info, fallback_nodes);
   ASSERT_TRUE(
       checkSegmentedBlockNumber(segmented_blocks, torch_tensorrt::core::partitioning::SegmentedBlock::kTensorRT, 2));
   ASSERT_TRUE(
@@ -109,8 +110,9 @@ TEST(Partitioning, SegmentSequentialModelWithMinBlockSizeCorrectly) {
   torch_tensorrt::core::partitioning::PartitionInfo partition_info;
   partition_info.enabled = true;
   partition_info.min_block_size = 3;
+  std::unordered_map<torch::jit::Node*, int> fallback_nodes;
   std::vector<torch_tensorrt::core::partitioning::SegmentedBlock> segmented_blocks =
-      torch_tensorrt::core::partitioning::segment_graph(g->block(), partition_info);
+      torch_tensorrt::core::partitioning::segment_graph(g->block(), partition_info, fallback_nodes);
   ASSERT_TRUE(
       checkSegmentedBlockNumber(segmented_blocks, torch_tensorrt::core::partitioning::SegmentedBlock::kTensorRT, 1));
   ASSERT_TRUE(
@@ -144,8 +146,9 @@ TEST(Partitioning, SegmentSequentialModelWithForcedOPCorrectly) {
   torch_tensorrt::core::partitioning::PartitionInfo partition_info;
   partition_info.enabled = true;
   partition_info.forced_fallback_operators.push_back("aten::relu");
+  std::unordered_map<torch::jit::Node*, int> fallback_nodes;
   std::vector<torch_tensorrt::core::partitioning::SegmentedBlock> segmented_blocks =
-      torch_tensorrt::core::partitioning::segment_graph(g->block(), partition_info);
+      torch_tensorrt::core::partitioning::segment_graph(g->block(), partition_info,fallback_nodes);
   ASSERT_TRUE(
       checkSegmentedBlockNumber(segmented_blocks, torch_tensorrt::core::partitioning::SegmentedBlock::kTensorRT, 3));
   ASSERT_TRUE(
@@ -179,8 +182,9 @@ TEST(Partitioning, SegmentBranchModelCorrectly) {
 
   torch_tensorrt::core::partitioning::PartitionInfo partition_info;
   partition_info.enabled = true;
+  std::unordered_map<torch::jit::Node*, int> fallback_nodes;
   std::vector<torch_tensorrt::core::partitioning::SegmentedBlock> segmented_blocks =
-      torch_tensorrt::core::partitioning::segment_graph(g->block(), partition_info);
+      torch_tensorrt::core::partitioning::segment_graph(g->block(), partition_info, fallback_nodes);
   ASSERT_TRUE(
       checkSegmentedBlockNumber(segmented_blocks, torch_tensorrt::core::partitioning::SegmentedBlock::kTensorRT, 2));
   ASSERT_TRUE(
@@ -215,8 +219,9 @@ TEST(Partitioning, SegmentBranchModelWithMinBlockSizeCorrectly) {
   torch_tensorrt::core::partitioning::PartitionInfo partition_info;
   partition_info.enabled = true;
   partition_info.min_block_size = 3;
+  std::unordered_map<torch::jit::Node*, int> fallback_nodes;
   std::vector<torch_tensorrt::core::partitioning::SegmentedBlock> segmented_blocks =
-      torch_tensorrt::core::partitioning::segment_graph(g->block(), partition_info);
+      torch_tensorrt::core::partitioning::segment_graph(g->block(), partition_info, fallback_nodes);
   ASSERT_TRUE(
       checkSegmentedBlockNumber(segmented_blocks, torch_tensorrt::core::partitioning::SegmentedBlock::kTensorRT, 1));
   ASSERT_TRUE(
@@ -255,8 +260,9 @@ TEST(Partitioning, SegmentBranchModelWithForcedFallbackOPCorrectly) {
   torch_tensorrt::core::partitioning::PartitionInfo partition_info;
   partition_info.enabled = true;
   partition_info.forced_fallback_operators.push_back("aten::relu");
+  std::unordered_map<torch::jit::Node*, int> fallback_nodes;
   torch_tensorrt::core::partitioning::PartitionedGraph segmented_blocks =
-      torch_tensorrt::core::partitioning::segment_graph(g->block(), partition_info);
+      torch_tensorrt::core::partitioning::segment_graph(g->block(), partition_info, fallback_nodes);
   ASSERT_TRUE(
       checkSegmentedBlockNumber(segmented_blocks, torch_tensorrt::core::partitioning::SegmentedBlock::kTensorRT, 3));
   ASSERT_TRUE(

diff --git a/tests/core/partitioning/test_shape_analysis.cpp b/tests/core/partitioning/test_shape_analysis.cpp
@@ -66,8 +66,9 @@ TEST(Partitioning, InferSequentialModelSegmentedBlockShapeCorrectly) {
     input_types.insert({g->inputs()[i], {at::kFloat}});
   }
   auto input_ivalues_map = torch_tensorrt::core::partitioning::generateRandomInputs(inputs_map, input_types);
+  std::unordered_map<torch::jit::Node*, int> fallback_nodes;
   std::vector<torch_tensorrt::core::partitioning::SegmentedBlock> segmented_blocks =
-      torch_tensorrt::core::partitioning::Partition(g->block(), input_ivalues_map, partition_info);
+      torch_tensorrt::core::partitioning::Partition(g->block(), input_ivalues_map, partition_info, fallback_nodes);
 
   ASSERT_TRUE(checkSegmentedBlockInputShape(
       segmented_blocks,
@@ -116,8 +117,9 @@ TEST(Partitioning, InferBranchModelSegmentedBlockShapeCorrectly) {
     input_types.insert({g->inputs()[i], {at::kFloat}});
   }
   auto input_ivalues_map = torch_tensorrt::core::partitioning::generateRandomInputs(inputs_map, input_types);
+  std::unordered_map<torch::jit::Node*, int> fallback_nodes;
   std::vector<torch_tensorrt::core::partitioning::SegmentedBlock> segmented_blocks =
-      torch_tensorrt::core::partitioning::Partition(g->block(), input_ivalues_map, partition_info);
+      torch_tensorrt::core::partitioning::Partition(g->block(), input_ivalues_map, partition_info, fallback_nodes);
 
   ASSERT_TRUE(checkSegmentedBlockInputShape(
       segmented_blocks,

diff --git a/tests/modules/custom_models.py b/tests/modules/custom_models.py
@@ -87,6 +87,18 @@ def forward(self, x):
         return x
 
 
+# Sample Inplace OP in Conditional Block Model
+class FallbackInplaceOPIf(nn.Module):
+    def __init__(self):
+        super(FallbackInplaceOPIf, self).__init__()
+    def forward(self, x, y):
+        mod_list = [x]
+        if x.sum() > y.sum():
+            mod_list.append(y)
+        z = torch.cat(mod_list)
+        return z
+
+
 def BertModule():
     model_name = "bert-base-uncased"
     enc = BertTokenizer.from_pretrained(model_name)

diff --git a/tests/modules/hub.py b/tests/modules/hub.py
@@ -96,6 +96,10 @@
         "model": cm.FallbackIf(),
         "path": "script"
     },
+    "inplace_op_if": {
+        "model": cm.FallbackInplaceOPIf(),
+        "path": "script"
+    },
     "bert-base-uncased": {
         "model": cm.BertModule(),
         "path": "trace"