#pragma once
#include <c10/util/safe_numerics.h>
#include <executorch/runtime/core/data_loader.h>
#include <executorch/runtime/core/error.h>
#include <executorch/runtime/core/result.h>
#include <executorch/runtime/platform/log.h>
#include <cstring>
namespace executorch {
namespace extension {
class BufferDataLoader final : public executorch::runtime::DataLoader {
public:
BufferDataLoader(const void* data, size_t size)
: data_(reinterpret_cast<const uint8_t*>(data)), size_(size) {}
ET_NODISCARD
executorch::runtime::Result<executorch::runtime::FreeableBuffer> load(
size_t offset,
size_t size,
ET_UNUSED const DataLoader::SegmentInfo& segment_info) const override {
size_t total_size;
bool overflow = c10::add_overflows(offset, size, &total_size);
ET_CHECK_OR_RETURN_ERROR(
!overflow && total_size <= size_,
InvalidArgument,
"offset %zu + size %zu > size_ %zu, or overflow detected",
offset,
size,
size_);
return executorch::runtime::FreeableBuffer(
data_ + offset, size, nullptr);
}
ET_NODISCARD executorch::runtime::Result<size_t> size() const override {
return size_;
}
ET_NODISCARD executorch::runtime::Error load_into(
size_t offset,
size_t size,
ET_UNUSED const SegmentInfo& segment_info,
void* buffer) const override {
ET_CHECK_OR_RETURN_ERROR(
buffer != nullptr,
InvalidArgument,
"Destination buffer cannot be null");
auto result = load(offset, size, segment_info);
if (!result.ok()) {
return result.error();
}
std::memcpy(buffer, result->data(), size);
return executorch::runtime::Error::Ok;
}
private:
const uint8_t* const data_; const size_t size_;
};
} }
namespace torch {
namespace executor {
namespace util {
using ::executorch::extension::BufferDataLoader;
} } }