test: add test case for rewrite input with params pass

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

Author: bowang007

core/lowering/passes/rewrite_inputs_with_params.cpp

@@ -1,13 +1,11 @@
-#include "torch/csrc/jit/ir/constants.h"
 #include "core/util/prelude.h"
-
+#include "torch/csrc/jit/ir/constants.h"
 
 namespace torch_tensorrt {
 namespace core {
 namespace lowering {
 namespace passes {
 
-
 void RewriteInputsWithParams(std::shared_ptr<torch::jit::Graph>& g, std::vector<torch::jit::IValue>& params) {
   auto input_size = g->inputs().size();
   auto param_it = params.rbegin();
@@ -33,6 +31,7 @@ void RewriteInputsWithParams(std::shared_ptr<torch::jit::Graph>& g, std::vector<
       }
     }
   }
+  LOG_GRAPH("After RewriteInputsWithParams: " << *g);
 }
 
 } // namespace passes

tests/core/lowering/BUILD

@@ -87,6 +87,10 @@ lowering_test(
     name = "test_unpack_reduce_ops",
 )
 
+lowering_test(
+    name = "test_rewrite_inputs_with_params",
+)
+
 test_suite(
     name = "lowering_tests",
     tests = [
@@ -102,6 +106,7 @@ test_suite(
         ":test_remove_detach_pass",
         ":test_remove_dropout_pass",
         ":test_remove_unnecessary_casts",
+        ":test_rewrite_inputs_with_params",
         ":test_unpack_hardsigmoid",
         ":test_unpack_hardswish",
         ":test_unpack_reduce_ops",

tests/core/lowering/test_rewrite_inputs_with_params.cpp

@@ -0,0 +1,33 @@
+#include <string>
+#include "core/compiler.h"
+#include "core/lowering/passes/passes.h"
+#include "gtest/gtest.h"
+#include "tests/util/util.h"
+#include "torch/csrc/jit/ir/irparser.h"
+#include "torch/csrc/jit/ir/subgraph_matcher.h"
+
+TEST(LoweringPasses, RewriteInputsWithParamsCorrectly) {
+  std::string source_graph = R"IR(
+    graph(%x: Tensor, %y: Tensor, %1 : Int(1)):
+        %out: Tensor = aten::sub(%x, %y, %1)
+        return (%out))IR";
+  std::string target_graph = R"IR(
+    graph(%x: Tensor,  %y : Tensor):
+        %2 : int = prim::Constant[value=0]()
+        %out: Tensor = aten::sub(%x, %y, %2)
+        return (%out))IR";
+
+  torch_tensorrt::core::util::logging::get_logger().set_reportable_log_level(
+      torch_tensorrt::core::util::logging::LogLevel::kGRAPH);
+  auto sg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(source_graph, &*sg);
+
+  torch::jit::IValue param0 = torch::jit::IValue(0);
+  std::vector<torch::jit::IValue> params{param0};
+  torch_tensorrt::core::lowering::passes::RewriteInputsWithParams(sg, params);
+  
+  auto tg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(target_graph, &*tg);
+
+  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
+}