Intorduce XNNPACKHeaderto manage flatbuffer data and constant data (#1523)

mcr229 · facebook-github-bot · commit 842c246b6ec5 · 2024-01-11T15:10:24.000-08:00
Summary:

Introducing the XNNPACKHeader to manage the flatbuffer data and constant data.

Previously, we have serialized constant data along with flatbuffer. However, with large weights and large tensors in general, this takes a large amount of time and memory converting our dataclass --&gt; json --&gt; flatbuffer. This has become a blocker on some larger models

To fix, we circumvent serializing constant tensors via flatbuffer, by appending the constant data after the flatbuffer payload. In order to do this, we need an XNNPACKHeader which will give us the flatbuffer offset, flatbuffer size, constant data offset, and constant data sizes.

It will look something like this:
	
```
             ┌───────────────────────────────────┐
             │XNNPACK Header                     │
             ├───────────────────────────────────┤
             │Padding for 16 byte alignment      │
             ├───────────────────────────────────┤
             │Flatbuffer-serialized payload data │
             │                                   │
             │                                   │
             ├───────────────────────────────────┤
             │Padding for 16 byte alignment      │
             ├───────────────────────────────────┤
             │Constant Data                      │
             │                                   │
             │                                   │
             └───────────────────────────────────┘
```

Within the XNNPACK Header, we hold the following:
- 4 bytes to offset the header magic
- 4 bytes for the header magic
- 2 bytes for the header length
- 4 bytes for the flatbuffer offset
- 4 bytes for the flatbuffer size
- 4 bytes for constant data offset
- 8 bytes for constant data size

Reviewed By: digantdesai

Differential Revision: D52497977
diff --git a/backends/xnnpack/serialization/xnnpack_graph_serialize.py b/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)
diff --git a/backends/xnnpack/test/TARGETS b/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",
+    ],
+)
diff --git a/backends/xnnpack/test/serialization/test_xnnheader.py b/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)
