Restore program segments from a view of the input (pytorch#23340)

## What is wrong today

`deserialize_pte_binary` reads a `.pte` file back into a Python
`Program`. Large data, such as delegate payloads and constant tensors,
is stored in segments at the end of the file.

To restore them, it first takes everything after the program data as a
new `bytes` object, then cuts each segment out of that. Slicing `bytes`
makes a copy, so while it runs the process holds:

1. the input data,
2. a copy of all the segments together,
3. a second copy of each segment as it is restored.

For a program with several large segments, for example a delegate plus
its constants, that is gigabytes of extra memory just to read the file.

## What this change does

`deserialize_pte_binary` passes a `memoryview` of the input to
`_restore_segments` instead of a sliced copy. A `memoryview` slice
points at the same memory and copies nothing.

```python
return _restore_segments(
    program=program,
    segment_data=memoryview(program_data)[segment_base_offset:],
)
```

`_restore_segments` still copies each segment into its own `bytes`, so:

- the restored program carries the same types as before (`bytes` in
delegate data, constant buffers and named data),
- nothing in the result keeps the whole input alive.

Each segment is now copied once instead of twice. The output does not
change.

When a file has only one segment and it runs to the end of the file,
Python does not copy a slice that covers the whole object, so main made
only one copy there too, and this change saves nothing. It helps every
program with more than one segment, for example a delegate plus
constants, or several delegates.

## What was tested

- New unit test in `exir/_serialize/test/test_program.py`:
`_restore_segments` receives a `memoryview` whose underlying object is
the input itself, and the restored delegate data is `bytes` equal to the
original blob. It fails before this change (it receives a `bytes` copy)
and passes after it.
- The rest of `test_program.py` and
`backends/cuda/tests/test_merge_ptes.py` (which reads merged programs
back with this function) pass before and after.
- Read back a 2 GiB program with two 1 GiB delegate segments on Linux
aarch64. Peak memory added by `deserialize_pte_binary` fell from 4.0 GiB
to 2.0 GiB, and the restored delegates were identical.

Author: shoumikhin

exir/_serialize/_program.py

@@ -670,7 +670,7 @@ def _restore_named_data(
     return named_data_store.get_named_data_store_output()
 
 
-def _restore_segments(program: Program, segment_data: bytes) -> PTEFile:
+def _restore_segments(program: Program, segment_data: memoryview) -> PTEFile:
     """Moves segments from `segment_data` into `program`.
 
     This should recreate the original Program that the segments were extracted
@@ -693,7 +693,9 @@ def _restore_segments(program: Program, segment_data: bytes) -> PTEFile:
             raise ValueError(
                 f"Segment {i} {segment} overflows data length {len(segment_data)}"
             )
-        segments.append(segment_data[segment.offset : segment.offset + segment.size])
+        segments.append(
+            bytes(segment_data[segment.offset : segment.offset + segment.size])
+        )
 
     # Restore delegate segments that weren't inlined previously.
     program = _restore_delegates(program, segments)
@@ -754,9 +756,11 @@ def deserialize_pte_binary(program_data: bytes) -> PTEFile:
     program: Program = _flatbuffer_to_program(program_data[:program_size])
 
     if segment_base_offset != 0:
-        # Move segment data back into the Program.
+        # A view, so the segment data is copied only once, when each segment is
+        # restored, rather than also as a whole before being split.
         return _restore_segments(
-            program=program, segment_data=program_data[segment_base_offset:]
+            program=program,
+            segment_data=memoryview(program_data)[segment_base_offset:],
         )
 
     return PTEFile(program=program, mutable_data=None, named_data=None)

exir/_serialize/test/test_program.py

@@ -16,13 +16,15 @@
 import unittest
 
 from typing import Dict, List, Sequence
+from unittest.mock import patch
 
 from executorch.exir._serialize._flatbuffer_program import _flatbuffer_to_program
 from executorch.exir._serialize._named_data_store import NamedDataStoreOutput
 from executorch.exir._serialize._program import (
     _ExtendedHeader,
     _get_extended_header,
     _program_to_json,
+    _restore_segments,
     deserialize_pte_binary,
     PTEFile,
     serialize_pte_binary,
@@ -749,6 +751,32 @@ def test_round_trip_with_segments(self) -> None:
         self.assertEqual(deserialized.mutable_data, None)
         self.assertEqual(deserialized.named_data, None)
 
+    def test_deserialize_restores_segments_from_a_view_of_the_input(self) -> None:
+        program = get_test_program()
+        blob = self.gen_blob_data(SEGMENT_ALIGNMENT * 4, b"\x10\x11\x01")
+        add_delegate_data(program, program.execution_plan[0], [blob])
+        pte_data = bytes(
+            serialize_pte_binary(
+                PTEFile(program=program),
+                extract_delegate_segments=True,
+                segment_alignment=SEGMENT_ALIGNMENT,
+            )
+        )
+
+        with patch(
+            "executorch.exir._serialize._program._restore_segments",
+            wraps=_restore_segments,
+        ) as restore_segments:
+            deserialized = deserialize_pte_binary(pte_data)
+
+        # A copy of the segment data is as large as all the delegates together.
+        segment_data = restore_segments.call_args.kwargs["segment_data"]
+        self.assertIsInstance(segment_data, memoryview)
+        self.assertIs(segment_data.obj, pte_data)
+        restored = deserialized.program.backend_delegate_data[-1].data
+        self.assertIsInstance(restored, bytes)
+        self.assertEqual(restored, blob)
+
     def test_no_constants(self) -> None:
         program = get_test_program()
         # Insert placeholder for non-const tensors.