tests(//tests/core/lowering): Adding a lowering testing framework

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

tests/core/BUILD

@@ -2,5 +2,6 @@ test_suite(
     name = "core_tests",
     tests = [
         "//tests/core/conversion:conversion_tests",
+        "//tests/core/lowering:lowering_tests",
     ],
 )

tests/core/conversion/evaluators/test_prim_evaluators.cpp

@@ -4,7 +4,7 @@
 #include "tests/util/util.h"
 #include "torch/csrc/jit/ir/irparser.h"
 
-TEST(Evaluators, PrimConstantConvertsCorrectly) {
+TEST(Evaluators, PrimConstantEvaluatesCorrectly) {
   const auto graph = R"IR(
       graph():
         %0 : int = prim::Constant[value=1]()

tests/core/lowering/BUILD

@@ -0,0 +1,19 @@
+load("//tests/core/lowering:lowering_test.bzl", "lowering_test")
+
+config_setting(
+    name = "use_pre_cxx11_abi",
+    values = {
+        "define": "abi=pre_cxx11_abi",
+    }
+)
+
+lowering_test(
+  name = "test_remove_contiguous_pass",
+)
+
+test_suite(
+    name = "lowering_tests",
+    tests = [
+        ":test_remove_contiguous_pass"
+    ]
+)

tests/core/lowering/lowering_test.bzl

@@ -0,0 +1,15 @@
+def lowering_test(name, visibility=None):
+    native.cc_test(
+        name = name,
+        srcs = [name + ".cpp"],
+        visibility = visibility,
+        deps = [
+            "//tests/util",
+            "//core",
+            "@googletest//:gtest_main",
+        ] + select({
+            ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
+            "//conditions:default":  ["@libtorch//:libtorch"],
+        }),
+        timeout="short"
+    )

tests/core/lowering/test_remove_contiguous_pass.cpp

@@ -0,0 +1,28 @@
+#include <string>
+#include "gtest/gtest.h"
+#include "tests/util/util.h"
+#include "core/compiler.h"
+#include "core/lowering/passes/passes.h"
+#include "torch/csrc/jit/ir/irparser.h"
+#include "torch/csrc/jit/ir/subgraph_matcher.h"
+
+TEST(LoweringPasses, RemoveContiguousLowersCorrectly) {
+  std::string source_graph = R"IR(
+    graph(%input, %1):
+      %2 = aten::contiguous(%input, %1)
+      %3 = foo::bar(%2)
+      return (%3))IR";
+  std::string target_graph = R"IR(
+    graph(%input, %1):
+      %3 = foo::bar(%input)
+      return (%3))IR";
+
+  auto sg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(source_graph, &*sg);
+  trtorch::core::lowering::passes::RemoveContiguous(sg);
+
+  auto tg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(target_graph, &*tg);
+
+  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
+}