Skip to content

Commit e9bf102

Browse files
lucylqlucylq
committed
[executorch][flat_tensor] DataMap implementation
Pull Request resolved: #7900 DataMap implementation that * Loads a flat_tensor file * Makes tensor information available via the named_data_map.h interface. TODO: in a later diff, update the ET runtime to hold onto the FreeableBuffers returned by the NDM. Then, the NDM will not persist the segment. T214294528 ghstack-source-id: 264501132 Differential Revision: [D67064580](https://our.internmc.facebook.com/intern/diff/D67064580/)
1 parent 15c8bdf commit e9bf102

File tree

8 files changed

+531
-2
lines changed

8 files changed

+531
-2
lines changed

‎extension/flat_tensor/TARGETS‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
2+
load(":targets.bzl", "define_common_targets")
3+
4+
oncall("executorch")
5+
6+
define_common_targets()

‎extension/flat_tensor/data_map.cpp‎

Lines changed: 249 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,249 @@
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/extension/flat_tensor/data_map.h>
10+
11+
#include <executorch/extension/flat_tensor/serialize/flat_tensor_header.h>
12+
#include <executorch/extension/flat_tensor/serialize/schema_generated.h>
13+
14+
#include <executorch/runtime/core/error.h>
15+
#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
16+
#include <executorch/runtime/core/freeable_buffer.h>
17+
#include <executorch/runtime/core/result.h>
18+
#include <executorch/runtime/core/span.h>
19+
#include <executorch/runtime/platform/compiler.h>
20+
21+
using executorch::runtime::Error;
22+
using executorch::runtime::FreeableBuffer;
23+
using executorch::runtime::Result;
24+
using executorch::runtime::Span;
25+
26+
using executorch::aten::ScalarType;
27+
using executorch::runtime::DataLoader;
28+
using executorch::runtime::TensorLayout;
29+
30+
namespace executorch {
31+
namespace extension {
32+
33+
namespace {
34+
/**
35+
* FlatTensor data must be aligned to this value to properly parse it. Must be a
36+
* power of 2. Note that max_align_t is the alignment that malloc() and new
37+
* guarantee.
38+
*/
39+
constexpr size_t kMinimumAlignment = alignof(std::max_align_t);
40+
41+
bool is_aligned(const void* data) {
42+
uintptr_t addr = reinterpret_cast<uintptr_t>(data);
43+
return addr % kMinimumAlignment == 0;
44+
}
45+
46+
Result<const flat_tensor_flatbuffer::TensorMetadata*> get_flat_tensor_metadata(
47+
const char* key,
48+
const flatbuffers::Vector<
49+
flatbuffers::Offset<flat_tensor_flatbuffer::TensorMetadata>>* tensors) {
50+
// Linear search by name.
51+
for (int i = 0; i < tensors->size(); i++) {
52+
if (std::strcmp(tensors->Get(i)->fully_qualified_name()->c_str(), key) ==
53+
0) {
54+
// TODO(T214294528): Support multiple segments in FlatTensor.
55+
if (tensors->Get(i)->segment_index() != 0) {
56+
return Error::InvalidExternalData;
57+
}
58+
return tensors->Get(i);
59+
}
60+
}
61+
return Error::NotFound;
62+
}
63+
64+
Result<const TensorLayout> create_tensor_layout(
65+
const flat_tensor_flatbuffer::TensorMetadata* tensor_metadata) {
66+
ScalarType scalar_type =
67+
static_cast<ScalarType>(tensor_metadata->scalar_type());
68+
const int dim = tensor_metadata->sizes()->size();
69+
const auto serialized_sizes = tensor_metadata->sizes()->data();
70+
const auto serialized_dim_order = tensor_metadata->dim_order()->data();
71+
return TensorLayout::create(
72+
Span<const int32_t>(serialized_sizes, dim),
73+
Span<const uint8_t>(serialized_dim_order, dim),
74+
scalar_type);
75+
}
76+
77+
} // namespace
78+
79+
ET_NODISCARD Result<const TensorLayout> DataMap::get_metadata(
80+
const char* key) const {
81+
Result<const flat_tensor_flatbuffer::TensorMetadata*> metadata_res =
82+
get_flat_tensor_metadata(key, flat_tensor_->tensors());
83+
if (!metadata_res.ok()) {
84+
return metadata_res.error();
85+
}
86+
return create_tensor_layout(metadata_res.get());
87+
}
88+
89+
ET_NODISCARD Result<FreeableBuffer> DataMap::get_data(const char* key) const {
90+
auto tensor_metadata = flat_tensor_->tensors();
91+
92+
Result<const flat_tensor_flatbuffer::TensorMetadata*> metadata_res =
93+
get_flat_tensor_metadata(key, tensor_metadata);
94+
if (!metadata_res.ok()) {
95+
return metadata_res.error();
96+
}
97+
const auto metadata = metadata_res.get();
98+
if (metadata->segment_index() == -1 || metadata->offset() == -1) {
99+
// Key doesn't exist.
100+
return Error::NotFound;
101+
}
102+
103+
Result<const TensorLayout> tensor_layout_res = create_tensor_layout(metadata);
104+
if (!tensor_layout_res.ok()) {
105+
return tensor_layout_res.error();
106+
}
107+
108+
return FreeableBuffer(
109+
static_cast<const uint8_t*>(data_ro_.data()) + metadata->offset(),
110+
tensor_layout_res.get().nbytes(),
111+
nullptr);
112+
}
113+
114+
ET_NODISCARD Result<size_t> DataMap::load_data_into(
115+
ET_UNUSED const char* key,
116+
ET_UNUSED void* buffer,
117+
ET_UNUSED size_t size) const {
118+
return Error::NotImplemented;
119+
}
120+
121+
ET_NODISCARD Result<size_t> DataMap::get_num_keys() const {
122+
return flat_tensor_->tensors()->size();
123+
}
124+
125+
ET_NODISCARD Result<const char*> DataMap::get_key(size_t index) const {
126+
if (index < 0 || index >= flat_tensor_->tensors()->size()) {
127+
return Error::InvalidArgument;
128+
}
129+
return flat_tensor_->tensors()->Get(index)->fully_qualified_name()->c_str();
130+
}
131+
132+
/* static */ Result<DataMap> DataMap::load(DataLoader* loader) {
133+
// Load data map.
134+
size_t flatbuffer_offset = 0;
135+
size_t flatbuffer_size = 0;
136+
size_t segment_base_offset = 0;
137+
size_t segment_data_size = 0;
138+
{
139+
// Check header.
140+
Result<FreeableBuffer> header = loader->load(
141+
/*offset=*/0,
142+
FlatTensorHeader::kNumHeadBytes,
143+
DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External));
144+
if (!header.ok()) {
145+
return header.error();
146+
}
147+
Result<FlatTensorHeader> fh =
148+
FlatTensorHeader::Parse(header->data(), header->size());
149+
if (fh.ok()) {
150+
// The header has the data map size.
151+
flatbuffer_offset = fh->flatbuffer_offset;
152+
flatbuffer_size = fh->flatbuffer_size;
153+
segment_base_offset = fh->segment_base_offset;
154+
segment_data_size = fh->segment_data_size;
155+
} else if (fh.error() == Error::NotFound) {
156+
// No header, throw error.
157+
ET_LOG(Error, "No FlatTensorHeader found.");
158+
return fh.error();
159+
} else {
160+
// corruption, throw error.
161+
ET_LOG(Error, "Flat tensor header may be corrupt.");
162+
return fh.error();
163+
}
164+
}
165+
166+
// Load flatbuffer data as a segment.
167+
Result<FreeableBuffer> flat_tensor_data = loader->load(
168+
/*offset=*/0,
169+
flatbuffer_offset + flatbuffer_size,
170+
DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External));
171+
if (!flat_tensor_data.ok()) {
172+
return flat_tensor_data.error();
173+
}
174+
175+
// Make sure magic matches.
176+
if (!flat_tensor_flatbuffer::FlatTensorBufferHasIdentifier(
177+
flat_tensor_data->data())) {
178+
ET_LOG(
179+
Error,
180+
"FlatTensor identifier '%.4s' != expected '%.4s'",
181+
flatbuffers::GetBufferIdentifier(flat_tensor_data->data()),
182+
flat_tensor_flatbuffer::FlatTensorIdentifier());
183+
return Error::InvalidExternalData;
184+
}
185+
186+
// The flatbuffer data must start at an aligned address to ensure internal
187+
// alignment of flatbuffer fields.
188+
ET_CHECK_OR_RETURN_ERROR(
189+
is_aligned(flat_tensor_data->data()),
190+
InvalidArgument,
191+
"FlatTensor data 0x%p must be aligned to %zu",
192+
flat_tensor_data->data(),
193+
kMinimumAlignment);
194+
195+
// Get pointer to root of flatbuffer table.
196+
const flat_tensor_flatbuffer::FlatTensor* flat_tensor =
197+
flat_tensor_flatbuffer::GetFlatTensor(flat_tensor_data->data());
198+
199+
// Get pointer to tensor metadata.
200+
const auto* s_tensor_metadata = flat_tensor->tensors();
201+
if (s_tensor_metadata == nullptr) {
202+
ET_LOG(Error, "FlatTensor has no tensor metadata.");
203+
return Error::InvalidExternalData;
204+
}
205+
206+
// Load constant data.
207+
const auto* s_data_segment = flat_tensor->segments();
208+
209+
// TODO(T214294528): Support multiple segments in FlatTensor.
210+
if (s_data_segment->size() != 1) {
211+
ET_LOG(
212+
Error,
213+
"FlatTensor has %u segments, only 1 supported.",
214+
s_data_segment->size());
215+
}
216+
// First segment offset should be 0.
217+
int segment_offset = s_data_segment->Get(0)->offset();
218+
if (segment_offset != 0) {
219+
ET_LOG(Error, "FlatTensor segment offset %d != 0", segment_offset);
220+
}
221+
// First segment size should be <= the total segment data size.
222+
int segment_size = s_data_segment->Get(0)->size();
223+
if (segment_size > segment_data_size) {
224+
ET_LOG(
225+
Error,
226+
"FlatTensor segment size %d > segment data size %zu",
227+
segment_size,
228+
segment_data_size);
229+
}
230+
231+
Result<FreeableBuffer> data_ro = loader->load(
232+
/*offset=*/segment_base_offset + segment_offset,
233+
segment_size,
234+
DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External));
235+
if (!data_ro.ok()) {
236+
return data_ro.error();
237+
}
238+
239+
return DataMap(
240+
std::move(flat_tensor_data.get()), flat_tensor, std::move(data_ro.get()));
241+
}
242+
243+
DataMap::~DataMap() {
244+
flat_tensor_data_.Free();
245+
data_ro_.Free();
246+
}
247+
248+
} // namespace extension
249+
} // namespace executorch

‎extension/flat_tensor/data_map.h‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
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+
#pragma once
10+
11+
#include <executorch/runtime/core/named_data_map.h>
12+
13+
#include <executorch/runtime/core/data_loader.h>
14+
#include <executorch/runtime/core/exec_aten/exec_aten.h>
15+
#include <executorch/runtime/core/result.h>
16+
#include <executorch/runtime/core/tensor_layout.h>
17+
#include <executorch/runtime/platform/compiler.h>
18+
19+
#include <utility>
20+
21+
// Forward declare flatbuffer types. This is a public header and must not
22+
// include the generated flatbuffer header.
23+
namespace flat_tensor_flatbuffer {
24+
struct FlatTensor;
25+
} // namespace flat_tensor_flatbuffer
26+
27+
namespace executorch {
28+
namespace extension {
29+
30+
/**
31+
* A NamedDataMap implementation for FlatTensor-serialized data.
32+
*/
33+
class DataMap final : public executorch::runtime::NamedDataMap {
34+
public:
35+
/**
36+
* Creates a new DataMap that wraps FlatTensor data.
37+
*
38+
* @param[in] loader The DataLoader that wraps the FlatTensor file.
39+
*/
40+
static executorch::runtime::Result<DataMap> load(
41+
executorch::runtime::DataLoader* loader);
42+
43+
ET_NODISCARD
44+
executorch::runtime::Result<const executorch::runtime::TensorLayout>
45+
get_metadata(const char* key) const override;
46+
ET_NODISCARD
47+
executorch::runtime::Result<executorch::runtime::FreeableBuffer> get_data(
48+
const char* key) const override;
49+
ET_NODISCARD executorch::runtime::Result<size_t>
50+
load_data_into(const char* key, void* buffer, size_t size) const override;
51+
52+
ET_NODISCARD executorch::runtime::Result<size_t> get_num_keys()
53+
const override;
54+
ET_NODISCARD executorch::runtime::Result<const char*> get_key(
55+
size_t index) const override;
56+
57+
DataMap(DataMap&&) noexcept = default;
58+
~DataMap() override;
59+
60+
private:
61+
DataMap(
62+
executorch::runtime::FreeableBuffer&& flat_tensor_data,
63+
const flat_tensor_flatbuffer::FlatTensor* flat_tensor,
64+
executorch::runtime::FreeableBuffer&& data_ro)
65+
: flat_tensor_data_(std::move(flat_tensor_data)),
66+
flat_tensor_(flat_tensor),
67+
data_ro_(std::move(data_ro)){};
68+
69+
// Not copyable or assignable.
70+
DataMap(const DataMap& rhs) = delete;
71+
DataMap& operator=(DataMap&& rhs) noexcept = delete;
72+
DataMap& operator=(const DataMap& rhs) = delete;
73+
74+
// Serialized flat_tensor flatbuffer data.
75+
executorch::runtime::FreeableBuffer flat_tensor_data_;
76+
77+
// Flatbuffer representation of the flat_tensor.
78+
const flat_tensor_flatbuffer::FlatTensor* flat_tensor_;
79+
80+
// Loaded read-only tensor data.
81+
executorch::runtime::FreeableBuffer data_ro_;
82+
};
83+
84+
} // namespace extension
85+
} // namespace executorch

‎extension/flat_tensor/targets.bzl‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
2+
3+
def define_common_targets():
4+
runtime.cxx_library(
5+
name = "data_map",
6+
srcs = [
7+
"data_map.cpp",
8+
],
9+
exported_headers = ["data_map.h"],
10+
deps = [
11+
"//executorch/extension/flat_tensor/serialize:schema",
12+
"//executorch/extension/flat_tensor/serialize:serialize",
13+
"//executorch/extension/flat_tensor/serialize:generated_headers",
14+
"//executorch/extension/flat_tensor/serialize:flat_tensor_header",
15+
"//executorch/runtime/core:core",
16+
"//executorch/runtime/core:evalue",
17+
"//executorch/runtime/core:named_data_map",
18+
"//executorch/runtime/core/exec_aten:lib",
19+
"//executorch/runtime/core/exec_aten/util:tensor_util",
20+
],
21+
visibility = [
22+
"//executorch/...",
23+
],
24+
)

‎extension/flat_tensor/test/TARGETS‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ load(":targets.bzl", "define_common_targets")
66

77
oncall("executorch")
88

9-
define_common_targets()
9+
define_common_targets(is_fbcode=True)
1010

1111
python_unittest(
1212
name = "serialize",

0 commit comments

Comments
 (0)