// Copyright 2020 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include <stddef.h>

#include "iree/base/api.h"
#include "iree/hal/allocator.h"
#include "iree/hal/buffer.h"
#include "iree/hal/buffer_heap_impl.h"
#include "iree/hal/resource.h"

typedef struct iree_hal_heap_allocator_t {
  iree_hal_resource_t resource;
  iree_allocator_t host_allocator;
  iree_allocator_t data_allocator;
  iree_string_view_t identifier;
  IREE_STATISTICS(iree_hal_heap_allocator_statistics_t statistics;)
} iree_hal_heap_allocator_t;

static const iree_hal_allocator_vtable_t iree_hal_heap_allocator_vtable;

iree_hal_heap_allocator_t* iree_hal_heap_allocator_cast(
    iree_hal_allocator_t* IREE_RESTRICT base_value) {
  return (iree_hal_heap_allocator_t*)base_value;
}

IREE_API_EXPORT iree_status_t iree_hal_allocator_create_heap(
    iree_string_view_t identifier, iree_allocator_t data_allocator,
    iree_allocator_t host_allocator, iree_hal_allocator_t** out_allocator) {
  IREE_ASSERT_ARGUMENT(out_allocator);
  *out_allocator = NULL;
  IREE_TRACE_ZONE_BEGIN(z0);

  iree_hal_heap_allocator_t* allocator = NULL;
  iree_host_size_t total_size =
      iree_sizeof_struct(*allocator) + identifier.size;
  iree_status_t status =
      iree_allocator_malloc(host_allocator, total_size, (void**)&allocator);
  if (iree_status_is_ok(status)) {
    iree_hal_resource_initialize(&iree_hal_heap_allocator_vtable,
                                 &allocator->resource);
    allocator->host_allocator = host_allocator;
    allocator->data_allocator = data_allocator;
    iree_string_view_append_to_buffer(
        identifier, &allocator->identifier,
        (char*)allocator + iree_sizeof_struct(*allocator));

    IREE_STATISTICS({
      // All start initialized to zero.
      iree_slim_mutex_initialize(&allocator->statistics.mutex);
    });

    *out_allocator = (iree_hal_allocator_t*)allocator;
  }

  IREE_TRACE_ZONE_END(z0);
  return status;
}

static void iree_hal_heap_allocator_destroy(
    iree_hal_allocator_t* IREE_RESTRICT base_allocator) {
  iree_hal_heap_allocator_t* allocator =
      iree_hal_heap_allocator_cast(base_allocator);
  iree_allocator_t host_allocator = allocator->host_allocator;
  IREE_TRACE_ZONE_BEGIN(z0);

  IREE_STATISTICS(iree_slim_mutex_deinitialize(&allocator->statistics.mutex));

  iree_allocator_free(host_allocator, allocator);

  IREE_TRACE_ZONE_END(z0);
}

static iree_allocator_t iree_hal_heap_allocator_host_allocator(
    const iree_hal_allocator_t* IREE_RESTRICT base_allocator) {
  iree_hal_heap_allocator_t* allocator =
      (iree_hal_heap_allocator_t*)base_allocator;
  return allocator->host_allocator;
}

static iree_status_t iree_hal_heap_allocator_trim(
    iree_hal_allocator_t* IREE_RESTRICT base_allocator) {
  return iree_ok_status();
}

static void iree_hal_heap_allocator_query_statistics(
    iree_hal_allocator_t* IREE_RESTRICT base_allocator,
    iree_hal_allocator_statistics_t* IREE_RESTRICT out_statistics) {
  IREE_STATISTICS({
    iree_hal_heap_allocator_t* allocator =
        iree_hal_heap_allocator_cast(base_allocator);
    iree_slim_mutex_lock(&allocator->statistics.mutex);
    memcpy(out_statistics, &allocator->statistics.base,
           sizeof(*out_statistics));
    iree_slim_mutex_unlock(&allocator->statistics.mutex);
  });
}

static iree_status_t iree_hal_heap_allocator_query_memory_heaps(
    iree_hal_allocator_t* IREE_RESTRICT base_allocator,
    iree_host_size_t capacity,
    iree_hal_allocator_memory_heap_t* IREE_RESTRICT heaps,
    iree_host_size_t* IREE_RESTRICT out_count) {
  const iree_host_size_t count = 1;
  if (out_count) *out_count = count;
  if (capacity < count) {
    // NOTE: lightweight as this is hit in normal pre-sizing usage.
    return iree_status_from_code(IREE_STATUS_OUT_OF_RANGE);
  }
  heaps[0] = (iree_hal_allocator_memory_heap_t){
      .type = IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL |
              IREE_HAL_MEMORY_TYPE_HOST_VISIBLE |
              IREE_HAL_MEMORY_TYPE_HOST_COHERENT,
      .allowed_usage = IREE_HAL_BUFFER_USAGE_TRANSFER |
                       IREE_HAL_BUFFER_USAGE_DISPATCH |
                       IREE_HAL_BUFFER_USAGE_SHARING_EXPORT |
                       IREE_HAL_BUFFER_USAGE_SHARING_REPLICATE |
                       IREE_HAL_BUFFER_USAGE_SHARING_CONCURRENT |
                       IREE_HAL_BUFFER_USAGE_SHARING_IMMUTABLE |
                       IREE_HAL_BUFFER_USAGE_MAPPING_SCOPED |
                       IREE_HAL_BUFFER_USAGE_MAPPING_PERSISTENT |
                       IREE_HAL_BUFFER_USAGE_MAPPING_OPTIONAL |
                       IREE_HAL_BUFFER_USAGE_MAPPING_ACCESS_RANDOM |
                       IREE_HAL_BUFFER_USAGE_MAPPING_ACCESS_SEQUENTIAL_WRITE,
      .max_allocation_size = ~(iree_device_size_t)0,
      .min_alignment = IREE_HAL_HEAP_BUFFER_ALIGNMENT,
  };
  return iree_ok_status();
}

static iree_hal_buffer_compatibility_t
iree_hal_heap_allocator_query_buffer_compatibility(
    iree_hal_allocator_t* IREE_RESTRICT base_allocator,
    iree_hal_buffer_params_t* IREE_RESTRICT params,
    iree_device_size_t* IREE_RESTRICT allocation_size) {
  // All buffers can be allocated on the heap and all heap-accessible buffers
  // can be imported/exported.
  iree_hal_buffer_compatibility_t compatibility =
      IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE |
      IREE_HAL_BUFFER_COMPATIBILITY_IMPORTABLE |
      IREE_HAL_BUFFER_COMPATIBILITY_EXPORTABLE;

  // Buffers can only be used on the queue if they are device visible.
  // This is not a strict requirement of heap buffers but matches devices that
  // have discrete memory spaces (remoting/sandboxed, GPUs, etc) and makes it
  // much easier to find issues of buffer definition with local devices that
  // will cause issues when used with real devices.
  if (iree_all_bits_set(params->type, IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE)) {
    if (iree_any_bit_set(params->usage, IREE_HAL_BUFFER_USAGE_TRANSFER)) {
      compatibility |= IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_TRANSFER;
    }
    if (iree_any_bit_set(params->usage,
                         IREE_HAL_BUFFER_USAGE_DISPATCH_STORAGE)) {
      compatibility |= IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_DISPATCH;
    }
  }

  // Always ensure we are host-visible.
  params->type |= IREE_HAL_MEMORY_TYPE_HOST_VISIBLE;

  // We are now optimal.
  params->type &= ~IREE_HAL_MEMORY_TYPE_OPTIMAL;

  // Host currently uses mapping to copy buffers, which is done a lot.
  // We could probably remove this mutation by preventing copies in those cases.
  // TODO(benvanik): check if transfer is still required for DMA copy source.
  params->usage |= IREE_HAL_BUFFER_USAGE_MAPPING_SCOPED |
                   IREE_HAL_BUFFER_USAGE_MAPPING_PERSISTENT |
                   IREE_HAL_BUFFER_USAGE_MAPPING_ACCESS_RANDOM |
                   IREE_HAL_BUFFER_USAGE_TRANSFER;

  return compatibility;
}

static iree_status_t iree_hal_heap_allocator_allocate_buffer(
    iree_hal_allocator_t* IREE_RESTRICT base_allocator,
    const iree_hal_buffer_params_t* IREE_RESTRICT params,
    iree_device_size_t allocation_size,
    iree_hal_buffer_t** IREE_RESTRICT out_buffer) {
  iree_hal_heap_allocator_t* allocator =
      iree_hal_heap_allocator_cast(base_allocator);

  // Coerce options into those required for use by heap-based devices.
  iree_hal_buffer_params_t compat_params = *params;
  if (!iree_all_bits_set(iree_hal_heap_allocator_query_buffer_compatibility(
                             base_allocator, &compat_params, &allocation_size),
                         IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE)) {
    return iree_make_status(
        IREE_STATUS_INVALID_ARGUMENT,
        "allocator cannot allocate a buffer with the given parameters");
  }

  // Allocate the buffer (both the wrapper and the contents).
  iree_hal_heap_allocator_statistics_t* statistics = NULL;
  IREE_STATISTICS(statistics = &allocator->statistics);
  iree_hal_buffer_t* buffer = NULL;
  IREE_RETURN_IF_ERROR(iree_hal_heap_buffer_create(
      statistics, &compat_params, allocation_size, allocator->data_allocator,
      allocator->host_allocator, &buffer));

  *out_buffer = buffer;
  return iree_ok_status();
}

static void iree_hal_heap_allocator_deallocate_buffer(
    iree_hal_allocator_t* IREE_RESTRICT base_allocator,
    iree_hal_buffer_t* IREE_RESTRICT base_buffer) {
  // We don't do any pooling yet.
  // TODO(benvanik): move stats tracking here.
  iree_hal_buffer_destroy(base_buffer);
}

static iree_status_t iree_hal_heap_allocator_import_buffer(
    iree_hal_allocator_t* IREE_RESTRICT base_allocator,
    const iree_hal_buffer_params_t* IREE_RESTRICT params,
    iree_hal_external_buffer_t* IREE_RESTRICT external_buffer,
    iree_hal_buffer_release_callback_t release_callback,
    iree_hal_buffer_t** IREE_RESTRICT out_buffer) {
  iree_hal_heap_allocator_t* allocator =
      iree_hal_heap_allocator_cast(base_allocator);

  // Coerce options into those required for use by heap-based devices.
  iree_hal_buffer_params_t compat_params = *params;
  iree_device_size_t allocation_size = external_buffer->size;
  if (!iree_all_bits_set(iree_hal_heap_allocator_query_buffer_compatibility(
                             base_allocator, &compat_params, &allocation_size),
                         IREE_HAL_BUFFER_COMPATIBILITY_IMPORTABLE)) {
    return iree_make_status(
        IREE_STATUS_INVALID_ARGUMENT,
        "allocator cannot import a buffer with the given parameters");
  }

  void* ptr = NULL;
  switch (external_buffer->type) {
    case IREE_HAL_EXTERNAL_BUFFER_TYPE_HOST_ALLOCATION:
      ptr = external_buffer->handle.host_allocation.ptr;
      break;
    case IREE_HAL_EXTERNAL_BUFFER_TYPE_DEVICE_ALLOCATION:
      ptr = (void*)((intptr_t)external_buffer->handle.device_allocation.ptr);
      break;
    default:
      return iree_make_status(IREE_STATUS_UNAVAILABLE,
                              "external buffer type not supported");
  }

  const iree_hal_buffer_placement_t placement = {
      .device = NULL,
      .queue_affinity = compat_params.queue_affinity
                            ? compat_params.queue_affinity
                            : IREE_HAL_QUEUE_AFFINITY_ANY,
      .flags = IREE_HAL_BUFFER_PLACEMENT_FLAG_NONE,
  };
  return iree_hal_heap_buffer_wrap(
      placement, compat_params.type, compat_params.access, compat_params.usage,
      external_buffer->size, iree_make_byte_span(ptr, external_buffer->size),
      release_callback, allocator->host_allocator, out_buffer);
}

static iree_status_t iree_hal_heap_allocator_export_buffer(
    iree_hal_allocator_t* IREE_RESTRICT base_allocator,
    iree_hal_buffer_t* IREE_RESTRICT buffer,
    iree_hal_external_buffer_type_t requested_type,
    iree_hal_external_buffer_flags_t requested_flags,
    iree_hal_external_buffer_t* IREE_RESTRICT out_external_buffer) {
  // For the heap allocator on CPUs, we can directly export requests for
  // both a host allocation and a device allocation as they are one and the
  // same. We switch below to store the pointer in the right part of the
  // union.
  if (requested_type != IREE_HAL_EXTERNAL_BUFFER_TYPE_HOST_ALLOCATION &&
      requested_type != IREE_HAL_EXTERNAL_BUFFER_TYPE_DEVICE_ALLOCATION) {
    return iree_make_status(IREE_STATUS_UNAVAILABLE,
                            "external buffer type not supported");
  }

  // Map the entire buffer persistently, if possible.
  iree_hal_buffer_mapping_t mapping;
  IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range(
      buffer, IREE_HAL_MAPPING_MODE_PERSISTENT,
      iree_hal_buffer_allowed_access(buffer), 0, IREE_WHOLE_BUFFER, &mapping));

  // Note that the returned pointer is unowned.
  out_external_buffer->type = requested_type;
  out_external_buffer->flags = requested_flags;
  out_external_buffer->size = mapping.contents.data_length;
  if (requested_type == IREE_HAL_EXTERNAL_BUFFER_TYPE_HOST_ALLOCATION) {
    out_external_buffer->handle.host_allocation.ptr = mapping.contents.data;
  } else if (requested_type ==
             IREE_HAL_EXTERNAL_BUFFER_TYPE_DEVICE_ALLOCATION) {
    out_external_buffer->handle.device_allocation.ptr =
        (uint64_t)(uintptr_t)mapping.contents.data;
  }
  return iree_ok_status();
}

static const iree_hal_allocator_vtable_t iree_hal_heap_allocator_vtable = {
    .destroy = iree_hal_heap_allocator_destroy,
    .host_allocator = iree_hal_heap_allocator_host_allocator,
    .trim = iree_hal_heap_allocator_trim,
    .query_statistics = iree_hal_heap_allocator_query_statistics,
    .query_memory_heaps = iree_hal_heap_allocator_query_memory_heaps,
    .query_buffer_compatibility =
        iree_hal_heap_allocator_query_buffer_compatibility,
    .allocate_buffer = iree_hal_heap_allocator_allocate_buffer,
    .deallocate_buffer = iree_hal_heap_allocator_deallocate_buffer,
    .import_buffer = iree_hal_heap_allocator_import_buffer,
    .export_buffer = iree_hal_heap_allocator_export_buffer,
};
