Merge pull request #719 from guoruoqian/support_aten_format

feat: support aten::format evaluator

Author: narendasan

core/conversion/evaluators/aten.cpp

@@ -706,7 +706,23 @@ auto aten_registrations TORCHTRT_UNUSED =
                     },
                     EvalOptions().validSchemas({
                         R"SIG(aten::copy_(Tensor(a!) self, Tensor src, bool non_blocking=False) -> (Tensor(a!)))SIG",
-                    })});
+                    })})
+        .evaluator({c10::Symbol::fromQualString("aten::format"),
+                    [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+                      int64_t input_num = n->inputs().size();
+                      std::vector<torch::jit::IValue> stack;
+                      for (auto v : n->inputs()) {
+                        stack.push_back(*args.at(v).IValue());
+                      }
+                      stack.push_back(input_num);
+                      auto& ops = torch::jit::getAllOperatorsFor(c10::Symbol::fromQualString("aten::format"));
+                      auto& aten_format = ops.front();
+                      aten_format->getOperation()(stack);
+                      std::string output;
+                      torch::jit::pop(stack, output);
+                      return output;
+                    },
+                    EvalOptions().validSchemas({"aten::format(str self, ...) -> (str)"})});
 } // namespace
 } // namespace evaluators
 } // namespace conversion

tests/core/conversion/evaluators/test_aten_evaluators.cpp

@@ -3,6 +3,7 @@
 #include "gtest/gtest.h"
 #include "tests/util/util.h"
 #include "torch/csrc/jit/ir/irparser.h"
+#include "torch/csrc/jit/runtime/jit_exception.h"
 #include "torch/torch.h"
 
 TEST(Evaluators, DivIntEvaluatesCorrectly) {
@@ -579,4 +580,89 @@ TEST(Evaluators, AndBoolResultIsFalseEvaluatesCorrectly) {
   auto trt_results = torch_tensorrt::tests::util::EvaluateGraph(g->block(), {});
 
   ASSERT_TRUE(jit_results[0] == trt_results[0]);
+}
+
+TEST(Evaluators, AtenFormatEvaluatesCorrectly) {
+  const auto graph = R"IR(
+      graph(%x_1 : Tensor, %x_2 : Tensor):
+        %0 : int = prim::Constant[value=1]()
+        %1 : str = prim::Constant[value="res{}_{}_"]()
+        %2 : int = prim::Constant[value=5]()
+        %2.1 : int = prim::Constant[value=2]()
+        %3 : str = prim::Constant[value="res5_2_"]()
+        %4 : str = aten::format(%1, %2, %2.1)
+        %5 : bool = aten::eq(%3, %4)
+        %y : Tensor = prim::If(%5)
+            block0():
+                %194 : Tensor = aten::add(%x_1, %x_2, %0)
+                -> (%194)
+            block1():
+                %195 : Tensor = aten::sub(%x_1, %x_2, %0)
+                -> (%195)
+        return (%y))IR";
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto in0 = at::randint(1, 10, {3, 4}, {at::kCUDA});
+  auto in1 = in0.clone();
+
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in0, in1});
+
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in0, in1});
+
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}
+
+TEST(Evaluators, AtenFormatRaiseExceptionEvaluatesCorrectly) {
+  const auto graph = R"IR(
+      graph(%x_1 : Tensor, %x_2 : Tensor):
+        %0 : int = prim::Constant[value=1]()
+        %1 : str = prim::Constant[value="res5_1"]()
+        %2 : str = prim::Constant[value="{} is not equal to {}"]()
+        %3 : str = prim::Constant[value="res5_2"]()
+        %5713 : Tensor = prim::Uninitialized()
+        %4 : str = aten::format(%2, %1, %3)
+        %5 : bool = aten::eq(%1, %3)
+        %y : Tensor = prim::If(%5)
+            block0():
+                %194 : Tensor = aten::add(%x_1, %x_2, %0)
+                -> (%194)
+            block1():
+                prim::RaiseException(%4)
+                -> (%5713)
+        return (%y))IR";
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto in0 = at::randint(1, 10, {3, 4}, {at::kCUDA});
+  auto in1 = in0.clone();
+
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  std::vector<at::Tensor> jit_results, trt_results;
+  std::string error_jit, error_torch_trt;
+  try {
+    jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in0, in1});
+  } catch (const torch::jit::JITException& error) {
+    error_jit = error.what();
+  }
+
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  try {
+    trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in0, in1});
+  } catch (const torch_tensorrt::Error& error) {
+    error_torch_trt = error.what();
+  }
+
+  auto position1 = error_jit.find("RuntimeError:");
+  auto position2 = error_torch_trt.find("Error from TorchScript:");
+  std::string jit_msg = error_jit.substr(position1 + 13);
+  std::string torch_trt_msg = error_torch_trt.substr(position2 + 23);
+  if (jit_msg == torch_trt_msg) {
+    ASSERT_TRUE(true);
+  } else {
+    ASSERT_TRUE(false);
+  }
 }