refactor: Split elementwise tests

tests/core/conversion/converters/BUILD

@@ -11,6 +11,14 @@ converter_test(
     name = "test_activation",
 )
 
+converter_test(
+    name = "test_add_sub_mul",
+)
+
+converter_test(
+    name = "test_atan2",
+)
+
 converter_test(
     name = "test_batch_norm",
 )
@@ -31,6 +39,10 @@ converter_test(
     name = "test_clone",
 )
 
+converter_test(
+    name = "test_clamp",
+)
+
 converter_test(
     name = "test_concat",
 )
@@ -47,10 +59,18 @@ converter_test(
     name = "test_copy",
 )
 
+converter_test(
+    name = "test_comparators",
+)
+
 converter_test(
     name = "test_cumsum",
 )
 
+converter_test(
+    name = "test_div",
+)
+
 converter_test(
     name = "test_einsum",
 )
@@ -147,17 +167,21 @@ test_suite(
     name = "converter_tests",
     tests = [
         ":test_activation",
+        ":test_add_sub_mul",
+        ":test_atan2",
         ":test_batch_norm",
         ":test_bitwise",
         ":test_cast",
+        ":test_clamp",
         ":test_clone",
+        ":test_comparators",
         ":test_concat",
         ":test_constant_pad",
         ":test_conv_deconv",
         ":test_copy",
         ":test_cumsum",
+        ":test_div",
         ":test_einsum",
-        ":test_element_wise",
         ":test_expand",
         ":test_instance_norm",
         ":test_interpolate",

tests/core/conversion/converters/test_add_sub_mul.cpp

@@ -0,0 +1,180 @@
+#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/torch.h"
+
+using torch_tensorrt::tests::util::pointwise_test_helper;
+
+TEST(Converters, ATenAddConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor, %1 : Tensor):
+        %2 : int = prim::Constant[value=1]()
+        %3 : Tensor = aten::add(%0, %1, %2)
+        return (%3))IR";
+  pointwise_test_helper(graph, false);
+  pointwise_test_helper(graph, false, false, {3, 4}, {4});
+  pointwise_test_helper(graph, false, false, {4}, {3, 4});
+  pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3});
+  pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3});
+}
+
+TEST(Converters, ATenAddWithAlphaConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor, %1 : Tensor):
+        %2 : float = prim::Constant[value=3.2]()
+        %3 : Tensor = aten::add(%0, %1, %2)
+        return (%3))IR";
+  pointwise_test_helper(graph, false);
+  pointwise_test_helper(graph, false, false, {3, 4}, {4});
+  pointwise_test_helper(graph, false, false, {4}, {3, 4});
+  pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3});
+  pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3});
+  pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kFloat, at::kInt);
+  pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kInt, at::kFloat);
+}
+
+TEST(Converters, ATenAddInplaceWithAlphaConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor, %1 : Tensor):
+        %2 : float = prim::Constant[value=7.6]()
+        %3 : Tensor = aten::add_(%0, %1, %2)
+        return (%3))IR";
+  pointwise_test_helper(graph, false);
+  pointwise_test_helper(graph, false, false, {3, 4}, {4});
+  pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3});
+  pointwise_test_helper(graph, false, false, {3, 4, 3}, {4, 3}, false, at::kFloat, at::kInt);
+}
+
+TEST(Converters, ATenAddImplicitWithIntAlphaConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor, %1 : Tensor):
+        %2 : int = prim::Constant[value=42]()
+        %3 : Tensor = aten::add_(%0, %1, %2)
+        return (%3))IR";
+  pointwise_test_helper(graph, false, false, {2, 2}, {2, 2}, false, at::kInt, at::kInt);
+  pointwise_test_helper(graph, false, false, {3, 4, 3}, {4, 3}, false, at::kInt, at::kInt);
+}
+
+TEST(Converters, ATenAddWithScalarConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=1]()
+        %scalar : float = prim::Constant[value=2.4]()
+        %3 : Tensor = aten::add(%0, %scalar, %2)
+        return (%3))IR";
+  pointwise_test_helper(graph, true);
+}
+
+TEST(Converters, ATenSubConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor, %1 : Tensor):
+        %2 : int = prim::Constant[value=2.3]()
+        %3 : Tensor = aten::sub(%0, %1, %2)
+        return (%3))IR";
+  pointwise_test_helper(graph, false);
+  pointwise_test_helper(graph, false, false, {3, 4}, {4});
+  pointwise_test_helper(graph, false, false, {4}, {3, 4});
+  pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3});
+  pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3});
+  pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kFloat, at::kInt);
+  pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kInt, at::kFloat);
+}
+
+TEST(Converters, ATenRsubWithTensorConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor, %1 : Tensor):
+        %2 : int = prim::Constant[value=2]()
+        %3 : Tensor = aten::rsub(%0, %1, %2)
+        return (%3))IR";
+  pointwise_test_helper(graph, false, false, {3, 4}, {4});
+  pointwise_test_helper(graph, false, false, {4}, {3, 4});
+  pointwise_test_helper(graph, false, true, {4, 3, 3, 3}, {4, 3, 3, 3});
+  pointwise_test_helper(graph, false, false, {4, 3, 3, 3}, {4, 3, 3, 3}, false, at::kInt, at::kFloat);
+  pointwise_test_helper(graph, false, false, {4, 3, 3, 3}, {4, 3, 3, 3}, false, at::kInt, at::kInt);
+}
+
+TEST(Converters, ATenRsubWithScalarConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=2]()
+        %scalar : float = prim::Constant[value=2.4]()
+        %3 : Tensor = aten::rsub(%0, %scalar, %2)
+        return (%3))IR";
+  pointwise_test_helper(graph, true, false, {4, 3, 3, 3});
+}
+
+TEST(Converters, ATenRsubWithIntScalarConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=2]()
+        %scalar : int = prim::Constant[value=8]()
+        %3 : Tensor = aten::rsub(%0, %scalar, %2)
+        return (%3))IR";
+  pointwise_test_helper(graph, true, false, {4, 3, 3, 3}, {}, false, at::kInt);
+}
+
+TEST(Converters, ATenMulConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor, %1 : Tensor):
+        %2 : Tensor = aten::mul(%0, %1)
+        return (%2))IR";
+  pointwise_test_helper(graph, false);
+  pointwise_test_helper(graph, false, false, {3, 4}, {4});
+  pointwise_test_helper(graph, false, false, {4}, {3, 4});
+  pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3});
+  pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3});
+  pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kFloat, at::kInt);
+  pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kInt, at::kFloat);
+}
+
+TEST(Converters, ATenSquareConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %1 : Tensor = aten::square(%0)
+        return (%1))IR";
+  pointwise_test_helper(graph, true);
+}
+
+TEST(Converters, ATenMulWithScalarConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %scalar : float = prim::Constant[value=2.4]()
+        %1 : Tensor = aten::mul(%0, %scalar)
+        return (%1))IR";
+  pointwise_test_helper(graph, true);
+}
+
+TEST(Converters, ATenMulWithIntScalarConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %scalar : int = prim::Constant[value=2]()
+        %1 : Tensor = aten::mul(%0, %scalar)
+        return (%1))IR";
+  pointwise_test_helper(graph, true, false, {5}, {5}, false, at::kInt);
+}
+
+TEST(Converters, ATenPowTensorConvertsCorrectly) {
+  const auto graph = R"IR(
+       graph(%x.1 : Tensor, %x2.1 : Tensor):
+          %3 : Tensor = aten::pow(%x.1, %x2.1)
+          return (%3))IR";
+  pointwise_test_helper(graph, false);
+  pointwise_test_helper(graph, false, false, {3, 4}, {4});
+  pointwise_test_helper(graph, false, false, {4}, {3, 4});
+  pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3});
+  pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3});
+  pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kFloat, at::kInt);
+  pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kInt, at::kFloat);
+}
+
+TEST(Converters, ATenPowScalarConvertsCorrectly) {
+  const auto graph = R"IR(
+        graph(%x.1 : Tensor):
+          %2 : int = prim::Constant[value=2]()
+          %3 : Tensor = aten::pow(%x.1, %2)
+          return (%3))IR";
+  pointwise_test_helper(graph, true);
+}

tests/core/conversion/converters/test_atan2.cpp

@@ -0,0 +1,76 @@
+#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/torch.h"
+
+TEST(Converters, ATenAtan2ConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.0 : Tensor, %x.1 : Tensor):
+        %2 : Tensor = aten::atan2(%x.0, %x.1)
+        return (%2))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  // Resize range to [-1, 1] to span multiple quadrants
+  auto in_0 = -2 * at::rand({2, 3, 5, 5}, {at::kCUDA}) + 1;
+  auto in_1 = -2 * at::rand({2, 3, 5, 5}, {at::kCUDA}) + 1;
+
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in_0, in_1});
+
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in_0, in_1});
+
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}
+
+TEST(Converters, ATenAtan2ManagesPosInfCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.0 : Tensor, %x.1 : Tensor):
+        %2 : Tensor = aten::atan2(%x.0, %x.1)
+        return (%2))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  // Expecting PI/2
+  auto in_0 = at::ones({4, 1, 7, 8}, {at::kCUDA});
+  auto in_1 = at::zeros({4, 1, 7, 8}, {at::kCUDA});
+
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in_0, in_1});
+
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in_0, in_1});
+
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}
+
+TEST(Converters, ATenAtan2ManagesNegInfCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.0 : Tensor, %x.1 : Tensor):
+        %2 : Tensor = aten::atan2(%x.0, %x.1)
+        return (%2))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  // Expecting -PI/2
+  auto in_0 = -1 * at::ones({4, 1, 7, 8}, {at::kCUDA});
+  auto in_1 = at::zeros({4, 1, 7, 8}, {at::kCUDA});
+
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in_0, in_1});
+
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in_0, in_1});
+
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}

tests/core/conversion/converters/test_clamp.cpp

@@ -0,0 +1,57 @@
+#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/torch.h"
+
+using torch_tensorrt::tests::util::pointwise_test_helper;
+
+TEST(Converters, ATenClampMinConvertsCorrectly) {
+  const auto graph = R"IR(
+  graph(%x.1 : Tensor):
+          %2 : float = prim::Constant[value=1.5]()
+          %3 : None = prim::Constant()
+          %4 : Tensor = aten::clamp(%x.1, %2, %3)
+          return (%4))IR";
+  pointwise_test_helper(graph, true);
+}
+
+TEST(Converters, ATenClampMaxConvertsCorrectly) {
+  const auto graph = R"IR(
+  graph(%x.1 : Tensor):
+          %2 : float = prim::Constant[value=3.5]()
+          %3 : None = prim::Constant()
+          %4 : Tensor = aten::clamp(%x.1, %3, %2)
+          return (%4))IR";
+  pointwise_test_helper(graph, true);
+}
+
+TEST(Converters, ATenClampMinMaxConvertsCorrectly) {
+  const auto graph = R"IR(
+  graph(%x.1 : Tensor):
+          %2 : float = prim::Constant[value=3.5]()
+          %3 : float = prim::Constant[value=1.5]()
+          %4 : Tensor = aten::clamp(%x.1, %3, %2)
+          return (%4))IR";
+  pointwise_test_helper(graph, true);
+}
+
+TEST(Converters, ATenClampMinimumConvertsCorrectly) {
+  const auto graph = R"IR(
+  graph(%x.1 : Tensor):
+          %2 : float = prim::Constant[value=2.5]()
+          %4 : Tensor = aten::clamp_min(%x.1, %2)
+          return (%4))IR";
+  pointwise_test_helper(graph, true);
+}
+
+TEST(Converters, ATenClampMaximumConvertsCorrectly) {
+  const auto graph = R"IR(
+  graph(%x.1 : Tensor):
+          %2 : float = prim::Constant[value=2.5]()
+          %4 : Tensor = aten::clamp_max(%x.1, %2)
+          return (%4))IR";
+  pointwise_test_helper(graph, true);
+}