[SR] Native implementation for IntImplicit

mikeiovine · mikeiovine · commit a9bc04f77e9f · 2022-03-22T11:06:19.000-07:00
Add a native implementation for `aten::IntImplicit`, which is similar to `aten::Int` except for a few extra checks it must do Differential Revision: [D35052997](https://our.internmc.facebook.com/intern/diff/D35052997/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D35052997/)! ghstack-source-id: 151910134 Pull Request resolved: #74562
diff --git a/benchmarks/static_runtime/test_static_runtime.cc b/benchmarks/static_runtime/test_static_runtime.cc
@@ -3015,3 +3015,27 @@ TEST(StaticRuntime, ConcatEmpty) {
   torch::jit::StaticModule smod(mod);
   EXPECT_THROW(smod({}), c10::Error);
 }
+
+TEST(StaticRuntime, IntImplicit) {
+  const auto src = R"IR(
+    graph(%a: Tensor):
+        %y: int = aten::IntImplicit(%a)
+        return (%y)
+  )IR";
+  testStaticRuntime(src, {at::tensor({1}, at::kInt).squeeze()});
+}
+
+TEST(StaticRuntime, IntImplicit_ThrowOnBadInputs) {
+  const auto src = R"IR(
+    graph(%a: Tensor):
+        %y: int = aten::IntImplicit(%a)
+        return (%y)
+  )IR";
+  auto graph = getGraphFromIR(src);
+  torch::jit::StaticModule smod(graph);
+  // Not 0D tensor
+  EXPECT_THROW(smod({at::tensor({1, 2}, at::kInt)}), std::runtime_error);
+  // Wrong dtype
+  EXPECT_THROW(
+      smod({at::tensor({1}, at::kFloat).squeeze()}), std::runtime_error);
+}
diff --git a/torch/csrc/jit/runtime/static/native_ops.cpp b/torch/csrc/jit/runtime/static/native_ops.cpp
@@ -1004,5 +1004,31 @@ REGISTER_NATIVE_OPERATOR_FUNCTOR(
       return nullptr;
     });
 
+REGISTER_NATIVE_OPERATOR_FUNCTOR(
+    aten::IntImplicit,
+    aten_IntImplicit,
+    [](Node* n) -> SROperator {
+      if (!n->matches(torch::schema("aten::IntImplicit(Tensor a) -> int"))) {
+        LogAndDumpSchema(n);
+        return nullptr;
+      }
+      return [](ProcessedNode* pnode) {
+        auto& tensor = pnode->Input(0).toTensor();
+        // JIT does a check for requires_grad, but we skip it here since SR is
+        // inference only
+        if (tensor.sizes().size() != 0) {
+          throw std::runtime_error(
+              "Cannot convert a tensor of dimension > 0 to scalar");
+        }
+        if (!isIntegralType(tensor.scalar_type())) {
+          std::stringstream ss;
+          ss << "Cannot input a tensor of type " << tensor.scalar_type()
+             << " as an integral argument";
+          throw std::runtime_error(ss.str());
+        }
+        pnode->Output(0) = at::native::item(tensor).toInt();
+      };
+    });
+
 } // namespace jit
 } // namespace torch
