Add manifest extension AoT

Summary:
Add some infra for us to optionally add some key structured data to the end of a pte. This diff is around enabling users to easily tag their model with a cryptographic signature. Has room to expand later.

A key design motivation is it would be ideal if this is transparent to the rest of the extensions we have today. Im claiming the prime footer real estate for this which is unused today by anything in tree. This should let it be composable with other formats like bundledProgram too.

Differential Revision: D82052721

Author: JacobSzwejbka

extension/manifest/TARGETS

@@ -0,0 +1,15 @@
+load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
+
+oncall("executorch")
+
+python_library(
+    name = "_manifest",
+    srcs = [
+        "_manifest.py",
+    ],
+    deps = [
+        "//executorch/exir/_serialize:lib",
+        "//executorch/exir:_warnings",
+    ],
+    visibility = ["PUBLIC"],
+)

extension/manifest/__init__.py

@@ -0,0 +1,12 @@
+# 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.
+
+from executorch.extension.manifest._manifest import Manifest, append_manifest
+
+__all__ = [
+    "Manifest",
+    "append_manifest",
+]

extension/manifest/_manifest.py

@@ -0,0 +1,154 @@
+from dataclasses import dataclass
+from typing import ClassVar, Literal
+
+from executorch.exir._serialize._cord import Cord
+from executorch.exir._serialize.padding import padding_required
+from executorch.exir._warnings import experimental
+
+# Byte order of numbers written to the manifest. Always little-endian
+# regardless of the host system, since all commonly-used modern CPUs are little
+# endian.
+_MANIFEST_BYTEORDER: Literal["little"] = "little"
+
+@dataclass 
+class _ManifestLayout:
+ """Python class mirroring the binary layout of the manifest.
+ separate from the Manifest class, which is the user facing 
+ representation.
+ """
+ EXPECTED_MAGIC: ClassVar[bytes] = b"em00"
+
+ EXPECTED_LENGTH: ClassVar[int] = (
+   # Header magic
+   4
+   # Header length
+   + 4
+   # Padding size
+   + 4 
+   # program offset
+   + 8
+   # signature
+   + 8
+ )
+
+ signature: int
+
+ # Size in bytes between the top of the manifest and the end of the data it was appended to
+ padding_size: int = 0
+
+ # Size in bytes between the top of the manifest and the start of the data it was appended to.
+ program_offset: int = 0
+
+ # The manifest length, in bytes, read from or to be written to the binary
+ # footer.
+ length: int = EXPECTED_LENGTH
+
+ # The magic bytes read from or to be written to the binary footer.
+ magic: bytes = EXPECTED_MAGIC
+
+ def is_valid(self) -> bool:
+     """Returns true if the manifest appears to be well-formed."""
+     return (
+         self.magic == _ManifestLayout.EXPECTED_MAGIC
+         and self.length >= _ManifestLayout.EXPECTED_LENGTH
+     )
+
+ def to_bytes(self) -> bytes:
+     """Returns the binary representation of the Manifest. Written
+     bottom up.
+
+     Note that this will ignore self.magic and self.length and will always
+     write the proper magic/length.
+     """
+     data: bytes = (
+         
+         # uint64_t: Signature unique ID for the data the manifest was appended to.
+         self.signature.to_bytes(8, byteorder=_MANIFEST_BYTEORDER)
+         # uint64_t: Size in bytes between the manifest and the data it was appended to.
+         +self.program_offset.to_bytes(8, byteorder=_MANIFEST_BYTEORDER)
+         # uint32_t: Size in bytes between the manifest and the data it was appended to.
+         +self.padding_size.to_bytes(4, byteorder=_MANIFEST_BYTEORDER)
+         # uint32_t: Size of this manifest. This makes it easier to add new
+         # fields to this header in the future. Always use the proper size
+         # (i.e., ignore self.length) since there's no reason to create an
+         # invalid footer.
+         + self.EXPECTED_LENGTH.to_bytes(4, byteorder=_MANIFEST_BYTEORDER)
+         # Manifest magic. This lets consumers detect whether the
+         # manifest was inserted or not. Always use the proper magic value
+         # (i.e., ignore self.magic) since there's no reason to create an
+         # invalid manifest.
+         + self.EXPECTED_MAGIC
+     )
+     return data
+
+ @staticmethod
+ def from_bytes(data: bytes) -> "_ManifestLayout":
+     """Tries to read a manifest from the provided data.
+
+     Does not validate that the header is well-formed. Callers should
+     use is_valid().
+
+     Args:
+         data: The data to read from.
+     Returns:
+         The contents of the serialized manifest.
+     Raises:
+         ValueError: If not enough data is provided.
+     """
+     if len(data) < _ManifestLayout.EXPECTED_LENGTH:
+         raise ValueError(
+             f"Not enough data for the manifest: {len(data)} "
+             + f"< {_ManifestLayout.EXPECTED_LENGTH}"
+         )
+
+     return _ManifestLayout(
+         magic=data[-4:],
+         length=int.from_bytes(data[-8:-4], byteorder=_MANIFEST_BYTEORDER),
+         padding_size=int.from_bytes(data[-12:-8], byteorder=_MANIFEST_BYTEORDER),
+         program_offset=int.from_bytes(data[-20:-12], byteorder=_MANIFEST_BYTEORDER),
+         signature=int.from_bytes(
+             data[-28:-20], byteorder=_MANIFEST_BYTEORDER
+         ),
+     )
+
+ @staticmethod
+ def from_manifest(manifest: "Manifest") -> "_ManifestLayout":
+     return _ManifestLayout(
+      signature = manifest.signature,
+      #padding_size will be properly initialized in the append_manifest function
+     )
+
+@experimental("This API is experimental and subject to change without notice.")
+@dataclass
+class Manifest:
+ """A manifest that can be appended to a binary blob. The manifest contains
+ meta information about the binary blob. You must know who created the manifest
+ to be able to interpret the data in the manifest."""
+
+ # Unique ID for the data the manifest was appended to. Often this might contain
+ # a crytographic signature for the data.
+ signature: int
+
+ @staticmethod
+ def _from_manifest_layout(layout: _ManifestLayout) -> "Manifest":
+   return Manifest(
+       signature=layout.signature,
+   )
+  
+ @staticmethod
+ def from_bytes(data: bytes) -> "Manifest":
+   """Tries to read a manifest from the provided data."""
+   layout = _ManifestLayout.from_bytes(data)
+   if not layout.is_valid():
+     raise ValueError("Cannot parse manifest from bytes")
+   return Manifest._from_manifest_layout(layout)
+
+@experimental("This API is experimental and subject to change without notice.")
+def append_manifest(pte_data: Cord, manifest: Manifest, alignment:int=16):
+   """Appends a manifest to the provided data."""
+   manifest_layout = _ManifestLayout.from_manifest(manifest)
+   manifest_layout.padding_size = padding_required(len(pte_data), alignment)
+   manifest_layout.program_offset = len(pte_data) + manifest_layout.padding_size
+   pte_data.append((b"\x00" * manifest_layout.padding_size))
+   pte_data.append(manifest_layout.to_bytes())
+   

extension/manifest/test/TARGETS

@@ -0,0 +1,16 @@
+load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest")
+
+oncall("executorch")
+
+python_unittest(
+    name = "test_manifest",
+    srcs = [
+        "test_manifest.py",
+    ],
+    deps = [
+        "//executorch/extension/manifest:_manifest",
+        "//executorch/extension/pybindings:portable_lib",
+        "//executorch/exir:lib",
+        "//caffe2:torch",
+    ],
+)