[MPS] Add support for flatbuffer serialization > 4GB

Differential Revision: D60876064

Pull Request resolved: #4574

Author: DenisVieriu97

backends/apple/mps/CMakeLists.txt

@@ -69,7 +69,7 @@ add_library(mpsdelegate ${_mps_backend__srcs})
 find_library(FOUNDATION_FRAMEWORK Foundation)
 find_library(METAL_FRAMEWORK Metal)
 find_library(MPS_FRAMEWORK MetalPerformanceShaders)
-find_library(MPS_GRAPG_FRAMEWORK MetalPerformanceShadersGraph)
+find_library(MPS_GRAPH_FRAMEWORK MetalPerformanceShadersGraph)
 
 target_link_libraries(
   mpsdelegate
@@ -79,7 +79,7 @@ target_link_libraries(
           ${FOUNDATION_FRAMEWORK}
           ${METAL_FRAMEWORK}
           ${MPS_FRAMEWORK}
-          ${MPS_GRAPG_FRAMEWORK}
+          ${MPS_GRAPH_FRAMEWORK}
 )
 
 target_link_options_shared_lib(mpsdelegate)

backends/apple/mps/mps_preprocess.py

@@ -2,9 +2,8 @@
 #  Copyright (c) 2023 Apple Inc. All rights reserved.
 #  Provided subject to the LICENSE file in the top level directory.
 #
-
 import logging
-from typing import Dict, final, List
+from typing import ClassVar, Dict, final, List, Tuple
 
 import torch
 
@@ -16,6 +15,8 @@
 )
 
 from executorch.backends.apple.mps.serialization.mps_graph_schema import (
+    Buffer,
+    DataSegment,
     MPSGraph,
     MPSTensor,
     OpType,
@@ -25,6 +26,7 @@
     convert_to_flatbuffer,
 )
 from executorch.backends.apple.mps.utils.mps_utils import is_parameter
+from executorch.exir._serialize._program import Cord
 
 from executorch.exir.backend.backend_details import (
     BackendDetails,
@@ -39,6 +41,29 @@
 
 @final
 class MPSBackend(BackendDetails):
+    @staticmethod
+    def slice_len_max(s):
+        assert s.start is not None
+        assert s.stop is not None
+        step = 1
+        if s.step is not None:
+            step = s.step
+        return max((s.stop - s.start) // step, 1)
+
+    MAGIC_IX: ClassVar[slice] = slice(4, 8)
+    DATA_SEGMENT_OFFSET_IX: ClassVar[slice] = slice(8, 16)
+    DATA_SEGMENT_SIZE_IX: ClassVar[slice] = slice(16, 24)
+
+    # magic bytes that should be at the beginning of the header
+    EXPECTED_MAGIC: ClassVar[bytes] = b"MP00"
+    # The length of the header in bytes
+    EXPECTED_LENGTH: ClassVar[int] = (
+        4
+        + slice_len_max(MAGIC_IX)
+        + slice_len_max(DATA_SEGMENT_OFFSET_IX)
+        + slice_len_max(DATA_SEGMENT_SIZE_IX)
+    )
+
     @staticmethod
     def preprocess(
         edge_program: ExportedProgram,
@@ -67,6 +92,7 @@ def preprocess(
             output_ids=[],
             constant_ids=[],
             graph_type=OpType.mps_graph,
+            constant_segment=DataSegment(0, 0),
         )
 
         convert_model_to_fp16 = True
@@ -100,10 +126,44 @@ def preprocess(
             else:
                 op_handler[node.op](edge_program, node_visitors, node, mps_graph)
 
+        segment_data, mps_graph = _extract_constant_segment(mps_graph)
+
+        # Add to aggregate segments cord with padding.
+        padding_length = _padding_required(len(segment_data), 16)
+        if padding_length > 0:
+            segment_data.append(b"\x00" * padding_length)
+
+        # Combine mps_graph with segment data
+        combined = Cord()
+        graph_bytes = convert_to_flatbuffer(mps_graph)
+
+        data_segment_offset: int = MPSBackend.EXPECTED_LENGTH
+        data_segment_offset = data_segment_offset + len(graph_bytes)
+
+        graph_padding_length = _padding_required(data_segment_offset, 16)
+        data_segment_offset = data_segment_offset + graph_padding_length
+        data_segment_size = len(segment_data)
+
+        data: bytes = (
+            b"\x00\x00\x00\x00"
+            + MPSBackend.EXPECTED_MAGIC
+            + data_segment_offset.to_bytes(8, byteorder="little")
+            + data_segment_size.to_bytes(8, byteorder="little")
+        )
+        assert len(data) == MPSBackend.EXPECTED_LENGTH
+
+        combined.append(data)
+        combined.append(graph_bytes)
+
+        if graph_padding_length > 0:
+            combined.append(b"\x00" * graph_padding_length)
+        # Append the segment data to the end of the mps graph
+        combined.append(segment_data)
+
         if logging.DEBUG >= logging.root.level:
             pretty_print(mps_graph)
 
-        return PreprocessResult(processed_bytes=convert_to_flatbuffer(mps_graph))
+        return PreprocessResult(processed_bytes=bytes(combined))
 
     @staticmethod
     def handle_call_function(
@@ -164,12 +224,42 @@ def handle_get_attr(
         pass
 
 
+def _padding_required(offset: int, alignment: int) -> int:
+    """Returns the padding required to align `offset` to `alignment`."""
+    remainder: int = offset % alignment
+    if remainder != 0:
+        return alignment - remainder
+    return 0
+
+
+def _extract_constant_segment(mps_graph: MPSGraph) -> Tuple[Cord, MPSGraph]:
+    """Extracts the constant segment from the MPSGraph and returns the updated MPSGraph along with the segment data."""
+    # Note that the beginning of the segment data is not aligned. Need to handle out of this call.
+    segment_data = Cord()
+    offset = 0
+    for i in range(len(mps_graph.mps_values)):
+        tensor = mps_graph.mps_values[i]
+        if tensor.constant_buffer_size > 0:
+            # Notice that buffer is already force aligned so we don't need to pad it
+            segment_data.append(tensor.constant_buffer.storage)
+
+            # Reset buffer to empty
+            tensor.constant_buffer = Buffer(storage=b"")
+            # Update segment offset
+            tensor.segment_offset = offset
+            offset += tensor.constant_buffer_size
+
+    return segment_data, mps_graph
+
+
 def tensor_to_str(mps_tensor: MPSTensor):
     tensor_str = "MPSTensor("
     tensor_str += "datatype=" + str(mps_tensor.datatype) + ", "
     tensor_str += "num_dims=" + str(mps_tensor.num_dims) + ", "
     tensor_str += "dims=" + str(mps_tensor.dims) + ", "
-    tensor_str += "constant_buffer_size=" + str(mps_tensor.constant_buffer_size)
+    tensor_str += "constant_buffer=" + str(mps_tensor.constant_buffer) + ", "
+    tensor_str += "constant_buffer_size=" + str(mps_tensor.constant_buffer_size) + ", "
+    tensor_str += "segment_offset=" + str(mps_tensor.segment_offset)
     tensor_str += ")"
 
     return tensor_str
@@ -193,3 +283,4 @@ def pretty_print(mps_graph: MPSGraph):
     logging.info(" Output ids:")
     for out_id in mps_graph.output_ids:
         logging.info(f"   {out_id}")
+    logging.info(f" Constant segment: {mps_graph.constant_segment}")

backends/apple/mps/runtime/MPSCompiler.mm

@@ -43,7 +43,7 @@
   Error err = Error::Ok;
 
   std::unique_ptr<MPSGraphBuilder> mpsGraphBuilder(
-    new MPSGraphBuilder(buffer_pointer, executor->_mpsGraphTensorToId));
+    new MPSGraphBuilder(buffer_pointer, num_bytes, executor->_mpsGraphTensorToId));
   err = mpsGraphBuilder->compileModel();
   ET_CHECK_OR_RETURN_ERROR(
     err == Error::Ok, Internal, "Failed to construct the MPS graph object");

backends/apple/mps/runtime/MPSDelegateHeader.h

@@ -0,0 +1,113 @@
+//
+//  Copyright (c) 2024 Apple Inc. All rights reserved.
+//  Provided subject to the LICENSE file in the top level directory.
+//
+
+#pragma once
+
+#include <executorch/runtime/core/result.h>
+
+namespace torch {
+namespace executor {
+namespace mps {
+namespace delegate {
+
+/**
+ * MPS-header that is embedded before the flatbuffer payload
+ *
+ */
+struct MPSDelegateHeader {
+  /**
+   * The minimum size of the MPSDelegateHeader. The caller should provide at
+   * least this many bytes of the head of the serialized MPS Data
+   */
+  static constexpr size_t kMinSize = 30;
+
+  /**
+   * 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] = {'M', 'P', '0', '0'};
+
+  /**
+   * The size in bytes of the header length. We store 2 bytes for the header
+   * length
+   */
+  static constexpr size_t kHeaderLengthSize = 2;
+
+  /**
+   * The expected location of the header length field relative to the beginning
+   * of the header.
+   */
+  static constexpr size_t kHeaderLengthOffset =
+      MPSDelegateHeader::kMagicOffset + MPSDelegateHeader::kMagicSize;
+
+  /*
+   * The expected location of the constant data offset field relative to the
+   * beginning of the header.
+   */
+  static constexpr size_t kConstantDataSegmentOffset = kHeaderLengthOffset;
+
+  /*
+   * The expected location of the constant data size field relative to the
+   * beginning of the header.
+   */
+  static constexpr size_t kConstantDataSizeOffset =
+      kConstantDataSegmentOffset + sizeof(uint64_t);
+
+  /**
+   * The expected location of the flatbuffer data offset field relative to the
+   * beginning of the header.
+   */
+  static constexpr size_t kFlatbufferDataOffsetOffset =
+      kConstantDataSizeOffset + sizeof(uint64_t);
+
+  /**
+   * 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 MPSHeader 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<MPSDelegateHeader> Parse(const void* data, size_t size);
+
+  /**
+   * The offset in bytes to the beginning of the constant data.
+   */
+  uint64_t constant_data_offset;
+  /**
+   * The size in bytes of the constant data.
+   */
+  uint64_t constant_data_size;
+  /**
+   * The offset in bytes to the beginning of the flatbuffer data.
+   */
+  uint64_t flatbuffer_offset;
+  /**
+   * The size in bytes of the flatbuffer data.
+   */
+  uint64_t flatbuffer_size;
+};
+
+} // namespace delegate
+} // namespace mps
+} // namespace executor
+} // namespace torch

backends/apple/mps/runtime/MPSDelegateHeader.mm

@@ -0,0 +1,53 @@
+//
+//  Copyright (c) 2024 Apple Inc. All rights reserved.
+//  Provided subject to the LICENSE file in the top level directory.
+//
+
+#include <executorch/backends/apple/mps/runtime/MPSDelegateHeader.h>
+
+#include <cstring>
+
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/result.h>
+
+namespace torch {
+namespace executor {
+namespace mps {
+namespace delegate {
+
+/// Interprets the 8 bytes at `data` as a little-endian uint64_t.
+uint64_t getUInt64LE(const uint8_t* data) {
+  return (uint64_t)data[0] | ((uint64_t)data[1] << 8) |
+      ((uint64_t)data[2] << 16) | ((uint64_t)data[3] << 24) |
+      ((uint64_t)data[4] << 32) | ((uint64_t)data[5] << 40) |
+      ((uint64_t)data[6] << 48) | ((uint64_t)data[7] << 56);
+}
+
+Result<MPSDelegateHeader> MPSDelegateHeader::Parse(const void* data, size_t size) {
+  const uint8_t* header_data = (const uint8_t*)data;
+
+  if (size < MPSDelegateHeader::kMinSize) {
+    return Error::InvalidArgument;
+  }
+
+  const uint8_t* magic_start = header_data + MPSDelegateHeader::kMagicOffset;
+  if (std::memcmp(magic_start, MPSDelegateHeader::kMagic, MPSDelegateHeader::kMagicSize) != 0) {
+    return Error::NotFound;
+  }
+
+  uint64_t constant_data_offset = getUInt64LE(header_data + MPSDelegateHeader::kConstantDataSegmentOffset);
+  uint64_t constant_data_size = getUInt64LE(header_data + MPSDelegateHeader::kConstantDataSizeOffset);
+  uint64_t flatbuffer_offset = MPSDelegateHeader::kFlatbufferDataOffsetOffset;
+  uint64_t flatbuffer_size = size - flatbuffer_offset;
+
+  return MPSDelegateHeader{
+      constant_data_offset,
+      constant_data_size,
+      flatbuffer_offset,
+      flatbuffer_size};
+}
+
+} // namespace delegate
+} // namespace mps
+} // namespace executor
+} // namespace torch

backends/apple/mps/runtime/MPSGraphBuilder.h

@@ -40,7 +40,8 @@ using NodePtr = const mpsgraph::MPSNode *;
  */
 class MPSGraphBuilder {
 public:
-  MPSGraphBuilder(const void *buffer_pointer, std::unordered_map<MPSGraphTensor *, int32_t> &mpsGraphTensorToId);
+  MPSGraphBuilder(const void *buffer_pointer, size_t num_bytes,
+                  std::unordered_map<MPSGraphTensor *, int32_t> &mpsGraphTensorToId);
   ~MPSGraphBuilder() = default;
 
   Error compileModel();
@@ -178,12 +179,15 @@ class MPSGraphBuilder {
   const mpsgraph::MPSGraph *_flatBufferGraph;
   // FlatBuffer raw bytes of the serialized MPS model.
   const void *_buffer_pointer;
+  size_t _num_bytes;
 
   bool _metal_kernel;
   MPSGraph *_mpsGraph;
   MPSGraphExecutable *_mpsGraphExecutable;
   NSMutableDictionary<MPSGraphTensor *, MPSGraphShapedType *> *_feeds;
   NSMutableArray<MPSGraphTensor *> *_targetTensors;
+
+  const uint8_t *_constant_data_ptr;
 };
 
 #undef _DEFINE_MPS_OP