Update runtime to read new schema (#1369)

lucylq · facebook-github-bot · commit 0d0dd1dae213 · 2023-12-07T11:02:46.000-08:00
Summary: Pull Request resolved: #1369 * Runtime branches on `constant_buffer` / `constant_segment`. * Throw error if both contain items. Reviewed By: dbort Differential Revision: D51638620 fbshipit-source-id: 03bbe2f502ab77e06c92ad55aa0ff368babe9b68
diff --git a/runtime/executor/program.cpp b/runtime/executor/program.cpp
@@ -144,13 +144,56 @@ Result<executorch_flatbuffer::ExecutionPlan*> get_execution_plan(
   const executorch_flatbuffer::Program* flatbuffer_program =
       executorch_flatbuffer::GetProgram(program_data->data());
 
-  // The FreeableBuffer owns the data that flatbuffer_program points into. Also
-  // keep a pointer to the loader so it can load more segments when necessary.
-  return Program(
-      loader,
-      segment_base_offset,
-      std::move(program_data.get()),
-      flatbuffer_program);
+  // Constant data may live inside the flatbuffer data (constant_buffer) or in a
+  // separate segment (constant_segment). It should not be in both.
+  const auto& constant_buffer = flatbuffer_program->constant_buffer();
+  const auto& constant_segment = flatbuffer_program->constant_segment();
+
+  // Check if the constant data is inside a separate segment.
+  if (constant_segment != nullptr && constant_segment->offsets()->size() > 0) {
+    ET_CHECK_OR_RETURN_ERROR(
+        constant_buffer->size() == 0,
+        InvalidState,
+        "constant_buffer contains %u items, constant_segment.offsets contains %u items. Only one should be used.",
+        constant_buffer->size(),
+        constant_segment->offsets()->size());
+
+    // Load constant segment.
+    // TODO(T171839323): Add test for segment_index > num available segments.
+    ET_CHECK_OR_RETURN_ERROR(
+        constant_segment->segment_index() <
+            flatbuffer_program->segments()->size(),
+        InvalidArgument,
+        "Constant segment index %d invalid for program segments range %d",
+        constant_segment->segment_index(),
+        flatbuffer_program->segments()->size());
+
+    const executorch_flatbuffer::DataSegment* data_segment =
+        flatbuffer_program->segments()->Get(constant_segment->segment_index());
+    Result<FreeableBuffer> constant_segment_data = loader->Load(
+        segment_base_offset + data_segment->offset(), data_segment->size());
+    if (!constant_segment_data.ok()) {
+      return constant_segment_data.error();
+    }
+    // The FreeableBuffer owns the data that flatbuffer_program points into.
+    // Also keep a pointer to the loader so it can load more segments when
+    // necessary.
+    return Program(
+        loader,
+        segment_base_offset,
+        std::move(program_data.get()),
+        flatbuffer_program,
+        std::move(constant_segment_data.get()));
+  } else {
+    // The constant data is stored inside the flatbuffer, so this program does
+    // not contain a separate segment for it.
+    return Program(
+        loader,
+        segment_base_offset,
+        std::move(program_data.get()),
+        flatbuffer_program,
+        /*constant_segment_data=*/FreeableBuffer{});
+  }
 }
 
 size_t Program::num_methods() const {
@@ -192,21 +235,67 @@ Result<MethodMeta> Program::method_meta(const char* method_name) const {
 }
 
 Result<const void*> Program::get_constant_buffer_data(
-    size_t buffer_index) const {
+    size_t buffer_index,
+    size_t nbytes) const {
   auto internal_program =
       static_cast<const executorch_flatbuffer::Program*>(internal_program_);
-  size_t size = internal_program->constant_buffer()->size();
-  ET_CHECK_OR_RETURN_ERROR(
-      buffer_index < size,
-      InvalidArgument,
-      "Constant buffer %zu out of program buffer range %zu",
-      buffer_index,
-      size);
 
-  const auto& constant_buffer = *internal_program->constant_buffer();
+  // Constant data is either in a separate segment (constant_segment_data) and
+  // loaded during Program::load, or stored inside the flatbuffer data
+  // (constant_buffer).
+  if (constant_segment_data_.data() != nullptr) {
+    size_t num_elems = internal_program->constant_segment()->offsets()->size();
+    ET_CHECK_OR_RETURN_ERROR(
+        buffer_index < num_elems,
+        InvalidArgument,
+        "Constant segment buffer index %zu invalid for program constant segment range %zu",
+        buffer_index,
+        num_elems);
+
+    // All constant data is stored in one segment, with each tensor aligned to
+    // @executorch_tensor_alignment. Tensor offsets are stored in the flatbuffer
+    // data in Program.constant_segment.offsets.
+    // The constant data at buffer_index is located at: base address of the
+    // constant segment + offset for tensor at buffer_index.
+    uint64_t offset = static_cast<uint64_t>(
+        (*internal_program->constant_segment()->offsets())[buffer_index]);
 
-  return static_cast<const void*>(
-      constant_buffer[buffer_index]->storage()->data());
+    size_t size = constant_segment_data_.size();
+    ET_CHECK_OR_RETURN_ERROR(
+        offset + nbytes <= size,
+        InvalidArgument,
+        "Constant segment offset %" PRIu64
+        " + size_bytes %zu invalid for program constant segment size %zu",
+        offset,
+        nbytes,
+        size);
+
+    // Offset is wrt the beginning of the constant segment.
+    return static_cast<const void*>(
+        static_cast<const unsigned char*>(constant_segment_data_.data()) +
+        offset);
+  } else {
+    // Otherwise, the constant data is stored inside Program.constant_buffer.
+    size_t num_elems = internal_program->constant_buffer()->size();
+    ET_CHECK_OR_RETURN_ERROR(
+        buffer_index < num_elems,
+        InvalidArgument,
+        "Constant buffer index %zu invalid for program constant buffer range %zu",
+        buffer_index,
+        num_elems);
+
+    const auto& constant_buffer = *internal_program->constant_buffer();
+
+    ET_CHECK_OR_RETURN_ERROR(
+        constant_buffer[buffer_index]->storage()->size() <= nbytes,
+        InvalidArgument,
+        "Constant buffer size %u larger than allocated nbytes %zu",
+        constant_buffer[buffer_index]->storage()->size(),
+        nbytes);
+
+    return static_cast<const void*>(
+        constant_buffer[buffer_index]->storage()->data());
+  }
 }
 
 Result<int64_t> Program::get_non_const_buffer_size(
diff --git a/runtime/executor/program.h b/runtime/executor/program.h
@@ -89,11 +89,13 @@ class Program final {
   ~Program() = default;
 
   /**
-   * Get the constant buffer inside Program with index buffer_idx
-   * @param[in] buffer_idx the index of the buffer in the constant_buffer
-   * @return The buffer with corresponding index
+   * Get the constant buffer inside Program with index buffer_idx.
+   * @param[in] buffer_idx the index of the buffer in the constant_buffer.
+   * @param[in] nbytes the number of bytes to read from the buffer.
+   * @return The buffer with corresponding index.
    */
-  Result<const void*> get_constant_buffer_data(size_t buffer_idx) const;
+  Result<const void*> get_constant_buffer_data(size_t buffer_idx, size_t nbytes)
+      const;
 
   /**
    * Returns the number of methods in the program.
@@ -254,12 +256,14 @@ class Program final {
       DataLoader* loader,
       size_t segment_base_offset,
       FreeableBuffer&& program_data,
-      const executorch_flatbuffer::Program* internal_program)
+      const executorch_flatbuffer::Program* internal_program,
+      FreeableBuffer&& constant_segment_data)
       : program_data_(std::move(program_data)),
         // Don't need the loader if there are no segments.
         loader_(segment_base_offset > 0 ? loader : nullptr),
         internal_program_(internal_program),
-        segment_base_offset_(segment_base_offset) {}
+        segment_base_offset_(segment_base_offset),
+        constant_segment_data_(std::move(constant_segment_data)) {}
 
   // Not copyable or assignable.
   Program(const Program& rhs) = delete;
@@ -279,6 +283,9 @@ class Program final {
   /// The offset to the first segment, in bytes. If zero, no segments should
   /// be present in internal_program_.
   size_t segment_base_offset_;
+
+  /// Constant segment data.
+  FreeableBuffer constant_segment_data_;
 };
 
 } // namespace executor
diff --git a/runtime/executor/tensor_parser_exec_aten.cpp b/runtime/executor/tensor_parser_exec_aten.cpp
@@ -55,8 +55,8 @@ __ET_NODISCARD Result<void*> getTensorDataPtr(
     size_t nbytes,
     HierarchicalAllocator* allocator) {
   if (s_tensor->constant_buffer_idx() > 0) {
-    auto data =
-        program->get_constant_buffer_data(s_tensor->constant_buffer_idx());
+    auto data = program->get_constant_buffer_data(
+        s_tensor->constant_buffer_idx(), nbytes);
     if (!data.ok()) {
       return data.error();
     }
