onnxruntime/winml/lib/Api.Image/DisjointBufferHelpers.cpp

#include "lib/Api.Image/pch.h"
#include "inc/DisjointBufferHelpers.h"

namespace _winml {

static void LoadOrStoreDisjointBuffers(
    bool should_load_buffer,
    size_t num_buffers,
    std::function<gsl::span<byte>(size_t)> get_buffer,
    gsl::span<byte>& buffer_span) {
  auto size_in_bytes = buffer_span.size_bytes();
  auto buffer = buffer_span.data();
  size_t offset_in_bytes = 0;
  for (size_t i = 0; i < num_buffers && size_in_bytes > offset_in_bytes; i++) {
    auto span = get_buffer(i);
    auto current_size_in_bytes = span.size_bytes();
    auto current_buffer = span.data();

    if (size_in_bytes - offset_in_bytes < current_size_in_bytes) {
      current_size_in_bytes = size_in_bytes - offset_in_bytes;
    }

    auto offset_buffer = buffer + offset_in_bytes;
    if (should_load_buffer) {
      memcpy(offset_buffer, current_buffer, current_size_in_bytes);
    } else {
      memcpy(current_buffer, offset_buffer, current_size_in_bytes);
    }

    offset_in_bytes += current_size_in_bytes;
  }
}

void LoadSpanFromDisjointBuffers(
    size_t num_buffers,
    std::function<gsl::span<byte>(size_t)> get_buffer,
    gsl::span<byte>& buffer_span) {
  LoadOrStoreDisjointBuffers(true /*load into the span*/, num_buffers, get_buffer, buffer_span);
}

void StoreSpanIntoDisjointBuffers(
    size_t num_buffers,
    std::function<gsl::span<byte>(size_t)> get_buffer,
    gsl::span<byte>& buffer_span) {
  LoadOrStoreDisjointBuffers(false /*store into buffers*/, num_buffers, get_buffer, buffer_span);
}

} // namespace _winml