Add testing for NumToTensor, Full, and Masked Fill lowering passes

- Generalize a case for NumToTensor in lowering pass

Author: gs-olive

core/lowering/passes/device_casting.cpp

@@ -34,7 +34,7 @@ void UnpackAndCastMaskedFill(std::shared_ptr<torch::jit::Graph>& graph) {
 
 void UnpackAndCastNumToTensor(std::shared_ptr<torch::jit::Graph>& graph) {
   std::string num_to_tensor_cast_pattern = R"IR(
-    graph(%1: int):
+    graph(%1: Scalar):
       %2: Tensor = prim::NumToTensor(%1)
       return (%2))IR";
 

tests/core/lowering/BUILD

@@ -31,6 +31,10 @@ lowering_test(
     name = "test_conv1d_pass",
 )
 
+lowering_test(
+    name = "test_device_casting",
+)
+
 lowering_test(
     name = "test_exception_elimination_pass",
 )
@@ -91,6 +95,7 @@ test_suite(
     name = "lowering_tests",
     tests = [
         ":test_conv1d_pass",
+        ":test_device_casting",
         ":test_exception_elimination_pass",
         ":test_linear_to_addmm",
         ":test_module_fallback_passes",

tests/core/lowering/test_device_casting.cpp

@@ -0,0 +1,118 @@
+#include <string>
+#include "core/compiler.h"
+#include "core/lowering/passes/passes.h"
+#include "core/util/prelude.h"
+#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(LoweringPasses, UnpackAndCastMaskedFillLowersCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.1: Tensor, %x.2: Tensor, %x.3: float):
+        %2 : Tensor = aten::masked_fill_(%x.1, %x.2, %x.3)
+        return (%2))IR";
+
+  auto in = at::rand({2, 3, 5, 7}, {at::kCUDA});
+  auto in2 = at::rand({2, 3, 5, 7}, {at::kCUDA}).to(torch::kBool);
+  auto in3 = 7.3;
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto jit_pre_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in, in2, in3});
+  torch_tensorrt::core::lowering::passes::UnpackAndCastMaskedFill(g);
+  torch::jit::EliminateCommonSubexpression(g);
+  auto jit_post_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in, in2, in3});
+
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_pre_results[0].toTensor(), jit_post_results[0].toTensor(), 2e-6));
+}
+
+TEST(LoweringPasses, UnpackAndCastNumToTensorLowersIntCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.1: int):
+        %2 : Tensor = prim::NumToTensor(%x.1)
+        return (%2))IR";
+
+  // Make range [0.01, 1.01] to ensure positives / avoid NaN with negative sqrt
+  auto in = 1;
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto jit_pre_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in});
+  torch_tensorrt::core::lowering::passes::UnpackAndCastNumToTensor(g);
+  torch::jit::EliminateCommonSubexpression(g);
+  auto jit_post_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in});
+
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_pre_results[0].toTensor(), jit_post_results[0].toTensor(), 2e-6));
+}
+
+TEST(LoweringPasses, UnpackAndCastNumToTensorLowersFloatCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.1: float):
+        %2 : Tensor = prim::NumToTensor(%x.1)
+        return (%2))IR";
+
+  auto in = 78.1;
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto jit_pre_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in});
+  torch_tensorrt::core::lowering::passes::UnpackAndCastNumToTensor(g);
+  torch::jit::EliminateCommonSubexpression(g);
+  auto jit_post_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in});
+
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_pre_results[0].toTensor(), jit_post_results[0].toTensor(), 2e-6));
+}
+
+TEST(LoweringPasses, UnpackAndCastFullIntLowersCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.1: int):
+        %5 : NoneType = prim::Constant()
+        %2 : int = prim::Constant[value=3]()
+        %10 : int[] = prim::ListConstruct(%2, %2)
+        %out : Tensor = aten::full(%10, %x.1, %5, %5, %5, %5)
+        return (%out))IR";
+
+  auto in = 4;
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto jit_pre_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in});
+  torch_tensorrt::core::lowering::passes::UnpackAndCastFull(g);
+  torch::jit::EliminateCommonSubexpression(g);
+  auto jit_post_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in});
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(
+      jit_pre_results[0].toTensor(), jit_post_results[0].toTensor().cpu(), 2e-6));
+}
+
+TEST(LoweringPasses, UnpackAndCastFullFloatLowersCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.1: float):
+        %5 : NoneType = prim::Constant()
+        %2 : int = prim::Constant[value=5]()
+        %3 : int = prim::Constant[value=4]()
+        %10 : int[] = prim::ListConstruct(%2, %3)
+        %out : Tensor = aten::full(%10, %x.1, %5, %5, %5, %5)
+        return (%out))IR";
+
+  auto in = 54.1;
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto jit_pre_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in});
+  torch_tensorrt::core::lowering::passes::UnpackAndCastFull(g);
+  torch::jit::EliminateCommonSubexpression(g);
+  auto jit_post_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in});
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(
+      jit_pre_results[0].toTensor(), jit_post_results[0].toTensor().cpu(), 2e-6));
+}