diff --git a/backends/qnnpack/test/TARGETS b/backends/qnnpack/test/TARGETS index d0a25fa9ce4..99097c58687 100644 --- a/backends/qnnpack/test/TARGETS +++ b/backends/qnnpack/test/TARGETS @@ -16,7 +16,6 @@ python_unittest( "//executorch/backends/qnnpack/partition:qnnpack_partitioner", "//executorch/exir:lib", "//executorch/exir/backend:backend_api", - "//executorch/exir/serialize:lib", "//executorch/extension/pybindings:portable", # @manual "//executorch/extension/pytree:pylib", ], diff --git a/backends/qnnpack/test/test_qnnpack.py b/backends/qnnpack/test/test_qnnpack.py index b6c65bdbae7..71aa47e03c2 100644 --- a/backends/qnnpack/test/test_qnnpack.py +++ b/backends/qnnpack/test/test_qnnpack.py @@ -19,7 +19,6 @@ from executorch.exir.backend.backend_api import to_backend, validation_disabled -# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`. from executorch.extension.pybindings.portable import ( # @manual _load_for_executorch_from_buffer, ) diff --git a/backends/vulkan/test/test_vulkan_delegate.py b/backends/vulkan/test/test_vulkan_delegate.py index 5ed35237bee..7b63edde711 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/backends/vulkan/test/test_vulkan_delegate.py @@ -19,7 +19,7 @@ ctypes.CDLL("libvulkan.so.1") -# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`. + from executorch.extension.pybindings.portable import ( # @manual _load_for_executorch_from_buffer, ) @@ -85,7 +85,6 @@ def forward(self, *args): ) # Test the model with executor - # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`. executorch_module = _load_for_executorch_from_buffer(executorch_program.buffer) # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`. inputs_flattened, _ = tree_flatten(sample_inputs) diff --git a/backends/xnnpack/test/TARGETS b/backends/xnnpack/test/TARGETS index 3305f259931..9b59a6eb9cd 100644 --- a/backends/xnnpack/test/TARGETS +++ b/backends/xnnpack/test/TARGETS @@ -28,7 +28,6 @@ python_unittest( "//executorch/exir:tracer", "//executorch/exir/backend:backend_api", "//executorch/exir/passes:spec_prop_pass", - "//executorch/exir/serialize:lib", "//executorch/extension/pybindings:portable", # @manual "//executorch/extension/pytree:pylib", ], @@ -59,7 +58,6 @@ python_unittest( "//executorch/exir/backend:backend_api", "//executorch/exir/dialects:lib", "//executorch/exir/passes:spec_prop_pass", - "//executorch/exir/serialize:lib", "//executorch/extension/pybindings:portable", # @manual "//executorch/extension/pytree:pylib", ], @@ -89,7 +87,6 @@ python_unittest( "//executorch/exir:tracer", "//executorch/exir/backend:backend_api", "//executorch/exir/passes:spec_prop_pass", - "//executorch/exir/serialize:lib", "//executorch/extension/pybindings:portable", # @manual "//executorch/extension/pytree:pylib", "//pytorch/vision:torchvision", @@ -133,6 +130,7 @@ python_unittest( "//caffe2:torch", "//executorch/backends/xnnpack/partition:xnnpack_partitioner", "//executorch/backends/xnnpack/test/tester:tester", + "//executorch/backends/xnnpack/utils:xnnpack_utils", "//pytorch/vision:torchvision", ], ) diff --git a/backends/xnnpack/test/test_xnnpack_utils.py b/backends/xnnpack/test/test_xnnpack_utils.py index 7af509c462a..9f2e1a817c9 100644 --- a/backends/xnnpack/test/test_xnnpack_utils.py +++ b/backends/xnnpack/test/test_xnnpack_utils.py @@ -37,7 +37,6 @@ from executorch.exir.passes.spec_prop_pass import SpecPropPass from executorch.exir.tracer import _default_decomposition_table -# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`. from executorch.extension.pybindings.portable import ( # @manual _load_for_executorch_from_buffer, ) @@ -230,7 +229,6 @@ def forward(self, *args): ) # Test the model with executor - # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`. executorch_module = _load_for_executorch_from_buffer(executorch_program.buffer) # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`. inputs_flattened, _ = tree_flatten(sample_inputs) diff --git a/examples/export/test/test_export.py b/examples/export/test/test_export.py index e4cb98bcffa..0f1135c20c8 100644 --- a/examples/export/test/test_export.py +++ b/examples/export/test/test_export.py @@ -13,7 +13,6 @@ from executorch.examples.export.utils import export_to_edge from executorch.examples.models import MODEL_NAME_TO_MODEL -# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`. from executorch.extension.pybindings.portable import ( # @manual _load_for_executorch_from_buffer, ) @@ -37,7 +36,7 @@ def _assert_eager_lowered_same_result( edge_model = export_to_edge(eager_model, example_inputs) executorch_prog = edge_model.to_executorch() - # pyre-ignore + pte_model = _load_for_executorch_from_buffer(executorch_prog.buffer) with torch.no_grad(): diff --git a/exir/backend/test/demos/rpc/test_rpc.py b/exir/backend/test/demos/rpc/test_rpc.py index f53d72b1b5e..fc754742bd3 100644 --- a/exir/backend/test/demos/rpc/test_rpc.py +++ b/exir/backend/test/demos/rpc/test_rpc.py @@ -17,7 +17,6 @@ ) from executorch.exir.backend.test.op_partitioner_demo import AddMulPartitionerDemo -# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`. from executorch.extension.pybindings.portable import ( # @manual _load_for_executorch_from_buffer, ) diff --git a/exir/backend/test/demos/test_delegate_aten_mode.py b/exir/backend/test/demos/test_delegate_aten_mode.py index 9198de3989b..7c80cf20cc9 100644 --- a/exir/backend/test/demos/test_delegate_aten_mode.py +++ b/exir/backend/test/demos/test_delegate_aten_mode.py @@ -15,7 +15,6 @@ BackendWithCompilerDemo, ) -# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`. from executorch.extension.pybindings.aten_mode_lib import ( # @manual _load_for_executorch_from_buffer, ) diff --git a/exir/backend/test/demos/test_xnnpack_qnnpack.py b/exir/backend/test/demos/test_xnnpack_qnnpack.py index f7df6c22605..4d8bbd4af0c 100644 --- a/exir/backend/test/demos/test_xnnpack_qnnpack.py +++ b/exir/backend/test/demos/test_xnnpack_qnnpack.py @@ -22,7 +22,6 @@ from executorch.exir.backend.backend_api import to_backend, validation_disabled from executorch.exir.passes.spec_prop_pass import SpecPropPass -# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`. from executorch.extension.pybindings.portable import ( # @manual _load_for_executorch_from_buffer, ) diff --git a/exir/backend/test/test_backends.py b/exir/backend/test/test_backends.py index 34f57a76c01..c79ccc728bb 100644 --- a/exir/backend/test/test_backends.py +++ b/exir/backend/test/test_backends.py @@ -46,7 +46,6 @@ Program, ) -# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`. from executorch.extension.pybindings.portable import ( # @manual _load_for_executorch_from_buffer, ) @@ -224,7 +223,6 @@ def forward(self, x): ) buff = exec_prog.buffer - # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`. executorch_module = _load_for_executorch_from_buffer(buff) model_inputs = torch.ones(1) model_outputs = executorch_module.forward([model_inputs]) @@ -281,7 +279,6 @@ def forward(self, a, x, b): ) buff = exec_prog.buffer - # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`. executorch_module = _load_for_executorch_from_buffer(buff) # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`. @@ -338,7 +335,6 @@ def forward(self, x): # This line should raise an exception like # RuntimeError: failed with error 0x12 - # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`. _load_for_executorch_from_buffer(buff) @vary_segments @@ -434,7 +430,6 @@ def forward(self, x): ) ) - # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`. executorch_module = _load_for_executorch_from_buffer(buff) model_inputs = torch.ones(1) @@ -561,7 +556,6 @@ def forward(self, x): ) flatbuffer = exec_prog.buffer - # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`. executorch_module = _load_for_executorch_from_buffer(flatbuffer) model_outputs = executorch_module.forward([*model_inputs]) @@ -858,7 +852,6 @@ def forward(self, a, x, b): # There should be 2 delegated modules self.assertEqual(counter, 2) - # pyre-ignore[16]: Module `executorch.extension.pybindings` has no attribute `portable`. executorch_module = _load_for_executorch_from_buffer(executorch_prog.buffer) # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`. inputs_flattened, _ = tree_flatten(inputs) diff --git a/exir/backend/test/test_backends_lifted.py b/exir/backend/test/test_backends_lifted.py index 0d52eb7af47..80b50315f21 100644 --- a/exir/backend/test/test_backends_lifted.py +++ b/exir/backend/test/test_backends_lifted.py @@ -49,7 +49,6 @@ Program, ) -# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`. from executorch.extension.pybindings.portable import ( # @manual _load_for_executorch_from_buffer, ) diff --git a/extension/pybindings/TARGETS b/extension/pybindings/TARGETS index 4e21936cdde..6d96f4dc90d 100644 --- a/extension/pybindings/TARGETS +++ b/extension/pybindings/TARGETS @@ -3,10 +3,25 @@ # targets.bzl. This file can contain fbcode-only targets. load("@fbcode//executorch/extension/pybindings:targets.bzl", "ATEN_MODULE_DEPS", "MODELS_ATEN_OPS_ATEN_MODE_GENERATED_LIB", "MODELS_ATEN_OPS_LEAN_MODE_GENERATED_LIB", "PORTABLE_MODULE_DEPS", "define_common_targets", "executorch_pybindings") +load("@fbcode_macros//build_defs:native_rules.bzl", "buck_genrule") load("@fbcode_macros//build_defs:python_library.bzl", "python_library") define_common_targets() +# In order to have pyre recognize the pybindings module, the name of the .pyi must exactly match the +# name of the lib. To avoid copy pasting the pyi file in tree a whole bunch of times we use genrules +# to do it for us +buck_genrule( + name = "pybindings_types_gen", + srcs = [":pybinding_types"], + outs = { + "aten_mode_lib.pyi": ["aten_mode_lib.pyi"], + "portable.pyi": ["portable.pyi"], + }, + cmd = "cp $(location :pybinding_types)/* $OUT/portable.pyi && cp $(location :pybinding_types)/* $OUT/aten_mode_lib.pyi", + visibility = ["//executorch/extension/pybindings/..."], +) + executorch_pybindings( srcs = [ "module.cpp", @@ -22,6 +37,7 @@ executorch_pybindings( ], cppdeps = PORTABLE_MODULE_DEPS + MODELS_ATEN_OPS_LEAN_MODE_GENERATED_LIB, python_module_name = "portable", + types = ["//executorch/extension/pybindings:pybindings_types_gen[portable.pyi]"], visibility = ["PUBLIC"], ) @@ -31,6 +47,7 @@ executorch_pybindings( ], cppdeps = ATEN_MODULE_DEPS + MODELS_ATEN_OPS_ATEN_MODE_GENERATED_LIB, python_module_name = "aten_mode_lib", + types = ["//executorch/extension/pybindings:pybindings_types_gen[aten_mode_lib.pyi]"], visibility = ["PUBLIC"], ) diff --git a/extension/pybindings/module.cpp b/extension/pybindings/module.cpp index 147750d36b1..dc9db2eaeaf 100644 --- a/extension/pybindings/module.cpp +++ b/extension/pybindings/module.cpp @@ -403,15 +403,17 @@ struct PyModule final { return std::make_unique(m.get_program_ptr(), m.get_program_len()); } - py::list run_method(const std::string& name, const py::sequence& pyinputs) { - std::vector inputs; - const auto inputs_size = py::len(pyinputs); - inputs.reserve(inputs_size); + py::list run_method( + const std::string& method_name, + const py::sequence& inputs) { + std::vector cpp_inputs; + const auto inputs_size = py::len(inputs); + cpp_inputs.reserve(inputs_size); for (size_t i = 0; i < inputs_size; ++i) { - inputs.emplace_back(pyToEValue(pyinputs[i], keep_alive_)); + cpp_inputs.emplace_back(pyToEValue(inputs[i], keep_alive_)); } - auto outputs = module_->run_method(name, inputs); + auto outputs = module_->run_method(method_name, cpp_inputs); const auto outputs_size = outputs.size(); py::list list(outputs_size); @@ -421,8 +423,8 @@ struct PyModule final { return list; } - py::list forward(const py::sequence& pyinputs) { - return run_method("forward", pyinputs); + py::list forward(const py::sequence& inputs) { + return run_method("forward", inputs); } private: @@ -461,7 +463,7 @@ void init_module_functions(py::module_& m) { m.def("_create_profile_block", &create_profile_block); m.def("_reset_profile_results", []() { EXECUTORCH_RESET_PROFILE_RESULTS(); }); - py::class_(m, "Module") + py::class_(m, "ExecutorchModule") .def("run_method", &PyModule::run_method) .def("forward", &PyModule::forward); diff --git a/extension/pybindings/pybindings.pyi b/extension/pybindings/pybindings.pyi new file mode 100644 index 00000000000..8d98e526fe3 --- /dev/null +++ b/extension/pybindings/pybindings.pyi @@ -0,0 +1,15 @@ +# 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. + +# pyre-strict +from typing import Any, Dict, List, Sequence, Tuple + +class ExecutorchModule: + def run_method(self, method_name: str, inputs: Sequence[Any]) -> List[Any]: ... + def forward(self, inputs: Sequence[Any]) -> List[Any]: ... + +def _load_for_executorch(path: str) -> ExecutorchModule: ... +def _load_for_executorch_from_buffer(buffer: bytes) -> ExecutorchModule: ... diff --git a/extension/pybindings/targets.bzl b/extension/pybindings/targets.bzl index 7f5eb2e533f..199e0c3fada 100644 --- a/extension/pybindings/targets.bzl +++ b/extension/pybindings/targets.bzl @@ -46,12 +46,13 @@ MODELS_ATEN_OPS_ATEN_MODE_GENERATED_LIB = [ "//executorch/kernels/aten:generated_lib_aten", ] -def executorch_pybindings(python_module_name, srcs = [], cppdeps = [], visibility = ["//executorch/..."]): +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, + types = types, base_module = "executorch.extension.pybindings", preprocessor_flags = [ "-DEXECUTORCH_PYTHON_MODULE_NAME={}".format(python_module_name), @@ -88,6 +89,15 @@ def define_common_targets(): visibility = ["//executorch/extension/pybindings/..."], ) + # cxx_python_extension kwarg 'types' cant take export_file rules directly and we need to rename the .pyi + # file to match the lib anyway, so we just expose the file like this and then have genrules consume and + # rename it before passing it to executorch pybindings. + runtime.filegroup( + name = "pybinding_types", + srcs = ["pybindings.pyi"], + visibility = ["//executorch/extension/pybindings/..."], + ) + executorch_pybindings( srcs = [ "module_stub.cpp", diff --git a/extension/pybindings/test/test.py b/extension/pybindings/test/test.py index 53985fee72a..24a46da6f81 100644 --- a/extension/pybindings/test/test.py +++ b/extension/pybindings/test/test.py @@ -15,14 +15,13 @@ from executorch.exir.scalar_type import ScalarType from executorch.exir.schema import Program -# executorch.extension.pybindings.portable is a cpp_python_extension target. -# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`. +# pyre-ignore from executorch.extension.pybindings.portable import ( - _get_io_metadata_for_program_operators, # @manual - _get_program_from_buffer, # @manual - _get_program_operators, # @manual - _load_for_executorch_from_buffer, # @manual - IOMetaData, # @manual + _get_io_metadata_for_program_operators, + _get_program_from_buffer, + _get_program_operators, + _load_for_executorch_from_buffer, + IOMetaData, ) diff --git a/shim/xplat/executorch/build/runtime_wrapper.bzl b/shim/xplat/executorch/build/runtime_wrapper.bzl index 78bd43a8cea..841463b551a 100644 --- a/shim/xplat/executorch/build/runtime_wrapper.bzl +++ b/shim/xplat/executorch/build/runtime_wrapper.bzl @@ -270,6 +270,8 @@ def _cxx_test(*args, **kwargs): def _cxx_python_extension(*args, **kwargs): _patch_kwargs_common(kwargs) kwargs["srcs"] = _patch_executorch_references(kwargs["srcs"]) + if "types" in kwargs: + kwargs["types"] = _patch_executorch_references(kwargs["types"]) env.cxx_python_extension(*args, **kwargs) def _export_file(*args, **kwargs):