add cpp serializer for flat tensor (ptd) (#7841)

Summary:

Leverage the flatbuffer builder apis to generate a .ptd on device. This will be used by training for checkpointing.

No other use cases really exist for generating a .ptd on device right now, so I didnt worry about making this more easily extensible by coming up with a cpp equivalent of a cord or by trying to think ahead of how this might integrate with delegates.

Later if we add support for delegates owning the weights under ET training we can revisit this.

Differential Revision: D67992901

Author: JacobSzwejbka

extension/flat_tensor/serialize/flat_tensor_header.h

@@ -38,6 +38,21 @@ struct FlatTensorHeader {
   // @lint-ignore CLANGTIDY facebook-hte-CArray
   static constexpr char kMagic[kMagicSize] = {'F', 'H', '0', '1'};
 
+  /// The expected length of the header, in bytes.
+  static constexpr uint32_t kHeaderExpectedLength =
+      // Header magic
+      4
+      // Header length
+      + 4
+      // Flatbuffer offset
+      + 8
+      // Flatbuffer data size
+      + 8
+      // Segment base offset
+      + 8
+      // Data size
+      + 8;
+
   /**
    * Look for and parse a FlatTensorHeader in the provided data.
    *

extension/flat_tensor/serialize/serialize.cpp

@@ -0,0 +1,176 @@
+/*
+ * 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/extension/flat_tensor/serialize/serialize.h>
+
+#include <executorch/extension/flat_tensor/serialize/flat_tensor_header.h>
+#include <executorch/extension/flat_tensor/serialize/scalar_type_generated.h>
+#include <executorch/extension/flat_tensor/serialize/schema_generated.h>
+
+#include <fstream>
+#include <string>
+
+namespace executorch {
+namespace extension {
+namespace flat_tensor {
+
+namespace {
+size_t padding_required(size_t offset, size_t alignment) {
+  // Returns the padding required to align `offset` to `alignment`.
+  size_t remainder = offset % alignment;
+  if (remainder != 0) {
+    return alignment - remainder;
+  }
+  return 0;
+}
+
+size_t aligned_size(size_t input_size, size_t alignment) {
+  // Returns input_size padded up to the next whole multiple of alignment.
+  return input_size + padding_required(input_size, alignment);
+}
+
+void write_nulls(std::ostream& out, size_t num_bytes) {
+  for (size_t i = 0; i < num_bytes; i++) {
+    out.write("\0", 1);
+  }
+}
+} // namespace
+
+runtime::Error save_ptd(
+    const std::string& path,
+    const std::map<std::string, exec_aten::Tensor>& tensor_map,
+    const size_t tensor_alignment) {
+  // Create File
+  std::ofstream file;
+  file.open(path);
+  runtime::Error e = save_ptd(file, tensor_map, tensor_alignment);
+  file.close();
+  return e;
+}
+
+runtime::Error save_ptd(
+    std::ostream& out,
+    const std::map<std::string, exec_aten::Tensor>& tensor_map,
+    const size_t tensor_alignment) {
+  // Assert the system is little endian. Since we are sending the data over
+  // the wire, we need to ensure that the data is always in the same format.
+  // for now we only support little endian.
+  int n = 1;
+  if (*(char*)&n != 1) {
+    ET_LOG(Error, "Cannot save_ptd on big endian system");
+    return runtime::Error::NotSupported;
+  }
+  // Create flatbuffer
+  flatbuffers::FlatBufferBuilder builder;
+
+  std::vector<flatbuffers::Offset<::flat_tensor_flatbuffer::TensorMetadata>>
+      tensors;
+  std::vector<flatbuffers::Offset<::flat_tensor_flatbuffer::DataSegment>>
+      buffers;
+
+  // Write the tensors.
+  size_t total_segment_size = 0;
+  size_t i = tensor_map.size();
+  for (const auto& [name, tensor] : tensor_map) {
+    auto name_offset = builder.CreateString(name);
+    // Write the tensor metadata.
+    auto tensor_metadata = ::flat_tensor_flatbuffer::CreateTensorMetadata(
+        builder,
+        name_offset,
+        static_cast<executorch_flatbuffer::ScalarType>(tensor.scalar_type()),
+        builder.CreateVector(tensor.sizes().data(), tensor.sizes().size()),
+        builder.CreateVector(
+            tensor.dim_order().data(), tensor.dim_order().size()),
+        0, // segment index
+        total_segment_size);
+
+    tensors.push_back(tensor_metadata);
+    // Don't pad last entry.
+    if (i != 1) {
+      // Precalculate the size of the data blob.
+      total_segment_size += aligned_size(tensor.nbytes(), tensor_alignment);
+    } else {
+      total_segment_size += tensor.nbytes();
+    }
+    i--;
+  }
+  // Only have one segment
+  buffers.push_back(::flat_tensor_flatbuffer::CreateDataSegment(
+      builder, 0, total_segment_size));
+
+  auto flat_tensor = CreateFlatTensor(
+      builder,
+      kSchemaVersion,
+      tensor_alignment,
+      builder.CreateVector(tensors),
+      builder.CreateVector(buffers));
+  builder.Finish(flat_tensor); // Our flatbuffer is created now.
+
+  // Calculate flatbuffer padding.
+  auto padded_flatbufer_size =
+      aligned_size(builder.GetSize(), tensor_alignment);
+  auto padded_header_size =
+      aligned_size(FlatTensorHeader::kHeaderExpectedLength, tensor_alignment);
+
+  // Write header
+  out.write(FlatTensorHeader::kMagic, sizeof(FlatTensorHeader::kMagic));
+  out.write(
+      reinterpret_cast<const char*>(&FlatTensorHeader::kHeaderExpectedLength),
+      sizeof(FlatTensorHeader::kHeaderExpectedLength));
+
+  FlatTensorHeader header = {
+      padded_header_size, // Offset to flatbuffer
+      builder.GetSize(), // flatbuffer size
+      padded_header_size + padded_flatbufer_size, // offset to segments
+      total_segment_size // segment data size
+  };
+
+  out.write(
+      reinterpret_cast<const char*>(&header.flatbuffer_offset),
+      sizeof(header.flatbuffer_offset));
+  out.write(
+      reinterpret_cast<const char*>(&header.flatbuffer_size),
+      sizeof(header.flatbuffer_size));
+  out.write(
+      reinterpret_cast<const char*>(&header.segment_base_offset),
+      sizeof(header.segment_base_offset));
+  out.write(
+      reinterpret_cast<const char*>(&header.segment_data_size),
+      sizeof(header.segment_data_size));
+
+  // Write header padding
+  write_nulls(
+      out,
+      padding_required(
+          FlatTensorHeader::kHeaderExpectedLength, tensor_alignment));
+
+  // Write flatbuffer
+  out.write(
+      reinterpret_cast<const char*>(builder.GetBufferPointer()),
+      builder.GetSize());
+
+  // Write flatbuffer padding
+  write_nulls(out, padding_required(builder.GetSize(), tensor_alignment));
+
+  // Write segment: buffers + tensor padding
+  i = tensor_map.size();
+  for (const auto& [name, tensor] : tensor_map) {
+    out.write(
+        reinterpret_cast<const char*>(tensor.data_ptr()), tensor.nbytes());
+    // Don't pad last entry.
+    if (i != 1) {
+      write_nulls(out, padding_required(tensor.nbytes(), tensor_alignment));
+    }
+    i--;
+  }
+  return runtime::Error::Ok;
+}
+
+} // namespace flat_tensor
+} // namespace extension
+} // namespace executorch

extension/flat_tensor/serialize/serialize.h

@@ -0,0 +1,53 @@
+/*
+ * 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/exec_aten/exec_aten.h>
+
+#include <map>
+#include <string>
+
+namespace executorch {
+namespace extension {
+namespace flat_tensor {
+
+/**
+ * Schema version of the .ptd format. Should be kept in sync with serialize.py
+ */
+constexpr uint32_t kSchemaVersion = 0;
+
+/**
+ * Creates a .ptd from the given tensor map.
+ *
+ * @param path The file path to save the .ptd to.
+ * @param tensor_map The map of tensor names to tensors to save.
+ * @param tensor_alignment The bytes tensor data should be aligned to.
+ * @return An error if the data could not be saved. Error::Ok for success.
+ */
+ET_EXPERIMENTAL runtime::Error save_ptd(
+    const std::string& path,
+    const std::map<std::string, exec_aten::Tensor>& tensor_map,
+    const size_t tensor_alignment);
+
+/**
+ * Creates a .ptd from the given tensor map.
+ *
+ * @param out The stream to write the .ptd data to.
+ * @param tensor_map The map of tensor names to tensors to save.
+ * @param tensor_alignment The bytes tensor data should be aligned to.
+ * @return An error if the data could not be saved. Error::Ok for success.
+ */
+ET_EXPERIMENTAL runtime::Error save_ptd(
+    std::ostream& out,
+    const std::map<std::string, exec_aten::Tensor>& tensor_map,
+    const size_t tensor_alignment);
+
+} // namespace flat_tensor
+} // namespace extension
+} // namespace executorch

extension/flat_tensor/serialize/serialize.py

@@ -266,7 +266,7 @@ def serialize(
         # Create FlatTensor, which describes of the contents of the file and
         # points to all the data segments. It will be serialized to flatbuffer.
         flat_tensor = FlatTensor(
-            version=0,
+            version=0,  # Keep in sync with c++ version number in serialize.h
             tensor_alignment=self.config.tensor_alignment,
             tensors=flat_tensor_metadata,
             segments=[DataSegment(offset=0, size=len(flat_tensor_data))],

extension/flat_tensor/serialize/targets.bzl

@@ -42,3 +42,18 @@ def define_common_targets():
         visibility = ["//executorch/..."],
         exported_deps = ["//executorch/runtime/core:core"],
     )
+
+    runtime.cxx_library(
+        name = "serialize_cpp",
+        srcs = ["serialize.cpp"],
+        deps = [
+            ":flat_tensor_header",
+            ":generated_headers",
+            "//executorch/runtime/core/exec_aten:lib",
+        ],
+        exported_headers = ["serialize.h"],
+        visibility = [
+            "//executorch/...",
+        ],
+        exported_external_deps = ["flatbuffers-api"],
+    )

extension/flat_tensor/test/targets.bzl

@@ -16,3 +16,16 @@ def define_common_targets():
             "//executorch/extension/flat_tensor/serialize:flat_tensor_header",
         ],
     )
+
+    runtime.cxx_test(
+        name = "serialize_cpp_test",
+        srcs = [
+            "test_serialize.cpp",
+        ],
+        deps = [
+            "//executorch/extension/flat_tensor/serialize:serialize_cpp",
+            "//executorch/extension/flat_tensor/serialize:generated_headers",
+            "//executorch/extension/flat_tensor/serialize:flat_tensor_header",
+            "//executorch/extension/tensor:tensor",
+        ],
+    )