move codegen pybindings out of public pybindings lib (#131)

Summary:
Pull Request resolved: #131

These are only used by codegen. Moving them out of the general pybindings lib and restricting visibility.

As a side note I dont think any of this stuff actually has to happen in c++. In fact it seems like itd be easier to do this on the Python Version of the schema object through deserialize_for_json.

Reviewed By: larryliu0820, dbort

Differential Revision: D48671946

fbshipit-source-id: d5493fcf27733fdab5df628f8d1d72ca3f201260

Author: JacobSzwejbka

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,

extension/pybindings/pybindings.cpp renamed to codegen/tools/selective_build.cpp

@@ -23,13 +23,14 @@ namespace {
 // Metadata for kernel call io variables.
 // dtype and dim_order will exist only if corresponding variable is Tensor.
 struct IOMetaData {
-  int type;
+  int kernel_type;
   int dtype;
   std::vector<unsigned int> dim_order;
 
   // Create tensor metadata. It records tensor's dtype and dim order.
   explicit IOMetaData(const executorch_flatbuffer::Tensor* t)
-      : type(static_cast<int>(executorch_flatbuffer::KernelTypes::Tensor)),
+      : kernel_type(
+            static_cast<int>(executorch_flatbuffer::KernelTypes::Tensor)),
         dtype(static_cast<int>(t->scalar_type())) {
     for (size_t i = 0; i < t->dim_order()->size(); i++) {
       dim_order.push_back(static_cast<unsigned int>(t->dim_order()->Get(i)));
@@ -38,7 +39,7 @@ struct IOMetaData {
 
   // Create metadata for non-tensor variable.
   explicit IOMetaData(executorch_flatbuffer::KernelTypes type)
-      : type(static_cast<int>(type)) {
+      : kernel_type(static_cast<int>(type)) {
     ET_CHECK(
         type != executorch_flatbuffer::KernelTypes::Tensor &&
         type != executorch_flatbuffer::KernelTypes::TensorList &&
@@ -54,10 +55,10 @@ struct KernelIOMetaDataComparsion {
       return lhs.size() < rhs.size();
     }
     for (size_t i = 0; i < lhs.size(); i++) {
-      if (lhs[i].type != rhs[i].type) {
-        return lhs[i].type < rhs[i].type;
+      if (lhs[i].kernel_type != rhs[i].kernel_type) {
+        return lhs[i].kernel_type < rhs[i].kernel_type;
       }
-      if (lhs[i].type !=
+      if (lhs[i].kernel_type !=
           static_cast<int>(executorch_flatbuffer::KernelTypes::Tensor)) {
         continue;
       }
@@ -241,12 +242,8 @@ py::dict _get_io_metadata_for_program_operators(
   return py_program_op_io_metadata;
 }
 
-void init_module_functions(py::module_&);
-
 PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
-  init_module_functions(m);
-
-  py::class_<executorch_flatbuffer::Program>(m, "Program");
+  py::class_<executorch_flatbuffer::Program>(m, "_Program");
 
   m.def(
       "_get_program_from_buffer",
@@ -263,8 +260,8 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
       &_get_io_metadata_for_program_operators,
       py::return_value_policy::copy);
 
-  py::class_<IOMetaData>(m, "IOMetaData")
-      .def_readwrite("type", &IOMetaData::type)
+  py::class_<IOMetaData>(m, "_IOMetaData")
+      .def_readwrite("kernel_type", &IOMetaData::kernel_type)
       .def_readwrite("dtype", &IOMetaData::dtype)
       .def_readwrite("dim_order", &IOMetaData::dim_order);
 }

codegen/tools/selective_build.pyi

@@ -0,0 +1,23 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Any, Dict, List
+
+class _Program: ...
+
+class _IOMetaData:
+    @property
+    def kernel_type(self) -> int: ...
+    @property
+    def dtype(self) -> int: ...
+    @property
+    def dim_order(self) -> List[int]: ...
+
+def _get_program_from_buffer(buffer: bytes) -> _Program: ...
+def _get_program_operators(program: _Program) -> List[str]: ...
+def _get_io_metadata_for_program_operators(
+    program: _Program,
+) -> Dict[str, Any]: ...

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": [":selective_build"],
         }),
     )
 
@@ -153,3 +153,25 @@ def define_common_targets(is_fbcode = False):
                 "//libfb/py:parutil",
             ],
         )
+
+    module_name = "selective_build"
+    runtime.cxx_python_extension(
+        name = module_name,
+        srcs = [
+            "selective_build.cpp",
+        ],
+        base_module = "executorch.codegen.tools",
+        types = ["{}.pyi".format(module_name)],
+        preprocessor_flags = [
+            "-DEXECUTORCH_PYTHON_MODULE_NAME={}".format(module_name),
+        ],
+        deps = [
+            "//executorch/schema:program",
+            "//executorch/util:read_file",
+        ],
+        external_deps = [
+            "pybind11",
+        ],
+        use_static_deps = True,
+        visibility = ["//executorch/codegen/..."],
+    )

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",

extension/pybindings/module_stub.cpp

@@ -13,7 +13,6 @@ PORTABLE_MODULE_DEPS = [
     "//executorch/runtime/kernel:operator_registry",
     "//executorch/runtime/executor:executor",
     "//executorch/schema:bundled_program_schema",
-    "//executorch/schema:program",
     "//executorch/extension/aten_util:aten_bridge",
     "//executorch/util:bundled_program_verification",
     "//executorch/extension/data_loader:buffer_data_loader",
@@ -29,7 +28,6 @@ ATEN_MODULE_DEPS = [
     "//executorch/runtime/executor:executor_aten",
     "//executorch/runtime/core/exec_aten:lib",
     "//executorch/schema:bundled_program_schema",
-    "//executorch/schema:program",
     "//executorch/extension/data_loader:buffer_data_loader",
     "//executorch/extension/data_loader:mmap_data_loader",
     "//executorch/extension/memory_allocator:malloc_memory_allocator",
@@ -49,17 +47,14 @@ MODELS_ATEN_OPS_ATEN_MODE_GENERATED_LIB = [
 def executorch_pybindings(python_module_name, srcs = [], cppdeps = [], visibility = ["//executorch/..."], types = []):
     runtime.cxx_python_extension(
         name = python_module_name,
-        srcs = [
-            "//executorch/extension/pybindings:pybindings.cpp",
-        ] + srcs,
+        srcs = srcs,
         types = types,
         base_module = "executorch.extension.pybindings",
         preprocessor_flags = [
             "-DEXECUTORCH_PYTHON_MODULE_NAME={}".format(python_module_name),
         ],
         deps = [
             "//executorch/runtime/core:core",
-            "//executorch/schema:program",
             "//executorch/util:read_file",
         ] + cppdeps,
         external_deps = [
@@ -77,13 +72,8 @@ def define_common_targets():
     TARGETS and BUCK files that call this function.
     """
 
-    # Export these so the internal fb/ subdir can create pybindings with custom internal deps
+    # Export this so the internal fb/ subdir can create pybindings with custom internal deps
     # without forking the pybinding source.
-    runtime.export_file(
-        name = "pybindings.cpp",
-        visibility = ["//executorch/extension/pybindings/..."],
-    )
-
     runtime.export_file(
         name = "module.cpp",
         visibility = ["//executorch/extension/pybindings/..."],
@@ -97,10 +87,3 @@ def define_common_targets():
         srcs = ["pybindings.pyi"],
         visibility = ["//executorch/extension/pybindings/..."],
     )
-
-    executorch_pybindings(
-        srcs = [
-            "module_stub.cpp",
-        ],
-        python_module_name = "operator",
-    )

extension/pybindings/targets.bzl

@@ -12,17 +12,9 @@
 import torch
 from executorch.exir import CaptureConfig
 from executorch.exir.print_program import pretty_print
-from executorch.exir.scalar_type import ScalarType
 from executorch.exir.schema import Program
 
-# pyre-ignore[21]
-from executorch.extension.pybindings.portable import (
-    _get_io_metadata_for_program_operators,
-    _get_program_from_buffer,
-    _get_program_operators,
-    _load_for_executorch_from_buffer,
-    IOMetaData,
-)
+from executorch.extension.pybindings.portable import _load_for_executorch_from_buffer
 
 
 class ModuleAdd(torch.nn.Module):
@@ -96,46 +88,6 @@ def test_e2e(self):
 
         self.assertEqual(str(expected), str(executorch_output))
 
-    def test_dump_operators(self):
-        # Create and serialize a program.
-        orig_program, _ = create_program()
-
-        # Deserialize the program and demonstrate that we could get its operator
-        # list.
-        program = _get_program_from_buffer(orig_program.buffer)
-        operators = _get_program_operators(program)
-        self.assertEqual(operators, ["aten::add.out"])
-
-    def test_get_op_io_meta(self):
-        # Checking whether get_op_io_meta returns the correct metadata for all its ios.
-        orig_program, inputs = create_program()
-
-        # Deserialize the program and demonstrate that we could get its operator
-        # list.
-        program = _get_program_from_buffer(orig_program.buffer)
-        program_op_io_metadata = _get_io_metadata_for_program_operators(program)
-
-        self.assertTrue(len(program_op_io_metadata) == 1)
-        self.assertTrue(isinstance(program_op_io_metadata, dict))
-
-        self.assertTrue("aten::add.out" in program_op_io_metadata)
-        self.assertTrue(isinstance(program_op_io_metadata["aten::add.out"], set))
-        self.assertTrue(len(program_op_io_metadata["aten::add.out"]) == 1)
-
-        for op_io_metadata in program_op_io_metadata["aten::add.out"]:
-            self.assertTrue(len(op_io_metadata) == 5)
-            self.assertTrue(isinstance(op_io_metadata, tuple))
-
-            for io_idx, io_metadata in enumerate(op_io_metadata):
-                self.assertTrue(isinstance(io_metadata, IOMetaData))
-                if io_idx == 2:
-                    # TODO(gasoonjia): Create a enum class to map KernelTypes to int, remove the hardcoded 2 and 5 below.
-                    self.assertEqual(io_metadata.type, 2)
-                else:
-                    self.assertEqual(io_metadata.type, 5)
-                    self.assertEqual(io_metadata.dtype, ScalarType.FLOAT)
-                    self.assertEqual(io_metadata.dim_order, [0, 1])
-
     def test_multiple_entry(self):
 
         program, inputs = create_program(ModuleMulti())

extension/pybindings/test/test.py

@@ -95,7 +95,7 @@ def define_common_targets():
             # Lock this down as tightly as possible to ensure that flatbuffers
             # are an implementation detail. Ideally this list would only include
             # //executorch/runtime/executor/...
-            "//executorch/extension/pybindings/...",
+            "//executorch/codegen/tools/...",
             "//executorch/runtime/executor/...",
             "//executorch/util/...",  # bundled_program_verification
         ],

schema/targets.bzl

@@ -110,6 +110,7 @@ def _remove_platform_specific_args(kwargs):
 def _remove_unsupported_kwargs(kwargs):
     """Removes environment unsupported kwargs
     """
+    kwargs.pop("types", None)  # will have to find a different way to handle .pyi files in oss
     return kwargs
 
 def _patch_headers(kwargs):

shim/xplat/executorch/build/env_interface.bzl

@@ -159,6 +159,7 @@ def _patch_kwargs_common(kwargs):
 
     Returns the possibly-modified `kwargs` parameter for chaining.
     """
+    env.remove_unsupported_kwargs(kwargs)
 
     # Be careful about dependencies on executorch targets for now, so that we
     # don't pick up unexpected clients while things are still in flux.
@@ -208,7 +209,6 @@ def _patch_kwargs_common(kwargs):
     return kwargs
 
 def _patch_kwargs_cxx(kwargs):
-    env.remove_unsupported_kwargs(kwargs)
     env.patch_platforms(kwargs)
     env.remove_platform_specific_args(kwargs)
     return _patch_kwargs_common(kwargs)