|
| 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/named_data_map/merged_data_map.h> |
| 10 | +#include <executorch/runtime/core/data_loader.h> |
| 11 | + |
| 12 | +#include <unordered_map> |
| 13 | +#include <vector> |
| 14 | + |
| 15 | +using executorch::aten::string_view; |
| 16 | +using executorch::ET_RUNTIME_NAMESPACE::NamedDataMap; |
| 17 | +using executorch::ET_RUNTIME_NAMESPACE::TensorLayout; |
| 18 | +using executorch::runtime::Error; |
| 19 | +using executorch::runtime::FreeableBuffer; |
| 20 | +using executorch::runtime::Result; |
| 21 | +using executorch::runtime::Span; |
| 22 | + |
| 23 | +namespace executorch::extension { |
| 24 | + |
| 25 | +/*static*/ Result<MergedDataMap> MergedDataMap::load( |
| 26 | + Span<const NamedDataMap*> named_data_maps) { |
| 27 | + std::vector<const NamedDataMap*> valid_data_maps; |
| 28 | + for (auto i : c10::irange(named_data_maps.size())) { |
| 29 | + if (named_data_maps[i] != nullptr && |
| 30 | + named_data_maps[i]->get_num_keys().get() > 0) { |
| 31 | + valid_data_maps.push_back(named_data_maps[i]); |
| 32 | + } |
| 33 | + } |
| 34 | + ET_CHECK_OR_RETURN_ERROR( |
| 35 | + !valid_data_maps.empty(), |
| 36 | + InvalidArgument, |
| 37 | + "No non-empty named data maps provided to merge"); |
| 38 | + |
| 39 | + // Check for duplicate keys. |
| 40 | + std::unordered_map<std::string, uint32_t> key_to_map_index; |
| 41 | + for (auto i : c10::irange(valid_data_maps.size())) { |
| 42 | + const auto cur_map = valid_data_maps[i]; |
| 43 | + uint32_t num_keys = cur_map->get_num_keys().get(); |
| 44 | + for (auto j : c10::irange(num_keys)) { |
| 45 | + const auto cur_key = cur_map->get_key(j).get(); |
| 46 | + const auto [it, inserted] = key_to_map_index.emplace(cur_key, i); |
| 47 | + ET_CHECK_OR_RETURN_ERROR( |
| 48 | + inserted, |
| 49 | + InvalidArgument, |
| 50 | + "Duplicate key %s in named data maps at index %u and %lu", |
| 51 | + cur_key, |
| 52 | + it->second, |
| 53 | + i); |
| 54 | + } |
| 55 | + } |
| 56 | + return MergedDataMap(std::move(valid_data_maps), std::move(key_to_map_index)); |
| 57 | +} |
| 58 | + |
| 59 | +ET_NODISCARD Result<const TensorLayout> MergedDataMap::get_tensor_layout( |
| 60 | + string_view key) const { |
| 61 | + const auto it = key_to_map_index_.find(key.data()); |
| 62 | + ET_CHECK_OR_RETURN_ERROR( |
| 63 | + it != key_to_map_index_.end(), |
| 64 | + NotFound, |
| 65 | + "Key %s not found in named data maps", |
| 66 | + key.data()); |
| 67 | + |
| 68 | + return named_data_maps_.at(it->second)->get_tensor_layout(key); |
| 69 | +} |
| 70 | + |
| 71 | +ET_NODISCARD |
| 72 | +Result<FreeableBuffer> MergedDataMap::get_data(string_view key) const { |
| 73 | + const auto it = key_to_map_index_.find(key.data()); |
| 74 | + ET_CHECK_OR_RETURN_ERROR( |
| 75 | + it != key_to_map_index_.end(), |
| 76 | + NotFound, |
| 77 | + "Key %s not found in named data maps", |
| 78 | + key.data()); |
| 79 | + return named_data_maps_.at(it->second)->get_data(key); |
| 80 | +} |
| 81 | + |
| 82 | +ET_NODISCARD Error MergedDataMap::load_data_into( |
| 83 | + string_view key, |
| 84 | + void* buffer, |
| 85 | + size_t size) const { |
| 86 | + const auto it = key_to_map_index_.find(key.data()); |
| 87 | + ET_CHECK_OR_RETURN_ERROR( |
| 88 | + it != key_to_map_index_.end(), |
| 89 | + NotFound, |
| 90 | + "Key %s not found in named data maps", |
| 91 | + key.data()); |
| 92 | + return named_data_maps_.at(it->second)->load_data_into(key, buffer, size); |
| 93 | +} |
| 94 | + |
| 95 | +ET_NODISCARD Result<uint32_t> MergedDataMap::get_num_keys() const { |
| 96 | + return key_to_map_index_.size(); |
| 97 | +} |
| 98 | + |
| 99 | +ET_NODISCARD Result<const char*> MergedDataMap::get_key(uint32_t index) const { |
| 100 | + uint32_t total_num_keys = get_num_keys().get(); |
| 101 | + ET_CHECK_OR_RETURN_ERROR( |
| 102 | + index < total_num_keys, |
| 103 | + InvalidArgument, |
| 104 | + "Index %u out of range of size %u", |
| 105 | + index, |
| 106 | + total_num_keys); |
| 107 | + for (auto i : c10::irange(named_data_maps_.size())) { |
| 108 | + auto num_keys = named_data_maps_[i]->get_num_keys().get(); |
| 109 | + if (index < num_keys) { |
| 110 | + return named_data_maps_[i]->get_key(index); |
| 111 | + } |
| 112 | + index -= num_keys; |
| 113 | + } |
| 114 | + // Shouldn't reach here. |
| 115 | + return Error::Internal; |
| 116 | +} |
| 117 | +} // namespace executorch::extension |
0 commit comments