Skip to content

Commit dc9909b

Browse files
lucylqlucylq
committed
[executorch][flat_tensor] DataMap implementation
Pull Request resolved: #7900 DataMap implementation that * Loads a flat_tensor file * Populates a map with {fqn: tensor} and {fqn: TensorLayout}. * Makes tensor information available via the named_data_map.h interface. For now, DataMap doesn't store the DataLoader. - If/when tensors are in their own segments, DataMap should also store a DataLoader. ghstack-source-id: 262759344 Differential Revision: [D67064580](https://our.internmc.facebook.com/intern/diff/D67064580/)
1 parent c4d9580 commit dc9909b

File tree

8 files changed

+455
-2
lines changed

8 files changed

+455
-2
lines changed

‎extension/flat_tensor/named_data_map/TARGETS

+6
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/named_data_map/data_map.cpp

+213
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,213 @@
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/named_data_map/data_map.h>
10+
#include <executorch/extension/flat_tensor/serialize/flat_tensor_header.h>
11+
#include <executorch/extension/flat_tensor/serialize/schema_generated.h>
12+
#include <executorch/runtime/core/error.h>
13+
#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
14+
#include <executorch/runtime/core/freeable_buffer.h>
15+
#include <executorch/runtime/core/result.h>
16+
#include <executorch/runtime/core/span.h>
17+
#include <executorch/runtime/platform/compiler.h>
18+
19+
#include <tuple>
20+
#include <unordered_map>
21+
22+
using executorch::runtime::Error;
23+
using executorch::runtime::FreeableBuffer;
24+
using executorch::runtime::Result;
25+
using executorch::runtime::Span;
26+
27+
using executorch::aten::ScalarType;
28+
using executorch::runtime::DataLoader;
29+
using executorch::runtime::TensorLayout;
30+
31+
namespace executorch {
32+
namespace extension {
33+
34+
namespace {
35+
/**
36+
* FlatTensor data must be aligned to this value to properly parse it. Must be a
37+
* power of 2. Note that max_align_t is the alignment that malloc() and new
38+
* guarantee.
39+
*/
40+
constexpr size_t kMinimumAlignment = alignof(std::max_align_t);
41+
42+
bool IsAligned(const void* data) {
43+
uintptr_t addr = reinterpret_cast<uintptr_t>(data);
44+
return addr % kMinimumAlignment == 0;
45+
}
46+
} // namespace
47+
48+
ET_NODISCARD Result<const TensorLayout> DataMap::get_metadata(
49+
const char* fqn) const {
50+
auto result = _name_to_tensor.find(fqn);
51+
if (result == _name_to_tensor.end()) {
52+
return Error::NotFound;
53+
}
54+
return std::get<2>(result->second);
55+
}
56+
57+
ET_NODISCARD Result<FreeableBuffer> DataMap::get_data(const char* fqn) const {
58+
auto result = _name_to_tensor.find(fqn);
59+
if (result == _name_to_tensor.end()) {
60+
return Error::NotFound;
61+
}
62+
int offset = std::get<1>(result->second);
63+
TensorLayout tensor = std::get<2>(result->second);
64+
65+
const uint8_t* data = static_cast<const uint8_t*>(_data_ro.data()) + offset;
66+
return FreeableBuffer(data, tensor.nbytes(), nullptr);
67+
}
68+
69+
ET_NODISCARD Error
70+
DataMap::load_data_into(const char* fqn, size_t size, void* buffer) const {
71+
return Error::NotImplemented;
72+
}
73+
74+
ET_NODISCARD Result<int> DataMap::get_num_keys() const {
75+
return _name_to_tensor.size();
76+
}
77+
78+
ET_NODISCARD Result<const char*> DataMap::get_key(int index) const {
79+
if (index <= 0 || index >= _name_to_tensor.size()) {
80+
return Error::InvalidArgument;
81+
}
82+
83+
auto iter = _name_to_tensor.begin();
84+
for (int i = 0; i < index; ++i) {
85+
++iter;
86+
}
87+
return iter->first.c_str();
88+
}
89+
90+
/* static */ Result<DataMap> DataMap::load(DataLoader* loader) {
91+
// Load data map.
92+
size_t flatbuffer_offset = 0;
93+
size_t flatbuffer_size = 0;
94+
size_t segment_base_offset = 0;
95+
size_t segment_data_size = 0;
96+
{
97+
// Check header.
98+
Result<FreeableBuffer> header = loader->load(
99+
/*offset=*/0,
100+
FlatTensorHeader::kNumHeadBytes,
101+
DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External));
102+
if (!header.ok()) {
103+
return header.error();
104+
}
105+
Result<FlatTensorHeader> fh =
106+
FlatTensorHeader::Parse(header->data(), header->size());
107+
if (fh.ok()) {
108+
// The header has the data map size.
109+
flatbuffer_offset = fh->flatbuffer_offset;
110+
flatbuffer_size = fh->flatbuffer_size;
111+
segment_base_offset = fh->segment_base_offset;
112+
segment_data_size = fh->segment_data_size;
113+
} else if (fh.error() == Error::NotFound) {
114+
// No header, throw error.
115+
ET_LOG(Error, "No FlatTensorHeader found.");
116+
return fh.error();
117+
} else {
118+
// corruption, throw error.
119+
ET_LOG(Error, "Flat tensor header may be corrupt.");
120+
return fh.error();
121+
}
122+
}
123+
124+
// Load flatbuffer data as a segment.
125+
Result<FreeableBuffer> flat_tensor_data = loader->load(
126+
/*offset=*/flatbuffer_offset,
127+
flatbuffer_size,
128+
DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External));
129+
if (!flat_tensor_data.ok()) {
130+
return flat_tensor_data.error();
131+
}
132+
133+
// Make sure magic matches.
134+
if (!flat_tensor::FlatTensorBufferHasIdentifier(flat_tensor_data->data())) {
135+
ET_LOG(
136+
Error,
137+
"FlatTensor identifier '%.4s' != expected '%.4s'",
138+
flatbuffers::GetBufferIdentifier(flat_tensor_data->data()),
139+
flat_tensor::FlatTensorIdentifier());
140+
return Error::InvalidExternalData;
141+
}
142+
143+
// The flatbuffer data must start at an aligned address to ensure internal
144+
// alignment of flatbuffer fields.
145+
ET_CHECK_OR_RETURN_ERROR(
146+
IsAligned(flat_tensor_data->data()),
147+
InvalidArgument,
148+
"FlatTensor data 0x%p must be aligned to %zu",
149+
flat_tensor_data->data(),
150+
kMinimumAlignment);
151+
152+
// Get pointer to root of flatbuffer table.
153+
const flat_tensor::FlatTensor* flat_tensor =
154+
flat_tensor::GetFlatTensor(flat_tensor_data->data());
155+
156+
// Get pointer to tensor metadata.
157+
const auto* s_tensor_metadata = flat_tensor->tensors();
158+
assert(s_tensor_metadata != nullptr);
159+
160+
std::unordered_map<std::string, std::tuple<int, int, TensorLayout>>
161+
fqn_to_tensor_layout = {};
162+
for (int i = 0; i < s_tensor_metadata->size(); i++) {
163+
// Create TensorLayouts.
164+
ScalarType scalar_type =
165+
static_cast<ScalarType>(s_tensor_metadata->Get(i)->scalar_type());
166+
const int dim = s_tensor_metadata->Get(i)->sizes()->size();
167+
168+
const auto serialized_sizes = s_tensor_metadata->Get(i)->sizes()->data();
169+
const auto serialized_dim_order =
170+
s_tensor_metadata->Get(i)->dim_order()->data();
171+
TensorLayout tensor_layout = TensorLayout(
172+
scalar_type,
173+
Span<const int32_t>(serialized_sizes, dim),
174+
Span<const uint8_t>(serialized_dim_order, dim));
175+
176+
int segment_index = s_tensor_metadata->Get(i)->segment_index();
177+
int offset = s_tensor_metadata->Get(i)->offset();
178+
std::string fqn = s_tensor_metadata->Get(i)->fully_qualified_name()->str();
179+
180+
auto val = std::make_tuple(segment_index, offset, tensor_layout);
181+
fqn_to_tensor_layout.insert({fqn, std::move(val)});
182+
}
183+
184+
// Load constant data.
185+
const auto* s_data_segment = flat_tensor->segments();
186+
187+
// Only support one segment for now.
188+
assert(s_data_segment->size() == 1);
189+
// First segment offset should be 0.
190+
int segment_offset = s_data_segment->Get(0)->offset();
191+
assert(segment_offset == 0);
192+
// First segment size should be <= the total segment data size.
193+
int segment_size = s_data_segment->Get(0)->size();
194+
assert(segment_size <= segment_data_size);
195+
196+
Result<FreeableBuffer> _data_ro = loader->load(
197+
/*offset=*/segment_base_offset + segment_offset,
198+
segment_size,
199+
DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External));
200+
if (!_data_ro.ok()) {
201+
return _data_ro.error();
202+
}
203+
204+
return DataMap(
205+
std::move(flat_tensor_data.get()),
206+
std::move(fqn_to_tensor_layout),
207+
std::move(_data_ro.get()));
208+
}
209+
210+
DataMap::~DataMap() {}
211+
212+
} // namespace extension
213+
} // namespace executorch

‎extension/flat_tensor/named_data_map/data_map.h

+83
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
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/data_loader.h>
12+
#include <executorch/runtime/core/exec_aten/exec_aten.h>
13+
#include <executorch/runtime/core/named_data_map.h>
14+
#include <executorch/runtime/core/result.h>
15+
#include <executorch/runtime/core/tensor_layout.h>
16+
#include <executorch/runtime/platform/compiler.h>
17+
18+
#include <unordered_map>
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 {
24+
struct FlatTensor;
25+
} // namespace flat_tensor
26+
27+
namespace executorch {
28+
namespace extension {
29+
30+
class DataMap final : public executorch::runtime::NamedDataMap {
31+
public:
32+
static executorch::runtime::Result<DataMap> load(
33+
executorch::runtime::DataLoader* loader);
34+
35+
ET_NODISCARD
36+
executorch::runtime::Result<const executorch::runtime::TensorLayout>
37+
get_metadata(const char* fqn) const override;
38+
ET_NODISCARD
39+
executorch::runtime::Result<executorch::runtime::FreeableBuffer> get_data(
40+
const char* fqn) const override;
41+
ET_NODISCARD runtime::Error
42+
load_data_into(const char* fqn, size_t size, void* buffer) const override;
43+
44+
ET_NODISCARD executorch::runtime::Result<int> get_num_keys() const override;
45+
ET_NODISCARD executorch::runtime::Result<const char*> get_key(
46+
int index) const override;
47+
48+
DataMap(DataMap&&) noexcept = default;
49+
~DataMap() override;
50+
51+
private:
52+
DataMap(
53+
executorch::runtime::FreeableBuffer&& flat_tensor_data,
54+
std::unordered_map<
55+
std::string,
56+
std::tuple<int, int, executorch::runtime::TensorLayout>>
57+
name_to_tensor,
58+
executorch::runtime::FreeableBuffer&& data_ro)
59+
: _flat_tensor_data(std::move(flat_tensor_data)),
60+
_name_to_tensor(std::move(name_to_tensor)),
61+
_data_ro(std::move(data_ro)) {}
62+
63+
// Not copyable or assignable.
64+
DataMap(const DataMap& rhs) = delete;
65+
DataMap& operator=(DataMap&& rhs) noexcept = delete;
66+
DataMap& operator=(const DataMap& rhs) = delete;
67+
68+
// FlatTensor flatbuffer data. Contains the data backing up
69+
// TensorLayout information in the _name_to_tensor map; must outlive it.
70+
executorch::runtime::FreeableBuffer _flat_tensor_data;
71+
72+
// Map of name to {segment index, offset, TensorLayout}.
73+
std::unordered_map<
74+
std::string,
75+
std::tuple<int, int, executorch::runtime::TensorLayout>>
76+
_name_to_tensor;
77+
78+
// Raw, read-only tensor data.
79+
executorch::runtime::FreeableBuffer _data_ro;
80+
};
81+
82+
} // namespace extension
83+
} // namespace executorch

‎extension/flat_tensor/named_data_map/targets.bzl

+23
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
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/exec_aten:lib",
18+
"//executorch/runtime/core/exec_aten/util:tensor_util",
19+
],
20+
visibility = [
21+
"//executorch/...",
22+
],
23+
)

‎extension/flat_tensor/test/TARGETS

+1-1
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)