Expand coverage of unary ops

Added support for:
aten::cos
aten::acos
aten::cosh
aten::sin
aten::asin
aten::sinh
aten::tan
aten::atan
aten::abs
aten::floor
aten::reciprocal
aten::ceil
aten::sqrt
aten::exp
aten::neg

Signed-off-by: Junjie Bai <[email protected]>

Author: bddppq

core/conversion/converters/impl/unary.cpp

@@ -8,28 +8,53 @@ namespace converters {
 namespace impl {
 namespace {
 
-auto unary_registrations = RegisterNodeConversionPatterns()
-    .pattern({
-        "aten::log(Tensor self) -> Tensor",
-        [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
-            auto in = args[0].ITensor();
-            auto log = ctx->net->addUnary(*in, nvinfer1::UnaryOperation::kLOG);
+#define convert(unary, trt_type)                                               \
+  auto unary##_registrations TRTORCH_UNUSED =                                  \
+      RegisterNodeConversionPatterns().pattern(                                \
+          {"aten::" #unary "(Tensor self) -> Tensor",                          \
+           [](ConversionCtx *ctx, const torch::jit::Node *n,                   \
+              args &args) -> bool {                                            \
+             auto in = args[0].ITensor();                                      \
+             auto unary =                                                      \
+                 ctx->net->addUnary(*in, nvinfer1::UnaryOperation::trt_type);  \
+                                                                               \
+             TRTORCH_CHECK(                                                    \
+                 unary,                                                        \
+                 "Unable to create " #unary " layer from node: " << *n);       \
+                                                                               \
+             unary->setName(util::node_info(n).c_str());                       \
+             auto out_value = n->outputs()[0];                                 \
+             auto out_tensor = unary->getOutput(0);                            \
+             out_tensor->setName(out_value->debugName().c_str());              \
+             ctx->value_tensor_map[out_value] = out_tensor;                    \
+             LOG_DEBUG(                                                        \
+                 "Output tensor shape: " << out_tensor->getDimensions());      \
+                                                                               \
+             return true;                                                      \
+           }});
 
-            TRTORCH_CHECK(log, "Unable to create log layer from node: " << *n);
+convert(cos, kCOS);
+convert(acos, kACOS);
+convert(cosh, kCOSH);
+convert(sin, kSIN);
+convert(asin, kASIN);
+convert(sinh, kSINH);
+convert(tan, kTAN);
+convert(atan, kATAN);
+convert(abs, kABS);
+convert(floor, kFLOOR);
+convert(reciprocal, kRECIP);
+convert(log, kLOG);
+convert(ceil, kCEIL);
+convert(sqrt, kSQRT);
+convert(exp, kEXP);
+convert(neg, kNEG);
+
+#undef convert
 
-            log->setName(util::node_info(n).c_str());
-            auto out_value = n->outputs()[0];
-            auto out_tensor = log->getOutput(0);
-            out_tensor->setName(out_value->debugName().c_str());
-            ctx->value_tensor_map[out_value] = out_tensor;
-            LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
-            
-            return true;
-        }
-    });
 } // namespace
 } // namespace impl
 } // namespace converters
 } // namespace conversion
 } // namespace core
-} // namespace trtorch 
+} // namespace trtorch

tests/core/converters/test_unary.cpp

@@ -4,22 +4,51 @@
 #include "tests/util/util.h"
 #include "core/compiler.h"
 
-TEST(Converters, ATenLogConvertsCorrectly) {
-    const auto graph = R"IR(
+namespace {
+std::string gen_test_graph(const std::string &unary) {
+  return R"IR(
       graph(%0 : Tensor):
-        %3 : Tensor = aten::log(%0)
+        %3 : Tensor = aten::)IR" +
+         unary + R"IR((%0)
         return (%3))IR";
+}
+} // namespace
 
-    auto g = std::make_shared<torch::jit::Graph>();
-    torch::jit::script::parseIR(graph, &*g);
-    
-    auto in = at::randint(1, 5, {5}, {at::kCUDA});
-    auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
-    auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
+#define test_unary(unary, name)                                                \
+  TEST(Converters, ATen##name##ConvertsCorrectly) {                            \
+    const auto graph = gen_test_graph(#unary);                                 \
+                                                                               \
+    auto g = std::make_shared<torch::jit::Graph>();                            \
+    torch::jit::script::parseIR(graph, &*g);                                   \
+                                                                               \
+    auto in = at::empty({10}, {at::kCUDA}).uniform_(0, 0.5);                   \
+    auto params =                                                              \
+        trtorch::core::conversion::get_named_params(g->inputs(), {});          \
+    auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});        \
+                                                                               \
+    in = at::clone(in);                                                        \
+    params = trtorch::core::conversion::get_named_params(g->inputs(), {});     \
+    auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});  \
+                                                                               \
+    ASSERT_TRUE(                                                               \
+        trtorch::tests::util::almostEqual(jit_results[0], trt_results[0]));    \
+  }
 
-    in = at::clone(in);
-    params = trtorch::core::conversion::get_named_params(g->inputs(), {});
-    auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
+test_unary(cos, Cos);
+test_unary(acos, Acos);
+test_unary(cosh, Cosh);
+test_unary(sin, Sin);
+test_unary(asin, Asin);
+test_unary(sinh, Sinh);
+test_unary(tan, Tan);
+test_unary(atan, Atan);
+test_unary(abs, Abs);
+test_unary(floor, Floor);
+test_unary(reciprocal, Reciprocal);
+test_unary(log, Log);
+test_unary(ceil, Ceil);
+test_unary(sqrt, Sqrt);
+test_unary(exp, Exp);
+test_unary(neg, Neg);
 
-    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0]));
-}
+#undef test_unary