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,160 @@
+/*
+ * 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/flat_tensor_header.h>
+#include <executorch/extension/flat_tensor/serialize/scalar_type_generated.h>
+#include <executorch/extension/flat_tensor/serialize/schema_generated.h>
+#include <executorch/extension/flat_tensor/serialize/serialize.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::TensorMetadata>> tensors;
+  std::vector<flatbuffers::Offset<::flat_tensor::DataSegment>> buffers;
+
+  // Write the tensors.
+  size_t total_segment_size = 0;
+  for (const auto& [name, tensor] : tensor_map) {
+    auto name_offset = builder.CreateString(name);
+    // Write the tensor metadata.
+    auto tensor_metadata = ::flat_tensor::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); // Precalculate the size of the data blob
+    total_segment_size += aligned_size(tensor.nbytes(), tensor_alignment);
+  }
+  // Only have one segment
+  buffers.push_back(
+      ::flat_tensor::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
+  for (const auto& [name, tensor] : tensor_map) {
+    out.write(
+        reinterpret_cast<const char*>(tensor.data_ptr()), tensor.nbytes());
+    write_nulls(out, padding_required(tensor.nbytes(), tensor_alignment));
+  }
+  return runtime::Error::Ok;
+}
+
+} // namespace flat_tensor
+} // namespace extension
+} // namespace executorch

extension/flat_tensor/serialize/serialize.h

@@ -0,0 +1,51 @@
+/*
+ * 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 Python.
+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/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",
+        ],
+    )

extension/flat_tensor/test/test_serialize.cpp

@@ -0,0 +1,115 @@
+/*
+ * 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/flat_tensor_header.h>
+#include <executorch/extension/flat_tensor/serialize/scalar_type_generated.h>
+#include <executorch/extension/flat_tensor/serialize/schema_generated.h>
+#include <executorch/extension/flat_tensor/serialize/serialize.h>
+
+#include <gtest/gtest.h>
+
+#include <executorch/extension/tensor/tensor_ptr.h>
+#include <executorch/runtime/core/result.h>
+#include <executorch/runtime/platform/runtime.h>
+#include <map>
+
+#include <sstream>
+
+using namespace ::testing;
+using executorch::runtime::Error;
+using executorch::runtime::Result;
+
+class FlatTensorSerializeTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    // Since these tests cause ET_LOG to be called, the PAL must be initialized
+    // first.
+    executorch::runtime::runtime_init();
+  }
+};
+
+TEST_F(FlatTensorSerializeTest, ValidFlatTensorSerialized) {
+  const size_t kTensorAlignment = 16;
+  std::map<std::string, exec_aten::Tensor> flat_tensor_map;
+
+  float linear_weight = 3.14f;
+  auto weight = executorch::extension::make_tensor_ptr({1}, &linear_weight);
+
+  float linear_bias = 2.0f;
+  auto bias = executorch::extension::make_tensor_ptr({1}, &linear_bias);
+
+  flat_tensor_map.insert({"linear.weight", *weight.get()});
+  flat_tensor_map.insert({"linear.bias", *bias.get()});
+
+  std::ostringstream buf;
+  auto err = executorch::extension::flat_tensor::save_ptd(
+      buf, flat_tensor_map, kTensorAlignment);
+  ASSERT_EQ(err, Error::Ok);
+  auto x = buf.str();
+  const char* byte_buffer = x.c_str();
+
+  // Check Magic
+  ASSERT_EQ(byte_buffer[0], 'F');
+  ASSERT_EQ(byte_buffer[1], 'H');
+  ASSERT_EQ(byte_buffer[2], '0');
+  ASSERT_EQ(byte_buffer[3], '1');
+
+  // Check Header
+  ASSERT_EQ(
+      *(uint32_t*)(byte_buffer + 4),
+      executorch::extension::FlatTensorHeader::kHeaderExpectedLength);
+  ASSERT_EQ(
+      *(uint64_t*)(byte_buffer + 8),
+      48); // Header is 40 bytes + 8 bytes of padding today, and then the
+           // flatbuffer starts.
+  ASSERT_EQ(
+      *(uint64_t*)(byte_buffer + 16),
+      224); // This is fragile, and depends on the schema, the builder, and the
+            // padding needed.
+  const uint64_t segment_offset = 48 +
+      224; // Segment offset, depends on the padded header and flatbuffer sizes.
+  ASSERT_EQ(*(uint64_t*)(byte_buffer + 24), segment_offset);
+  ASSERT_EQ(
+      *(uint64_t*)(byte_buffer + 32),
+      32); // Segment total size, 8 bytes of data (2 floats), 24 bytes of
+           // padding.
+
+  // Check Flatbuffer
+  auto flat_tensor = ::flat_tensor::GetFlatTensor(byte_buffer + 48);
+
+  ASSERT_EQ(
+      flat_tensor->version(),
+      executorch::extension::flat_tensor::kSchemaVersion);
+  ASSERT_EQ(flat_tensor->tensor_alignment(), 16);
+  ASSERT_EQ(flat_tensor->tensors()->size(), 2);
+  ASSERT_EQ(flat_tensor->segments()->size(), 1);
+
+  auto tensor0 = flat_tensor->tensors()->Get(0);
+  ASSERT_EQ(strcmp(tensor0->fully_qualified_name()->c_str(), "linear.bias"), 0);
+  ASSERT_EQ(tensor0->scalar_type(), executorch_flatbuffer::ScalarType::FLOAT);
+  ASSERT_EQ(tensor0->sizes()->size(), 1);
+  ASSERT_EQ(tensor0->segment_index(), 0);
+  ASSERT_EQ(tensor0->offset(), 0);
+
+  auto tensor1 = flat_tensor->tensors()->Get(1);
+  ASSERT_EQ(
+      strcmp(tensor1->fully_qualified_name()->c_str(), "linear.weight"), 0);
+  ASSERT_EQ(tensor1->scalar_type(), executorch_flatbuffer::ScalarType::FLOAT);
+  ASSERT_EQ(tensor1->sizes()->size(), 1);
+  ASSERT_EQ(tensor1->segment_index(), 0);
+  ASSERT_EQ(tensor1->offset(), 16);
+
+  // Test Segments
+  auto segment = flat_tensor->segments()->Get(0);
+
+  ASSERT_EQ(segment->offset(), 0);
+  ASSERT_EQ(segment->size(), 32);
+  uint8_t* data = (uint8_t*)(byte_buffer + segment_offset);
+  ASSERT_EQ(*(float*)(data + 0), linear_bias);
+  ASSERT_EQ(*(float*)(data + 16), linear_weight);
+}