Serialize a program without copying its delegate blobs (pytorch#24458)

### What is wrong today

`to_executorch()` serializes the program right away. Before it moves
each delegate blob out of the program and into a segment,
`serialize_pte_binary` deep-copies the program, so that the caller's
program does not change.

`copy.deepcopy` shares an object whose type is exactly `bytes`. But it
copies an object of any other bytes type, for example a subclass of
`bytes`. The Core ML backend returns such a type for models over 1 GiB:
it maps the model files into memory instead of reading them. So for
these models, `to_executorch()` briefly held 2 more copies of the
delegate blob, and the serialized program kept 1 of them for as long as
the `ExecutorchProgramManager` lived.

With a 4 GiB delegate blob of a `bytes` subclass, peak memory in
`to_executorch()` grew by 8 GiB, and 4 GiB stayed in use after it
returned.

### What this change does

It seeds the `copy.deepcopy` memo with the delegate blobs, so the copied
program shares them instead of copying them. The blobs are immutable, so
sharing them is safe. The rest of the program is still copied, and the
output does not change.

### What was tested

- New test `test_serialize_shares_delegate_blobs` serializes a program
whose delegate blob is a `bytes` subclass, and checks that the output
holds that same object. It fails without this change.
- The `exir/_serialize` tests pass, and the other `exir` tests give the
same results with and without this change.
- A test backend whose `preprocess` returns a 4 GiB blob, on Linux
x86_64. With a `bytes` subclass, the extra peak memory in
`to_executorch()` went from 8 GiB to none, and the `.pte` file is byte
for byte the same as before. With plain `bytes`, nothing changed.

Author: shoumikhin

exir/_serialize/_program.py

@@ -413,10 +413,17 @@ def serialize_pte_binary(
     if constant_tensor_alignment is None:
         constant_tensor_alignment = ALIGNMENT
 
-    # Don't modify the original program.
-    # TODO(T144120904): Could avoid yet more huge copies with a more shallow
-    # copy, reusing the actual data blobs.
-    program = copy.deepcopy(pte_file.program)
+    # Don't modify the original program, but share its delegate blobs rather
+    # than copy them: deepcopy shares only exact bytes, and Core ML returns
+    # another bytes type for models over 1 GiB. The memo lives only for this
+    # call, because deepcopy fills it with every object it copies.
+    program = copy.deepcopy(
+        pte_file.program,
+        memo={
+            id(entry.data): entry.data
+            for entry in pte_file.program.backend_delegate_data
+        },
+    )
 
     # Store extracted segment data, with any buffer-specific alignment.
     # This may be constant data, delegate data or named data.

exir/_serialize/test/test_program.py

@@ -777,6 +777,24 @@ def test_deserialize_restores_segments_from_a_view_of_the_input(self) -> None:
         self.assertIsInstance(restored, bytes)
         self.assertEqual(restored, blob)
 
+    def test_serialize_shares_delegate_blobs(self) -> None:
+        # Not exactly bytes, like the blob Core ML returns for models over
+        # 1 GiB, which deepcopy would copy.
+        class Blob(bytes):
+            pass
+
+        program = get_test_program()
+        blob = Blob(self.gen_blob_data(SEGMENT_ALIGNMENT, b"\x10\x11\x01"))
+        add_delegate_data(program, program.execution_plan[0], [blob])
+
+        pte_data = serialize_pte_binary(
+            PTEFile(program=program),
+            extract_delegate_segments=True,
+            segment_alignment=SEGMENT_ALIGNMENT,
+        )
+
+        self.assertTrue(any(buffer is blob for buffer in pte_data._buffers))
+
     def test_no_constants(self) -> None:
         program = get_test_program()
         # Insert placeholder for non-const tensors.