Support ExecutorchModule as a callable in pybindings (#1312)

Summary:
Pull Request resolved: #1312

Users would be able to like to do the following, which is more consistent with how eager mode models and graph modules are callables.
```
model = _load_for_executorch(pte_file)
model(*inputs)
```
instead of
```
model = _load_for_executorch(pte_file)
model.forward(*inputs)
```

Reviewed By: chakriu, JacobSzwejbka

Differential Revision: D51673732

fbshipit-source-id: 2d754a0282929e28ff75cd34c95645087384f9f3

Author: tarun292

extension/pybindings/pybindings.cpp

@@ -480,6 +480,12 @@ struct PyModule final {
     return run_method("forward", inputs);
   }
 
+  py::list forward_single_input(const torch::Tensor& inputTensor) {
+    py::list py_list;
+    py_list.append(py::cast(inputTensor));
+    return run_method("forward", py_list);
+  }
+
   bool has_etdump() {
     return module_->has_etdump();
   }
@@ -589,8 +595,9 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
       .def("forward", &PyModule::forward)
       .def("has_etdump", &PyModule::has_etdump)
       .def(
-          "write_etdump_result_to_file",
-          &PyModule::write_etdump_result_to_file);
+          "write_etdump_result_to_file", &PyModule::write_etdump_result_to_file)
+      .def("__call__", &PyModule::forward)
+      .def("__call__", &PyModule::forward_single_input);
 
   py::class_<PyBundledModule>(m, "BundledModule");
 }

extension/pybindings/pybindings.pyi

@@ -8,6 +8,7 @@
 from typing import Any, Dict, List, Sequence, Tuple
 
 class ExecutorchModule:
+    def __call__(self, inputs: Any) -> List[Any]: ...
     def run_method(self, method_name: str, inputs: Sequence[Any]) -> List[Any]: ...
     def forward(self, inputs: Sequence[Any]) -> List[Any]: ...
     # Bundled program methods.

extension/pybindings/test/make_test.py

@@ -40,6 +40,9 @@ def forward(self, x, y):
             def get_methods_to_export(self):
                 return ("forward",)
 
+            def get_inputs(self):
+                return (torch.ones(2, 2), torch.ones(2, 2))
+
         class ModuleMulti(torch.nn.Module):
             """The module to serialize and execute."""
 
@@ -55,13 +58,31 @@ def forward2(self, x, y):
             def get_methods_to_export(self):
                 return ("forward", "forward2")
 
+            def get_inputs(self):
+                return (torch.ones(2, 2), torch.ones(2, 2))
+
+        class ModuleAddSingleInput(torch.nn.Module):
+            """The module to serialize and execute."""
+
+            def __init__(self):
+                super(ModuleAddSingleInput, self).__init__()
+
+            def forward(self, x):
+                return x + x
+
+            def get_methods_to_export(self):
+                return ("forward",)
+
+            def get_inputs(self):
+                return (torch.ones(2, 2),)
+
         def create_program(
             eager_module: torch.nn.Module,
         ) -> Tuple[Program, Tuple[Any, ...]]:
             """Returns an executorch program based on ModuleAdd, along with inputs."""
 
             # Trace the test module and create a serialized ExecuTorch program.
-            inputs = (torch.ones(2, 2), torch.ones(2, 2))
+            inputs = eager_module.get_inputs()
             input_map = {}
             for method in eager_module.get_methods_to_export():
                 input_map[method] = inputs
@@ -116,8 +137,39 @@ def lower_function_call():
             outputs = lower_function_call()
             tester.assertTrue(torch.allclose(outputs[0], torch.ones(2, 2) * 2))
 
+        def test_module_callable(tester):
+            # Create an ExecuTorch program from ModuleAdd.
+            exported_program, inputs = create_program(ModuleAdd())
+
+            # Use pybindings to load and execute the program.
+            executorch_module = load_fn(exported_program.buffer)
+            # Invoke the callable on executorch_module instead of calling module.forward.
+            executorch_output = executorch_module(inputs)[0]
+
+            # The test module adds the two inputs, so its output should be the same
+            # as adding them directly.
+            expected = inputs[0] + inputs[1]
+            tester.assertEqual(str(expected), str(executorch_output))
+
+        def test_module_single_input(tester):
+            # Create an ExecuTorch program from ModuleAdd.
+            exported_program, inputs = create_program(ModuleAddSingleInput())
+
+            # Use pybindings to load and execute the program.
+            executorch_module = load_fn(exported_program.buffer)
+            # Inovke the callable on executorch_module instead of calling module.forward.
+            # Use only one input to test this case.
+            executorch_output = executorch_module(inputs[0])[0]
+
+            # The test module adds the two inputs, so its output should be the same
+            # as adding them directly.
+            expected = inputs[0] + inputs[0]
+            tester.assertEqual(str(expected), str(executorch_output))
+
         test_e2e(tester)
         test_multiple_entry(tester)
         test_output_lifespan(tester)
+        test_module_callable(tester)
+        test_module_single_input(tester)
 
     return wrapper

extension/pybindings/test/test_pybindings.py

@@ -35,4 +35,4 @@
 
 class PybindingsTest(unittest.TestCase):
     def test(self):
-        make_test(self, _load_for_executorch_from_buffer)
+        make_test(self, _load_for_executorch_from_buffer)(self)