Merge pull request #310 from inocsin/div_scalar

support div.Scalar(Tensor self, Scalar other)

Author: narendasan

core/conversion/converters/impl/element_wise.cpp

@@ -213,6 +213,20 @@ auto element_wise_registrations TRTORCH_UNUSED =
                     div->setName(util::node_info(n).c_str());
                     auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0));
 
+                    LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+                    return true;
+                  }})
+        .pattern({"aten::div.Scalar(Tensor self, Scalar other) -> (Tensor)",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    auto self = args[0].ITensorOrFreeze(ctx);
+                    auto otherScalar = args[1].unwrapToScalar().to<float>();
+                    auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+                    auto div =
+                        add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n));
+                    TRTORCH_CHECK(div, "Unable to create div layer from node: " << *n);
+
+                    div->setName(util::node_info(n).c_str());
+                    auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0));
                     LOG_DEBUG("Output tensor shape: " << out->getDimensions());
                     return true;
                   }})
@@ -229,6 +243,20 @@ auto element_wise_registrations TRTORCH_UNUSED =
                     div->setName(util::node_info(n).c_str());
                     auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0));
 
+                    LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+                    return true;
+                  }})
+        .pattern({"aten::div_.Scalar(Tensor self, Scalar other) -> (Tensor)",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    auto self = args[0].ITensorOrFreeze(ctx);
+                    auto otherScalar = args[1].unwrapToScalar().to<float>();
+                    auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+                    auto div =
+                        add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n));
+                    TRTORCH_CHECK(div, "Unable to create div layer from node: " << *n);
+
+                    div->setName(util::node_info(n).c_str());
+                    auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0));
                     LOG_DEBUG("Output tensor shape: " << out->getDimensions());
                     return true;
                   }})

core/conversion/evaluators/aten.cpp

@@ -426,7 +426,8 @@ auto aten_registrations TRTORCH_UNUSED =
                       }
                     },
                     EvalOptions().validSchemas({
-                        "aten::div.Scalar(Scalar a, Scalar b) -> (float)",
+                        "aten::div.float(float a, float b) -> (float)",
+                        "aten::div.int(int a, int b) -> (float)",
                     })})
         .evaluator({c10::Symbol::fromQualString("aten::floordiv"),
                     [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {

tests/core/conversion/converters/test_element_wise.cpp

@@ -123,6 +123,15 @@ TEST(Converters, ATenDivConvertsCorrectly) {
   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, ATenPowTensorConvertsCorrectly) {
   const auto graph = R"IR(
        graph(%x.1 : Tensor, %x2.1 : Tensor):

tests/core/conversion/evaluators/BUILD

@@ -11,9 +11,14 @@ evaluator_test(
   name = "test_prim_evaluators",
 )
 
+evaluator_test(
+  name = "test_aten_evaluators",
+)
+
 test_suite(
     name = "evaluator_tests",
     tests = [
-        ":test_prim_evaluators"
+        ":test_prim_evaluators",
+        ":test_aten_evaluators"
     ]
 )

tests/core/conversion/evaluators/test_aten_evaluators.cpp

@@ -0,0 +1,39 @@
+#include <string>
+#include "core/compiler.h"
+#include "gtest/gtest.h"
+#include "tests/util/util.h"
+#include "torch/csrc/jit/ir/irparser.h"
+
+TEST(Evaluators, DivIntEvaluatesCorrectly) {
+  const auto graph = R"IR(
+      graph():
+        %1 : int = prim::Constant[value=9]()
+        %2 : int = prim::Constant[value=4]()
+        %3 : float = aten::div(%1, %2)
+        return (%3))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto jit_results = trtorch::tests::util::EvaluateGraphJIT(g, {});
+  auto trt_results = trtorch::tests::util::EvaluateGraph(g->block(), {});
+
+  ASSERT_TRUE(jit_results[0] == trt_results[0]);
+}
+
+TEST(Evaluators, DivFloatEvaluatesCorrectly) {
+  const auto graph = R"IR(
+      graph():
+        %1 : float = prim::Constant[value=9.1]()
+        %2 : float = prim::Constant[value=4.2]()
+        %3 : float = aten::div(%1, %2)
+        return (%3))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto jit_results = trtorch::tests::util::EvaluateGraphJIT(g, {});
+  auto trt_results = trtorch::tests::util::EvaluateGraph(g->block(), {});
+
+  ASSERT_TRUE(jit_results[0] == trt_results[0]);
+}