[pytorch/executorch][diff_train] Serialize NamedData in PTE file (#8847)

1. Serialize NamedData in PTE file
2. Add NamedDataStore to EdgeProgramManager
 ---
Serializing NamedData is slightly different to constant/delegate data as each segment comes with its own alignment.

**An example:**
Given NamedData = {"key": data}. Data is 250 bytes.

- BackendA requires data with alignment=3
- BackendB requires data with alignment=4

Then, data0 should be serialized with alignment of lcm(3, 4) = 12

At serialization, ExecuTorch has a 'segment_alignment' that defaults to 128. Data is now serialized to lcm(12, 128) = 384.

Inside the DataSegment, we want to store the original size of the data (250). The offset of the subsequent DataSegment would be 384 bytes after the start of this one.

**Design**
Introduce a new dataclass 'AlignedData' that stores the buffer and any alignment that's required. This is used when assembling Program.segments to ensure we get lcm(buffer_alignment, segment_alignment).


Note: The default segment_alignment can be overridden inside 'ExecutorchBackendConfig'.

Pull Request resolved: #8835


Internal:
<< DO NOT EDIT BELOW THIS LINE >>

**GitHub Author**: pytorchbot <[email protected]>
**GitHub Repo**: [pytorch/executorch](https://github.com/pytorch/executorch)
**GitHub Pull Request**: [#8835](#8835)

Initially generated by: https://www.internalfb.com/intern/sandcastle/job/18014400207802303/

This was imported as part of a Diff Train.
Please review this as soon as possible. Since it is a direct copy of a commit on
GitHub, there shouldn't be much to do.

diff-train-source-id: 314ab18

Co-authored-by: lucylq <[email protected]>

Differential Revision: [D69764150](https://our.internmc.facebook.com/intern/diff/D69764150/)
ghstack-source-id: 269082799

Author: lucylq

exir/_serialize/_program.py

@@ -8,10 +8,11 @@
 
 import copy
 import json
+import math
 import re
 
 from dataclasses import dataclass
-from typing import ClassVar, List, Literal, Optional, Tuple
+from typing import ClassVar, Dict, List, Literal, Optional, Tuple
 
 from executorch.exir._serialize._cord import Cord
 from executorch.exir._serialize._dataclass import _DataclassEncoder, _json_to_dataclass
@@ -20,6 +21,10 @@
     _program_flatbuffer_to_json,
     _program_json_to_flatbuffer,
 )
+from executorch.exir._serialize._named_data_store import (
+    BufferEntry,
+    NamedDataStoreOutput,
+)
 
 from executorch.exir._serialize.padding import aligned_size, pad_to, padding_required
 
@@ -29,6 +34,7 @@
     Buffer,
     DataLocation,
     DataSegment,
+    NamedData,
     Program,
     SubsegmentOffsets,
 )
@@ -41,6 +47,24 @@
 _HEADER_BYTEORDER: Literal["little"] = "little"
 
 
+@dataclass
+class AlignedData:
+    """
+    Holds data that should be aligned, for serialization.
+
+    Attributes:
+        data: The data to serialize, as a cord.
+        alignment: The alignment required for the data.
+    """
+
+    data: Cord
+    alignment: int
+
+    def __init__(self, data: Cord, alignment: Optional[int] = None) -> None:
+        self.data = data
+        self.alignment = alignment or 1
+
+
 def _program_to_json(program: Program) -> str:
     """Returns the JSON representation of the given Program."""
     return json.dumps(program, cls=_DataclassEncoder)
@@ -213,7 +237,7 @@ def _get_extended_header(program_data: bytes) -> Optional[_ExtendedHeader]:
 
 def _extract_delegate_segments(
     program: Program,
-    segments: List[Cord],
+    segments: List[AlignedData],
 ) -> None:
     """Extracts the delegate segments inlined in the program into a list of buffers.
         The program is modified in-place to remove the delegate data.
@@ -253,7 +277,7 @@ def _extract_delegate_segments(
                 segment_index = segment_index_map.get(inline.data)
                 if segment_index is None:
                     segment_index = len(segments)
-                    segments.append(Cord(inline.data))
+                    segments.append(AlignedData(Cord(inline.data)))
                     segment_index_map[inline.data] = segment_index
                 delegate.processed = BackendDelegateDataReference(
                     location=DataLocation.SEGMENT,
@@ -316,6 +340,44 @@ def _extract_constant_segment(
     return constant_segment_data, constant_segment_offsets
 
 
+def _extract_named_data(
+    program: Program,
+    segments: List[AlignedData],
+    buffers: List[BufferEntry],
+    name_to_buffer_idx: Dict[str, int],
+) -> None:
+    """Modifies the program in-place to add references to the named data
+        segments.
+
+    Args:
+        program: The program to extract segments from. Modified in-place.
+        segments: A list of buffers to append extracted segments to. Modified in-place.
+        buffers: A list of unique buffers and the information required to
+            serialize them. Not modified.
+        name_to_buffer_idx: A map from the name of a blob to the index in buffers.
+            Not modified.
+    """
+    if program.named_data is not None and len(program.named_data) > 0:
+        raise ValueError("Program already has named data.")
+
+    # Map from buffer_idx to segment_idx.
+    segment_index_map: Dict[int, int] = {}
+
+    named_data: List[NamedData] = []
+    for name, buffer_idx in name_to_buffer_idx.items():
+        segment_index = segment_index_map.get(buffer_idx, None)
+        if segment_index is None:
+            segment_index = len(segments)
+            segment_index_map[buffer_idx] = segment_index
+            segments.append(
+                AlignedData(
+                    Cord(buffers[buffer_idx].buffer), buffers[buffer_idx].alignment
+                )
+            )
+        named_data.append(NamedData(key=name, segment_index=segment_index))
+    program.named_data = named_data
+
+
 def serialize_pte_binary(
     program: Program,
     *,
@@ -324,6 +386,7 @@ def serialize_pte_binary(
     segment_alignment: int = 128,
     constant_tensor_alignment: Optional[int] = None,
     delegate_alignment: Optional[int] = None,
+    named_data: Optional[NamedDataStoreOutput] = None,
 ) -> Cord:
     """Returns the runtime binary representation of the given Program.
 
@@ -343,6 +406,8 @@ def serialize_pte_binary(
         delegate_alignment: If provided, the minimum alignment of delegate data
             in the program. Must be a power of 2. If not provided, uses the
             value in the schema file.
+        named_data: If provided, named blobs to be stored in segments
+            after the PTE file.
     Returns:
         The serialized form of the Program, ready for execution by the runtime.
     """
@@ -355,8 +420,9 @@ def serialize_pte_binary(
     # copy, reusing the actual data blobs.
     program = copy.deepcopy(program)
 
-    # Store extracted segment data; this may be constant data or delegate data.
-    segments: List[Cord] = []
+    # Store extracted segment data, with any buffer-specific alignment.
+    # This may be constant data, delegate data or named data.
+    segments: List[AlignedData] = []
 
     constant_segment_data, constant_segment_offsets = _extract_constant_segment(
         program.constant_buffer, tensor_alignment=constant_tensor_alignment
@@ -374,7 +440,7 @@ def serialize_pte_binary(
         # Clear the constant buffer, as constant data will be stored in segments.
         program.constant_buffer = []
         # Add to the aggregate segments cord.
-        segments.append(constant_segment_data)
+        segments.append(AlignedData(constant_segment_data))
 
     if mutable_data is not None:
         mutable_segment_data, mutable_segment_offsets = _extract_constant_segment(
@@ -389,31 +455,34 @@ def serialize_pte_binary(
                 ),
             ]
             # Add to the aggregate segments cord.
-            segments.append(mutable_segment_data)
+            segments.append(AlignedData(mutable_segment_data))
 
     if extract_delegate_segments:
         _extract_delegate_segments(program, segments)
+    if named_data is not None:
+        _extract_named_data(program, segments, named_data.buffers, named_data.pte_data)
 
     # Append all segments into a single Cord, adding any necessary padding to ensure that
     # each segment begins at the required alignment.
     # Update program.segments with the offsets to each segment.
     segments_data = Cord()
-    for data in segments:
+    for segment in segments:
         prev_end = (
             (program.segments[-1].offset + program.segments[-1].size)
             if program.segments
             else 0
         )
+        alignment = math.lcm(segment_alignment, segment.alignment)
         program.segments.append(
             DataSegment(
-                offset=aligned_size(prev_end, segment_alignment), size=len(data)
+                offset=aligned_size(prev_end, alignment), size=len(segment.data)
             )
         )
         # Add to aggregate segments cord with padding.
-        padding_length = padding_required(len(segments_data), segment_alignment)
+        padding_length = padding_required(len(segments_data), alignment)
         if padding_length > 0:
             segments_data.append(b"\x00" * padding_length)
-        segments_data.append(data)
+        segments_data.append(segment.data)
 
     # Convert to a standard flatbuffer binary.
     result: _FlatbufferResult = _program_json_to_flatbuffer(

exir/_serialize/_serialize.py

@@ -6,12 +6,12 @@
 
 # pyre-strict
 
-
-from typing import Dict, Tuple
+from typing import Dict, Optional, Tuple
 
 from executorch.exir._serialize import _serialize_pte_binary
 
 from executorch.exir._serialize._cord import Cord
+from executorch.exir._serialize._named_data_store import NamedDataStoreOutput
 from executorch.exir._serialize.data_serializer import (
     DataPayload,
     DataSerializer,
@@ -28,17 +28,32 @@ def serialize_for_executorch(
     emitter_output: EmitterOutput,
     config: ExecutorchBackendConfig,
     data_serializer: DataSerializer,
+    named_data: Optional[NamedDataStoreOutput] = None,
 ) -> Tuple[Cord, Dict[str, Cord]]:
     """Serialize the output from Emitter into ExecuTorch artifacts; PTE and PTD files."""
 
     # Serialize PTE file.
+    pte_named_data = None
+    if (
+        named_data is not None
+        and len(named_data.buffers) > 0
+        and len(named_data.pte_data) > 0
+    ):
+        # Create a separate NamedDataStoreOutput with only pte_data; exclude
+        # external_data, which shouldn't be serialized with the PTE file.
+        pte_named_data = NamedDataStoreOutput(
+            buffers=named_data.buffers,
+            pte_data=named_data.pte_data,
+            external_data={},
+        )
     pte: Cord = _serialize_pte_binary(
         program=emitter_output.program,
         mutable_data=emitter_output.mutable_data,
         extract_delegate_segments=config.extract_delegate_segments,
         segment_alignment=config.segment_alignment,
         constant_tensor_alignment=config.constant_tensor_alignment,
         delegate_alignment=config.delegate_alignment,
+        named_data=pte_named_data,
     )
 
     # Serialize PTD files.
@@ -88,4 +103,10 @@ def serialize_for_executorch(
                 )
             )
 
+    if named_data is None or len(named_data.external_data) == 0:
+        return pte, ptd_files
+
+    if len(named_data.buffers) == 0:
+        raise RuntimeError("External data exists, but there are no buffers provided.")
+
     return pte, ptd_files