test: add test suite for conditionals

Signed-off-by: Bo Wang <[email protected]>

Author: bowang007

tests/core/partitioning/BUILD

@@ -7,6 +7,13 @@ config_setting(
     }
 )
 
+filegroup(
+    name = "jit_models",
+    srcs = ["//tests/modules:resnet50_traced.jit.pt",
+            "//tests/modules:mobilenet_v2_traced.jit.pt",
+            "//tests/modules:conditional_scripted.jit.pt"]
+)
+
 partitioning_test(
   name = "test_segmentation",
 )
@@ -35,17 +42,34 @@ cc_test(
       "//conditions:default":  ["@libtorch//:libtorch"],
   }),
   data = [
-      "//tests/modules:jit_models"
+      ":jit_models"
+  ]
+)
+
+cc_test(
+  name = "test_conditionals",
+  srcs = ["test_conditionals.cpp"],
+  deps = [
+      "//tests/util",
+      "//core",
+      "@googletest//:gtest_main",
+  ] + select({
+      ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
+      "//conditions:default":  ["@libtorch//:libtorch"],
+  }),
+  data = [
+      ":jit_models"
   ]
 )
 
 test_suite(
-    name = "partitioning_tests",
+    name = "partitioning_test",
     tests = [
         ":test_segmentation",
         ":test_shape_analysis",
         ":test_tensorrt_conversion",
         ":test_stitched_graph",
-        ":test_fallback_graph_output"
+        ":test_fallback_graph_output",
+        ":test_conditionals"
     ]
 )

tests/core/partitioning/test_conditionals.cpp

@@ -0,0 +1,43 @@
+#include <string>
+#include <unordered_set>
+#include "core/compiler.h"
+#include "gtest/gtest.h"
+#include "tests/util/util.h"
+#include "torch/script.h"
+
+size_t count_trt_engines_in_conditionals(std::shared_ptr<torch::jit::Graph> g) {
+  size_t count = 0;
+  for (auto n : g->nodes()) {
+    if (n->kind() == torch::jit::prim::If) {
+      std::vector<torch::jit::Block*> blocks{n->blocks()[0], n->blocks()[1]};
+      for (auto cur_block : blocks) {
+        for (auto n : cur_block->nodes()) {
+          if (n->kind().toQualString() == std::string("tensorrt::execute_engine")) {
+            ++count;
+          }
+        }
+      }
+    }
+  }
+  return count;
+}
+
+TEST(Partitioning, FallbackOnConditionalsCorrectly) {
+  torch::jit::script::Module mod;
+  try {
+    mod = torch::jit::load("tests/modules/conditional_scripted.jit.pt");
+  } catch (const c10::Error& e) {
+    std::cerr << "error loading the model\n";
+    return;
+  }
+
+  std::vector<trtorch::core::ir::InputRange> input_ranges{trtorch::core::ir::InputRange({3, 3, 16, 16})};
+  trtorch::core::CompileSpec cfg(input_ranges);
+  cfg.partition_info.enabled = true;
+  torch::jit::script::Module new_mod = trtorch::core::CompileGraph(mod, cfg);
+  auto g = new_mod.get_method("forward").graph();
+
+  auto conditional_engines_count = count_trt_engines_in_conditionals(g);
+
+  ASSERT_TRUE(conditional_engines_count == 2);
+}

tests/modules/hub.py

@@ -95,3 +95,31 @@ def forward(self, x):
 
 trace_model = torch.jit.trace(model, x)
 torch.jit.save(trace_model, "pooling_traced.jit.pt")
+
+
+# Sample Conditional Model (for testing partitioning and fallback in conditionals)
+class FallbackIf(torch.nn.Module):
+    def __init__(self):
+        super(FallbackIf, self).__init__()
+        self.relu1 = torch.nn.ReLU()
+        self.conv1 = torch.nn.Conv2d(3, 32, 3, 1, 1)
+        self.log_sig = torch.nn.LogSigmoid()
+        self.conv2 = torch.nn.Conv2d(32, 32, 3, 1, 1)
+        self.conv3 = torch.nn.Conv2d(32, 3, 3, 1, 1)
+
+    def forward(self, x):
+        x = self.relu1(x)
+        x_first = x[0][0][0][0].item()
+        if x_first > 0:
+            x = self.conv1(x)
+            x1 = self.log_sig(x)
+            x2 = self.conv2(x)
+            x = self.conv3(x1 + x2)
+        else:
+            x = self.log_sig(x)
+        x = self.conv1(x)
+        return x
+
+conditional_model = FallbackIf().eval().cuda()
+conditional_script_model = torch.jit.script(conditional_model)
+torch.jit.save(conditional_script_model, "conditional_scripted.jit.pt")