From c73ce483cb40f20340b35bc23c862197a0e75bb2 Mon Sep 17 00:00:00 2001 From: Dheeraj Peri Date: Thu, 1 Dec 2022 01:26:15 -0800 Subject: [PATCH 1/4] refactor: Split elementwise tests Signed-off-by: Dheeraj Peri --- tests/core/conversion/converters/BUILD | 26 +- .../converters/test_add_sub_mul.cpp | 180 +++++ .../core/conversion/converters/test_atan2.cpp | 76 +++ .../core/conversion/converters/test_clamp.cpp | 57 ++ .../converters/test_comparators.cpp | 137 ++++ tests/core/conversion/converters/test_div.cpp | 169 +++++ .../converters/test_element_wise.cpp | 634 ------------------ tests/util/BUILD | 1 + tests/util/util.cpp | 53 ++ tests/util/util.h | 11 + 10 files changed, 709 insertions(+), 635 deletions(-) create mode 100644 tests/core/conversion/converters/test_add_sub_mul.cpp create mode 100644 tests/core/conversion/converters/test_atan2.cpp create mode 100644 tests/core/conversion/converters/test_clamp.cpp create mode 100644 tests/core/conversion/converters/test_comparators.cpp create mode 100644 tests/core/conversion/converters/test_div.cpp delete mode 100644 tests/core/conversion/converters/test_element_wise.cpp diff --git a/tests/core/conversion/converters/BUILD b/tests/core/conversion/converters/BUILD index 98630b6dc1..901ca94998 100644 --- a/tests/core/conversion/converters/BUILD +++ b/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", diff --git a/tests/core/conversion/converters/test_add_sub_mul.cpp b/tests/core/conversion/converters/test_add_sub_mul.cpp new file mode 100644 index 0000000000..3631a92a44 --- /dev/null +++ b/tests/core/conversion/converters/test_add_sub_mul.cpp @@ -0,0 +1,180 @@ +#include +#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); +} diff --git a/tests/core/conversion/converters/test_atan2.cpp b/tests/core/conversion/converters/test_atan2.cpp new file mode 100644 index 0000000000..d899dff759 --- /dev/null +++ b/tests/core/conversion/converters/test_atan2.cpp @@ -0,0 +1,76 @@ +#include +#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::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::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::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)); +} \ No newline at end of file diff --git a/tests/core/conversion/converters/test_clamp.cpp b/tests/core/conversion/converters/test_clamp.cpp new file mode 100644 index 0000000000..2ae7f1d510 --- /dev/null +++ b/tests/core/conversion/converters/test_clamp.cpp @@ -0,0 +1,57 @@ +#include +#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); +} \ No newline at end of file diff --git a/tests/core/conversion/converters/test_comparators.cpp b/tests/core/conversion/converters/test_comparators.cpp new file mode 100644 index 0000000000..0107fa6837 --- /dev/null +++ b/tests/core/conversion/converters/test_comparators.cpp @@ -0,0 +1,137 @@ +#include +#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, ATenGreaterThanConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %2 : Tensor = aten::gt(%0, %1) + return (%2))IR"; + pointwise_test_helper(graph, false, false, {5, 5}, {5, 5}); +} + +TEST(Converters, ATenGreaterThanScalarConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %scalar : float = prim::Constant[value=3]() + %2 : Tensor = aten::gt(%0, %scalar) + return (%2))IR"; + pointwise_test_helper(graph, true, false, {5, 5}); +} + +TEST(Converters, ATenLessThanConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %2 : Tensor = aten::lt(%0, %1) + return (%2))IR"; + pointwise_test_helper(graph, false, false, {5, 5}, {5, 5}); +} + +TEST(Converters, ATenLessThanScalarConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %scalar : float = prim::Constant[value=3]() + %2 : Tensor = aten::lt(%0, %scalar) + return (%2))IR"; + pointwise_test_helper(graph, true, false, {5, 5}); +} + +TEST(Converters, ATenEqualConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %2 : Tensor = aten::eq(%0, %1) + return (%2))IR"; + pointwise_test_helper(graph, false, false, {5, 5}, {5, 5}); +} + +TEST(Converters, ATenEqualScalarConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %scalar : float = prim::Constant[value=3]() + %2 : Tensor = aten::eq(%0, %scalar) + return (%2))IR"; + pointwise_test_helper(graph, true, false, {5, 5}); +} + +TEST(Converters, ATenGEConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %2 : Tensor = aten::ge(%0, %1) + return (%2))IR"; + pointwise_test_helper(graph, false, false, {5, 5}, {5, 5}); +} + +TEST(Converters, ATenGEScalarConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %scalar : float = prim::Constant[value=3]() + %2 : Tensor = aten::ge(%0, %scalar) + return (%2))IR"; + pointwise_test_helper(graph, true, false, {5, 5}); +} + +TEST(Converters, ATenLEConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %2 : Tensor = aten::le(%0, %1) + return (%2))IR"; + pointwise_test_helper(graph, false, false, {5, 5}, {5, 5}); +} + +TEST(Converters, ATenLEScalarConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %scalar : float = prim::Constant[value=3]() + %2 : Tensor = aten::le(%0, %scalar) + return (%2))IR"; + pointwise_test_helper(graph, true, false, {5, 5}); +} + +TEST(Converters, ATenNeTensorConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor, + %y.1 : Tensor): + %3 : Tensor = aten::ne(%x.1, %y.1) + return (%3))IR"; + pointwise_test_helper(graph, false, false, {3, 4}, {3, 4}); + pointwise_test_helper(graph, false, true, {3, 4}, {3, 4}); +} + +TEST(Converters, ATenNeScalarConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int = prim::Constant[value=2]() + %3 : Tensor = aten::ne(%x.1, %2) + return (%3))IR"; + pointwise_test_helper(graph, true, false, {3, 4, 2}); +} + +TEST(Converters, ATenMaxConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %2 : Tensor = aten::max(%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}); +} + +TEST(Converters, ATenMinConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %2 : Tensor = aten::min(%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}); +} \ No newline at end of file diff --git a/tests/core/conversion/converters/test_div.cpp b/tests/core/conversion/converters/test_div.cpp new file mode 100644 index 0000000000..7670132be4 --- /dev/null +++ b/tests/core/conversion/converters/test_div.cpp @@ -0,0 +1,169 @@ +#include +#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, ATenDivConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %2 : Tensor = aten::div(%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}); +} + +TEST(Converters, ATenDivWithScalarConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %scalar : float = prim::Constant[value=2.4]() + %1 : Tensor = aten::div(%0, %scalar) + return (%1))IR"; + pointwise_test_helper(graph, true); +} + +TEST(Converters, ATenDivRoundingFloorConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %3 : str = prim::Constant[value="floor"]() + %2 : Tensor = aten::div(%0, %1, %3) + return (%2))IR"; + pointwise_test_helper(graph, false, false, {5}, {5}, true); + pointwise_test_helper(graph, false, false, {3, 4}, {4}, true); + pointwise_test_helper(graph, false, false, {4}, {3, 4}, true); + pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}, true); + pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}, true); + 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, ATenDivRoundingTruncConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %3 : str = prim::Constant[value="trunc"]() + %2 : Tensor = aten::div(%0, %1, %3) + return (%2))IR"; + pointwise_test_helper(graph, false, false, {5}, {5}, true); + pointwise_test_helper(graph, false, false, {3, 4}, {4}, true); + pointwise_test_helper(graph, false, false, {4}, {3, 4}, true); + pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}, true); + pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}, true); + 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, ATenDivRoundingNoneConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %3 : None = prim::Constant() + %2 : Tensor = aten::div(%0, %1, %3) + return (%2))IR"; + pointwise_test_helper(graph, false, false, {5}, {5}, true); + pointwise_test_helper(graph, false, false, {3, 4}, {4}, true); + pointwise_test_helper(graph, false, false, {4}, {3, 4}, true); + pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}, true); + pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}, true); +} + +TEST(Converters, ATenDivRoundingTruncWithIntsConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %trunc : str = prim::Constant[value="trunc"]() + %out : Tensor = aten::div(%0, %1, %trunc) + return (%out))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + // Avoid divide-by-zero issues by making denominator >= 1 + auto in_0 = at::randint(-5, 5, {4, 1, 7, 8}, {at::kCUDA}).to(torch::kInt32); + auto in_1 = at::randint(1, 10, {4, 1, 7, 8}, {at::kCUDA}).to(torch::kInt32); + + 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::exactlyEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]))); +} + +TEST(Converters, ATenFloorDivideConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %2 : Tensor = aten::floor_divide(%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, ATenFloorDivideWithScalarConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %scalar : float = prim::Constant[value=2.4]() + %1 : Tensor = aten::floor_divide(%0, %scalar) + return (%1))IR"; + pointwise_test_helper(graph, true); + pointwise_test_helper(graph, true, false, {5}, {5}, false, at::kInt); +} + +TEST(Converters, ATenRemainderConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %2 : Tensor = aten::remainder(%0, %1) + return (%2))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, &*g); + + auto input1 = at::randint(-5, 5, {4, 5}, {at::kCUDA}); + auto input2 = at::randint(1, 5, {5}, {at::kCUDA}); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {input1, input2}); + + torch_tensorrt::core::lowering::passes::ReduceRemainder(g); + + input1 = at::clone(input1); + input2 = at::clone(input2); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {input1, input2}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 5e-2)); +} + +TEST(Converters, ATenRemainderWithScalarConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %scalar : float = prim::Constant[value=2.4]() + %1 : Tensor = aten::remainder(%0, %scalar) + return (%1))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, &*g); + + auto in = at::randint(-5, 5, {5}, {at::kCUDA}); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); + + torch_tensorrt::core::lowering::passes::ReduceRemainder(g); + + in = at::clone(in); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} \ No newline at end of file diff --git a/tests/core/conversion/converters/test_element_wise.cpp b/tests/core/conversion/converters/test_element_wise.cpp deleted file mode 100644 index 82e959dfcd..0000000000 --- a/tests/core/conversion/converters/test_element_wise.cpp +++ /dev/null @@ -1,634 +0,0 @@ -#include -#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" - -void pointwise_test_helper( - std::string graph_ir, - bool singleInput, - bool dynamicInput = false, - std::vector shape1 = {5}, - std::vector shape2 = {5}, - bool negative_input = false, - at::ScalarType type1 = at::kFloat, - at::ScalarType type2 = at::kFloat) { - auto g = std::make_shared(); - torch::jit::parseIR(graph_ir, g.get()); - - // singleInput case is enabled when elementwise operation is performed - // with an input and a constant embedded in graph - std::vector torch_inputs; - int first_min = negative_input ? -5 : 1; - int first_max = 5; - int second_min = 1; - int second_max = 5; - if (type1 == at::kBool) { - first_min = 0; - first_max = 1; - } - if (type2 == at::kBool) { - second_min = 0; - second_max = 1; - } - torch_inputs.push_back(at::randint(first_min, first_max, shape1, at::TensorOptions(at::kCUDA).dtype(type1))); - if (!singleInput) { - torch_inputs.push_back(at::randint(second_min, second_max, shape2, at::TensorOptions(at::kCUDA).dtype(type2))); - } - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, torch_inputs); - - std::vector trt_inputs; - for (auto in : torch_inputs) { - trt_inputs.push_back(at::clone(in)); - } - - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - std::vector trt_results; - if (dynamicInput) { - trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, trt_inputs); - } else { - trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, trt_inputs); - } - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} - -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, 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, ATenDivConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %1 : Tensor): - %2 : Tensor = aten::div(%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}); -} - -TEST(Converters, ATenDivWithScalarConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor): - %scalar : float = prim::Constant[value=2.4]() - %1 : Tensor = aten::div(%0, %scalar) - return (%1))IR"; - pointwise_test_helper(graph, true); -} - -TEST(Converters, ATenDivRoundingFloorConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %1 : Tensor): - %3 : str = prim::Constant[value="floor"]() - %2 : Tensor = aten::div(%0, %1, %3) - return (%2))IR"; - pointwise_test_helper(graph, false, false, {5}, {5}, true); - pointwise_test_helper(graph, false, false, {3, 4}, {4}, true); - pointwise_test_helper(graph, false, false, {4}, {3, 4}, true); - pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}, true); - pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}, true); - 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, ATenDivRoundingTruncConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %1 : Tensor): - %3 : str = prim::Constant[value="trunc"]() - %2 : Tensor = aten::div(%0, %1, %3) - return (%2))IR"; - pointwise_test_helper(graph, false, false, {5}, {5}, true); - pointwise_test_helper(graph, false, false, {3, 4}, {4}, true); - pointwise_test_helper(graph, false, false, {4}, {3, 4}, true); - pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}, true); - pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}, true); - 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, ATenDivRoundingNoneConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %1 : Tensor): - %3 : None = prim::Constant() - %2 : Tensor = aten::div(%0, %1, %3) - return (%2))IR"; - pointwise_test_helper(graph, false, false, {5}, {5}, true); - pointwise_test_helper(graph, false, false, {3, 4}, {4}, true); - pointwise_test_helper(graph, false, false, {4}, {3, 4}, true); - pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}, true); - pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}, true); -} - -TEST(Converters, ATenDivRoundingTruncWithIntsConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %1 : Tensor): - %trunc : str = prim::Constant[value="trunc"]() - %out : Tensor = aten::div(%0, %1, %trunc) - return (%out))IR"; - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - // Avoid divide-by-zero issues by making denominator >= 1 - auto in_0 = at::randint(-5, 5, {4, 1, 7, 8}, {at::kCUDA}).to(torch::kInt32); - auto in_1 = at::randint(1, 10, {4, 1, 7, 8}, {at::kCUDA}).to(torch::kInt32); - - 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::exactlyEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]))); -} - -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); -} - -TEST(Converters, ATenNeTensorConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor, - %y.1 : Tensor): - %3 : Tensor = aten::ne(%x.1, %y.1) - return (%3))IR"; - pointwise_test_helper(graph, false, false, {3, 4}, {3, 4}); - pointwise_test_helper(graph, false, true, {3, 4}, {3, 4}); -} - -TEST(Converters, ATenNeScalarConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : int = prim::Constant[value=2]() - %3 : Tensor = aten::ne(%x.1, %2) - return (%3))IR"; - pointwise_test_helper(graph, true, false, {3, 4, 2}); -} - -TEST(Converters, ATenFloorDivideConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %1 : Tensor): - %2 : Tensor = aten::floor_divide(%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, ATenFloorDivideWithScalarConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor): - %scalar : float = prim::Constant[value=2.4]() - %1 : Tensor = aten::floor_divide(%0, %scalar) - return (%1))IR"; - pointwise_test_helper(graph, true); - pointwise_test_helper(graph, true, false, {5}, {5}, false, at::kInt); -} - -TEST(Converters, ATenMaxConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %1 : Tensor): - %2 : Tensor = aten::max(%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}); -} - -TEST(Converters, ATenMinConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %1 : Tensor): - %2 : Tensor = aten::min(%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}); -} - -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, 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); -} - -TEST(Converters, ATenGreaterThanConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %1 : Tensor): - %2 : Tensor = aten::gt(%0, %1) - return (%2))IR"; - pointwise_test_helper(graph, false, false, {5, 5}, {5, 5}); -} - -TEST(Converters, ATenGreaterThanScalarConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor): - %scalar : float = prim::Constant[value=3]() - %2 : Tensor = aten::gt(%0, %scalar) - return (%2))IR"; - pointwise_test_helper(graph, true, false, {5, 5}); -} - -TEST(Converters, ATenLessThanConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %1 : Tensor): - %2 : Tensor = aten::lt(%0, %1) - return (%2))IR"; - pointwise_test_helper(graph, false, false, {5, 5}, {5, 5}); -} - -TEST(Converters, ATenLessThanScalarConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor): - %scalar : float = prim::Constant[value=3]() - %2 : Tensor = aten::lt(%0, %scalar) - return (%2))IR"; - pointwise_test_helper(graph, true, false, {5, 5}); -} - -TEST(Converters, ATenEqualConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %1 : Tensor): - %2 : Tensor = aten::eq(%0, %1) - return (%2))IR"; - pointwise_test_helper(graph, false, false, {5, 5}, {5, 5}); -} - -TEST(Converters, ATenEqualScalarConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor): - %scalar : float = prim::Constant[value=3]() - %2 : Tensor = aten::eq(%0, %scalar) - return (%2))IR"; - pointwise_test_helper(graph, true, false, {5, 5}); -} - -TEST(Converters, ATenGEConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %1 : Tensor): - %2 : Tensor = aten::ge(%0, %1) - return (%2))IR"; - pointwise_test_helper(graph, false, false, {5, 5}, {5, 5}); -} - -TEST(Converters, ATenGEScalarConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor): - %scalar : float = prim::Constant[value=3]() - %2 : Tensor = aten::ge(%0, %scalar) - return (%2))IR"; - pointwise_test_helper(graph, true, false, {5, 5}); -} - -TEST(Converters, ATenLEConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %1 : Tensor): - %2 : Tensor = aten::le(%0, %1) - return (%2))IR"; - pointwise_test_helper(graph, false, false, {5, 5}, {5, 5}); -} - -TEST(Converters, ATenLEScalarConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor): - %scalar : float = prim::Constant[value=3]() - %2 : Tensor = aten::le(%0, %scalar) - return (%2))IR"; - pointwise_test_helper(graph, true, false, {5, 5}); -} - -TEST(Converters, ATenRemainderConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %1 : Tensor): - %2 : Tensor = aten::remainder(%0, %1) - return (%2))IR"; - - auto g = std::make_shared(); - torch::jit::parseIR(graph, &*g); - - auto input1 = at::randint(-5, 5, {4, 5}, {at::kCUDA}); - auto input2 = at::randint(1, 5, {5}, {at::kCUDA}); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {input1, input2}); - - torch_tensorrt::core::lowering::passes::ReduceRemainder(g); - - input1 = at::clone(input1); - input2 = at::clone(input2); - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {input1, input2}); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 5e-2)); -} - -TEST(Converters, ATenRemainderWithScalarConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor): - %scalar : float = prim::Constant[value=2.4]() - %1 : Tensor = aten::remainder(%0, %scalar) - return (%1))IR"; - - auto g = std::make_shared(); - torch::jit::parseIR(graph, &*g); - - auto in = at::randint(-5, 5, {5}, {at::kCUDA}); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); - - torch_tensorrt::core::lowering::passes::ReduceRemainder(g); - - in = at::clone(in); - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} - -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::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::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::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)); -} diff --git a/tests/util/BUILD b/tests/util/BUILD index c69fda8a4a..a31cda26f1 100644 --- a/tests/util/BUILD +++ b/tests/util/BUILD @@ -30,6 +30,7 @@ cc_library( ], deps = [ "@tensorrt//:nvinfer", + "@googletest//:gtest_main", ] + select({ ":use_pre_cxx11_abi": [ "@libtorch_pre_cxx11_abi//:libtorch", diff --git a/tests/util/util.cpp b/tests/util/util.cpp index 39d1eb14b1..4f87a8a0f8 100644 --- a/tests/util/util.cpp +++ b/tests/util/util.cpp @@ -1,5 +1,7 @@ #include "util.h" +#include #include "core/util/prelude.h" +#include "gtest/gtest.h" #include "torch/script.h" #include "torch/torch.h" @@ -48,6 +50,57 @@ bool exactlyEqual(const at::Tensor& computed_tensor, const at::Tensor& gt_tensor return (computed_tensor - gt_tensor).abs().max().item() == 0.f; } +void pointwise_test_helper( + std::string graph_ir, + bool singleInput, + bool dynamicInput, + std::vector shape1, + std::vector shape2, + bool negative_input, + at::ScalarType type1, + at::ScalarType type2) { + auto g = std::make_shared(); + torch::jit::parseIR(graph_ir, g.get()); + + // singleInput case is enabled when elementwise operation is performed + // with an input and a constant embedded in graph + std::vector torch_inputs; + int first_min = negative_input ? -5 : 1; + int first_max = 5; + int second_min = 1; + int second_max = 5; + if (type1 == at::kBool) { + first_min = 0; + first_max = 1; + } + if (type2 == at::kBool) { + second_min = 0; + second_max = 1; + } + torch_inputs.push_back(at::randint(first_min, first_max, shape1, at::TensorOptions(at::kCUDA).dtype(type1))); + if (!singleInput) { + torch_inputs.push_back(at::randint(second_min, second_max, shape2, at::TensorOptions(at::kCUDA).dtype(type2))); + } + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, torch_inputs); + + std::vector trt_inputs; + for (auto in : torch_inputs) { + trt_inputs.push_back(at::clone(in)); + } + + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + std::vector trt_results; + if (dynamicInput) { + trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, trt_inputs); + } else { + trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, trt_inputs); + } + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + } // namespace util } // namespace tests } // namespace torch_tensorrt diff --git a/tests/util/util.h b/tests/util/util.h index 23207ef02c..d7cf7accee 100644 --- a/tests/util/util.h +++ b/tests/util/util.h @@ -6,6 +6,7 @@ #include "ATen/Tensor.h" #include "core/ir/ir.h" #include "core/util/prelude.h" +#include "torch/csrc/jit/ir/irparser.h" namespace torch_tensorrt { namespace tests { @@ -21,6 +22,16 @@ bool almostEqual(const at::Tensor& computed_tensor, const at::Tensor& gt_tensor, bool exactlyEqual(const at::Tensor& a, const at::Tensor& b); +void pointwise_test_helper( + std::string graph_ir, + bool singleInput, + bool dynamicInput = false, + std::vector shape1 = {5}, + std::vector shape2 = {5}, + bool negative_input = false, + at::ScalarType type1 = at::kFloat, + at::ScalarType type2 = at::kFloat); + std::vector RunEngine(std::string& eng, std::vector inputs); // Runs an arbitrary JIT graph and returns results From fd75dd31913c3e070ec03c9263b2ac0a794c017c Mon Sep 17 00:00:00 2001 From: Dheeraj Peri Date: Thu, 20 Apr 2023 01:52:06 -0700 Subject: [PATCH 2/4] chore: Split tests to reduce execution time Signed-off-by: Dheeraj Peri --- tests/core/conversion/converters/BUILD | 45 + .../core/conversion/converters/test_index.cpp | 294 +++++ .../converters/test_masked_fill.cpp | 99 ++ .../conversion/converters/test_reduce.cpp | 237 ---- .../core/conversion/converters/test_roll.cpp | 84 ++ .../conversion/converters/test_scatter.cpp | 79 ++ .../conversion/converters/test_select.cpp | 1170 ----------------- .../core/conversion/converters/test_slice.cpp | 332 +++++ .../core/conversion/converters/test_split.cpp | 174 +++ .../conversion/converters/test_unbind.cpp | 88 ++ .../conversion/converters/test_unpack.cpp | 243 ++++ .../core/conversion/converters/test_where.cpp | 68 + 12 files changed, 1506 insertions(+), 1407 deletions(-) create mode 100644 tests/core/conversion/converters/test_index.cpp create mode 100644 tests/core/conversion/converters/test_masked_fill.cpp create mode 100644 tests/core/conversion/converters/test_roll.cpp create mode 100644 tests/core/conversion/converters/test_scatter.cpp create mode 100644 tests/core/conversion/converters/test_slice.cpp create mode 100644 tests/core/conversion/converters/test_split.cpp create mode 100644 tests/core/conversion/converters/test_unbind.cpp create mode 100644 tests/core/conversion/converters/test_unpack.cpp create mode 100644 tests/core/conversion/converters/test_where.cpp diff --git a/tests/core/conversion/converters/BUILD b/tests/core/conversion/converters/BUILD index 901ca94998..a8c57b1b41 100644 --- a/tests/core/conversion/converters/BUILD +++ b/tests/core/conversion/converters/BUILD @@ -95,6 +95,10 @@ converter_test( name = "test_matrix_multiply", ) +converter_test( + name = "test_masked_fill", +) + converter_test( name = "test_max", ) @@ -115,6 +119,10 @@ converter_test( name = "test_reduce", ) +converter_test( + name = "test_roll", +) + converter_test( name = "test_reflection_pad", ) @@ -123,6 +131,10 @@ converter_test( name = "test_replication_pad", ) +converter_test( + name = "test_scatter", +) + converter_test( name = "test_shuffle", ) @@ -139,6 +151,10 @@ converter_test( name = "test_interpolate", ) +converter_test( + name = "test_index", +) + converter_test( name = "test_select", ) @@ -147,6 +163,14 @@ converter_test( name = "test_stack", ) +converter_test( + name = "test_slice", +) + +converter_test( + name = "test_split", +) + converter_test( name = "test_topk", ) @@ -159,10 +183,22 @@ converter_test( name = "test_unsqueeze", ) +converter_test( + name = "test_unbind", +) + +converter_test( + name = "test_unpack", +) + converter_test( name = "test_squeeze", ) +converter_test( + name = "test_where", +) + test_suite( name = "converter_tests", tests = [ @@ -185,22 +221,31 @@ test_suite( ":test_expand", ":test_instance_norm", ":test_interpolate", + ":test_index", ":test_layer_norm", ":test_linear", ":test_lstm_cell", ":test_matrix_multiply", + ":test_masked_fill", ":test_max", ":test_normalize", ":test_pooling", ":test_reduce", + ":test_roll", ":test_replication_pad", + ":test_scatter", ":test_select", ":test_shuffle", ":test_softmax", ":test_squeeze", ":test_stack", + ":test_split", + ":test_slice", ":test_topk", ":test_unary", ":test_unsqueeze", + ":test_unbind", + ":test_unpack", + ":test_where", ], ) diff --git a/tests/core/conversion/converters/test_index.cpp b/tests/core/conversion/converters/test_index.cpp new file mode 100644 index 0000000000..34e50f2abd --- /dev/null +++ b/tests/core/conversion/converters/test_index.cpp @@ -0,0 +1,294 @@ +#include +#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" + +TEST(Converters, ATenIndexSelectConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %index : Int (2)): + %2 : int = prim::Constant[value=0]() + %3 : Tensor = aten::index_select(%0, %2, %index) + return (%3))IR"; + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + auto in = at::randint(1, 10, {4, 4, 4}, {at::kCUDA}); + auto index = at::randint(0, 4, {2}, {at::kCUDA}).to(torch::kI32); + + auto jit_in = at::clone(in); + auto jit_index = at::clone(index); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {jit_index}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_index = at::clone(index); + auto trt_params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {trt_index}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, trt_params, {trt_in}); + + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} + +TEST(Converters, ATenIndexSelectNegativeDimConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %index : Int (5)): + %2 : int = prim::Constant[value=-1]() + %3 : Tensor = aten::index_select(%0, %2, %index) + return (%3))IR"; + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {5, 3, 9}, {at::kCUDA}); + auto index = at::randint(0, 9, {5}, {at::kCUDA}).to(torch::kI32); + + auto jit_in = at::clone(in); + auto jit_index = at::clone(index); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {jit_index}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_index = at::clone(index); + auto trt_params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {trt_index}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, trt_params, {trt_in}); + + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} + +TEST(Converters, ATenIndexTensorOneIndiceConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor, + %index : Tensor): + %18 : Tensor?[] = prim::ListConstruct(%index) + %19 : Tensor = aten::index(%x.1, %18) + return (%19))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto in1 = at::randint(1, 10, {5, 10}, {at::kCUDA}); + auto in2 = at::full({2}, 4, {at::kCUDA}); + auto options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + auto in2_trt = at::full({2}, 4, {options}); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, in2}); + + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, in2_trt}); + + ASSERT_TRUE( + torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); +} + +TEST(Converters, ATenIndexTensorFullIndicesConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor, + %index0 : Tensor, + %index1 : Tensor, + %index2 : Tensor): + %18 : Tensor?[] = prim::ListConstruct(%index0, %index1, %index2) + %19 : Tensor = aten::index(%x.1, %18) + return (%19))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto in1 = at::randint(1, 10, {5, 10, 4}, {at::kCUDA}); + auto index0 = at::tensor({0, 1, 2, 3}, {at::kCUDA}).to(torch::kLong); + auto index1 = at::tensor({1, 3, 4, 6}, {at::kCUDA}).to(torch::kLong); + auto index2 = at::tensor({3, 2, 1, 0}, {at::kCUDA}).to(torch::kLong); + auto index0_trt = index0.to(torch::kInt32); + auto index1_trt = index1.to(torch::kInt32); + auto index2_trt = index2.to(torch::kInt32); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, index0, index1, index2}); + + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, index0_trt, index1_trt, index2_trt}); + + ASSERT_TRUE( + torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); +} + +TEST(Converters, ATenIndexTensorRepeatedFullIndicesConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor, + %index0 : Tensor, + %index1 : Tensor, + %index2 : Tensor): + %18 : Tensor?[] = prim::ListConstruct(%index0, %index1, %index2) + %19 : Tensor = aten::index(%x.1, %18) + %20 : Tensor = aten::index(%x.1, %18) + return (%19, %20))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto in1 = at::randint(1, 10, {5, 10, 4}, {at::kCUDA}); + auto index0 = at::tensor({0, 1, 2, 3}, {at::kCUDA}).to(torch::kLong); + auto index1 = at::tensor({1, 3, 4, 6}, {at::kCUDA}).to(torch::kLong); + auto index2 = at::tensor({3, 2, 1, 0}, {at::kCUDA}).to(torch::kLong); + auto index0_trt = index0.to(torch::kInt32); + auto index1_trt = index1.to(torch::kInt32); + auto index2_trt = index2.to(torch::kInt32); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, index0, index1, index2}); + + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, index0_trt, index1_trt, index2_trt}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[1], trt_results[1], 2e-6)); +} + +TEST(Converters, ATenIndexTensorIdx0Idx1NoneConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor, + %index0 : Tensor, + %index1 : Tensor): + %5 : NoneType = prim::Constant() + %18 : Tensor?[] = prim::ListConstruct(%index0, %index1, %5) + %19 : Tensor = aten::index(%x.1, %18) + return (%19))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto in1 = at::randint(1, 10, {5, 10, 4}, {at::kCUDA}); + auto index0 = at::tensor({0, 1, 2, 3}, {at::kCUDA}).to(torch::kLong); + auto index1 = at::tensor({1, 3, 4, 6}, {at::kCUDA}).to(torch::kLong); + auto index0_trt = index0.to(torch::kInt32); + auto index1_trt = index1.to(torch::kInt32); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, index0, index1}); + + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, index0_trt, index1_trt}); + LOG_DEBUG(trt_results); + + ASSERT_TRUE( + torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); +} + +TEST(Converters, ATenIndexTensorIdx0NoneIdx1ConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor, + %index0 : Tensor, + %index1 : Tensor): + %5 : NoneType = prim::Constant() + %18 : Tensor?[] = prim::ListConstruct(%index0, %5, %index1) + %19 : Tensor = aten::index(%x.1, %18) + return (%19))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto in1 = at::randint(1, 10, {5, 10, 4}, {at::kCUDA}); + auto index0 = at::tensor({0, 1, 2, 3}, {at::kCUDA}).to(torch::kLong); + auto index1 = at::tensor({3, 2, 1, 0}, {at::kCUDA}).to(torch::kLong); + auto index0_trt = index0.to(torch::kInt32); + auto index1_trt = index1.to(torch::kInt32); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, index0, index1}); + + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, index0_trt, index1_trt}); + + ASSERT_TRUE( + torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); +} + +TEST(Converters, ATenIndexTensorNoneIdx0Idx1ConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor, + %index0 : Tensor, + %index1 : Tensor): + %5 : NoneType = prim::Constant() + %18 : Tensor?[] = prim::ListConstruct(%5, %index0, %index1) + %19 : Tensor = aten::index(%x.1, %18) + return (%19))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto in1 = at::randint(1, 10, {5, 10, 4}, {at::kCUDA}); + auto index0 = at::tensor({0, 1, 2, 3}, {at::kCUDA}).to(torch::kLong); + auto index1 = at::tensor({3, 2, 1, 0}, {at::kCUDA}).to(torch::kLong); + auto index0_trt = index0.to(torch::kInt32); + auto index1_trt = index1.to(torch::kInt32); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, index0, index1}); + + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, index0_trt, index1_trt}); + + ASSERT_TRUE( + torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); +} + +TEST(Converters, ATenIndexTensorIdxsNoneConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor, + %index0 : Tensor, + %index1 : Tensor, + %index2 : Tensor): + %5 : NoneType = prim::Constant() + %18 : Tensor?[] = prim::ListConstruct(%index0, %index1, %index2, %5) + %19 : Tensor = aten::index(%x.1, %18) + return (%19))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto in1 = at::randint(1, 10, {4, 8, 8, 4}, {at::kCUDA}); + auto index0 = at::full({4, 13, 1}, 1, {at::kCUDA}).to(torch::kLong); + auto index1 = at::full({4, 13, 1}, 2, {at::kCUDA}).to(torch::kLong); + auto index2 = at::full({4, 13, 1}, 3, {at::kCUDA}).to(torch::kLong); + auto index0_trt = index0.to(torch::kInt32); + auto index1_trt = index1.to(torch::kInt32); + auto index2_trt = index2.to(torch::kInt32); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, index0, index1, index2}); + + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, index0_trt, index1_trt, index2_trt}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, ATenIndexTensorNoneIdx1ConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor, + %index0 : Tensor): + %5 : NoneType = prim::Constant() + %18 : Tensor?[] = prim::ListConstruct(%5, %index0) + %19 : Tensor = aten::index(%x.1, %18) + return (%19))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto in1 = at::randint(1, 10, {1, 3, 480, 928}, {at::kCUDA}); + auto index0 = at::tensor({2, 1, 0}, {at::kCUDA}).to(torch::kLong); + + auto index0_trt = index0.to(torch::kInt32); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, index0}); + + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, index0_trt}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} \ No newline at end of file diff --git a/tests/core/conversion/converters/test_masked_fill.cpp b/tests/core/conversion/converters/test_masked_fill.cpp new file mode 100644 index 0000000000..518b31dc02 --- /dev/null +++ b/tests/core/conversion/converters/test_masked_fill.cpp @@ -0,0 +1,99 @@ +#include +#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" + +TEST(Converters, ATenMaskedFillZerosConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %44 : Device = prim::Constant[value="cuda"]() + %8 : bool = prim::Constant[value=0]() + %7 : None = prim::Constant() + %f32_dtype: int = prim::Constant[value=11]() + %1 : int = prim::Constant[value=0]() # bert.py:5:26 + %2 : int = prim::Constant[value=1]() # bert.py:5:32 + %33 : int = prim::Constant[value=2]() # bert.py:6:31 + %3 : int[] = prim::ListConstruct(%1, %1, %2) + %4 : int[] = prim::ListConstruct(%2, %2, %1) + %5 : int[][] = prim::ListConstruct(%3, %4) + %9 : Tensor = aten::tensor(%5, %f32_dtype, %7, %8) # bert.py:5:11 + %mask.1 : Tensor = aten::to(%9, %44, %7, %8, %8) # bert.py:5:11 + %mask.2 : Tensor = trt::const(%mask.1) + %34 : Tensor = aten::masked_fill(%x.1, %mask.1, %33) # bert.py:6:11 + return (%34, %mask.2))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, &*g); + + auto in = at::zeros({1, 2, 3}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + torch_tensorrt::core::lowering::passes::RemoveNOPs(g); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + ASSERT_TRUE( + torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); +} + +TEST(Converters, ATenMaskedFillMixedTypesFloatIntConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor, %x.2 : Tensor): + %val : float = prim::Constant[value=4.0]() + %out : Tensor = aten::masked_fill(%x.1, %x.2, %val) + return (%out))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, &*g); + + // Input is a float tensor, filled with an int --> expecting float tensor out + auto in1 = at::rand({2, 3, 5, 7}, {at::kCUDA}).to(torch::kFloat32); + auto in2 = (2 * at::rand({2, 3, 5, 7}, {at::kCUDA})).to(torch::kBool); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, in2}); + + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, in2}); + + ASSERT_TRUE( + torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); + + // Ensure data types match in outputs + ASSERT_TRUE(jit_results[0].dtype() == trt_results[0].dtype()); +} + +TEST(Converters, ATenMaskedFillMixedTypesIntFloatConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor, %x.2 : Tensor): + %val : int = prim::Constant[value=4]() + %out : Tensor = aten::masked_fill(%x.1, %x.2, %val) + return (%out))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, &*g); + + // Input is an integer tensor, filled with a float --> expecting integer tensor out + auto in1 = at::rand({1, 3, 5, 7}, {at::kCUDA}).to(torch::kInt32); + auto in2 = (2 * at::rand({1, 3, 5, 7}, {at::kCUDA})).to(torch::kBool); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, in2}); + + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, in2}); + + ASSERT_TRUE( + torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); + + // Ensure data types match in outputs + ASSERT_TRUE(jit_results[0].dtype() == trt_results[0].dtype()); +} \ No newline at end of file diff --git a/tests/core/conversion/converters/test_reduce.cpp b/tests/core/conversion/converters/test_reduce.cpp index 47e8b8d154..3cdb2d3b84 100644 --- a/tests/core/conversion/converters/test_reduce.cpp +++ b/tests/core/conversion/converters/test_reduce.cpp @@ -392,240 +392,3 @@ TEST(Converters, ATenAllDimDynamicConvertsCorrectly) { auto in = at::randint(0, 2, {64, 2}, at::kCUDA).to(torch::kHalf); test_body(graph, in, true); } - -TEST(Converters, UnpackVarLowersCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 - %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 - %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 - %6 : int[] = prim::ListConstruct(%3) - %7 : Tensor = aten::var(%x.1, %6, %5, %4) # test_zeros.py:10:26 - return (%7))IR"; - - auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); - - in = at::clone(in); - torch_tensorrt::core::lowering::passes::UnpackVar(g); - torch::jit::EliminateCommonSubexpression(g); - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} - -TEST(Converters, UnpackVarKeepDimsLowersCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 - %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 - %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 - %6 : int[] = prim::ListConstruct(%3) - %7 : Tensor = aten::var(%x.1, %6, %5, %5) # test_zeros.py:10:26 - return (%7))IR"; - - auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); - - in = at::clone(in); - torch_tensorrt::core::lowering::passes::UnpackVar(g); - torch::jit::EliminateCommonSubexpression(g); - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} - -TEST(Converters, UnpackVarUnbiasedLowersCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 - %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 - %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 - %6 : int[] = prim::ListConstruct(%3) - %7 : Tensor = aten::var(%x.1, %6, %4, %4) # test_zeros.py:10:26 - return (%7))IR"; - - auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); - - in = at::clone(in); - torch_tensorrt::core::lowering::passes::UnpackVar(g); - torch::jit::EliminateCommonSubexpression(g); - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} - -TEST(Converters, UnpackVarUnbiasedKeepDimsLowersCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 - %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 - %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 - %6 : int[] = prim::ListConstruct(%3) - %7 : Tensor = aten::var(%x.1, %6, %4, %5) # test_zeros.py:10:26 - return (%7))IR"; - - auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); - - in = at::clone(in); - torch_tensorrt::core::lowering::passes::UnpackVar(g); - torch::jit::EliminateCommonSubexpression(g); - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} - -TEST(Converters, UnpackStdLowersCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 - %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 - %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 - %6 : int[] = prim::ListConstruct(%3) - %7 : Tensor = aten::std(%x.1, %6, %5, %4) # test_zeros.py:10:26 - return (%7))IR"; - - auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); - - in = at::clone(in); - torch_tensorrt::core::lowering::passes::UnpackStd(g); - torch_tensorrt::core::lowering::passes::UnpackVar(g); - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} - -TEST(Converters, UnpackStdKeepDimsLowersCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 - %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 - %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 - %6 : int[] = prim::ListConstruct(%3) - %7 : Tensor = aten::std(%x.1, %6, %5, %5) # test_zeros.py:10:26 - return (%7))IR"; - - auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); - - in = at::clone(in); - torch_tensorrt::core::lowering::passes::UnpackStd(g); - torch_tensorrt::core::lowering::passes::UnpackVar(g); - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} - -TEST(Converters, UnpackStdUnbiasedLowersCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 - %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 - %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 - %6 : int[] = prim::ListConstruct(%3) - %7 : Tensor = aten::std(%x.1, %6, %4, %4) # test_zeros.py:10:26 - return (%7))IR"; - - auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); - - in = at::clone(in); - torch_tensorrt::core::lowering::passes::UnpackStd(g); - torch_tensorrt::core::lowering::passes::UnpackVar(g); - torch::jit::EliminateCommonSubexpression(g); - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} - -TEST(Converters, UnpackStdUnbiasedKeepDimsLowersCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 - %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 - %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 - %one : int = prim::Constant[value=1]() - %6 : int[] = prim::ListConstruct(%3, %one) - %7 : Tensor = aten::std(%x.1, %6, %4, %5) # test_zeros.py:10:26 - return (%7))IR"; - - auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); - - in = at::clone(in); - torch_tensorrt::core::lowering::passes::UnpackStd(g); - torch_tensorrt::core::lowering::passes::UnpackVar(g); - torch::jit::EliminateCommonSubexpression(g); - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} - -TEST(Converters, UnpackVarUnbiasedNegAxisLowersCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %37 : bool = prim::Constant[value=1]() - %53 : int[] = prim::Constant[value=[-1]]() - %69 : Tensor = aten::var(%x.1, %53, %37, %37) - return (%69))IR"; - - auto in = at::randint(-5, 5, {2, 20, 768}, at::kCUDA).to(at::kFloat); - - auto jit_in = at::clone(in); - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - in = at::clone(in); - torch_tensorrt::core::lowering::passes::UnpackVar(g); - torch::jit::EliminateCommonSubexpression(g); - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {jit_in}); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} diff --git a/tests/core/conversion/converters/test_roll.cpp b/tests/core/conversion/converters/test_roll.cpp new file mode 100644 index 0000000000..693fd47aef --- /dev/null +++ b/tests/core/conversion/converters/test_roll.cpp @@ -0,0 +1,84 @@ +#include +#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" + +TEST(Converters, ATenRollConvertsCorrectly) { + const auto graph = R"IR( + graph(%1 : Tensor): + %2 : int[] = prim::Constant[value=[1, 0, 3, 7]]() + %3 : int[] = prim::Constant[value=[0, 1, 2, 3]]() + %4 : Tensor = aten::roll(%1, %2, %3) + return (%4))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + // Run Pytorch + auto in = at::randint(1, 10, {2, 3, 4, 5}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} + +TEST(Converters, ATenRollShiftsNegativeConvertsCorrectly) { + const auto graph = R"IR( + graph(%1 : Tensor): + %2 : int[] = prim::Constant[value=[0, -3, -3]]() + %3 : int[] = prim::Constant[value=[1, 2, 3]]() + %4 : Tensor = aten::roll(%1, %2, %3) + return (%4))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + // Run Pytorch + auto in = at::randint(1, 10, {1, 3, 4, 5}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} + +TEST(Converters, ATenRollDimsNegativeConvertsCorrectly) { + const auto graph = R"IR( + graph(%1 : Tensor): + %2 : int[] = prim::Constant[value=[0, -3, -3]]() + %3 : int[] = prim::Constant[value=[1, 2, -1]]() + %4 : Tensor = aten::roll(%1, %2, %3) + return (%4))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + // Run Pytorch + auto in = at::randint(1, 10, {1, 3, 4, 5}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} \ No newline at end of file diff --git a/tests/core/conversion/converters/test_scatter.cpp b/tests/core/conversion/converters/test_scatter.cpp new file mode 100644 index 0000000000..b7d0883249 --- /dev/null +++ b/tests/core/conversion/converters/test_scatter.cpp @@ -0,0 +1,79 @@ +#include +#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" + +TEST(Converters, ScatterValueConvertsCorrectly) { + const auto graph = R"IR( + graph(%data : Tensor, + %index.1 : Tensor): + %value : int = prim::Constant[value=100]() + %dim : int = prim::Constant[value=1]() + %5 : NoneType = prim::Constant() + %6 : bool = prim::Constant[value=0]() + %7 : int = prim::Constant[value=4]() + %index : Tensor = aten::to(%index.1, %7, %6, %6, %5) + %10 : Tensor = aten::scatter(%data, %dim, %index, %value) + return (%10))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto index = at::randint(0, 5, {2, 2}, {at::kCUDA}); + auto data = at::randn({5, 5}, {at::kCUDA}); + + auto jit_index = at::clone(index); + auto jit_data = at::clone(data); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_data, jit_index}); + + auto trt_index = at::clone(index); + auto trt_data = at::clone(data); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_data, trt_index}); + + for (size_t i = 0; i < jit_results.size(); i++) { + auto trt = trt_results[i].reshape(jit_results[i].sizes()); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); + } +} + +TEST(Converters, ScatterSrcConvertsCorrectly) { + const auto graph = R"IR( + graph(%data : Tensor, + %src : Tensor, + %index.1 : Tensor): + %dim : int = prim::Constant[value=1]() + %5 : NoneType = prim::Constant() + %6 : bool = prim::Constant[value=0]() + %7 : int = prim::Constant[value=4]() + %index : Tensor = aten::to(%index.1, %7, %6, %6, %5) + %10 : Tensor = aten::scatter(%data, %dim, %index, %src) + return (%10))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto index = at::randint(0, 4, {2, 2}, {at::kCUDA}); + auto data = at::randn({5, 5}, {at::kCUDA}); + auto src = at::randn({2, 2}, {at::kCUDA}); + + auto jit_index = at::clone(index); + auto jit_data = at::clone(data); + auto jit_src = at::clone(src); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_data, jit_src, jit_index}); + + auto trt_index = at::clone(index); + auto trt_data = at::clone(data); + auto trt_src = at::clone(src); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_data, trt_src, trt_index}); + + for (size_t i = 0; i < jit_results.size(); i++) { + auto trt = trt_results[i].reshape(jit_results[i].sizes()); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); + } +} \ No newline at end of file diff --git a/tests/core/conversion/converters/test_select.cpp b/tests/core/conversion/converters/test_select.cpp index d93dd5b2c5..d2af33f099 100644 --- a/tests/core/conversion/converters/test_select.cpp +++ b/tests/core/conversion/converters/test_select.cpp @@ -165,60 +165,6 @@ TEST(Converters, ATenSelectEmptyTensorConvertsCorrectly) { ASSERT_TRUE(torch_tensorrt::tests::util::sameShape(jit_results[0], trt_results[0])); } -TEST(Converters, ATenIndexSelectConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %index : Int (2)): - %2 : int = prim::Constant[value=0]() - %3 : Tensor = aten::index_select(%0, %2, %index) - return (%3))IR"; - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - auto in = at::randint(1, 10, {4, 4, 4}, {at::kCUDA}); - auto index = at::randint(0, 4, {2}, {at::kCUDA}).to(torch::kI32); - - auto jit_in = at::clone(in); - auto jit_index = at::clone(index); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {jit_index}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_index = at::clone(index); - auto trt_params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {trt_index}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, trt_params, {trt_in}); - - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - -TEST(Converters, ATenIndexSelectNegativeDimConvertsCorrectly) { - const auto graph = R"IR( - graph(%0 : Tensor, %index : Int (5)): - %2 : int = prim::Constant[value=-1]() - %3 : Tensor = aten::index_select(%0, %2, %index) - return (%3))IR"; - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {5, 3, 9}, {at::kCUDA}); - auto index = at::randint(0, 9, {5}, {at::kCUDA}).to(torch::kI32); - - auto jit_in = at::clone(in); - auto jit_index = at::clone(index); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {jit_index}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_index = at::clone(index); - auto trt_params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {trt_index}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, trt_params, {trt_in}); - - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - TEST(Converters, ATenNarrowStartScalarConvertsCorrectly) { const auto graph = R"IR( graph(%x.1 : Tensor): @@ -273,1119 +219,3 @@ TEST(Converters, ATenEmbeddingConvertsCorrectly) { ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); } - -TEST(Converters, ATenRollConvertsCorrectly) { - const auto graph = R"IR( - graph(%1 : Tensor): - %2 : int[] = prim::Constant[value=[1, 0, 3, 7]]() - %3 : int[] = prim::Constant[value=[0, 1, 2, 3]]() - %4 : Tensor = aten::roll(%1, %2, %3) - return (%4))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - // Run Pytorch - auto in = at::randint(1, 10, {2, 3, 4, 5}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - -TEST(Converters, ATenRollShiftsNegativeConvertsCorrectly) { - const auto graph = R"IR( - graph(%1 : Tensor): - %2 : int[] = prim::Constant[value=[0, -3, -3]]() - %3 : int[] = prim::Constant[value=[1, 2, 3]]() - %4 : Tensor = aten::roll(%1, %2, %3) - return (%4))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - // Run Pytorch - auto in = at::randint(1, 10, {1, 3, 4, 5}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - -TEST(Converters, ATenRollDimsNegativeConvertsCorrectly) { - const auto graph = R"IR( - graph(%1 : Tensor): - %2 : int[] = prim::Constant[value=[0, -3, -3]]() - %3 : int[] = prim::Constant[value=[1, 2, -1]]() - %4 : Tensor = aten::roll(%1, %2, %3) - return (%4))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - // Run Pytorch - auto in = at::randint(1, 10, {1, 3, 4, 5}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - -TEST(Converters, ATenSliceConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : None = prim::Constant() - %3 : int = prim::Constant[value=2]() - %4 : int = prim::Constant[value=4]() - %5 : int = prim::Constant[value=1]() - %6 : int = prim::Constant[value=0]() - %7 : Tensor = aten::select(%x.1, %6, %6) - %8 : Tensor = aten::select(%7, %6, %5) - %9 : Tensor = aten::slice(%8, %6, %5, %4, %3) - %10 : Tensor = aten::slice(%9, %5, %2, %2, %5) - return (%10))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {1, 3, 5, 5}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - -TEST(Converters, ATenSliceNegStartIndexConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : int = prim::Constant[value=1]() - %3 : int = prim::Constant[value=9223372036854775807]() - %4 : int = prim::Constant[value=-2]() - %5 : int = prim::Constant[value=0]() - %6 : Tensor = aten::slice(%x.1, %5, %4, %3, %2) - %7 : Tensor = aten::slice(%6, %2, %5, %3, %2) - return (%7))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {6, 3}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - -TEST(Converters, ATenSliceNegEndIndexConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : int = prim::Constant[value=3]() - %3 : int = prim::Constant[value=9223372036854775807]() - %4 : int = prim::Constant[value=2]() - %5 : int = prim::Constant[value=-3]() - %6 : int = prim::Constant[value=1]() - %7 : int = prim::Constant[value=-2]() - %8 : int = prim::Constant[value=0]() - %9 : Tensor = aten::slice(%x.1, %8, %8, %7, %6) - %10 : Tensor = aten::slice(%9, %6, %8, %5, %6) - %11 : Tensor = aten::slice(%10, %4, %8, %3, %6) - %12 : Tensor = aten::slice(%11, %2, %8, %3, %6) - return (%12))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {6, 5, 3, 3}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - -TEST(Converters, ATenSliceListConvertsCorrectly) { - const auto graph = R"IR( - graph(%x : Tensor): - %1 : NoneType = prim::Constant() - %2 : int = prim::Constant[value=2]() - %3 : int = prim::Constant[value=1]() - %4 : int = prim::Constant[value=3]() - %list : Tensor[] = aten::unbind(%x, %4) - %slice : Tensor[] = aten::slice(%list, %1, %2, %3) - %out.1 : Tensor, %out.2 : Tensor = prim::ListUnpack(%slice) - return (%out.1, %out.2))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in_x = at::randint(1, 10, {6, 5, 3, 3}, {at::kCUDA}); - - auto jit_in_x = at::clone(in_x); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in_x}); - - auto trt_in_x = at::clone(in_x); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in_x}); - - for (size_t i = 0; i < jit_results.size(); i++) { - auto trt = trt_results[i].reshape(jit_results[i].sizes()); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); - } -} - -TEST(Converters, ATenSliceDynamicBatchConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : None = prim::Constant() - %dim : int = prim::Constant[value=0]() - %start : int = prim::Constant[value=1]() - %end : int = prim::Constant[value=15]() - %step : int = prim::Constant[value=2]() - %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step) - return (%9))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {16, 32}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - // dynamic shape in batch - auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, true); - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - -TEST(Converters, ATenSliceDynamicBatchLargeEndConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : None = prim::Constant() - %dim : int = prim::Constant[value=0]() - %start : int = prim::Constant[value=1]() - %end : int = prim::Constant[value=9223372036854775807]() - %step : int = prim::Constant[value=2]() - %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step) - return (%9))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {16, 32}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - // dynamic shape in batch - auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, true); - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - -TEST(Converters, ATenSliceDynamicNegStartBatchConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : None = prim::Constant() - %dim : int = prim::Constant[value=0]() - %start : int = prim::Constant[value=-15]() - %end : int = prim::Constant[value=15]() - %step : int = prim::Constant[value=2]() - %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step) - return (%9))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {16, 32}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - // dynamic shape in batch - auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, true); - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - -TEST(Converters, ATenSliceDynamicNegEndBatchConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : None = prim::Constant() - %dim : int = prim::Constant[value=0]() - %start : int = prim::Constant[value=1]() - %end : int = prim::Constant[value=-2]() - %step : int = prim::Constant[value=3]() - %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step) - return (%9))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {16, 32}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - // dynamic shape in batch - auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, true); - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - -TEST(Converters, ATenSliceDynamicNoneBatchConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %dim : int = prim::Constant[value=0]() - %start : None = prim::Constant() - %end : None = prim::Constant() - %step : int = prim::Constant[value=3]() - %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step) - return (%9))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {16, 32}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - // dynamic shape in batch - auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, true); - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - -TEST(Converters, ATenSliceDynamicConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : None = prim::Constant() - %dim : int = prim::Constant[value=1]() - %start : int = prim::Constant[value=3]() - %end : int = prim::Constant[value=32]() - %step : int = prim::Constant[value=3]() - %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step) - return (%9))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {16, 32}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - // dynamic shape in dim 1, slice in dim 1 - auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, false); - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - -TEST(Converters, ATenSliceDynamic2ConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : None = prim::Constant() - %dim : int = prim::Constant[value=1]() - %start : int = prim::Constant[value=3]() - %end : int = prim::Constant[value=17]() - %step : int = prim::Constant[value=3]() - %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step) - return (%9))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {16, 32}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - // dynamic shape in batch, slice in dim 1 - auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, true); - auto trt = trt_results[0].reshape(jit_results[0].sizes()); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); -} - -TEST(Converters, ATenSplitSizesInScriptingConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : int[] = prim::Constant[value=[1, 2]]() - %3 : int = prim::Constant[value=1]() - %4 : Tensor[] = aten::split(%x.1, %2, %3) - %x1.1 : Tensor, %x2.1 : Tensor = prim::ListUnpack(%4) - return (%x1.1, %x2.1))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {1, 3, 4, 4}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - - for (size_t i = 0; i < jit_results.size(); i++) { - auto trt = trt_results[i].reshape(jit_results[i].sizes()); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); - } -} - -TEST(Converters, ATenSplitSizesinTracingConvertsCorrectly) { - const auto graph = R"IR( - graph(%argument_1.1 : Tensor): - %2 : int[] = prim::Constant[value=[1, 2]]() - %3 : int = prim::Constant[value=1]() - %4 : Tensor[] = aten::split_with_sizes(%argument_1.1, %2, %3) - %5 : Tensor, %6 : Tensor = prim::ListUnpack(%4) - return (%5, %6))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {1, 3, 4, 4}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - - for (size_t i = 0; i < jit_results.size(); i++) { - auto trt = trt_results[i].reshape(jit_results[i].sizes()); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); - } -} - -TEST(Converters, ATenSplitFixedConvertsCorrectly) { - const auto graph = R"IR( - graph(%argument_1.1 : Tensor): - %2 : int = prim::Constant[value=1]() - %3 : Tensor[] = aten::split(%argument_1.1, %2, %2) - %4 : Tensor, %5 : Tensor, %6 : Tensor = prim::ListUnpack(%3) - return (%4, %5, %6))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {1, 3, 4, 4}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - - for (size_t i = 0; i < jit_results.size(); i++) { - auto trt = trt_results[i].reshape(jit_results[i].sizes()); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); - } -} - -TEST(Converters, ATenSplitFixedHasRemainderConvertsCorrectly) { - const auto graph = R"IR( - graph(%argument_1.1 : Tensor): - %2 : int = prim::Constant[value=2]() - %2.1 : int = prim::Constant[value=1]() - %3 : Tensor[] = aten::split(%argument_1.1, %2, %2.1) - %4 : Tensor, %5 : Tensor, %6 : Tensor = prim::ListUnpack(%3) - return (%4, %5, %6))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, &*g); - - auto in = at::randint(1, 10, {1, 5, 4, 4}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - - for (size_t i = 0; i < jit_results.size(); i++) { - auto trt = trt_results[i].reshape(jit_results[i].sizes()); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); - } -} - -TEST(Converters, ATenSplitAndAddConvertsCorrectly) { - const auto graph = R"IR( - graph(%argument_1.1 : Tensor): - %2 : int = prim::Constant[value=2]() - %2.1 : int = prim::Constant[value=1]() - %3 : Tensor[] = aten::split(%argument_1.1, %2, %2.1) - %4 : Tensor, %5 : Tensor = prim::ListUnpack(%3) - %6 : Tensor = aten::add(%4, %5, %2.1) - return (%6))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, &*g); - - auto in = at::randint(1, 10, {1, 4, 4, 4}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - - for (size_t i = 0; i < jit_results.size(); i++) { - auto trt = trt_results[i].reshape(jit_results[i].sizes()); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); - } -} - -TEST(Converters, ATenSplitNegativeDimsConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : int = prim::Constant[value=1]() - %n1 : int = prim::Constant[value=-1]() - %3 : Tensor[] = aten::split(%x.1, %2, %n1) - %4 : Tensor, %5 : Tensor, %6 : Tensor, %7 : Tensor = prim::ListUnpack(%3) - return (%4, %5, %6, %7))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {1, 3, 4, 4}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - - for (size_t i = 0; i < jit_results.size(); i++) { - auto trt = trt_results[i].reshape(jit_results[i].sizes()); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); - } -} - -TEST(Converters, ATenMaskedFillZerosConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %44 : Device = prim::Constant[value="cuda"]() - %8 : bool = prim::Constant[value=0]() - %7 : None = prim::Constant() - %f32_dtype: int = prim::Constant[value=11]() - %1 : int = prim::Constant[value=0]() # bert.py:5:26 - %2 : int = prim::Constant[value=1]() # bert.py:5:32 - %33 : int = prim::Constant[value=2]() # bert.py:6:31 - %3 : int[] = prim::ListConstruct(%1, %1, %2) - %4 : int[] = prim::ListConstruct(%2, %2, %1) - %5 : int[][] = prim::ListConstruct(%3, %4) - %9 : Tensor = aten::tensor(%5, %f32_dtype, %7, %8) # bert.py:5:11 - %mask.1 : Tensor = aten::to(%9, %44, %7, %8, %8) # bert.py:5:11 - %mask.2 : Tensor = trt::const(%mask.1) - %34 : Tensor = aten::masked_fill(%x.1, %mask.1, %33) # bert.py:6:11 - return (%34, %mask.2))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, &*g); - - auto in = at::zeros({1, 2, 3}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - torch_tensorrt::core::lowering::passes::RemoveNOPs(g); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - - ASSERT_TRUE( - torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); -} - -TEST(Converters, ATenMaskedFillMixedTypesFloatIntConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor, %x.2 : Tensor): - %val : float = prim::Constant[value=4.0]() - %out : Tensor = aten::masked_fill(%x.1, %x.2, %val) - return (%out))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, &*g); - - // Input is a float tensor, filled with an int --> expecting float tensor out - auto in1 = at::rand({2, 3, 5, 7}, {at::kCUDA}).to(torch::kFloat32); - auto in2 = (2 * at::rand({2, 3, 5, 7}, {at::kCUDA})).to(torch::kBool); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, in2}); - - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, in2}); - - ASSERT_TRUE( - torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); - - // Ensure data types match in outputs - ASSERT_TRUE(jit_results[0].dtype() == trt_results[0].dtype()); -} - -TEST(Converters, ATenMaskedFillMixedTypesIntFloatConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor, %x.2 : Tensor): - %val : int = prim::Constant[value=4]() - %out : Tensor = aten::masked_fill(%x.1, %x.2, %val) - return (%out))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, &*g); - - // Input is an integer tensor, filled with a float --> expecting integer tensor out - auto in1 = at::rand({1, 3, 5, 7}, {at::kCUDA}).to(torch::kInt32); - auto in2 = (2 * at::rand({1, 3, 5, 7}, {at::kCUDA})).to(torch::kBool); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, in2}); - - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, in2}); - - ASSERT_TRUE( - torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); - - // Ensure data types match in outputs - ASSERT_TRUE(jit_results[0].dtype() == trt_results[0].dtype()); -} - -TEST(Converters, ATenIndexTensorOneIndiceConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor, - %index : Tensor): - %18 : Tensor?[] = prim::ListConstruct(%index) - %19 : Tensor = aten::index(%x.1, %18) - return (%19))IR"; - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto in1 = at::randint(1, 10, {5, 10}, {at::kCUDA}); - auto in2 = at::full({2}, 4, {at::kCUDA}); - auto options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); - auto in2_trt = at::full({2}, 4, {options}); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, in2}); - - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, in2_trt}); - - ASSERT_TRUE( - torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); -} - -TEST(Converters, ATenIndexTensorFullIndicesConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor, - %index0 : Tensor, - %index1 : Tensor, - %index2 : Tensor): - %18 : Tensor?[] = prim::ListConstruct(%index0, %index1, %index2) - %19 : Tensor = aten::index(%x.1, %18) - return (%19))IR"; - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto in1 = at::randint(1, 10, {5, 10, 4}, {at::kCUDA}); - auto index0 = at::tensor({0, 1, 2, 3}, {at::kCUDA}).to(torch::kLong); - auto index1 = at::tensor({1, 3, 4, 6}, {at::kCUDA}).to(torch::kLong); - auto index2 = at::tensor({3, 2, 1, 0}, {at::kCUDA}).to(torch::kLong); - auto index0_trt = index0.to(torch::kInt32); - auto index1_trt = index1.to(torch::kInt32); - auto index2_trt = index2.to(torch::kInt32); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, index0, index1, index2}); - - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, index0_trt, index1_trt, index2_trt}); - - ASSERT_TRUE( - torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); -} - -TEST(Converters, ATenIndexTensorRepeatedFullIndicesConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor, - %index0 : Tensor, - %index1 : Tensor, - %index2 : Tensor): - %18 : Tensor?[] = prim::ListConstruct(%index0, %index1, %index2) - %19 : Tensor = aten::index(%x.1, %18) - %20 : Tensor = aten::index(%x.1, %18) - return (%19, %20))IR"; - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto in1 = at::randint(1, 10, {5, 10, 4}, {at::kCUDA}); - auto index0 = at::tensor({0, 1, 2, 3}, {at::kCUDA}).to(torch::kLong); - auto index1 = at::tensor({1, 3, 4, 6}, {at::kCUDA}).to(torch::kLong); - auto index2 = at::tensor({3, 2, 1, 0}, {at::kCUDA}).to(torch::kLong); - auto index0_trt = index0.to(torch::kInt32); - auto index1_trt = index1.to(torch::kInt32); - auto index2_trt = index2.to(torch::kInt32); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, index0, index1, index2}); - - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, index0_trt, index1_trt, index2_trt}); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[1], trt_results[1], 2e-6)); -} - -TEST(Converters, ATenIndexTensorIdx0Idx1NoneConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor, - %index0 : Tensor, - %index1 : Tensor): - %5 : NoneType = prim::Constant() - %18 : Tensor?[] = prim::ListConstruct(%index0, %index1, %5) - %19 : Tensor = aten::index(%x.1, %18) - return (%19))IR"; - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto in1 = at::randint(1, 10, {5, 10, 4}, {at::kCUDA}); - auto index0 = at::tensor({0, 1, 2, 3}, {at::kCUDA}).to(torch::kLong); - auto index1 = at::tensor({1, 3, 4, 6}, {at::kCUDA}).to(torch::kLong); - auto index0_trt = index0.to(torch::kInt32); - auto index1_trt = index1.to(torch::kInt32); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, index0, index1}); - - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, index0_trt, index1_trt}); - LOG_DEBUG(trt_results); - - ASSERT_TRUE( - torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); -} - -TEST(Converters, ATenIndexTensorIdx0NoneIdx1ConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor, - %index0 : Tensor, - %index1 : Tensor): - %5 : NoneType = prim::Constant() - %18 : Tensor?[] = prim::ListConstruct(%index0, %5, %index1) - %19 : Tensor = aten::index(%x.1, %18) - return (%19))IR"; - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto in1 = at::randint(1, 10, {5, 10, 4}, {at::kCUDA}); - auto index0 = at::tensor({0, 1, 2, 3}, {at::kCUDA}).to(torch::kLong); - auto index1 = at::tensor({3, 2, 1, 0}, {at::kCUDA}).to(torch::kLong); - auto index0_trt = index0.to(torch::kInt32); - auto index1_trt = index1.to(torch::kInt32); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, index0, index1}); - - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, index0_trt, index1_trt}); - - ASSERT_TRUE( - torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); -} - -TEST(Converters, ATenIndexTensorNoneIdx0Idx1ConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor, - %index0 : Tensor, - %index1 : Tensor): - %5 : NoneType = prim::Constant() - %18 : Tensor?[] = prim::ListConstruct(%5, %index0, %index1) - %19 : Tensor = aten::index(%x.1, %18) - return (%19))IR"; - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto in1 = at::randint(1, 10, {5, 10, 4}, {at::kCUDA}); - auto index0 = at::tensor({0, 1, 2, 3}, {at::kCUDA}).to(torch::kLong); - auto index1 = at::tensor({3, 2, 1, 0}, {at::kCUDA}).to(torch::kLong); - auto index0_trt = index0.to(torch::kInt32); - auto index1_trt = index1.to(torch::kInt32); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, index0, index1}); - - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, index0_trt, index1_trt}); - - ASSERT_TRUE( - torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6)); -} - -TEST(Converters, ATenIndexTensorIdxsNoneConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor, - %index0 : Tensor, - %index1 : Tensor, - %index2 : Tensor): - %5 : NoneType = prim::Constant() - %18 : Tensor?[] = prim::ListConstruct(%index0, %index1, %index2, %5) - %19 : Tensor = aten::index(%x.1, %18) - return (%19))IR"; - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto in1 = at::randint(1, 10, {4, 8, 8, 4}, {at::kCUDA}); - auto index0 = at::full({4, 13, 1}, 1, {at::kCUDA}).to(torch::kLong); - auto index1 = at::full({4, 13, 1}, 2, {at::kCUDA}).to(torch::kLong); - auto index2 = at::full({4, 13, 1}, 3, {at::kCUDA}).to(torch::kLong); - auto index0_trt = index0.to(torch::kInt32); - auto index1_trt = index1.to(torch::kInt32); - auto index2_trt = index2.to(torch::kInt32); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, index0, index1, index2}); - - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, index0_trt, index1_trt, index2_trt}); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} - -TEST(Converters, ATenIndexTensorNoneIdx1ConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor, - %index0 : Tensor): - %5 : NoneType = prim::Constant() - %18 : Tensor?[] = prim::ListConstruct(%5, %index0) - %19 : Tensor = aten::index(%x.1, %18) - return (%19))IR"; - - auto g = std::make_shared(); - torch::jit::parseIR(graph, g.get()); - - auto in1 = at::randint(1, 10, {1, 3, 480, 928}, {at::kCUDA}); - auto index0 = at::tensor({2, 1, 0}, {at::kCUDA}).to(torch::kLong); - - auto index0_trt = index0.to(torch::kInt32); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, index0}); - - params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, index0_trt}); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} - -TEST(Converters, ATenUnbindConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : int = prim::Constant[value=0]() - %3 : Tensor[] = aten::unbind(%x.1, %2) - %o1.1 : Tensor, %o2.1 : Tensor = prim::ListUnpack(%3) - return (%o1.1, %o2.1))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {2, 3, 4, 4}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - - for (size_t i = 0; i < jit_results.size(); i++) { - auto trt = trt_results[i]; - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); - } -} - -TEST(Converters, ATenUnbindNegativeAxisConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : int = prim::Constant[value=-1]() - %3 : Tensor[] = aten::unbind(%x.1, %2) - %o1.1 : Tensor, %o2.1 : Tensor = prim::ListUnpack(%3) - return (%o1.1, %o2.1))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto in = at::randint(1, 10, {5, 2}, {at::kCUDA}); - - auto jit_in = at::clone(in); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); - - auto trt_in = at::clone(in); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); - - for (size_t i = 0; i < jit_results.size(); i++) { - auto trt = trt_results[i]; - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); - } -} - -TEST(Converters, ATenUnbindEvaluatedTensor) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : None = prim::Constant() - %3 : int[] = aten::size(%x.1) - %z.1 : Tensor = aten::zeros(%3, %2, %2, %2, %2) - %5 : int = prim::Constant[value=-1]() - %6 : Tensor[] = aten::unbind(%z.1, %5) - %o1.1 : Tensor, %o2.1 : Tensor = prim::ListUnpack(%6) - return (%o1.1, %o2.1))IR"; - - auto in = at::randint(1, 10, {2}, {at::kCUDA}); - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); - - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); - - for (size_t i = 0; i < jit_results.size(); i++) { - auto trt = trt_results[i]; - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i].cuda(), trt, 2e-6)); - } -} - -TEST(Converters, ScatterValueConvertsCorrectly) { - const auto graph = R"IR( - graph(%data : Tensor, - %index.1 : Tensor): - %value : int = prim::Constant[value=100]() - %dim : int = prim::Constant[value=1]() - %5 : NoneType = prim::Constant() - %6 : bool = prim::Constant[value=0]() - %7 : int = prim::Constant[value=4]() - %index : Tensor = aten::to(%index.1, %7, %6, %6, %5) - %10 : Tensor = aten::scatter(%data, %dim, %index, %value) - return (%10))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto index = at::randint(0, 5, {2, 2}, {at::kCUDA}); - auto data = at::randn({5, 5}, {at::kCUDA}); - - auto jit_index = at::clone(index); - auto jit_data = at::clone(data); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_data, jit_index}); - - auto trt_index = at::clone(index); - auto trt_data = at::clone(data); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_data, trt_index}); - - for (size_t i = 0; i < jit_results.size(); i++) { - auto trt = trt_results[i].reshape(jit_results[i].sizes()); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); - } -} - -TEST(Converters, ScatterSrcConvertsCorrectly) { - const auto graph = R"IR( - graph(%data : Tensor, - %src : Tensor, - %index.1 : Tensor): - %dim : int = prim::Constant[value=1]() - %5 : NoneType = prim::Constant() - %6 : bool = prim::Constant[value=0]() - %7 : int = prim::Constant[value=4]() - %index : Tensor = aten::to(%index.1, %7, %6, %6, %5) - %10 : Tensor = aten::scatter(%data, %dim, %index, %src) - return (%10))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto index = at::randint(0, 4, {2, 2}, {at::kCUDA}); - auto data = at::randn({5, 5}, {at::kCUDA}); - auto src = at::randn({2, 2}, {at::kCUDA}); - - auto jit_index = at::clone(index); - auto jit_data = at::clone(data); - auto jit_src = at::clone(src); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_data, jit_src, jit_index}); - - auto trt_index = at::clone(index); - auto trt_data = at::clone(data); - auto trt_src = at::clone(src); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_data, trt_src, trt_index}); - - for (size_t i = 0; i < jit_results.size(); i++) { - auto trt = trt_results[i].reshape(jit_results[i].sizes()); - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); - } -} - -TEST(Converters, WhereConvertsCorrectly) { - const auto graph = R"IR( - graph(%condition : Tensor, - %x : Tensor, - %y : Tensor): - %out : Tensor = aten::where(%condition, %x, %y) - return (%out))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - auto condition = at::randint(0, 2, {5, 5}, {at::kCUDA}).to(torch::kBool); - auto x = at::randn({5, 5}, {at::kCUDA}); - auto y = at::randn({5, 5}, {at::kCUDA}); - - auto jit_condition = at::clone(condition); - auto jit_x = at::clone(x); - auto jit_y = at::clone(y); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_condition, jit_x, jit_y}); - - auto trt_condition = at::clone(condition); - auto trt_x = at::clone(x); - auto trt_y = at::clone(y); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_condition, trt_x, trt_y}); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} - -TEST(Converters, WhereConvertsMismatchedShapesCorrectly) { - const auto graph = R"IR( - graph(%condition : Tensor, - %x : Tensor, - %y : Tensor): - %out : Tensor = aten::where(%condition, %x, %y) - return (%out))IR"; - - auto g = std::make_shared(); - - torch::jit::parseIR(graph, g.get()); - - // As per Torch behavior, the input Tensors are expected to be broadcasted - // along their respective dimension in the largest-rank Tensor provided - auto condition = at::randint(0, 2, {7, 5}, {at::kCUDA}).to(torch::kBool); - auto x = at::randn({2, 7, 5}, {at::kCUDA}); - auto y = at::randn({5}, {at::kCUDA}); - - auto jit_condition = at::clone(condition); - auto jit_x = at::clone(x); - auto jit_y = at::clone(y); - auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); - auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_condition, jit_x, jit_y}); - - auto trt_condition = at::clone(condition); - auto trt_x = at::clone(x); - auto trt_y = at::clone(y); - auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_condition, trt_x, trt_y}); - - ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); -} diff --git a/tests/core/conversion/converters/test_slice.cpp b/tests/core/conversion/converters/test_slice.cpp new file mode 100644 index 0000000000..83ba879291 --- /dev/null +++ b/tests/core/conversion/converters/test_slice.cpp @@ -0,0 +1,332 @@ +#include +#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" + +TEST(Converters, ATenSliceConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : None = prim::Constant() + %3 : int = prim::Constant[value=2]() + %4 : int = prim::Constant[value=4]() + %5 : int = prim::Constant[value=1]() + %6 : int = prim::Constant[value=0]() + %7 : Tensor = aten::select(%x.1, %6, %6) + %8 : Tensor = aten::select(%7, %6, %5) + %9 : Tensor = aten::slice(%8, %6, %5, %4, %3) + %10 : Tensor = aten::slice(%9, %5, %2, %2, %5) + return (%10))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {1, 3, 5, 5}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} + +TEST(Converters, ATenSliceNegStartIndexConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int = prim::Constant[value=1]() + %3 : int = prim::Constant[value=9223372036854775807]() + %4 : int = prim::Constant[value=-2]() + %5 : int = prim::Constant[value=0]() + %6 : Tensor = aten::slice(%x.1, %5, %4, %3, %2) + %7 : Tensor = aten::slice(%6, %2, %5, %3, %2) + return (%7))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {6, 3}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} + +TEST(Converters, ATenSliceNegEndIndexConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int = prim::Constant[value=3]() + %3 : int = prim::Constant[value=9223372036854775807]() + %4 : int = prim::Constant[value=2]() + %5 : int = prim::Constant[value=-3]() + %6 : int = prim::Constant[value=1]() + %7 : int = prim::Constant[value=-2]() + %8 : int = prim::Constant[value=0]() + %9 : Tensor = aten::slice(%x.1, %8, %8, %7, %6) + %10 : Tensor = aten::slice(%9, %6, %8, %5, %6) + %11 : Tensor = aten::slice(%10, %4, %8, %3, %6) + %12 : Tensor = aten::slice(%11, %2, %8, %3, %6) + return (%12))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {6, 5, 3, 3}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} + +TEST(Converters, ATenSliceListConvertsCorrectly) { + const auto graph = R"IR( + graph(%x : Tensor): + %1 : NoneType = prim::Constant() + %2 : int = prim::Constant[value=2]() + %3 : int = prim::Constant[value=1]() + %4 : int = prim::Constant[value=3]() + %list : Tensor[] = aten::unbind(%x, %4) + %slice : Tensor[] = aten::slice(%list, %1, %2, %3) + %out.1 : Tensor, %out.2 : Tensor = prim::ListUnpack(%slice) + return (%out.1, %out.2))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in_x = at::randint(1, 10, {6, 5, 3, 3}, {at::kCUDA}); + + auto jit_in_x = at::clone(in_x); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in_x}); + + auto trt_in_x = at::clone(in_x); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in_x}); + + for (size_t i = 0; i < jit_results.size(); i++) { + auto trt = trt_results[i].reshape(jit_results[i].sizes()); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); + } +} + +TEST(Converters, ATenSliceDynamicBatchConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : None = prim::Constant() + %dim : int = prim::Constant[value=0]() + %start : int = prim::Constant[value=1]() + %end : int = prim::Constant[value=15]() + %step : int = prim::Constant[value=2]() + %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step) + return (%9))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {16, 32}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + // dynamic shape in batch + auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, true); + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} + +TEST(Converters, ATenSliceDynamicBatchLargeEndConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : None = prim::Constant() + %dim : int = prim::Constant[value=0]() + %start : int = prim::Constant[value=1]() + %end : int = prim::Constant[value=9223372036854775807]() + %step : int = prim::Constant[value=2]() + %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step) + return (%9))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {16, 32}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + // dynamic shape in batch + auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, true); + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} + +TEST(Converters, ATenSliceDynamicNegStartBatchConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : None = prim::Constant() + %dim : int = prim::Constant[value=0]() + %start : int = prim::Constant[value=-15]() + %end : int = prim::Constant[value=15]() + %step : int = prim::Constant[value=2]() + %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step) + return (%9))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {16, 32}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + // dynamic shape in batch + auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, true); + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} + +TEST(Converters, ATenSliceDynamicNegEndBatchConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : None = prim::Constant() + %dim : int = prim::Constant[value=0]() + %start : int = prim::Constant[value=1]() + %end : int = prim::Constant[value=-2]() + %step : int = prim::Constant[value=3]() + %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step) + return (%9))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {16, 32}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + // dynamic shape in batch + auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, true); + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} + +TEST(Converters, ATenSliceDynamicNoneBatchConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %dim : int = prim::Constant[value=0]() + %start : None = prim::Constant() + %end : None = prim::Constant() + %step : int = prim::Constant[value=3]() + %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step) + return (%9))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {16, 32}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + // dynamic shape in batch + auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, true); + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} + +TEST(Converters, ATenSliceDynamicConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : None = prim::Constant() + %dim : int = prim::Constant[value=1]() + %start : int = prim::Constant[value=3]() + %end : int = prim::Constant[value=32]() + %step : int = prim::Constant[value=3]() + %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step) + return (%9))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {16, 32}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + // dynamic shape in dim 1, slice in dim 1 + auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, false); + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} + +TEST(Converters, ATenSliceDynamic2ConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : None = prim::Constant() + %dim : int = prim::Constant[value=1]() + %start : int = prim::Constant[value=3]() + %end : int = prim::Constant[value=17]() + %step : int = prim::Constant[value=3]() + %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step) + return (%9))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {16, 32}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + // dynamic shape in batch, slice in dim 1 + auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, true); + auto trt = trt_results[0].reshape(jit_results[0].sizes()); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); +} \ No newline at end of file diff --git a/tests/core/conversion/converters/test_split.cpp b/tests/core/conversion/converters/test_split.cpp new file mode 100644 index 0000000000..87bd5a16e0 --- /dev/null +++ b/tests/core/conversion/converters/test_split.cpp @@ -0,0 +1,174 @@ +#include +#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" + +TEST(Converters, ATenSplitSizesInScriptingConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int[] = prim::Constant[value=[1, 2]]() + %3 : int = prim::Constant[value=1]() + %4 : Tensor[] = aten::split(%x.1, %2, %3) + %x1.1 : Tensor, %x2.1 : Tensor = prim::ListUnpack(%4) + return (%x1.1, %x2.1))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {1, 3, 4, 4}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + for (size_t i = 0; i < jit_results.size(); i++) { + auto trt = trt_results[i].reshape(jit_results[i].sizes()); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); + } +} + +TEST(Converters, ATenSplitSizesinTracingConvertsCorrectly) { + const auto graph = R"IR( + graph(%argument_1.1 : Tensor): + %2 : int[] = prim::Constant[value=[1, 2]]() + %3 : int = prim::Constant[value=1]() + %4 : Tensor[] = aten::split_with_sizes(%argument_1.1, %2, %3) + %5 : Tensor, %6 : Tensor = prim::ListUnpack(%4) + return (%5, %6))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {1, 3, 4, 4}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + for (size_t i = 0; i < jit_results.size(); i++) { + auto trt = trt_results[i].reshape(jit_results[i].sizes()); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); + } +} + +TEST(Converters, ATenSplitFixedConvertsCorrectly) { + const auto graph = R"IR( + graph(%argument_1.1 : Tensor): + %2 : int = prim::Constant[value=1]() + %3 : Tensor[] = aten::split(%argument_1.1, %2, %2) + %4 : Tensor, %5 : Tensor, %6 : Tensor = prim::ListUnpack(%3) + return (%4, %5, %6))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {1, 3, 4, 4}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + for (size_t i = 0; i < jit_results.size(); i++) { + auto trt = trt_results[i].reshape(jit_results[i].sizes()); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); + } +} + +TEST(Converters, ATenSplitFixedHasRemainderConvertsCorrectly) { + const auto graph = R"IR( + graph(%argument_1.1 : Tensor): + %2 : int = prim::Constant[value=2]() + %2.1 : int = prim::Constant[value=1]() + %3 : Tensor[] = aten::split(%argument_1.1, %2, %2.1) + %4 : Tensor, %5 : Tensor, %6 : Tensor = prim::ListUnpack(%3) + return (%4, %5, %6))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, &*g); + + auto in = at::randint(1, 10, {1, 5, 4, 4}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + for (size_t i = 0; i < jit_results.size(); i++) { + auto trt = trt_results[i].reshape(jit_results[i].sizes()); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); + } +} + +TEST(Converters, ATenSplitAndAddConvertsCorrectly) { + const auto graph = R"IR( + graph(%argument_1.1 : Tensor): + %2 : int = prim::Constant[value=2]() + %2.1 : int = prim::Constant[value=1]() + %3 : Tensor[] = aten::split(%argument_1.1, %2, %2.1) + %4 : Tensor, %5 : Tensor = prim::ListUnpack(%3) + %6 : Tensor = aten::add(%4, %5, %2.1) + return (%6))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, &*g); + + auto in = at::randint(1, 10, {1, 4, 4, 4}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + for (size_t i = 0; i < jit_results.size(); i++) { + auto trt = trt_results[i].reshape(jit_results[i].sizes()); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); + } +} + +TEST(Converters, ATenSplitNegativeDimsConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int = prim::Constant[value=1]() + %n1 : int = prim::Constant[value=-1]() + %3 : Tensor[] = aten::split(%x.1, %2, %n1) + %4 : Tensor, %5 : Tensor, %6 : Tensor, %7 : Tensor = prim::ListUnpack(%3) + return (%4, %5, %6, %7))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {1, 3, 4, 4}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + for (size_t i = 0; i < jit_results.size(); i++) { + auto trt = trt_results[i].reshape(jit_results[i].sizes()); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); + } +} \ No newline at end of file diff --git a/tests/core/conversion/converters/test_unbind.cpp b/tests/core/conversion/converters/test_unbind.cpp new file mode 100644 index 0000000000..0062a055bb --- /dev/null +++ b/tests/core/conversion/converters/test_unbind.cpp @@ -0,0 +1,88 @@ +#include +#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" + +TEST(Converters, ATenUnbindConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int = prim::Constant[value=0]() + %3 : Tensor[] = aten::unbind(%x.1, %2) + %o1.1 : Tensor, %o2.1 : Tensor = prim::ListUnpack(%3) + return (%o1.1, %o2.1))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {2, 3, 4, 4}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + for (size_t i = 0; i < jit_results.size(); i++) { + auto trt = trt_results[i]; + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); + } +} + +TEST(Converters, ATenUnbindNegativeAxisConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : int = prim::Constant[value=-1]() + %3 : Tensor[] = aten::unbind(%x.1, %2) + %o1.1 : Tensor, %o2.1 : Tensor = prim::ListUnpack(%3) + return (%o1.1, %o2.1))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto in = at::randint(1, 10, {5, 2}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + for (size_t i = 0; i < jit_results.size(); i++) { + auto trt = trt_results[i]; + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt, 2e-6)); + } +} + +TEST(Converters, ATenUnbindEvaluatedTensor) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : None = prim::Constant() + %3 : int[] = aten::size(%x.1) + %z.1 : Tensor = aten::zeros(%3, %2, %2, %2, %2) + %5 : int = prim::Constant[value=-1]() + %6 : Tensor[] = aten::unbind(%z.1, %5) + %o1.1 : Tensor, %o2.1 : Tensor = prim::ListUnpack(%6) + return (%o1.1, %o2.1))IR"; + + auto in = at::randint(1, 10, {2}, {at::kCUDA}); + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); + + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); + + for (size_t i = 0; i < jit_results.size(); i++) { + auto trt = trt_results[i]; + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i].cuda(), trt, 2e-6)); + } +} \ No newline at end of file diff --git a/tests/core/conversion/converters/test_unpack.cpp b/tests/core/conversion/converters/test_unpack.cpp new file mode 100644 index 0000000000..9e540723fa --- /dev/null +++ b/tests/core/conversion/converters/test_unpack.cpp @@ -0,0 +1,243 @@ +#include +#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" + +TEST(Converters, UnpackVarLowersCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 + %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 + %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 + %6 : int[] = prim::ListConstruct(%3) + %7 : Tensor = aten::var(%x.1, %6, %5, %4) # test_zeros.py:10:26 + return (%7))IR"; + + auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); + + in = at::clone(in); + torch_tensorrt::core::lowering::passes::UnpackVar(g); + torch::jit::EliminateCommonSubexpression(g); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, UnpackVarKeepDimsLowersCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 + %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 + %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 + %6 : int[] = prim::ListConstruct(%3) + %7 : Tensor = aten::var(%x.1, %6, %5, %5) # test_zeros.py:10:26 + return (%7))IR"; + + auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); + + in = at::clone(in); + torch_tensorrt::core::lowering::passes::UnpackVar(g); + torch::jit::EliminateCommonSubexpression(g); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, UnpackVarUnbiasedLowersCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 + %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 + %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 + %6 : int[] = prim::ListConstruct(%3) + %7 : Tensor = aten::var(%x.1, %6, %4, %4) # test_zeros.py:10:26 + return (%7))IR"; + + auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); + + in = at::clone(in); + torch_tensorrt::core::lowering::passes::UnpackVar(g); + torch::jit::EliminateCommonSubexpression(g); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, UnpackVarUnbiasedKeepDimsLowersCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 + %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 + %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 + %6 : int[] = prim::ListConstruct(%3) + %7 : Tensor = aten::var(%x.1, %6, %4, %5) # test_zeros.py:10:26 + return (%7))IR"; + + auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); + + in = at::clone(in); + torch_tensorrt::core::lowering::passes::UnpackVar(g); + torch::jit::EliminateCommonSubexpression(g); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, UnpackStdLowersCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 + %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 + %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 + %6 : int[] = prim::ListConstruct(%3) + %7 : Tensor = aten::std(%x.1, %6, %5, %4) # test_zeros.py:10:26 + return (%7))IR"; + + auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); + + in = at::clone(in); + torch_tensorrt::core::lowering::passes::UnpackStd(g); + torch_tensorrt::core::lowering::passes::UnpackVar(g); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, UnpackStdKeepDimsLowersCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 + %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 + %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 + %6 : int[] = prim::ListConstruct(%3) + %7 : Tensor = aten::std(%x.1, %6, %5, %5) # test_zeros.py:10:26 + return (%7))IR"; + + auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); + + in = at::clone(in); + torch_tensorrt::core::lowering::passes::UnpackStd(g); + torch_tensorrt::core::lowering::passes::UnpackVar(g); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, UnpackStdUnbiasedLowersCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 + %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 + %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 + %6 : int[] = prim::ListConstruct(%3) + %7 : Tensor = aten::std(%x.1, %6, %4, %4) # test_zeros.py:10:26 + return (%7))IR"; + + auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); + + in = at::clone(in); + torch_tensorrt::core::lowering::passes::UnpackStd(g); + torch_tensorrt::core::lowering::passes::UnpackVar(g); + torch::jit::EliminateCommonSubexpression(g); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, UnpackStdUnbiasedKeepDimsLowersCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %5 : bool = prim::Constant[value=0]() # test_zeros.py:10:65 + %4 : bool = prim::Constant[value=1]() # test_zeros.py:10:50 + %3 : int = prim::Constant[value=0]() # test_zeros.py:10:39 + %one : int = prim::Constant[value=1]() + %6 : int[] = prim::ListConstruct(%3, %one) + %7 : Tensor = aten::std(%x.1, %6, %4, %5) # test_zeros.py:10:26 + return (%7))IR"; + + auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA); + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in}); + + in = at::clone(in); + torch_tensorrt::core::lowering::passes::UnpackStd(g); + torch_tensorrt::core::lowering::passes::UnpackVar(g); + torch::jit::EliminateCommonSubexpression(g); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in}); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, UnpackVarUnbiasedNegAxisLowersCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %37 : bool = prim::Constant[value=1]() + %53 : int[] = prim::Constant[value=[-1]]() + %69 : Tensor = aten::var(%x.1, %53, %37, %37) + return (%69))IR"; + + auto in = at::randint(-5, 5, {2, 20, 768}, at::kCUDA).to(at::kFloat); + + auto jit_in = at::clone(in); + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + in = at::clone(in); + torch_tensorrt::core::lowering::passes::UnpackVar(g); + torch::jit::EliminateCommonSubexpression(g); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {jit_in}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} \ No newline at end of file diff --git a/tests/core/conversion/converters/test_where.cpp b/tests/core/conversion/converters/test_where.cpp new file mode 100644 index 0000000000..23482662bf --- /dev/null +++ b/tests/core/conversion/converters/test_where.cpp @@ -0,0 +1,68 @@ +#include +#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" + +TEST(Converters, WhereConvertsCorrectly) { + const auto graph = R"IR( + graph(%condition : Tensor, + %x : Tensor, + %y : Tensor): + %out : Tensor = aten::where(%condition, %x, %y) + return (%out))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + auto condition = at::randint(0, 2, {5, 5}, {at::kCUDA}).to(torch::kBool); + auto x = at::randn({5, 5}, {at::kCUDA}); + auto y = at::randn({5, 5}, {at::kCUDA}); + + auto jit_condition = at::clone(condition); + auto jit_x = at::clone(x); + auto jit_y = at::clone(y); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_condition, jit_x, jit_y}); + + auto trt_condition = at::clone(condition); + auto trt_x = at::clone(x); + auto trt_y = at::clone(y); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_condition, trt_x, trt_y}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, WhereConvertsMismatchedShapesCorrectly) { + const auto graph = R"IR( + graph(%condition : Tensor, + %x : Tensor, + %y : Tensor): + %out : Tensor = aten::where(%condition, %x, %y) + return (%out))IR"; + + auto g = std::make_shared(); + + torch::jit::parseIR(graph, g.get()); + + // As per Torch behavior, the input Tensors are expected to be broadcasted + // along their respective dimension in the largest-rank Tensor provided + auto condition = at::randint(0, 2, {7, 5}, {at::kCUDA}).to(torch::kBool); + auto x = at::randn({2, 7, 5}, {at::kCUDA}); + auto y = at::randn({5}, {at::kCUDA}); + + auto jit_condition = at::clone(condition); + auto jit_x = at::clone(x); + auto jit_y = at::clone(y); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_condition, jit_x, jit_y}); + + auto trt_condition = at::clone(condition); + auto trt_x = at::clone(x); + auto trt_y = at::clone(y); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_condition, trt_x, trt_y}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} \ No newline at end of file From 90fd8c74431204b3ed289133916643cf324c65cc Mon Sep 17 00:00:00 2001 From: Dheeraj Peri Date: Thu, 20 Apr 2023 11:51:54 -0700 Subject: [PATCH 3/4] chore: Fix missing headers in tests Signed-off-by: Dheeraj Peri --- tests/core/conversion/converters/test_masked_fill.cpp | 1 + tests/core/conversion/converters/test_unpack.cpp | 2 ++ tests/core/conversion/converters/test_where.cpp | 1 + 3 files changed, 4 insertions(+) diff --git a/tests/core/conversion/converters/test_masked_fill.cpp b/tests/core/conversion/converters/test_masked_fill.cpp index 518b31dc02..2c375463e5 100644 --- a/tests/core/conversion/converters/test_masked_fill.cpp +++ b/tests/core/conversion/converters/test_masked_fill.cpp @@ -1,3 +1,4 @@ +#include #include #include "core/compiler.h" #include "core/lowering/passes/passes.h" diff --git a/tests/core/conversion/converters/test_unpack.cpp b/tests/core/conversion/converters/test_unpack.cpp index 9e540723fa..858462b003 100644 --- a/tests/core/conversion/converters/test_unpack.cpp +++ b/tests/core/conversion/converters/test_unpack.cpp @@ -4,6 +4,8 @@ #include "gtest/gtest.h" #include "tests/util/util.h" #include "torch/csrc/jit/ir/irparser.h" +#include "torch/csrc/jit/passes/common_subexpression_elimination.h" +#include "torch/torch.h" TEST(Converters, UnpackVarLowersCorrectly) { const auto graph = R"IR( diff --git a/tests/core/conversion/converters/test_where.cpp b/tests/core/conversion/converters/test_where.cpp index 23482662bf..34b3696582 100644 --- a/tests/core/conversion/converters/test_where.cpp +++ b/tests/core/conversion/converters/test_where.cpp @@ -1,3 +1,4 @@ +#include #include #include "core/compiler.h" #include "core/lowering/passes/passes.h" From 76f6a09a8ccb3c6ec80117fa4dbf12543bcc0a1d Mon Sep 17 00:00:00 2001 From: Dheeraj Peri Date: Thu, 20 Apr 2023 12:49:42 -0700 Subject: [PATCH 4/4] chore: fix headers Signed-off-by: Dheeraj Peri --- tests/core/conversion/converters/test_index.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/core/conversion/converters/test_index.cpp b/tests/core/conversion/converters/test_index.cpp index 34e50f2abd..b405d7a436 100644 --- a/tests/core/conversion/converters/test_index.cpp +++ b/tests/core/conversion/converters/test_index.cpp @@ -1,3 +1,4 @@ +#include #include #include "core/compiler.h" #include "core/lowering/passes/passes.h"