1 /*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9 #include <executorch/runtime/executor/tensor_parser.h>
10
11 #include <executorch/runtime/core/evalue.h>
12 #include <executorch/runtime/core/exec_aten/exec_aten.h>
13 #include <executorch/runtime/executor/memory_manager.h>
14 #include <executorch/runtime/executor/program.h>
15 #include <executorch/runtime/platform/profiler.h>
16 #include <executorch/schema/program_generated.h>
17
18 namespace executorch {
19 namespace runtime {
20 namespace deserialization {
21
22 // Provides access to private Program methods.
23 class TensorParser final {
24 public:
load_mutable_subsegment_into(const Program * program,size_t mutable_data_segments_index,size_t offset_index,size_t size,void * buffer)25 ET_NODISCARD static Error load_mutable_subsegment_into(
26 const Program* program,
27 size_t mutable_data_segments_index,
28 size_t offset_index,
29 size_t size,
30 void* buffer) {
31 return program->load_mutable_subsegment_into(
32 mutable_data_segments_index, offset_index, size, buffer);
33 }
34 };
35
36 namespace {
37
38 // Retrieve the buffer specified by the allocation_info
getMemPlannedPtr(const executorch_flatbuffer::AllocationDetails * allocation_info,size_t nbytes,HierarchicalAllocator * allocator)39 ET_NODISCARD Result<void*> getMemPlannedPtr(
40 const executorch_flatbuffer::AllocationDetails* allocation_info,
41 size_t nbytes,
42 HierarchicalAllocator* allocator) {
43 // Normal non-constant Tensor. Allocate data using mem_id and offset.
44
45 // TODO(T142455629): make the allocator actually id based and not indexed
46 // based. -1 is a hack to get the memory ids 0 aligned because previously
47 // 0 was reserved
48 const uint32_t memory_id = allocation_info->memory_id() - 1;
49
50 // Originally this field was a single uint32_t, but we need 64 bits for
51 // larger models. To preserve backwards compatibility, the high bits are
52 // managed in a separate uint32_t field.
53 const uint32_t memory_offset_low = allocation_info->memory_offset_low();
54 const uint32_t memory_offset_high = allocation_info->memory_offset_high();
55
56 size_t memory_offset = memory_offset_low;
57 if (memory_offset_high > 0) {
58 // The compiler should remove this always-true check on 64-bit systems.
59 ET_CHECK_OR_RETURN_ERROR(
60 sizeof(size_t) >= sizeof(uint64_t),
61 NotSupported,
62 "size_t cannot hold memory offset 0x%08" PRIx32 ".%08" PRIx32,
63 memory_offset_high,
64 memory_offset_low);
65 memory_offset |= static_cast<size_t>(memory_offset_high) << 32;
66 }
67 return allocator->get_offset_address(memory_id, memory_offset, nbytes);
68 }
69 } // namespace
70
parseTensorList(const flatbuffers::Vector<int32_t> * tensor_indices,EValue * values_,MemoryManager * memory_manager)71 ET_NODISCARD Result<BoxedEvalueList<exec_aten::Tensor>> parseTensorList(
72 const flatbuffers::Vector<int32_t>* tensor_indices,
73 EValue* values_,
74 MemoryManager* memory_manager) {
75 EXECUTORCH_SCOPE_PROF("TensorParser::parseTensorList");
76
77 auto* tensor_list =
78 memory_manager->method_allocator()->allocateList<exec_aten::Tensor>(
79 tensor_indices->size());
80 if (tensor_list == nullptr) {
81 return Error::MemoryAllocationFailed;
82 }
83 auto* evalp_list = memory_manager->method_allocator()->allocateList<EValue*>(
84 tensor_indices->size());
85 if (evalp_list == nullptr) {
86 return Error::MemoryAllocationFailed;
87 }
88
89 // For each tensor index look up the corresponding Tensor (which has been
90 // already allocated) and stick it in the list.
91 size_t output_idx = 0;
92 for (int32_t tensor_index : *tensor_indices) {
93 // Placement new as the list elements are not initialized, so calling
94 // copy assignment is not defined if its non trivial.
95 new (&tensor_list[output_idx]) exec_aten::Tensor(
96 values_[static_cast<size_t>(tensor_index)].toTensor());
97 evalp_list[output_idx] = &values_[static_cast<size_t>(tensor_index)];
98 output_idx++;
99 }
100
101 return BoxedEvalueList<exec_aten::Tensor>(
102 evalp_list, tensor_list, tensor_indices->size());
103 }
104
getTensorDataPtr(const executorch_flatbuffer::Tensor * s_tensor,const Program * program,size_t nbytes,HierarchicalAllocator * allocator)105 ET_NODISCARD Result<void*> getTensorDataPtr(
106 const executorch_flatbuffer::Tensor* s_tensor,
107 const Program* program,
108 size_t nbytes,
109 HierarchicalAllocator* allocator) {
110 auto data_buffer_idx = s_tensor->data_buffer_idx();
111 const executorch_flatbuffer::AllocationDetails* allocation_info =
112 s_tensor->allocation_info();
113
114 // Memory Planned, with initial state
115 if (data_buffer_idx > 0 && allocation_info != nullptr) {
116 auto planned_ptr = getMemPlannedPtr(allocation_info, nbytes, allocator);
117 if (!planned_ptr.ok()) {
118 return planned_ptr.error();
119 }
120 auto err = TensorParser::load_mutable_subsegment_into(
121 program, 0, s_tensor->data_buffer_idx(), nbytes, planned_ptr.get());
122
123 if (err != Error::Ok) {
124 return err;
125 }
126 return planned_ptr;
127
128 // Constant
129 } else if (data_buffer_idx > 0 && allocation_info == nullptr) {
130 auto const_data =
131 program->get_constant_buffer_data(data_buffer_idx, nbytes);
132 if (!const_data.ok()) {
133 return const_data.error();
134 }
135
136 // The const_cast is 'ok' here because the program and runtime should
137 // guarantee that this data is never modified.
138 return const_cast<void*>(const_data.get());
139
140 // Memory planned, no initial state
141 } else if (data_buffer_idx == 0 && allocation_info != nullptr) {
142 return getMemPlannedPtr(allocation_info, nbytes, allocator);
143
144 // Pointer recived at runtime
145 } else { // data_buffer_idx == 0 && allocation_info == nullptr,
146 return nullptr;
147 }
148 }
149
150 } // namespace deserialization
151 } // namespace runtime
152 } // namespace executorch
153