Prep TensorParser for Loading Mutable Segments (#4545)

Summary:
Pull Request resolved: #4545

Pretty soon itll be possible for a tensor to have both a memplanned pointer and an initial buffer.

Differential Revision: D60774447

fbshipit-source-id: aed6491c59f58f60e4a221ac8e2b2edae17942a7

Author: JacobSzwejbka

runtime/executor/tensor_parser_exec_aten.cpp

@@ -19,6 +19,41 @@ namespace torch {
 namespace executor {
 namespace deserialization {
 
+namespace {
+
+// Retrieve the buffer specified by the allocation_info
+__ET_NODISCARD Result<void*> getMemPlannedPtr(
+    const executorch_flatbuffer::AllocationDetails* allocation_info,
+    size_t nbytes,
+    HierarchicalAllocator* allocator) {
+  // Normal non-constant Tensor. Allocate data using mem_id and offset.
+
+  // TODO(T142455629): make the allocator actually id based and not indexed
+  // based. -1 is a hack to get the memory ids 0 aligned because previously
+  // 0 was reserved
+  const uint32_t memory_id = allocation_info->memory_id() - 1;
+
+  // Originally this field was a single uint32_t, but we need 64 bits for
+  // larger models. To preserve backwards compatibility, the high bits are
+  // managed in a separate uint32_t field.
+  const uint32_t memory_offset_low = allocation_info->memory_offset_low();
+  const uint32_t memory_offset_high = allocation_info->memory_offset_high();
+
+  size_t memory_offset = memory_offset_low;
+  if (memory_offset_high > 0) {
+    // The compiler should remove this always-true check on 64-bit systems.
+    ET_CHECK_OR_RETURN_ERROR(
+        sizeof(size_t) >= sizeof(uint64_t),
+        NotSupported,
+        "size_t cannot hold memory offset 0x%08" PRIx32 ".%08" PRIx32,
+        memory_offset_high,
+        memory_offset_low);
+    memory_offset |= static_cast<size_t>(memory_offset_high) << 32;
+  }
+  return allocator->get_offset_address(memory_id, memory_offset, nbytes);
+}
+} // namespace
+
 __ET_NODISCARD Result<BoxedEvalueList<exec_aten::Tensor>> parseTensorList(
     const flatbuffers::Vector<int32_t>* tensor_indices,
     EValue* values_,
@@ -53,49 +88,41 @@ __ET_NODISCARD Result<void*> getTensorDataPtr(
     const Program* program,
     size_t nbytes,
     HierarchicalAllocator* allocator) {
-  if (s_tensor->data_buffer_idx() > 0) {
-    auto data =
-        program->get_constant_buffer_data(s_tensor->data_buffer_idx(), nbytes);
-    if (!data.ok()) {
-      return data.error();
+  auto data_buffer_idx = s_tensor->data_buffer_idx();
+  const executorch_flatbuffer::AllocationDetails* allocation_info =
+      s_tensor->allocation_info();
+
+  // Memory Planned, with initial state
+  if (data_buffer_idx > 0 && allocation_info != nullptr) {
+    // Stub case for now.
+
+    // Get memory planned data pointer
+
+    // Call something like program.load_into_buffer(s_tensor->segment_idx,
+    // s_tensor->data_buffer_idx, mem_planned_buffer, nbytes)
+
+    return Error::NotImplemented;
+
+    // Constant
+  } else if (data_buffer_idx > 0 && allocation_info == nullptr) {
+    auto const_data =
+        program->get_constant_buffer_data(data_buffer_idx, nbytes);
+    if (!const_data.ok()) {
+      return const_data.error();
     }
+
     // The const_cast is 'ok' here because the program and runtime should
     // guarantee that this data is never modified.
-    return const_cast<void*>(data.get());
-  }
+    return const_cast<void*>(const_data.get());
 
-  const executorch_flatbuffer::AllocationDetails* allocation_info =
-      s_tensor->allocation_info();
-  if (allocation_info != nullptr) {
-    // Normal non-constant Tensor. Allocate data using mem_id and offset.
-
-    // TODO(T142455629): make the allocator actually id based and not indexed
-    // based. -1 is a hack to get the memory ids 0 aligned because previously
-    // 0 was reserved
-    const uint32_t memory_id = allocation_info->memory_id() - 1;
-
-    // Originally this field was a single uint32_t, but we need 64 bits for
-    // larger models. To preserve backwards compatibility, the high bits are
-    // managed in a separate uint32_t field.
-    const uint32_t memory_offset_low = allocation_info->memory_offset_low();
-    const uint32_t memory_offset_high = allocation_info->memory_offset_high();
-
-    size_t memory_offset = memory_offset_low;
-    if (memory_offset_high > 0) {
-      // The compiler should remove this always-true check on 64-bit systems.
-      ET_CHECK_OR_RETURN_ERROR(
-          sizeof(size_t) >= sizeof(uint64_t),
-          NotSupported,
-          "size_t cannot hold memory offset 0x%08" PRIx32 ".%08" PRIx32,
-          memory_offset_high,
-          memory_offset_low);
-      memory_offset |= static_cast<size_t>(memory_offset_high) << 32;
-    }
-    return allocator->get_offset_address(memory_id, memory_offset, nbytes);
-  }
+    // Memory planned, no initial state
+  } else if (data_buffer_idx == 0 && allocation_info != nullptr) {
+    return getMemPlannedPtr(allocation_info, nbytes, allocator);
 
-  // The tensor's data will be allocated as part of execution.
-  return nullptr;
+    // Pointer recived at runtime
+  } else { // data_buffer_idx == 0 && allocation_info == nullptr,
+    return nullptr;
+  }
 }
 
 } // namespace deserialization