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

Differential Revision: D48671946

fbshipit-source-id: af5d51d59253534c3d318fa01316a502e7ec0446

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": [] if runtime.is_oss else ["//executorch/codegen/tools:selective_build"],  # TODO(larryliu0820) :selective_build doesn't build in OSS yet
         }),
     )
 
@@ -153,3 +153,43 @@ 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",
+        ],
+        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/..."],
+    )
+
+    runtime.python_test(
+        # @autodeps-skip
+        name = "test_selective_build",
+        srcs = [
+            "test/test_selective_build.py",
+        ],
+        deps = [
+            "//caffe2:torch",
+            "//caffe2:torch_fx",
+            "//executorch/codegen/tools:selective_build",
+            "//executorch/exir:lib",
+            "//executorch/exir:pass_manager",
+            "//executorch/exir:scalar_type",
+            "//executorch/exir/emit:lib",
+            "//executorch/exir/passes:lib",
+            "//executorch/exir/serialize:lib",
+        ],
+    )

codegen/tools/test/test_selective_build.py

@@ -0,0 +1,121 @@
+# 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.
+
+import unittest
+from typing import Any, Optional, Tuple
+
+import executorch.exir as exir
+
+import torch
+
+from executorch.codegen.tools.selective_build import (
+    _get_io_metadata_for_program_operators,
+    _get_program_from_buffer,
+    _get_program_operators,
+    _IOMetaData,
+)
+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
+
+
+class ModuleAdd(torch.nn.Module):
+    """The module to serialize and execute."""
+
+    def __init__(self):
+        super(ModuleAdd, self).__init__()
+
+    def forward(self, x, y):
+        return x + y
+
+    def get_methods_to_export(self):
+        return ("forward",)
+
+
+class ModuleMulti(torch.nn.Module):
+    """The module to serialize and execute."""
+
+    def __init__(self):
+        super(ModuleMulti, self).__init__()
+
+    def forward(self, x, y):
+        return x + y
+
+    def forward2(self, x, y):
+        return x + y + 1
+
+    def get_methods_to_export(self):
+        return ("forward", "forward2")
+
+
+def create_program(
+    eager_module: Optional[torch.nn.Module] = None,
+) -> Tuple[Program, Tuple[Any, ...]]:
+    """Returns an executorch program based on ModuleAdd, along with inputs."""
+
+    if eager_module is None:
+        eager_module = ModuleAdd()
+
+    # Trace the test module and create a serialized Executorch program.
+    inputs = (torch.ones(2, 2), torch.ones(2, 2))
+    input_map = {}
+    for method in eager_module.get_methods_to_export():
+        input_map[method] = inputs
+
+    # These cleanup passes are required to convert the `add` op to its out
+    # variant, along with some other transformations.
+    exec_prog = (
+        exir.capture_multiple(eager_module, input_map, config=CaptureConfig())
+        .to_edge()
+        .to_executorch()
+    )
+
+    # Create the Executorch program from the graph.
+    pretty_print(exec_prog.program)
+    return (exec_prog, inputs)
+
+
+class PybindingsTest(unittest.TestCase):
+    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.kernel_type, 2)
+                else:
+                    self.assertEqual(io_metadata.kernel_type, 5)
+                    self.assertEqual(io_metadata.dtype, ScalarType.FLOAT)
+                    self.assertEqual(io_metadata.dim_order, [0, 1])

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