move codegen pybindings out of public pybindings lib

codegen/tools/gen_oplist.py

@@ -78,43 +78,38 @@ class KernelType(IntEnum):
 
 
 def _get_operators(model_file: str) -> List[str]:
-    # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.extension.pybindings.operator`.
-    # pyre-ignore: Undefined attribute [16]: Module `executorch.extension.pybindings` has no attribute `operator`
-    from executorch.extension.pybindings.operator import (
+    from executorch.codegen.tools.selective_build import (
         _get_program_from_buffer,
         _get_program_operators,
     )
 
     print("Processing model file: ", model_file)
     with open(model_file, "rb") as f:
         buf = f.read()
-    # pyre-ignore: Undefined attribute [16]: Module `executorch.extension.pybindings` has no attribute `operator`.
+
     program = _get_program_from_buffer(buf)
-    # pyre-ignore: Undefined attribute [16]: Module `executorch.extension.pybindings` has no attribute `operator`.
     operators = _get_program_operators(program)
     print(f"Model file loaded, operators are: {operators}")
     return operators
 
 
 def _get_kernel_metadata_for_model(model_file: str) -> Dict[str, List[str]]:
 
-    from executorch.extension.pybindings.operator import (
+    from executorch.codegen.tools.selective_build import (
         _get_io_metadata_for_program_operators,
         _get_program_from_buffer,
-        IOMetaData,
+        _IOMetaData,
     )
 
     with open(model_file, "rb") as f:
         buf = f.read()
-    # pyre-ignore: Undefined attribute [16]: Module `executorch.extension.pybindings` has no attribute `operator`.
+
     program = _get_program_from_buffer(buf)
-    # pyre-ignore: Undefined attribute [16]: Module `executorch.extension.pybindings` has no attribute `operator`.
     operators_with_io_metadata = _get_io_metadata_for_program_operators(program)
 
     op_kernel_key_list: Dict[str, List[str]] = {}
 
-    # pyre-ignore: Undefined or invalid type [11]: Annotation `IOMetaData` is not defined as a type.Pyre
-    specialized_kernels: Set[List[IOMetaData]]
+    specialized_kernels: Set[List[_IOMetaData]]
     for op_name, specialized_kernels in operators_with_io_metadata.items():
         print(op_name)
         if op_name not in op_kernel_key_list:
@@ -124,7 +119,7 @@ def _get_kernel_metadata_for_model(model_file: str) -> Dict[str, List[str]]:
             version = "v1"
             kernel_key = version + "/"
             for io_metadata in specialized_kernel:
-                if io_metadata.type in [
+                if io_metadata.kernel_type in [
                     KernelType.TENSOR,
                     KernelType.TENSOR_LIST,
                     KernelType.OPTIONAL_TENSOR_LIST,

codegen/tools/targets.bzl

@@ -18,7 +18,7 @@ def define_common_targets(is_fbcode = False):
         external_deps = ["torchgen"],
         deps = select({
             "DEFAULT": [],
-            "ovr_config//os:linux": [] if runtime.is_oss else ["//executorch/extension/pybindings:operator"],  # TODO(larryliu0820): pybindings:operator doesn't build in OSS yet
+            "ovr_config//os:linux": [] if runtime.is_oss else ["//executorch/codegen/tools/fb:selective_build"],  # TODO(larryliu0820) :selective_build doesn't build in OSS yet
         }),
     )
 

extension/pybindings/module.cpp

@@ -26,7 +26,6 @@
 #include <executorch/runtime/platform/profiler.h>
 #include <executorch/runtime/platform/runtime.h>
 #include <executorch/schema/bundled_program_schema_generated.h>
-#include <executorch/schema/program_generated.h>
 #include <executorch/util/TestMemoryConfig.h>
 #include <executorch/util/bundled_program_verification.h>
 #include <executorch/util/read_file.h>
@@ -463,7 +462,7 @@ void create_profile_block(const std::string& name) {
 
 } // namespace
 
-void init_module_functions(py::module_& m) {
+PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
   m.def("_load_for_executorch", PyModule::load_from_file, py::arg("path"));
   m.def(
       "_load_for_executorch_from_buffer",