Intorduce XNNPACKHeaderto manage flatbuffer data and constant data

backends/xnnpack/serialization/xnnpack_graph_serialize.py

@@ -8,14 +8,20 @@
 import os
 import tempfile
 
-from dataclasses import fields, is_dataclass
+from dataclasses import dataclass, fields, is_dataclass
+from typing import ClassVar, Literal
 
 import pkg_resources
 from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import XNNGraph
 from executorch.exir._serialize._dataclass import _DataclassEncoder
 
 from executorch.exir._serialize._flatbuffer import _flatc_compile
 
+# Byte order of numbers written to program headers. Always little-endian
+# regardless of the host system, since all commonly-used modern CPUs are little
+# endian.
+_HEADER_BYTEORDER: Literal["little"] = "little"
+
 
 def sanity_check_xnngraph_dataclass(table, name: str = ""):
     """
@@ -68,6 +74,168 @@ def check_for_sym(obj, name):
             check_for_sym(o, _name_field)
 
 
+@dataclass
+class XNNHeader:
+    # Class Constants
+    MAGIC_OFFSET: ClassVar[slice] = slice(4, 8)
+    HEADER_SIZE_OFFSET: ClassVar[slice] = slice(8, 10)
+    FLATBUFFER_OFFSET_OFFSET: ClassVar[slice] = slice(10, 14)
+    FLATBUFFER_SIZE_OFFSET: ClassVar[slice] = slice(14, 18)
+    CONSTANT_DATA_OFFSET_OFFSET: ClassVar[slice] = slice(18, 22)
+    CONSTANT_DATA_SIZE_OFFSET: ClassVar[slice] = slice(22, 30)
+
+    # magic bytes that should be at the beginning of the header
+    EXPECTED_MAGIC: ClassVar[bytes] = b"XH00"
+    # The length of the header in bytes.
+    EXPECTED_LENGTH: ClassVar[int] = (
+        # Zeros magic
+        # We offset the magic by 4 bytes so that it is in the same location
+        # as the flatbuffer payload's magic. This way we can dynamically
+        # choose between the XNNPACK Header and Flatbuffer Header
+        4
+        # Header magic
+        + 4
+        # Header Length
+        + 2
+        # Flatbuffer offset
+        + 4
+        # Flatbuffer size
+        + 4
+        # Constant Data offset
+        + 4
+        # Constant Data size
+        + 8
+    )
+
+    # Instance attributes. @dataclass will turn these into ctor args.
+
+    # offset to the flatbuffer data
+    flatbuffer_offset: int
+
+    # flatbuffer size
+    flatbuffer_size: int
+
+    # offset to the constant data
+    constant_data_offset: int
+
+    # constant data size
+    constant_data_size: int
+
+    @staticmethod
+    def from_bytes(data: bytes) -> "XNNHeader":
+        """
+        Converts the given bytes into an XNNHeader object.
+
+        We check that the magic and length is valid, but do not check that the offset and
+        size values are valid. We ensure here that the XNNHeader metadata is valid (magic and length)
+        but not the offsets and sizes themselves. Callers should use is_valid() to validate the
+        header contents
+
+        Args:
+            data: Data to read from
+        Returns:
+            XNNHeader object that contains the parsed data
+        Raises:
+            ValueError: if not enough data is provided, or if parsed length/magic are invalid
+        """
+        if len(data) > XNNHeader.EXPECTED_LENGTH:
+            raise ValueError(
+                f"Invalid XNNHeader: expected no more than {XNNHeader.EXPECTED_LENGTH} bytes, got {len(data)}"
+            )
+
+        magic: bytes = data[XNNHeader.MAGIC_OFFSET]
+        length_bytes: bytes = data[XNNHeader.HEADER_SIZE_OFFSET]
+        flatbuffer_offset_bytes: bytes = data[XNNHeader.FLATBUFFER_OFFSET_OFFSET]
+        flatbuffer_size_bytes: bytes = data[XNNHeader.FLATBUFFER_SIZE_OFFSET]
+        constant_data_offset_bytes: bytes = data[XNNHeader.CONSTANT_DATA_OFFSET_OFFSET]
+        constant_data_size_bytes: bytes = data[XNNHeader.CONSTANT_DATA_SIZE_OFFSET]
+
+        length = int.from_bytes(length_bytes, byteorder=_HEADER_BYTEORDER)
+
+        if magic != XNNHeader.EXPECTED_MAGIC:
+            raise ValueError(
+                f"Invalid XNNHeader: invalid magic bytes {magic}, expected {XNNHeader.EXPECTED_MAGIC}"
+            )
+        if length != len(data):
+            raise ValueError(
+                f"Invalid XNNHeader: Invalid parsed length: data given was {len(data)} bytes, parsed length was {length} bytes"
+            )
+
+        return XNNHeader(
+            flatbuffer_offset=int.from_bytes(
+                flatbuffer_offset_bytes, byteorder=_HEADER_BYTEORDER
+            ),
+            flatbuffer_size=int.from_bytes(
+                flatbuffer_size_bytes, byteorder=_HEADER_BYTEORDER
+            ),
+            constant_data_offset=int.from_bytes(
+                constant_data_offset_bytes, byteorder=_HEADER_BYTEORDER
+            ),
+            constant_data_size=int.from_bytes(
+                constant_data_size_bytes, byteorder=_HEADER_BYTEORDER
+            ),
+        )
+
+    def is_valid(self) -> bool:
+        """
+        Sanity checks the the XNNHeader.
+
+        We check that the flatbuffer size is non_zero and that the constant data offset
+        is after the flatbuffer payload. We check that the constant data size is non-negative.
+
+        Returns:
+            True if the XNNHeader is valid, False otherwise
+        """
+        # flatbuffer payload must have a non-zero size
+        valid_flatbuffer_size = self.flatbuffer_size > 0
+        # constant data offset is after flatbuffer payload
+        valid_const_data_offset = (
+            self.constant_data_offset >= self.flatbuffer_offset + self.flatbuffer_size
+        )
+        valid_const_data_size = self.constant_data_size >= 0
+
+        return (
+            valid_flatbuffer_size and valid_const_data_offset and valid_const_data_size
+        )
+
+    def to_bytes(self) -> bytes:
+        """
+        Converts XNNHeader to bytes for serialization.
+
+        Returns:
+            Returns the binary representation of the XNNPACK Header.
+        """
+
+        # We expect the given offsets and sizes to be valid
+        if not self.is_valid():
+            raise ValueError("Invalid XNNHeader: header failed is_valid() check")
+
+        data: bytes = (
+            # Padding for magic bytes. This is so that header magic is in the same position
+            # as the flatbuffer magic, and allows consumer to detect whether the header is
+            # being used or not
+            b"\x00\x00\x00\x00"
+            # XNNPACK Header's magic. This allows consumer to detect whether or not the header
+            # is being used or the flatbuffer header is being used
+            + self.EXPECTED_MAGIC
+            # uint16_t: Size of this header. This makes it easier to add new fields to the header
+            # in the future.
+            + self.EXPECTED_LENGTH.to_bytes(2, byteorder=_HEADER_BYTEORDER)
+            # uint32_t: Offset to the start of the flatbuffer data
+            + self.flatbuffer_offset.to_bytes(4, byteorder=_HEADER_BYTEORDER)
+            # uint32_t: Size of the flatbuffer data payload
+            + self.flatbuffer_size.to_bytes(4, byteorder=_HEADER_BYTEORDER)
+            # uint32_t: Offset to the start of the constant data
+            + self.constant_data_offset.to_bytes(4, byteorder=_HEADER_BYTEORDER)
+            # uint64_t: Size of the constant data
+            + self.constant_data_size.to_bytes(8, byteorder=_HEADER_BYTEORDER)
+        )
+
+        assert len(data) == XNNHeader.EXPECTED_LENGTH
+
+        return data
+
+
 def convert_to_flatbuffer(xnnpack_graph: XNNGraph) -> bytes:
     sanity_check_xnngraph_dataclass(xnnpack_graph)
     xnnpack_graph_json = json.dumps(xnnpack_graph, cls=_DataclassEncoder)

backends/xnnpack/test/TARGETS

@@ -1,6 +1,8 @@
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 load(":targets.bzl", "define_common_targets")
 
+oncall("executorch")
+
 define_common_targets()
 
 runtime.python_test(
@@ -60,3 +62,13 @@ runtime.python_test(
         "libtorch",
     ],
 )
+
+runtime.python_test(
+    name = "test_xnnpack_serialization",
+    srcs = glob([
+        "serialization/*.py",
+    ]),
+    deps = [
+        "//executorch/backends/xnnpack:xnnpack_preprocess",
+    ],
+)

backends/xnnpack/test/serialization/test_xnnheader.py

@@ -0,0 +1,104 @@
+# 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 unittest
+
+from executorch.backends.xnnpack.serialization.xnnpack_graph_serialize import XNNHeader
+
+EXAMPLE_FLATBUFFER_OFFSET: int = 0x11223344
+EXAMPLE_FLATBUFFER_SIZE: int = 0x55667788
+EXAMPLE_CONSTANT_DATA_OFFSET: int = EXAMPLE_FLATBUFFER_OFFSET + EXAMPLE_FLATBUFFER_SIZE
+EXAMPLE_CONSTANT_DATA_SIZE: int = 0x99AABBCC99AABBCC
+
+# If header layout or magic changes, this test must change too.
+# The layout of the header is a contract, not an implementation detail
+EXAMPLE_HEADER_DATA: bytes = (
+    # zeros
+    b"\x00\x00\x00\x00"
+    # magic
+    + b"XH00"
+    # All Values below are littl Endian
+    # header length
+    + b"\x1E\x00"
+    # Flatbuffer Offset
+    + b"\x44\x33\x22\x11"
+    # Flatbuffer Size
+    + b"\x88\x77\x66\x55"
+    # Constant Data Offset
+    + b"\xCC\xAA\x88\x66"
+    # Constant Data Size
+    + b"\xCC\xBB\xAA\x99\xCC\xBB\xAA\x99"
+)
+
+
+class TestXNNHeader(unittest.TestCase):
+    def test_to_bytes(self) -> None:
+        header = XNNHeader(
+            EXAMPLE_FLATBUFFER_OFFSET,
+            EXAMPLE_FLATBUFFER_SIZE,
+            EXAMPLE_CONSTANT_DATA_OFFSET,
+            EXAMPLE_CONSTANT_DATA_SIZE,
+        )
+        self.assertEqual(header.to_bytes(), EXAMPLE_HEADER_DATA)
+        self.assertTrue(header.is_valid())
+
+    def test_from_bytes(self) -> None:
+        header = XNNHeader.from_bytes(EXAMPLE_HEADER_DATA)
+        self.assertEqual(header.flatbuffer_offset, EXAMPLE_FLATBUFFER_OFFSET)
+        self.assertEqual(header.flatbuffer_size, EXAMPLE_FLATBUFFER_SIZE)
+        self.assertEqual(header.constant_data_offset, EXAMPLE_CONSTANT_DATA_OFFSET)
+        self.assertEqual(header.constant_data_size, EXAMPLE_CONSTANT_DATA_SIZE)
+
+    def test_invalid_metadata(self) -> None:
+        WRONG_MAGIC_DATA = EXAMPLE_HEADER_DATA[0:4] + b"YT01" + EXAMPLE_HEADER_DATA[8:]
+        with self.assertRaisesRegex(
+            ValueError,
+            "Invalid XNNHeader: invalid magic bytes b'YT01', expected b'XH00'",
+        ):
+            XNNHeader.from_bytes(WRONG_MAGIC_DATA)
+
+        WRONG_LENGTH_DATA = (
+            EXAMPLE_HEADER_DATA[0:8] + b"\x1D\x00" + EXAMPLE_HEADER_DATA[10:]
+        )
+        with self.assertRaisesRegex(
+            ValueError,
+            "Invalid XNNHeader: Invalid parsed length: data given was 30 bytes, parsed length was 29 bytes",
+        ):
+            XNNHeader.from_bytes(WRONG_LENGTH_DATA)
+
+        with self.assertRaisesRegex(
+            ValueError,
+            "Invalid XNNHeader: expected no more than 30 bytes, got 31",
+        ):
+            XNNHeader.from_bytes(EXAMPLE_HEADER_DATA + b"\x00")
+
+    def test_invalid_flatbuffer_size(self) -> None:
+        header = XNNHeader(
+            EXAMPLE_FLATBUFFER_OFFSET,
+            0,
+            EXAMPLE_CONSTANT_DATA_OFFSET,
+            EXAMPLE_CONSTANT_DATA_SIZE,
+        )
+
+        with self.assertRaises(ValueError):
+            header.to_bytes()
+
+    def test_invalid_constant_data_offset(self) -> None:
+        header = XNNHeader(
+            EXAMPLE_FLATBUFFER_OFFSET,
+            EXAMPLE_FLATBUFFER_SIZE,
+            EXAMPLE_FLATBUFFER_OFFSET + EXAMPLE_FLATBUFFER_SIZE - 1,
+            EXAMPLE_CONSTANT_DATA_SIZE,
+        )
+
+        with self.assertRaises(ValueError):
+            header.to_bytes()
+
+    def test_to_bytes_same_as_from_bytes(self) -> None:
+        header = XNNHeader.from_bytes(EXAMPLE_HEADER_DATA)
+
+        to_bytes = header.to_bytes()
+        self.assertEquals(EXAMPLE_HEADER_DATA, to_bytes)