Refactor NnapiCompilation registration into it's own file (#63183)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63183

Move registration of NnapiCompilation into it's own file, so that `nnapi_bind.cpp` (which contains the implementation of NnapiCompilation) can be moved to `aten_cpu`, while maintaining the selectiveness for registration.

`nnapi_bind.cpp` is moved to `aten_cpu` in https://github.com/pytorch/pytorch/pull/62919. See the PR for more details on why it's needed.

ghstack-source-id: 135900318

Test Plan: Nnapi unit tests: `python test/test_nnapi.py`

Reviewed By: iseeyuan

Differential Revision: D30288708

fbshipit-source-id: 6ed5967fa6bd018075469d18e68f844d413cf265

diff --git a/aten/src/ATen/nnapi/nnapi_bind.cpp b/aten/src/ATen/nnapi/nnapi_bind.cpp
index 7b11c8c..0ab8997 100644 (file)
--- a/aten/src/ATen/nnapi/nnapi_bind.cpp
+++ b/aten/src/ATen/nnapi/nnapi_bind.cpp
@@ -1,7 +1,6 @@
 #include <vector>
 
 #include <ATen/ATen.h>
-
 #include <ATen/nnapi/nnapi_bind.h>
 #include <ATen/nnapi/nnapi_wrapper.h>
 #include <ATen/nnapi/nnapi_model_loader.h>
@@ -178,32 +177,6 @@ void NnapiCompilation::get_operand_type(const at::Tensor& t, ANeuralNetworksOper
   CAFFE_THROW("Bad dtype");
 }
 
-// Set flag if running on ios
-#ifdef __APPLE__
-  #include <TargetConditionals.h>
-  #if TARGET_OS_IPHONE
-    #define IS_IOS_NNAPI_BIND
-  #endif
-#endif
-
-#ifndef IS_IOS_NNAPI_BIND
-// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-static auto register_NnapiCompilation = [](){
-  try {
-    return torch::jit::class_<NnapiCompilation>("_nnapi", "Compilation")
-        .def(torch::jit::init<>())
-        .def("init", &NnapiCompilation::init)
-        .def("run", &NnapiCompilation::run)
-        ;
-  } catch (std::exception& exn) {
-    LOG(ERROR) << "Failed to register class nnapi.Compilation: " << exn.what();
-    throw;
-  }
-}();
-#else
-  #undef IS_IOS_NNAPI_BIND
-#endif
-
 } // namespace bind
 } // namespace nnapi
 } // namespace torch

diff --git a/aten/src/ATen/nnapi/nnapi_register.cpp b/aten/src/ATen/nnapi/nnapi_register.cpp
new file mode 100644 (file)
index 0000000..25d5441
--- /dev/null
+++ b/aten/src/ATen/nnapi/nnapi_register.cpp
@@ -0,0 +1,21 @@
+#include <ATen/nnapi/nnapi_bind.h>
+
+// Set flag if running on ios
+#ifdef __APPLE__
+  #include <TargetConditionals.h>
+  #if TARGET_OS_IPHONE
+    #define IS_IOS_NNAPI_BIND
+  #endif
+#endif
+
+#ifndef IS_IOS_NNAPI_BIND
+TORCH_LIBRARY(_nnapi, m) {
+  m.class_<torch::nnapi::bind::NnapiCompilation>("Compilation")
+    .def(torch::jit::init<>())
+    .def("init", &torch::nnapi::bind::NnapiCompilation::init)
+    .def("run", &torch::nnapi::bind::NnapiCompilation::run)
+    ;
+}
+#else
+  #undef IS_IOS_NNAPI_BIND
+#endif

diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl
index bfcf55a..e63e6e7 100644 (file)
--- a/tools/build_variables.bzl
+++ b/tools/build_variables.bzl
@@ -1116,6 +1116,7 @@ aten_native_source_non_codegen_list = [
     # Files not in native, but depends on native symbols
     # "aten/src/ATen/TensorIndexing.cpp",
     "aten/src/ATen/TensorIterator.cpp",
+    "aten/src/ATen/nnapi/nnapi_register.cpp",
     "aten/src/ATen/nnapi/nnapi_bind.cpp",
     "aten/src/ATen/nnapi/nnapi_wrapper.cpp",
     "aten/src/ATen/nnapi/nnapi_model_loader.cpp",