data_loader.h

/*
 * 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 <cstddef>
#include <cstdint>
#include <limits>

#include <executorch/runtime/core/freeable_buffer.h>
#include <executorch/runtime/core/result.h>
#include <executorch/runtime/platform/compiler.h>

namespace executorch {
namespace runtime {

/**
 * Loads from a data source.
 *
 * See //executorch/extension/data_loader for common implementations.
 */
class DataLoader {
 public:
  /**
   * Describes the content of the segment.
   */
  struct SegmentInfo {
    /**
     * Represents the purpose of the segment.
     */
    enum class Type {
      /**
       * Data for the actual program.
       */
      Program,
      /**
       * Holds constant tensor data.
       */
      Constant,
      /**
       * Data used for initializing a backend.
       */
      Backend,
      /**
       * Data used for initializing mutable tensors.
       */
      Mutable,
      /**
       * Data used for initializing external tensors.
       */
      External,
    };

    /// Type of the segment.
    Type segment_type;

    /// Index of the segment within the segment list. Undefined for program
    /// segments.
    size_t segment_index;

    /// An optional, null-terminated string describing the segment. For
    /// `Backend` segments, this is the backend ID. Null for other segment
    /// types.
    const char* descriptor;

    SegmentInfo() = default;

    explicit SegmentInfo(
        Type segment_type_,
        size_t segment_index_ = 0,
        const char* descriptor_ = nullptr)
        : segment_type(segment_type_),
          segment_index(segment_index_),
          descriptor(descriptor_) {}
  };

  virtual ~DataLoader() = default;

  /**
   * Loads data from the underlying data source.
   *
   * NOTE: This must be thread-safe. If this call modifies common state, the
   * implementation must do its own locking.
   *
   * @param offset The byte offset in the data source to start loading from.
   * @param size The number of bytes to load.
   * @param segment_info Information about the segment being loaded.
   *
   * @returns a `FreeableBuffer` that owns the loaded data.
   */
  ET_NODISCARD virtual Result<FreeableBuffer>
  load(size_t offset, size_t size, const SegmentInfo& segment_info) const = 0;

  /**
   * Loads data from the underlying data source into the provided buffer.
   *
   * NOTE: This must be thread-safe. If this call modifies common state, the
   * implementation must do its own locking.
   *
   * @param offset The byte offset in the data source to start loading from.
   * @param size The number of bytes to load.
   * @param segment_info Information about the segment being loaded.
   * @param buffer The buffer to load data into. Must point to at least `size`
   * bytes of memory.
   *
   * @returns an Error indicating if the load was successful.
   */
  ET_NODISCARD virtual Error load_into(
      size_t offset,
      size_t size,
      const SegmentInfo& segment_info,
      void* buffer) const {
    // Using a stub implementation here instead of pure virtual to expand the
    // data_loader interface in a backwards compatible way.
    (void)buffer;
    (void)offset;
    (void)size;
    (void)segment_info;
    ET_LOG(Error, "load_into() not implemented for this data loader.");
    return Error::NotImplemented;
  }

  /**
   * Returns the length of the underlying data source, typically the file size.
   */
  ET_NODISCARD virtual Result<size_t> size() const = 0;

  /** Loads data from a 64-bit source offset. */
  ET_NODISCARD virtual Result<FreeableBuffer> load_at_offset(
      uint64_t offset,
      size_t size,
      const SegmentInfo& segment_info) const {
#if SIZE_MAX < UINT64_MAX
    if (offset > std::numeric_limits<size_t>::max()) {
      ET_LOG(
          Error,
          "load_at_offset() source range cannot be represented by this data loader.");
      return Error::NotSupported;
    }
#endif
    if (static_cast<uint64_t>(size) >
        static_cast<uint64_t>(std::numeric_limits<size_t>::max()) - offset) {
      ET_LOG(
          Error,
          "load_at_offset() source range cannot be represented by this data loader.");
      return Error::NotSupported;
    }
    return load(static_cast<size_t>(offset), size, segment_info);
  }

  /** Loads data from a 64-bit source offset into the provided buffer. */
  ET_NODISCARD virtual Error load_into_at_offset(
      uint64_t offset,
      size_t size,
      const SegmentInfo& segment_info,
      void* buffer) const {
#if SIZE_MAX < UINT64_MAX
    if (offset > std::numeric_limits<size_t>::max()) {
      ET_LOG(
          Error,
          "load_into_at_offset() source range cannot be represented by this data loader.");
      return Error::NotSupported;
    }
#endif
    if (static_cast<uint64_t>(size) >
        static_cast<uint64_t>(std::numeric_limits<size_t>::max()) - offset) {
      ET_LOG(
          Error,
          "load_into_at_offset() source range cannot be represented by this data loader.");
      return Error::NotSupported;
    }
    return load_into(static_cast<size_t>(offset), size, segment_info, buffer);
  }

  /**
   * Returns the 64-bit length of the underlying data source.
   *
   * The default implementation widens the result from `size()`. Implementations
   * that wrap sources larger than `size_t` can represent must override this
   * method.
   */
  ET_NODISCARD virtual Result<uint64_t> source_size() const {
    Result<size_t> result = size();
    if (!result.ok()) {
      return result.error();
    }
    return static_cast<uint64_t>(result.get());
  }
};

} // namespace runtime
} // namespace executorch

namespace torch {
namespace executor {
// TODO(T197294990): Remove these deprecated aliases once all users have moved
// to the new `::executorch` namespaces.
using ::executorch::runtime::DataLoader;
} // namespace executor
} // namespace torch