[CoreAI] Reject mutated user inputs at lowering (pytorch#23317)

ExecuTorch owns user inputs, so a Core AI delegate can never take over
their in-place mutation. `coreai-torch` lowers such a mutation to Core
AI state, which leaves the `.aimodel` boundary unable to match what
ExecuTorch passes in.
* `CoreAIBackend.preprocess` now raises before conversion when the
program mutates a user input: `Core AI delegates cannot mutate
ExecuTorch-owned user inputs: [...]`.
* Mutated buffers that the delegate owns are unaffected.
* Tests (`test_io_compat.py`):
* `test_mutated_user_input_is_rejected` lowers a model that mutates its
input and expects the error.
* The `MutableBufferBoundaryTest` docstring now describes the current
behavior.

Test plan: `python -m pytest backends/apple/coreai` gives 126 passed.

Author: metascroy

backends/apple/coreai/compiler/preprocess.py

@@ -603,6 +603,12 @@ def preprocess(
                     "build-time output directory (set it via coreai_sidecar_dir)"
                 )
 
+        mutated_inputs = edge_program.graph_signature.user_inputs_to_mutate
+        if mutated_inputs:
+            raise ValueError(
+                "Core AI delegates cannot mutate ExecuTorch-owned user inputs: "
+                f"{sorted(mutated_inputs.values())}"
+            )
         edge_inputs, edge_outputs = _edge_io(edge_program)
         input_names = [f"input_{i}" for i in range(len(edge_inputs))]
         output_names = [f"output_{i}" for i in range(len(edge_outputs))]

backends/apple/coreai/test/test_io_compat.py

@@ -26,7 +26,10 @@
     assert_io_compatible,
     io_mismatches,
 )
-from executorch.backends.apple.coreai.compiler.preprocess import _convert_to_aiprogram
+from executorch.backends.apple.coreai.compiler.preprocess import (
+    _convert_to_aiprogram,
+    CoreAIBackend,
+)
 from executorch.backends.apple.coreai.partition.partitioner import CoreAIPartitioner
 from executorch.exir import to_edge, to_edge_transform_and_lower
 from executorch.exir.lowered_backend_module import executorch_call_delegate
@@ -299,15 +302,23 @@ def forward(self, x):
 
 
 class MutableBufferBoundaryTest(unittest.TestCase):
-    """A mutated buffer crosses the boundary as an input as well as an output.
-
-    coreai gives it a graph argument and a result, so both sides of the
-    compatibility check have to count it.
-    """
+    """Mutated buffers the delegate owns stay off the ExecuTorch boundary."""
 
     def test_mutated_buffer_lowers(self):
         _lower(_MutatedBuffer(), (torch.randn(2, 8),))
 
+    def test_mutated_user_input_is_rejected(self):
+        class MutatedInput(nn.Module):
+            def forward(self, x):
+                x.add_(1.0)
+                return x * 2.0
+
+        edge = to_edge(
+            torch.export.export(MutatedInput(), (torch.randn(2, 8),))
+        ).exported_program()
+        with self.assertRaisesRegex(ValueError, "cannot mutate ExecuTorch-owned"):
+            CoreAIBackend.preprocess(edge, [])
+
 
 if __name__ == "__main__":
     unittest.main()