Serialize constant Data outside of flatbuffer (#1542)

Summary:

We introduce the `serialize_xnnpack_binary` method which serializees the constant data outside of the flatbuffer. It leverages the xnnheader introduced in the previous diff to store offsets and sizes for both the flatbuffer payload as well as the constant data payload.

Note here we have not yet switched the delegate to use the new `serialize_xnnpack_binary` function as this new serialization also requires changes on the runtime side. This will be tested in the diff which follows.

Differential Revision: D52498367

Author: mcr229

backends/xnnpack/serialization/schema.fbs

@@ -281,6 +281,17 @@ table XNNLeakyReLU {
   flags: uint;
 }
 
+// Describes data offsets for constant data
+table ConstantDataOffset {
+  // Constant data offsets are relative to the constant data base offset provided
+  // in the XNNPACKHeader.
+  offset: uint64;
+
+  // The size in bytes of valid data starting at the offset. The constant data
+  // may be followed by padding before the next piece of constant data
+  size: uint64;
+}
+
 table XNNGraph {
   // Schema version.
   version:string;
@@ -299,11 +310,16 @@ table XNNGraph {
   // Tables of constant data, used for constant Values (e.g.
   // data field of weight tensors). Each constant is assigned an index into the table
   // which are each individually aligned. 0 index is reserved to be pointed to by non-constant
-  // Tensors
+  // Tensors. Exactly one of constant_buffer and constant_data must be non-empty
   constant_buffer:[Buffer];
 
   // the list index is memory buffer id, the value is the memory buffer size.
   mem_buffer_sizes: [uint];
+
+  // List of the constant data that follows the XNNGraph in this file. Each constant data is assigned an index into
+  // the table. 0 index is reserved to be pointed to by non-constant Tensor. Exactly one of constant_buffer and
+  // constant_data must be non-empty
+  constant_data:[ConstantDataOffset];
 }
 
 root_type XNNGraph;

backends/xnnpack/serialization/xnnpack_graph_schema.py

@@ -417,6 +417,12 @@ class Buffer:
     storage: bytes
 
 
+@dataclass
+class ConstantDataOffset:
+    offset: int
+    size: int
+
+
 @dataclass
 class XNNGraph:
     version: str
@@ -429,3 +435,5 @@ class XNNGraph:
 
     constant_buffer: List[Buffer]
     mem_buffer_sizes: List[int]
+
+    constant_data: List[ConstantDataOffset]

backends/xnnpack/serialization/xnnpack_graph_serialize.py

@@ -9,10 +9,14 @@
 import tempfile
 
 from dataclasses import dataclass, fields, is_dataclass
-from typing import ClassVar, Literal
+from typing import ClassVar, List, Literal, Tuple
 
 import pkg_resources
-from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import XNNGraph
+from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import (
+    Buffer,
+    ConstantDataOffset,
+    XNNGraph,
+)
 from executorch.exir._serialize._dataclass import _DataclassEncoder
 
 from executorch.exir._serialize._flatbuffer import _flatc_compile
@@ -236,6 +240,72 @@ def to_bytes(self) -> bytes:
         return data
 
 
+def _padding_required(offset: int, alignment: int) -> int:
+    """Returns the padding required to align `offset` to `alignment`."""
+    remainder: int = offset % alignment
+    if remainder != 0:
+        return alignment - remainder
+    return 0
+
+
+def _aligned_size(input_size: int, alignment: int) -> int:
+    """Returns input_size padded up to the next whole multiple of alignment."""
+    return input_size + _padding_required(input_size, alignment)
+
+
+def _pad_to(data: bytes, length: int) -> bytes:
+    """Returns the input followed by enough zero bytes to become the requested length.
+
+    Args:
+        data: The data to pad.
+        length: The length of the returned data.
+    Returns:
+        The padded data.
+    Raises:
+        ValueError: If the requested length is less than the input length.
+    """
+    if length < len(data):
+        raise ValueError(f"Data length {len(data)} > padded length {length}")
+    if length > len(data):
+        data = data + b"\x00" * (length - len(data))
+    assert len(data) == length
+    return data
+
+
+def _extract_constant_data(
+    constant_buffer: List[Buffer],
+    tensor_alignment: int,
+) -> Tuple[bytes, List[int]]:
+    """Copies the tensors from the provided list into a single buffer and tracks the offsets
+    of each tensor.
+
+        constant_buffer: list of Buffers from which to extract constants from. Not modified.
+        tensor_alignment: Alignment in bytes. The starting offset of each tensor in the
+            constant segment will be aligned to this value. Default to 16.
+
+    Returns:
+        A tuple of (constant segment, list of offsets for each tensor in the segment)
+    """
+    constant_segment_data: bytearray = bytearray()
+    constant_segment_offsets: List[int] = []
+    current_offset: int = 0
+    for i in range(len(constant_buffer)):
+        buffer = constant_buffer[i]
+        buffer_length = len(buffer.storage)
+        pad_length = _padding_required(buffer_length, tensor_alignment)
+
+        # Append each constant buffer to the constant segment.
+        constant_segment_data += buffer.storage
+        # Add padding for all but the last tensor.
+        if i < len(constant_buffer) - 1:
+            constant_segment_data += b"\x00" * pad_length
+
+        # Append constant data offset.
+        constant_segment_offsets.append(current_offset)
+        current_offset += buffer_length + pad_length
+    return bytes(constant_segment_data), constant_segment_offsets
+
+
 def convert_to_flatbuffer(xnnpack_graph: XNNGraph) -> bytes:
     sanity_check_xnngraph_dataclass(xnnpack_graph)
     xnnpack_graph_json = json.dumps(xnnpack_graph, cls=_DataclassEncoder)
@@ -251,3 +321,67 @@ def convert_to_flatbuffer(xnnpack_graph: XNNGraph) -> bytes:
         output_path = os.path.join(d, "schema.bin")
         with open(output_path, "rb") as output_file:
             return output_file.read()
+
+
+def serialize_xnnpack_binary(xnnpack_graph: XNNGraph) -> bytes:
+    """Returns the runtime binary representation of the given XNNGraph.
+
+    Args:
+        xnnpack_graph: XNNGraph object to serialize.
+
+    Returns:
+        The serialized form of the XNNGraph, ready for execution by XNNPACK Backend
+    """
+    constant_tensor_alignment = 16
+
+    # Extract constant data from the graph
+    constant_data, constant_data_offsets = _extract_constant_data(
+        xnnpack_graph.constant_buffer, constant_tensor_alignment
+    )
+
+    assert len(constant_data_offsets) == len(xnnpack_graph.mem_buffer_sizes)
+
+    for offset_idx in range(len(constant_data_offsets)):
+        constant_data_offset = constant_data_offsets[offset_idx]
+        constant_data_size = xnnpack_graph.mem_buffer_sizes[offset_idx]
+        xnnpack_graph.constant_data.append(
+            ConstantDataOffset(constant_data_offset, constant_data_size)
+        )
+
+    # We are moving all constant data from the graph to the constant data section.
+    # So we remove all constant buffers
+    xnnpack_graph.constant_buffer = []
+    xnnpack_graph.mem_buffer_sizes = []
+
+    # Convert the XNNGraph to a flatbuffer
+    flatbuffer_payload = convert_to_flatbuffer(xnnpack_graph)
+
+    # size of flatbuffer data, padded to be 16 byte aligned
+    padded_flatbuffer_length: int = _aligned_size(
+        input_size=len(flatbuffer_payload),
+        alignment=constant_tensor_alignment,
+    )
+    # size of header to insert, padded to be `constant_tensor_alignment` byte aligned
+    padded_header_length: int = _aligned_size(
+        input_size=XNNHeader.EXPECTED_LENGTH,
+        alignment=constant_tensor_alignment,
+    )
+
+    # Create the XNNPACK Header
+    header: bytes = XNNHeader(
+        flatbuffer_offset=padded_header_length,
+        flatbuffer_size=len(flatbuffer_payload),
+        constant_data_offset=padded_header_length + padded_flatbuffer_length,
+        constant_data_size=len(constant_data),
+    ).to_bytes()
+
+    # Concatenate the header, flatbuffer data, and constant data
+    # Constant data does not need to be padded to alignment because nothing follows it
+
+    return b"".join(
+        [
+            _pad_to(header, padded_header_length),
+            _pad_to(flatbuffer_payload, padded_flatbuffer_length),
+            constant_data,
+        ]
+    )

backends/xnnpack/test/serialization/test_serialization.py

@@ -0,0 +1,98 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import random
+import unittest
+from typing import List, Tuple
+
+from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import (
+    Buffer,
+    XNNGraph,
+)
+
+from executorch.backends.xnnpack.serialization.xnnpack_graph_serialize import (
+    _HEADER_BYTEORDER,
+    serialize_xnnpack_binary,
+    XNNHeader,
+)
+
+
+class TestSerialization(unittest.TestCase):
+    def _generate_random_const_buffers(
+        self, num_tensors: int
+    ) -> Tuple[List[Buffer], List[int]]:
+        """
+        Helper function to generate `num_tensor` buffers of random sizes and random contents,
+        we return a tuple of (list_of_buffers, list_of_mem_sizes),
+        """
+        buffers = []
+        mem_sizes = []
+        for _ in range(num_tensors):
+            buffer_size = random.randint(1, 1000)
+            buffer = bytearray(os.urandom(buffer_size))
+            buffers.append(Buffer(storage=bytes(buffer)))
+            mem_sizes.append(buffer_size)
+
+        return buffers, mem_sizes
+
+    def test_serialize_xnnpack_binary(self):
+        xnn_graph = XNNGraph(
+            version="0",
+            xnodes=[],
+            xvalues=[],
+            num_externs=0,
+            input_ids=[],
+            output_ids=[],
+            constant_buffer=[Buffer(storage=b"")],
+            mem_buffer_sizes=[0],
+            constant_data=[],
+        )
+        buffers, sizes = self._generate_random_const_buffers(5)
+        xnn_graph.constant_buffer.extend(buffers)
+        xnn_graph.mem_buffer_sizes.extend(sizes)
+        buffers = xnn_graph.constant_buffer
+
+        serialized_binary = serialize_xnnpack_binary(xnn_graph)
+        offsets = xnn_graph.constant_data
+
+        # Check header
+        self.assertEqual(serialized_binary[0:4], b"\x00\x00\x00\x00")
+        self.assertEqual(serialized_binary[XNNHeader.MAGIC_OFFSET], b"XH00")
+        flatbuffer_offset_bytes = serialized_binary[XNNHeader.FLATBUFFER_OFFSET_OFFSET]
+        constant_data_offset_bytes = serialized_binary[
+            XNNHeader.CONSTANT_DATA_OFFSET_OFFSET
+        ]
+
+        # Check flatbuffer is at flatbuffer offset
+        flatbuffer_offset = int.from_bytes(
+            flatbuffer_offset_bytes, byteorder=_HEADER_BYTEORDER
+        )
+        # Flatbuffer magic should be in the same spot as the Header's magic
+        self.assertEqual(
+            serialized_binary[flatbuffer_offset:][XNNHeader.MAGIC_OFFSET], b"XN00"
+        )
+
+        # Check constant data
+        # Check that constant buffers have been moved to constant data
+        self.assertEqual(len(offsets), len(buffers))
+        self.assertEqual(len(xnn_graph.constant_buffer), 0)
+
+        constant_data_offset = int.from_bytes(
+            constant_data_offset_bytes, byteorder=_HEADER_BYTEORDER
+        )
+        constant_data_payload = serialized_binary[constant_data_offset:]
+
+        # We check that constant data indexes stored in the xnn_graph correctly index
+        # into the correct buffer in the constant data section
+        for idx in range(1, len(offsets)):
+            offset = offsets[idx].offset
+            size = offsets[idx].size
+
+            constant_data_bytes = constant_data_payload[offset : offset + size]
+            constant_buffer_bytes = buffers[idx].storage
+
+            self.assertEqual(constant_data_bytes, constant_buffer_bytes)

backends/xnnpack/xnnpack_preprocess.py

@@ -231,6 +231,7 @@ def preprocess(
             output_ids=[],
             constant_buffer=[Buffer(storage=b"")],
             mem_buffer_sizes=[0],
+            constant_data=[],
         )
 
         node_visitors = get_node_visitors(ep, node_to_external_map)