Handle XNNHeader in XNNPACK Runtime

Summary:
We introduce XNNHeader on runtime side to handle the newly introduced XNNHeader ahead of time. XNNHeader manages the offsets and sizes of the flatbuffer payload and the constant data payload so that it is accessible by the XNNCompiler

It is important to note that on serialization side, we have not yet switched our serialization method to `serialize_xnnpack_binary` so this does not yet use the new serialization format. However, passing tests on this illustrates BC as old models will still be able to run on this new runtime.

Passing tests here show that the Header Magic correctly works in discerning between using the XNNHeader and the Flatbuffer header

Differential Revision: D52556131

Author: mcr229

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -7,6 +7,7 @@
  */
 
 #include <executorch/backends/xnnpack/runtime/XNNCompiler.h>
+#include <executorch/backends/xnnpack/runtime/XNNHeader.h>
 #include <executorch/backends/xnnpack/schema_generated.h>
 #include <executorch/backends/xnnpack/threadpool/threadpool.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
@@ -113,6 +114,7 @@ Error defineTensor(
     std::unordered_map<uint32_t, uint32_t>& remapped_ids,
     ValuePtr value,
     GraphPtr flatbuffer_graph,
+    const uint8_t* constant_data_ptr,
     XNNExecutor* executor,
     MemoryAllocator* runtime_allocator) {
   const fb_xnnpack::XNNTensorValue* tensor_value = nullptr;
@@ -151,11 +153,20 @@ Error defineTensor(
 
   // Get Pointer to constant data from flatbuffer, if its non-constant
   // it is a nullptr
-  const auto& constant_buffer = *flatbuffer_graph->constant_buffer();
   auto buffer_idx = tensor_value->constant_buffer_idx();
-  const auto buffer_ptr = buffer_idx == 0
-      ? nullptr
-      : constant_buffer[buffer_idx]->storage()->data();
+  const uint8_t* buffer_ptr = nullptr;
+  if (buffer_idx) {
+    if (!constant_data_ptr) {
+      const auto& constant_buffer = *flatbuffer_graph->constant_buffer();
+      buffer_ptr = constant_buffer[buffer_idx]->storage()->data();
+    } else {
+      const auto& constant_data_offsets = *flatbuffer_graph->constant_data();
+      uint32_t constant_data_offset =
+          constant_data_offsets[buffer_idx]->offset();
+      buffer_ptr = constant_data_ptr + constant_data_offset;
+    }
+  }
+
   xnn_status status;
   // The type we might have to convert to
   auto dq_datatype = getDataType(tensor_value->dq_datatype());
@@ -1429,14 +1440,30 @@ __ET_NODISCARD Error XNNCompiler::compileModel(
     size_t num_bytes,
     XNNExecutor* executor,
     MemoryAllocator* runtime_allocator) {
+  Result<XNNHeader> header = XNNHeader::Parse(buffer_pointer, num_bytes);
+  const uint8_t* flatbuffer_data =
+      reinterpret_cast<const uint8_t*>(buffer_pointer);
+  const uint8_t* constant_data = nullptr;
+
+  // Header status can only either be Error::Ok or Error::NotFound
+  if (header.ok()) {
+    flatbuffer_data = reinterpret_cast<const uint8_t*>(buffer_pointer) +
+        header->flatbuffer_offset;
+    constant_data = reinterpret_cast<const uint8_t*>(buffer_pointer) +
+        header->constant_data_offset;
+  } else if (header.error() != Error::NotFound) {
+    ET_LOG(Error, "XNNHeader may be corrupt");
+    return header.error();
+  }
+
   ET_CHECK_OR_RETURN_ERROR(
-      fb_xnnpack::XNNGraphBufferHasIdentifier(buffer_pointer),
+      fb_xnnpack::XNNGraphBufferHasIdentifier(flatbuffer_data),
       DelegateInvalidCompatibility,
       "XNNPACK Delegate Serialization Format version identifier '%.4s' != expected '%.4s'",
-      flatbuffers::GetBufferIdentifier(buffer_pointer),
+      flatbuffers::GetBufferIdentifier(flatbuffer_data),
       fb_xnnpack::XNNGraphIdentifier());
 
-  auto flatbuffer_graph = fb_xnnpack::GetXNNGraph(buffer_pointer);
+  auto flatbuffer_graph = fb_xnnpack::GetXNNGraph(flatbuffer_data);
   // initialize xnnpack
   xnn_status status = xnn_initialize(/*allocator =*/nullptr);
   ET_CHECK_OR_RETURN_ERROR(
@@ -1476,6 +1503,7 @@ __ET_NODISCARD Error XNNCompiler::compileModel(
         remapped_ids,
         value,
         flatbuffer_graph,
+        constant_data,
         executor,
         runtime_allocator);
 

backends/xnnpack/runtime/XNNHeader.cpp

@@ -0,0 +1,92 @@
+/*
+ * 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.
+ */
+
+#include <executorch/backends/xnnpack/runtime/XNNHeader.h>
+
+#include <cinttypes>
+#include <cstring>
+
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/result.h>
+
+#pragma clang diagnostic ignored "-Wdeprecated"
+
+namespace torch {
+namespace executor {
+namespace xnnpack {
+namespace delegate {
+
+namespace {
+
+/// The expected location of the header length field relative to the beginning
+/// of the header.
+static constexpr size_t kHeaderLengthOffset =
+    XNNHeader::kMagicOffset + XNNHeader::kMagicSize;
+
+/// The expected location of the flatbuffer data offset field relative to the
+/// beginning of the header.
+static constexpr size_t kFlatbufferDataOffsetOffset =
+    kHeaderLengthOffset + sizeof(unsigned short);
+
+/// The expected location of the flatbuffer data size field relative to the
+/// beginning of the header.
+static constexpr size_t kFlatbufferDataSizeOffset =
+    kFlatbufferDataOffsetOffset + sizeof(uint32_t);
+
+/// The expected location of the constant data offset field relative to the
+/// beginning of the header.
+static constexpr size_t kConstantDataOffsetOffset =
+    kFlatbufferDataSizeOffset + sizeof(uint32_t);
+
+/// The expected location of the constant data size field relative to the
+/// beginning of the header.
+static constexpr size_t kConstantDataSizeOffset =
+    kConstantDataOffsetOffset + sizeof(uint32_t);
+
+/// Interprets the 4 bytes at `data` as a little-endian uint32_t.
+uint32_t GetUInt32LE(const uint8_t* data) {
+  return (uint32_t)data[0] | ((uint32_t)data[1] << 8) |
+      ((uint32_t)data[2] << 16) | ((uint32_t)data[3] << 24);
+}
+
+} // namespace
+
+Result<XNNHeader> XNNHeader::Parse(const void* data, size_t size) {
+  const uint8_t* header_data = (const uint8_t*)data;
+
+  const uint8_t* magic_start = header_data + XNNHeader::kMagicOffset;
+  if (std::memcmp(magic_start, XNNHeader::kMagic, XNNHeader::kMagicSize) != 0) {
+    return Error::NotFound;
+  }
+
+  uint32_t flatbuffer_offset =
+      GetUInt32LE(header_data + kFlatbufferDataOffsetOffset);
+
+  uint32_t flatbuffer_size =
+      GetUInt32LE(header_data + kFlatbufferDataSizeOffset);
+
+  uint32_t constant_data_offset =
+      GetUInt32LE(header_data + kConstantDataOffsetOffset);
+
+  uint32_t constant_data_size =
+      GetUInt32LE(header_data + kConstantDataSizeOffset);
+
+  return XNNHeader{
+      flatbuffer_offset,
+      flatbuffer_size,
+      constant_data_offset,
+      constant_data_size};
+}
+
+// Define storage for the static.
+constexpr char XNNHeader::kMagic[kMagicSize];
+
+} // namespace delegate
+} // namespace xnnpack
+} // namespace executor
+} // namespace torch

backends/xnnpack/runtime/XNNHeader.h

@@ -0,0 +1,84 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <executorch/runtime/core/result.h>
+
+namespace torch {
+namespace executor {
+namespace xnnpack {
+namespace delegate {
+
+/**
+ * An extended XNNPACK-header that is embeded before the flatbuffer payload
+ *
+ */
+struct XNNHeader {
+  /**
+   * The magic offset. This offset is the same as the offset for flatbuffer
+   * header so we will be able to check if the header is is either the
+   * flatbuffer head or the wrapper header we introduce here
+   */
+  static constexpr size_t kMagicOffset = 4;
+
+  /**
+   * The magic bytes that identify the header.
+   *
+   * This is the canonical definition of the expected value. If the header
+   * layout ever changes in a compatibility-breaking way, increment the digits
+   * in the magic. But, doing so will prevent older binaries from recognizing
+   * the presence of the header. The compatibility-preserving way to make
+   * changes is to increase the header's length field and add new fields at the
+   * end.
+   */
+  static constexpr size_t kMagicSize = 4;
+  static constexpr char kMagic[kMagicSize] = {'X', 'H', '0', '0'};
+
+  /**
+   * The size in bytes of the header length. We store 2 bytes for the header
+   * length
+   */
+  static constexpr size_t kHeaderLengthSize = 2;
+
+  /**
+   * Look for and parse an ExtendedHeader in the provided data.
+   *
+   * @param[in] data The contents of the beginning of the serialized binary
+   *     Program data, starting at offset 0 (i.e., the head of the file).
+   * @param[in] size Length of `data` in bytes.
+   *
+   * @returns an XNNHeader if the header was found and is valid. Returns an
+   *     error if size was too short, if the header was not found, or if the
+   *     header appeared to be corrupt.
+   */
+  static Result<XNNHeader> Parse(const void* data, size_t size);
+
+  /**
+   * The offset in bytes to the beginning of the flatbuffer data.
+   */
+  uint32_t flatbuffer_offset;
+  /**
+   * The size in bytes of the flatbuffer data.
+   */
+  uint32_t flatbuffer_size;
+
+  /**
+   * The offset in bytes to the beginning of the constant data.
+   */
+  uint32_t constant_data_offset;
+  /**
+   * The size in bytes of the constant data.
+   */
+  uint32_t constant_data_size;
+};
+
+} // namespace delegate
+} // namespace xnnpack
+} // namespace executor
+} // namespace torch