[runtime] Let Method.execute borrow outputs

Module.forward and the native method take clone_outputs, but the public
Method wrapper of executorch.runtime always copied every output, so a caller
of the recommended API could not avoid that copy.

Method.execute now takes clone_outputs, with the same default and meaning
as the bindings: False returns outputs that share the method's memory and
may change on the next execution.

Author: shoumikhin

runtime/__init__.py

@@ -136,16 +136,20 @@ class Method:
     def __init__(self, method: ExecuTorchMethod) -> None:
         self._method = method
 
-    def execute(self, inputs: Sequence[Any]) -> Sequence[Any]:
+    def execute(
+        self, inputs: Sequence[Any], clone_outputs: bool = True
+    ) -> Sequence[Any]:
         """Executes the method with the given inputs.
 
         Args:
             inputs: A sequence of input values, typically torch.Tensor objects.
+            clone_outputs: If False, the outputs share the method's memory and
+                a later execution may overwrite them, which avoids a copy.
 
         Returns:
             A list of output values, typically torch.Tensor objects.
         """
-        return self._method(inputs)
+        return self._method(inputs, clone_outputs=clone_outputs)
 
     @property
     def metadata(self) -> MethodMeta:
@@ -188,9 +192,11 @@ def load_method(self, name: str) -> Optional[Method]:
         program placed them. A method whose buffers are all on the host shares
         one set of arenas with every other host-only method of this program,
         sized to the largest of them, so two such methods overwrite each
-        other's intermediate values. Outputs are copied out on every call, so
-        that sharing does not change what execute returns. A method with any
-        buffer placed on an accelerator gets its own arenas instead, claimed
+        other's intermediate values. By default execute copies its outputs out,
+        so that sharing does not change what it returns. With
+        ``clone_outputs=False`` the outputs share that memory, so the next
+        execution of any method of this program may overwrite them. A method
+        with any buffer placed on an accelerator gets its own arenas instead, claimed
         when that method is first loaded and owned by that method until the
         method itself is released. This program caches every method it loads,
         so in normal use those arenas live as long as the program does. A

runtime/test/test_runtime.py

@@ -52,6 +52,23 @@ def test_smoke(self):
         outputs = method.execute(inputs)
         self.assertTrue(torch.allclose(outputs[0], inputs[0] + inputs[1]))
 
+    def test_execute_can_borrow_outputs(self):
+        ep, inputs = create_program(ModuleAdd())
+        method = Runtime.get().load_program(ep.buffer).load_method("forward")
+
+        def address(output):
+            return torch.from_dlpack(output).data_ptr()
+
+        # Borrowed outputs are views of the method's own memory, so each
+        # execution writes the same place. Cloned outputs are copies.
+        borrowed = [method.execute(inputs, clone_outputs=False)[0] for _ in range(2)]
+        self.assertEqual(address(borrowed[0]), address(borrowed[1]))
+        cloned = [method.execute(inputs)[0] for _ in range(2)]
+        self.assertNotEqual(address(cloned[0]), address(cloned[1]))
+        torch.testing.assert_close(
+            torch.from_dlpack(borrowed[1]), inputs[0] + inputs[1]
+        )
+
     def test_module_with_multiple_method_names(self):
         ep, inputs = create_program(ModuleMulti())
         runtime = Runtime.get()