xref: /aosp_15_r20/external/executorch/backends/xnnpack/test/serialization/test_xnnheader.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_serialize import XNNHeader
10
11EXAMPLE_FLATBUFFER_OFFSET: int = 0x11223344
12EXAMPLE_FLATBUFFER_SIZE: int = 0x55667788
13EXAMPLE_CONSTANT_DATA_OFFSET: int = EXAMPLE_FLATBUFFER_OFFSET + EXAMPLE_FLATBUFFER_SIZE
14EXAMPLE_CONSTANT_DATA_SIZE: int = 0x99AABBCC99AABBCC
15
16# If header layout or magic changes, this test must change too.
17# The layout of the header is a contract, not an implementation detail
18EXAMPLE_HEADER_DATA: bytes = (
19    # zeros
20    b"\x00\x00\x00\x00"
21    # magic
22    + b"XH00"
23    # All Values below are littl Endian
24    # header length
25    + b"\x1E\x00"
26    # Flatbuffer Offset
27    + b"\x44\x33\x22\x11"
28    # Flatbuffer Size
29    + b"\x88\x77\x66\x55"
30    # Constant Data Offset
31    + b"\xCC\xAA\x88\x66"
32    # Constant Data Size
33    + b"\xCC\xBB\xAA\x99\xCC\xBB\xAA\x99"
34)
35
36
37class TestXNNHeader(unittest.TestCase):
38    def test_to_bytes(self) -> None:
39        header = XNNHeader(
40            EXAMPLE_FLATBUFFER_OFFSET,
41            EXAMPLE_FLATBUFFER_SIZE,
42            EXAMPLE_CONSTANT_DATA_OFFSET,
43            EXAMPLE_CONSTANT_DATA_SIZE,
44        )
45        self.assertEqual(header.to_bytes(), EXAMPLE_HEADER_DATA)
46        self.assertTrue(header.is_valid())
47
48    def test_from_bytes(self) -> None:
49        header = XNNHeader.from_bytes(EXAMPLE_HEADER_DATA)
50        self.assertEqual(header.flatbuffer_offset, EXAMPLE_FLATBUFFER_OFFSET)
51        self.assertEqual(header.flatbuffer_size, EXAMPLE_FLATBUFFER_SIZE)
52        self.assertEqual(header.constant_data_offset, EXAMPLE_CONSTANT_DATA_OFFSET)
53        self.assertEqual(header.constant_data_size, EXAMPLE_CONSTANT_DATA_SIZE)
54
55    def test_invalid_metadata(self) -> None:
56        WRONG_MAGIC_DATA = EXAMPLE_HEADER_DATA[0:4] + b"YT01" + EXAMPLE_HEADER_DATA[8:]
57        with self.assertRaisesRegex(
58            ValueError,
59            "Invalid XNNHeader: invalid magic bytes b'YT01', expected b'XH00'",
60        ):
61            XNNHeader.from_bytes(WRONG_MAGIC_DATA)
62
63        WRONG_LENGTH_DATA = (
64            EXAMPLE_HEADER_DATA[0:8] + b"\x1D\x00" + EXAMPLE_HEADER_DATA[10:]
65        )
66        with self.assertRaisesRegex(
67            ValueError,
68            "Invalid XNNHeader: Invalid parsed length: data given was 30 bytes, parsed length was 29 bytes",
69        ):
70            XNNHeader.from_bytes(WRONG_LENGTH_DATA)
71
72        with self.assertRaisesRegex(
73            ValueError,
74            "Invalid XNNHeader: expected no more than 30 bytes, got 31",
75        ):
76            XNNHeader.from_bytes(EXAMPLE_HEADER_DATA + b"\x00")
77
78    def test_invalid_flatbuffer_size(self) -> None:
79        header = XNNHeader(
80            EXAMPLE_FLATBUFFER_OFFSET,
81            0,
82            EXAMPLE_CONSTANT_DATA_OFFSET,
83            EXAMPLE_CONSTANT_DATA_SIZE,
84        )
85
86        with self.assertRaises(ValueError):
87            header.to_bytes()
88
89    def test_invalid_constant_data_offset(self) -> None:
90        header = XNNHeader(
91            EXAMPLE_FLATBUFFER_OFFSET,
92            EXAMPLE_FLATBUFFER_SIZE,
93            EXAMPLE_FLATBUFFER_OFFSET + EXAMPLE_FLATBUFFER_SIZE - 1,
94            EXAMPLE_CONSTANT_DATA_SIZE,
95        )
96
97        with self.assertRaises(ValueError):
98            header.to_bytes()
99
100    def test_to_bytes_same_as_from_bytes(self) -> None:
101        header = XNNHeader.from_bytes(EXAMPLE_HEADER_DATA)
102
103        to_bytes = header.to_bytes()
104        self.assertEqual(EXAMPLE_HEADER_DATA, to_bytes)
105