diff --git a/runtime/core/named_data_map.h b/runtime/core/named_data_map.h new file mode 100644 index 00000000000..31a60abb89f --- /dev/null +++ b/runtime/core/named_data_map.h @@ -0,0 +1,77 @@ +/* + * 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 + +namespace executorch { +namespace runtime { + +/** + * Interface to access and retrieve data via name. + * See executorch/extension/flat_tensor/ for an example. + */ +class ET_EXPERIMENTAL NamedDataMap { + public: + virtual ~NamedDataMap() = default; + /** + * Get metadata by key. + * + * @param key The name of the tensor. + * @return Result containing TensorLayout with tensor metadata. + */ + ET_NODISCARD virtual Result + get_metadata(const char* key) const = 0; + /** + * Get data by key. + * + * @param key Name of the data. + * @return Result containing a FreeableBuffer with the tensor data. + */ + ET_NODISCARD virtual Result get_data( + const char* key) const = 0; + + /** + * Loads data corresponding to the key into the provided buffer. + * + * @param key The name of the data. + * @param size The number of bytes to load. Use `get_metadata` to retrieve the + * size of the data for a given key. + * @param buffer The buffer to load the data into. Must point to at least + * `size` bytes of memory. + * @return Result containing the number of bytes written on success. This will + * fail if the buffer is too small. + */ + ET_NODISCARD virtual Result + load_data_into(const char* key, void* buffer, size_t size) const = 0; + + /** + * Get the number of keys in the NamedDataMap. + * + * @return Result containing the number of keys. + */ + ET_NODISCARD virtual Result get_num_keys() const = 0; + + /** + * Get the key at the given index. + * + * @param index The index of the key to retrieve. + * @return Result containing the key at the given index. Note: the returned + * pointer is only valid for the lifetime of the DataMap. + */ + ET_NODISCARD virtual Result get_key(size_t index) const = 0; +}; + +} // namespace runtime +} // namespace executorch diff --git a/runtime/core/targets.bzl b/runtime/core/targets.bzl index d64d4579536..253e2526107 100644 --- a/runtime/core/targets.bzl +++ b/runtime/core/targets.bzl @@ -41,6 +41,7 @@ def define_common_targets(): "defines.h", "error.h", "freeable_buffer.h", + "named_data_map.h", "result.h", "span.h", "tensor_layout.h",