From 87eded30247019e5fe1d5cec5f6206b3d79d7cff Mon Sep 17 00:00:00 2001 From: lucylq Date: Thu, 23 Jan 2025 11:53:20 -0800 Subject: [PATCH 1/2] [executorch][flat_tensor] DataMap implementation 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. Differential Revision: [D67064580](https://our.internmc.facebook.com/intern/diff/D67064580/) [ghstack-poisoned] --- extension/flat_tensor/named_data_map/TARGETS | 6 + .../flat_tensor/named_data_map/data_map.cpp | 209 ++++++++++++++++++ .../flat_tensor/named_data_map/data_map.h | 83 +++++++ .../flat_tensor/named_data_map/targets.bzl | 23 ++ extension/flat_tensor/test/TARGETS | 2 +- extension/flat_tensor/test/data_map_test.cpp | 102 +++++++++ extension/flat_tensor/test/targets.bzl | 28 ++- runtime/core/data_loader.h | 4 + 8 files changed, 455 insertions(+), 2 deletions(-) create mode 100644 extension/flat_tensor/named_data_map/TARGETS create mode 100644 extension/flat_tensor/named_data_map/data_map.cpp create mode 100644 extension/flat_tensor/named_data_map/data_map.h create mode 100644 extension/flat_tensor/named_data_map/targets.bzl create mode 100644 extension/flat_tensor/test/data_map_test.cpp diff --git a/extension/flat_tensor/named_data_map/TARGETS b/extension/flat_tensor/named_data_map/TARGETS new file mode 100644 index 00000000000..08e83a5f3c4 --- /dev/null +++ b/extension/flat_tensor/named_data_map/TARGETS @@ -0,0 +1,6 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +define_common_targets() diff --git a/extension/flat_tensor/named_data_map/data_map.cpp b/extension/flat_tensor/named_data_map/data_map.cpp new file mode 100644 index 00000000000..605876ead18 --- /dev/null +++ b/extension/flat_tensor/named_data_map/data_map.cpp @@ -0,0 +1,209 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +#include +#include +#include + +#include +#include +#include +#include +#include + +#include + +#include +#include + +using executorch::runtime::Error; +using executorch::runtime::FreeableBuffer; +using executorch::runtime::Result; +using executorch::runtime::Span; + +using executorch::aten::ScalarType; +using executorch::runtime::DataLoader; +using executorch::runtime::TensorLayout; + +namespace executorch { +namespace extension { + +namespace { +/** + * FlatTensor data must be aligned to this value to properly parse it. Must be a + * power of 2. Note that max_align_t is the alignment that malloc() and new + * guarantee. + */ +constexpr size_t kMinimumAlignment = alignof(std::max_align_t); + +bool IsAligned(const void* data) { + uintptr_t addr = reinterpret_cast(data); + return addr % kMinimumAlignment == 0; +} +} // namespace + +ET_NODISCARD Result DataMap::get_metadata( + const char* fqn) const { + auto result = _tensor_map.find(fqn); + if (result == _tensor_map.end()) { + return Error::NotFound; + } + return std::get<2>(result->second); +} + +ET_NODISCARD Result DataMap::get_data(const char* fqn) const { + auto result = _tensor_map.find(fqn); + if (result == _tensor_map.end()) { + return Error::NotFound; + } + int offset = std::get<1>(result->second); + TensorLayout tensor = std::get<2>(result->second); + + const uint8_t* data = static_cast(_data_ro.data()) + offset; + return FreeableBuffer(data, tensor.nbytes(), nullptr); +} + +ET_NODISCARD Result DataMap::get_num_keys() const { + return _tensor_map.size(); +} + +ET_NODISCARD Result DataMap::get_key(int index) const { + if (index <= 0 || index >= _tensor_map.size()) { + return Error::InvalidArgument; + } + + auto iter = _tensor_map.begin(); + for (int i = 0; i < index; ++i) { + ++iter; + } + return iter->first.c_str(); +} + +/* static */ Result DataMap::load(DataLoader* loader) { + // Load data map. + size_t flatbuffer_offset = 0; + size_t flatbuffer_size = 0; + size_t segment_base_offset = 0; + size_t segment_data_size = 0; + { + // Check header. + Result header = loader->load( + /*offset=*/0, + FlatTensorHeader::kNumHeadBytes, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External)); + if (!header.ok()) { + return header.error(); + } + Result fh = + FlatTensorHeader::Parse(header->data(), header->size()); + if (fh.ok()) { + // The header has the data map size. + flatbuffer_offset = fh->flatbuffer_offset; + flatbuffer_size = fh->flatbuffer_size; + segment_base_offset = fh->segment_base_offset; + segment_data_size = fh->segment_data_size; + } else if (fh.error() == Error::NotFound) { + // No header, throw error. + ET_LOG(Error, "No FlatTensorHeader found."); + return fh.error(); + } else { + // corruption, throw error. + ET_LOG(Error, "Flat tensor header may be corrupt."); + return fh.error(); + } + } + + // Load flatbuffer data as a segment. + Result flat_tensor_data = loader->load( + /*offset=*/flatbuffer_offset, + flatbuffer_size, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External)); + if (!flat_tensor_data.ok()) { + return flat_tensor_data.error(); + } + + // Make sure magic matches. + if (!flat_tensor::FlatTensorBufferHasIdentifier(flat_tensor_data->data())) { + ET_LOG( + Error, + "FlatTensor identifier '%.4s' != expected '%.4s'", + flatbuffers::GetBufferIdentifier(flat_tensor_data->data()), + flat_tensor::FlatTensorIdentifier()); + return Error::InvalidExternalData; + } + + // The flatbuffer data must start at an aligned address to ensure internal + // alignment of flatbuffer fields. + ET_CHECK_OR_RETURN_ERROR( + IsAligned(flat_tensor_data->data()), + InvalidArgument, + "FlatTensor data 0x%p must be aligned to %zu", + flat_tensor_data->data(), + kMinimumAlignment); + + // Get pointer to root of flatbuffer table. + const flat_tensor::FlatTensor* flat_tensor = + flat_tensor::GetFlatTensor(flat_tensor_data->data()); + + // Get pointer to tensor metadata. + const auto* s_tensor_metadata = flat_tensor->tensors(); + assert(s_tensor_metadata != nullptr); + + std::unordered_map> + fqn_to_tensor_layout = {}; + for (int i = 0; i < s_tensor_metadata->size(); i++) { + // Create TensorLayouts. + ScalarType scalar_type = + static_cast(s_tensor_metadata->Get(i)->scalar_type()); + const int dim = s_tensor_metadata->Get(i)->sizes()->size(); + + const auto serialized_sizes = s_tensor_metadata->Get(i)->sizes()->data(); + const auto serialized_dim_order = + s_tensor_metadata->Get(i)->dim_order()->data(); + TensorLayout tensor_layout = TensorLayout( + scalar_type, + Span(serialized_sizes, dim), + Span(serialized_dim_order, dim)); + + int segment_index = s_tensor_metadata->Get(i)->segment_index(); + int offset = s_tensor_metadata->Get(i)->offset(); + std::string fqn = s_tensor_metadata->Get(i)->fully_qualified_name()->str(); + + auto val = std::make_tuple(segment_index, offset, tensor_layout); + fqn_to_tensor_layout.insert({fqn, std::move(val)}); + } + + // Load constant data. + const auto* s_data_segment = flat_tensor->segments(); + + // Only support one segment for now. + assert(s_data_segment->size() == 1); + // First segment offset should be 0. + int segment_offset = s_data_segment->Get(0)->offset(); + assert(segment_offset == 0); + // First segment size should be <= the total segment data size. + int segment_size = s_data_segment->Get(0)->size(); + assert(segment_size <= segment_data_size); + + Result _data_ro = loader->load( + /*offset=*/segment_base_offset + segment_offset, + segment_size, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External)); + if (!_data_ro.ok()) { + return _data_ro.error(); + } + + return DataMap( + std::move(flat_tensor_data.get()), + std::move(fqn_to_tensor_layout), + std::move(_data_ro.get())); +} + +DataMap::~DataMap() {} + +} // namespace extension +} // namespace executorch diff --git a/extension/flat_tensor/named_data_map/data_map.h b/extension/flat_tensor/named_data_map/data_map.h new file mode 100644 index 00000000000..4a7e531f663 --- /dev/null +++ b/extension/flat_tensor/named_data_map/data_map.h @@ -0,0 +1,83 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +#include +#include +#include +#include + +#include + +#include +#include +#include + +// Forward declare flatbuffer types. This is a public header and must not +// include the generated flatbuffer header. +namespace flat_tensor { +struct FlatTensor; +} // namespace flat_tensor + +namespace executorch { +namespace extension { + +class DataMap final : public executorch::runtime::NamedDataMap { + public: + static executorch::runtime::Result load( + executorch::runtime::DataLoader* loader); + + ET_NODISCARD + executorch::runtime::Result + get_metadata(const char* fqn) const override; + ET_NODISCARD + executorch::runtime::Result get_data( + const char* fqn) const override; + + ET_NODISCARD executorch::runtime::Result get_num_keys() const override; + ET_NODISCARD executorch::runtime::Result get_key( + int index) const override; + + DataMap(DataMap&&) noexcept = default; + ~DataMap() override; + + private: + DataMap( + executorch::runtime::FreeableBuffer&& flat_tensor_data, + std::unordered_map< + std::string, + std::tuple> tensor_map, + executorch::runtime::FreeableBuffer&& data_ro) + : _flat_tensor_data(std::move(flat_tensor_data)), + _tensor_map(std::move(tensor_map)), + _data_ro(std::move(data_ro)) {} + + // Not copyable or assignable. + DataMap(const DataMap& rhs) = delete; + DataMap& operator=(DataMap&& rhs) noexcept = delete; + DataMap& operator=(const DataMap& rhs) = delete; + + // FlatTensor flatbuffer data. Contains the data backing up + // TensorLayout information in the _tensor_map; must outlive it. + executorch::runtime::FreeableBuffer _flat_tensor_data; + + // Map of fqn to {segment index, offset, TensorLayout}. + std::unordered_map< + std::string, + std::tuple> + _tensor_map; + + // Raw, read-only tensor data. + executorch::runtime::FreeableBuffer _data_ro; +}; + +} // namespace extension +} // namespace executorch diff --git a/extension/flat_tensor/named_data_map/targets.bzl b/extension/flat_tensor/named_data_map/targets.bzl new file mode 100644 index 00000000000..25c0eec2f88 --- /dev/null +++ b/extension/flat_tensor/named_data_map/targets.bzl @@ -0,0 +1,23 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + runtime.cxx_library( + name = "data_map", + srcs = [ + "data_map.cpp", + ], + exported_headers = ["data_map.h"], + deps = [ + "//executorch/extension/flat_tensor/serialize:schema", + "//executorch/extension/flat_tensor/serialize:serialize", + "//executorch/extension/flat_tensor/serialize:generated_headers", + "//executorch/extension/flat_tensor/serialize:flat_tensor_header", + "//executorch/runtime/core:core", + "//executorch/runtime/core:evalue", + "//executorch/runtime/core/exec_aten:lib", + "//executorch/runtime/core/exec_aten/util:tensor_util", + ], + visibility = [ + "//executorch/...", + ], + ) diff --git a/extension/flat_tensor/test/TARGETS b/extension/flat_tensor/test/TARGETS index c9989b67554..0dc97974af8 100644 --- a/extension/flat_tensor/test/TARGETS +++ b/extension/flat_tensor/test/TARGETS @@ -6,7 +6,7 @@ load(":targets.bzl", "define_common_targets") oncall("executorch") -define_common_targets() +define_common_targets(is_fbcode=True) python_unittest( name = "serialize", diff --git a/extension/flat_tensor/test/data_map_test.cpp b/extension/flat_tensor/test/data_map_test.cpp new file mode 100644 index 00000000000..028e13eb454 --- /dev/null +++ b/extension/flat_tensor/test/data_map_test.cpp @@ -0,0 +1,102 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include + +using namespace ::testing; +using executorch::extension::DataMap; +using executorch::extension::FlatTensorHeader; +using executorch::runtime::DataLoader; +using executorch::runtime::Error; +using executorch::runtime::FreeableBuffer; +using executorch::runtime::Result; +using executorch::runtime::TensorLayout; +using torch::executor::util::FileDataLoader; + +class DataMapTest : public ::testing::Test { + protected: + void SetUp() override { + // Since these tests cause ET_LOG to be called, the PAL must be initialized + // first. + executorch::runtime::runtime_init(); + } +}; + +TEST_F(DataMapTest, LoadDataMap) { + const char* path = std::getenv("ET_MODULE_LINEAR_DATA"); + Result loader = FileDataLoader::from(path); + ASSERT_EQ(loader.error(), Error::Ok); + + Result header = loader->load( + /*offset=*/0, + FlatTensorHeader::kNumHeadBytes, + /*segment_info=*/ + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External)); + + ASSERT_EQ(header.error(), Error::Ok); + + auto data_map_loader_ = + std::make_unique(std::move(loader.get())); + + Result data_map = DataMap::load(data_map_loader_.get()); + EXPECT_EQ(data_map.error(), Error::Ok); + + // Check tensor metadata. + Result const_a_res = data_map->get_metadata("a"); + assert(const_a_res.ok()); + + const TensorLayout const_a = const_a_res.get(); + EXPECT_EQ(const_a.scalar_type(), executorch::aten::ScalarType::Float); + auto sizes_a = const_a.sizes(); + EXPECT_EQ(sizes_a.size(), 2); + EXPECT_EQ(sizes_a[0], 2); + EXPECT_EQ(sizes_a[1], 2); + auto dim_order_a = const_a.dim_order(); + EXPECT_EQ(dim_order_a.size(), 2); + EXPECT_EQ(dim_order_a[0], 0); + EXPECT_EQ(dim_order_a[1], 1); + + Result const_b_res = data_map->get_metadata("b"); + assert(const_b_res.ok()); + + const TensorLayout const_b = const_b_res.get(); + EXPECT_EQ(const_b.scalar_type(), executorch::aten::ScalarType::Float); + auto sizes_b = const_b.sizes(); + EXPECT_EQ(sizes_b.size(), 2); + EXPECT_EQ(sizes_b[0], 2); + EXPECT_EQ(sizes_b[1], 2); + auto dim_order_b = const_b.dim_order(); + EXPECT_EQ(dim_order_b.size(), 2); + EXPECT_EQ(dim_order_b[0], 0); + EXPECT_EQ(dim_order_b[1], 1); + + // Check tensor data. + Result data_a_res = data_map->get_data("a"); + assert(data_a_res.ok()); + // Check we have the correct tensor data. + FreeableBuffer data_a = std::move(data_a_res.get()); + EXPECT_EQ(data_a.size(), 16); + + Result data_b_res = data_map->get_data("b"); + assert(data_b_res.ok()); + // Check we have the correct tensor data. + FreeableBuffer data_b = std::move(data_b_res.get()); + EXPECT_EQ(data_b.size(), 16); +} diff --git a/extension/flat_tensor/test/targets.bzl b/extension/flat_tensor/test/targets.bzl index 0c08106a5c5..21226207578 100644 --- a/extension/flat_tensor/test/targets.bzl +++ b/extension/flat_tensor/test/targets.bzl @@ -1,6 +1,6 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") -def define_common_targets(): +def define_common_targets(is_fbcode=False): """Defines targets that should be shared between fbcode and xplat. The directory containing this targets.bzl file should also contain both @@ -16,3 +16,29 @@ def define_common_targets(): "//executorch/extension/flat_tensor/serialize:flat_tensor_header", ], ) + + if not runtime.is_oss and is_fbcode: + modules_env = { + # The tests use this var to find the program file to load. This uses + # an fbcode target path because the authoring/export tools + # intentionally don't work in xplat (since they're host-only tools). + "ET_MODULE_LINEAR_PROGRAM": "$(location fbcode//executorch/test/models:exported_programs_with_data_separated[ModuleLinear.pte])", + "ET_MODULE_LINEAR_DATA": "$(location fbcode//executorch/test/models:exported_programs_with_data_separated[ModuleLinear.ptd])", + } + + runtime.cxx_test( + name = "data_map", + srcs = [ + "data_map_test.cpp", + ], + deps = [ + "//executorch/extension/flat_tensor/serialize:schema", + "//executorch/extension/flat_tensor/serialize:generated_headers", + "//executorch/extension/flat_tensor/named_data_map:data_map", + "//executorch/extension/flat_tensor/serialize:flat_tensor_header", + "//executorch/extension/data_loader:buffer_data_loader", + "//executorch/extension/data_loader:file_data_loader", + "//executorch/runtime/core/exec_aten:lib", + ], + env = modules_env, + ) diff --git a/runtime/core/data_loader.h b/runtime/core/data_loader.h index 876212de8fb..45fd1bc8189 100644 --- a/runtime/core/data_loader.h +++ b/runtime/core/data_loader.h @@ -48,6 +48,10 @@ class DataLoader { * Data used for initializing mutable tensors. */ Mutable, + /** + * Data used for initializing external tensors. + */ + External, }; /// Type of the segment. From b8f4a784303b396e96c5552a529baad8959e69c6 Mon Sep 17 00:00:00 2001 From: lucylq Date: Thu, 23 Jan 2025 12:06:09 -0800 Subject: [PATCH 2/2] Update on "[executorch][flat_tensor] DataMap implementation" 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. Differential Revision: [D67064580](https://our.internmc.facebook.com/intern/diff/D67064580/) [ghstack-poisoned] --- .../flat_tensor/named_data_map/data_map.cpp | 19 +++++++++---------- .../flat_tensor/named_data_map/data_map.h | 18 ++++++++---------- extension/flat_tensor/test/data_map_test.cpp | 8 +++----- extension/flat_tensor/test/targets.bzl | 8 ++++---- 4 files changed, 24 insertions(+), 29 deletions(-) diff --git a/extension/flat_tensor/named_data_map/data_map.cpp b/extension/flat_tensor/named_data_map/data_map.cpp index 605876ead18..4f18b0001bd 100644 --- a/extension/flat_tensor/named_data_map/data_map.cpp +++ b/extension/flat_tensor/named_data_map/data_map.cpp @@ -5,18 +5,17 @@ * This source code is licensed under the BSD-style license found in the * LICENSE file in the root directory of this source tree. */ + #include #include #include - #include +#include #include #include #include #include -#include - #include #include @@ -48,16 +47,16 @@ bool IsAligned(const void* data) { ET_NODISCARD Result DataMap::get_metadata( const char* fqn) const { - auto result = _tensor_map.find(fqn); - if (result == _tensor_map.end()) { + auto result = _name_to_tensor.find(fqn); + if (result == _name_to_tensor.end()) { return Error::NotFound; } return std::get<2>(result->second); } ET_NODISCARD Result DataMap::get_data(const char* fqn) const { - auto result = _tensor_map.find(fqn); - if (result == _tensor_map.end()) { + auto result = _name_to_tensor.find(fqn); + if (result == _name_to_tensor.end()) { return Error::NotFound; } int offset = std::get<1>(result->second); @@ -68,15 +67,15 @@ ET_NODISCARD Result DataMap::get_data(const char* fqn) const { } ET_NODISCARD Result DataMap::get_num_keys() const { - return _tensor_map.size(); + return _name_to_tensor.size(); } ET_NODISCARD Result DataMap::get_key(int index) const { - if (index <= 0 || index >= _tensor_map.size()) { + if (index <= 0 || index >= _name_to_tensor.size()) { return Error::InvalidArgument; } - auto iter = _tensor_map.begin(); + auto iter = _name_to_tensor.begin(); for (int i = 0; i < index; ++i) { ++iter; } diff --git a/extension/flat_tensor/named_data_map/data_map.h b/extension/flat_tensor/named_data_map/data_map.h index 4a7e531f663..14ddc29c1b5 100644 --- a/extension/flat_tensor/named_data_map/data_map.h +++ b/extension/flat_tensor/named_data_map/data_map.h @@ -8,16 +8,13 @@ #pragma once -#include - #include +#include #include #include +#include #include -#include - -#include #include #include @@ -54,10 +51,11 @@ class DataMap final : public executorch::runtime::NamedDataMap { executorch::runtime::FreeableBuffer&& flat_tensor_data, std::unordered_map< std::string, - std::tuple> tensor_map, + std::tuple> + name_to_tensor, executorch::runtime::FreeableBuffer&& data_ro) : _flat_tensor_data(std::move(flat_tensor_data)), - _tensor_map(std::move(tensor_map)), + _name_to_tensor(std::move(name_to_tensor)), _data_ro(std::move(data_ro)) {} // Not copyable or assignable. @@ -66,14 +64,14 @@ class DataMap final : public executorch::runtime::NamedDataMap { DataMap& operator=(const DataMap& rhs) = delete; // FlatTensor flatbuffer data. Contains the data backing up - // TensorLayout information in the _tensor_map; must outlive it. + // TensorLayout information in the _name_to_tensor map; must outlive it. executorch::runtime::FreeableBuffer _flat_tensor_data; - // Map of fqn to {segment index, offset, TensorLayout}. + // Map of name to {segment index, offset, TensorLayout}. std::unordered_map< std::string, std::tuple> - _tensor_map; + _name_to_tensor; // Raw, read-only tensor data. executorch::runtime::FreeableBuffer _data_ro; diff --git a/extension/flat_tensor/test/data_map_test.cpp b/extension/flat_tensor/test/data_map_test.cpp index 028e13eb454..447809b9adb 100644 --- a/extension/flat_tensor/test/data_map_test.cpp +++ b/extension/flat_tensor/test/data_map_test.cpp @@ -6,18 +6,16 @@ * LICENSE file in the root directory of this source tree. */ -#include - #include #include +#include +#include +#include #include #include #include #include -#include -#include - #include using namespace ::testing; diff --git a/extension/flat_tensor/test/targets.bzl b/extension/flat_tensor/test/targets.bzl index 21226207578..0fd7c4cd923 100644 --- a/extension/flat_tensor/test/targets.bzl +++ b/extension/flat_tensor/test/targets.bzl @@ -32,12 +32,12 @@ def define_common_targets(is_fbcode=False): "data_map_test.cpp", ], deps = [ - "//executorch/extension/flat_tensor/serialize:schema", - "//executorch/extension/flat_tensor/serialize:generated_headers", - "//executorch/extension/flat_tensor/named_data_map:data_map", - "//executorch/extension/flat_tensor/serialize:flat_tensor_header", "//executorch/extension/data_loader:buffer_data_loader", "//executorch/extension/data_loader:file_data_loader", + "//executorch/extension/flat_tensor/named_data_map:data_map", + "//executorch/extension/flat_tensor/serialize:flat_tensor_header", + "//executorch/extension/flat_tensor/serialize:generated_headers", + "//executorch/extension/flat_tensor/serialize:schema", "//executorch/runtime/core/exec_aten:lib", ], env = modules_env,