add cpp serializer for flat tensor (ptd)

Differential Revision: D67992901

Pull Request resolved: #7841

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",
+        ],
+    )