xref: /aosp_15_r20/external/executorch/backends/xnnpack/test/serialization/test_serialization.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# All rights reserved.
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7import unittest
8
9from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import (
10    ConstantDataOffset,
11    XNNGraph,
12)
13
14from executorch.backends.xnnpack.serialization.xnnpack_graph_serialize import (
15    _HEADER_BYTEORDER,
16    serialize_xnnpack_binary,
17    XNNHeader,
18)
19
20
21class TestSerialization(unittest.TestCase):
22    def test_serialize_xnnpack_binary(self):
23        xnn_graph = XNNGraph(
24            version="0",
25            xnodes=[],
26            xvalues=[],
27            num_externs=0,
28            input_ids=[],
29            output_ids=[],
30            constant_data=[ConstantDataOffset(0, 0)],
31        )
32
33        constant_data_bytes = b"\x00" * 24
34        serialized_binary = serialize_xnnpack_binary(
35            xnn_graph, bytearray(constant_data_bytes)
36        )
37
38        # Check header
39        self.assertEqual(serialized_binary[0:4], b"\x00\x00\x00\x00")
40        self.assertEqual(serialized_binary[XNNHeader.MAGIC_OFFSET], b"XH00")
41        flatbuffer_offset_bytes = serialized_binary[XNNHeader.FLATBUFFER_OFFSET_OFFSET]
42
43        # Check flatbuffer is at flatbuffer offset
44        flatbuffer_offset = int.from_bytes(
45            flatbuffer_offset_bytes, byteorder=_HEADER_BYTEORDER
46        )
47        # Flatbuffer magic should be in the same spot as the Header's magic
48        self.assertEqual(
49            serialized_binary[flatbuffer_offset:][XNNHeader.MAGIC_OFFSET], b"XN01"
50        )
51