// Copyright 2023 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 "iree/hal/utils/caching_allocator.h"

#include "iree/base/internal/synchronization.h"

// Default capacity of a pool free list when not specified by the user.
#define IREE_HAL_CACHING_ALLOCATOR_DEFAULT_FREE_LIST_CAPACITY 64

//===----------------------------------------------------------------------===//
// iree_hal_caching_allocator_pool_t
//===----------------------------------------------------------------------===//

IREE_TRACE(
    static const char* IREE_HAL_CACHING_ALLOCATOR_ID = "Free Cached Memory");

void iree_hal_caching_allocator_pool_params_initialize(
    iree_hal_allocator_memory_heap_t heap,
    iree_hal_caching_allocator_pool_params_t* out_params) {
  IREE_ASSERT_ARGUMENT(out_params);
  memset(out_params, 0, sizeof(*out_params));
  out_params->heap = heap;
  out_params->max_allocation_size = heap.max_allocation_size;
  out_params->max_allocation_capacity = IREE_DEVICE_SIZE_MAX;
  out_params->max_free_allocation_count =
      IREE_HAL_CACHING_ALLOCATOR_DEFAULT_FREE_LIST_CAPACITY;
}

// Pool of arbitrarily-sized device allocations for a particular heap.
// This maintains a free list of blocks available for use but does not track
// outstanding allocations.
//
// Thread-safe. Pools can service requests from multiple threads concurrently by
// way of a pool-specific mutex. The mutex will not be held during underlying
// allocator operations such as when acquiring a new allocation as these can be
// extremely slow and the underlying allocator is also assumed thread-safe.
typedef iree_alignas(
    iree_max_align_t) struct iree_hal_caching_allocator_pool_t {
  // Defines which heap this pool allocates from and the pool limits.
  iree_hal_caching_allocator_pool_params_t params;

  // Underlying device allocator used to allocate storage buffers.
  // Unretained as the parent allocator retains it for us.
  iree_hal_allocator_t* device_allocator;

  // Guards access to the pool data structures as buffers can be
  // acquired/released from multiple threads if shared across user-visible
  // devices.
  //
  // Note that we keep the mutex per-pool so that if we do need to allocate or
  // free we can do so without holding the lock.
  iree_slim_mutex_t mutex;

  // Total size, in bytes, of all outstanding allocations made from this pool.
  // This only includes allocations we are able to pool as we otherwise cannot
  // observe imported/exported buffers.
  iree_device_size_t total_allocated_size;

  // Total size, in bytes, of all free buffers currently in this pool.
  iree_device_size_t free_allocated_size;

  // Flat MRU list of available buffers with max_free_allocation_count slots.
  // Sorted by ascending recency (the higher the index the more recent).
  // If we really cared about optimizing the interior removal then we'd want
  // a linked/skip list or some bucketing but that's really for the higher
  // level allocators to do.
  iree_host_size_t free_count;
  iree_hal_buffer_t* free_buffers[];
} iree_hal_caching_allocator_pool_t;

static void iree_hal_caching_allocator_pool_trim(
    iree_hal_caching_allocator_pool_t* pool);

// Initializes a buffer pool in |out_pool|.
// Buffer device storage will be allocated from |device_allocator|.
static void iree_hal_caching_allocator_pool_initialize(
    iree_hal_caching_allocator_pool_params_t params,
    iree_hal_allocator_t* device_allocator,
    iree_hal_caching_allocator_pool_t* out_pool) {
  IREE_TRACE_ZONE_BEGIN(z0);

  out_pool->params = params;
  out_pool->device_allocator = device_allocator;
  iree_slim_mutex_initialize(&out_pool->mutex);
  out_pool->total_allocated_size = 0;
  out_pool->free_allocated_size = 0;
  out_pool->free_count = 0;

  IREE_TRACE_SET_PLOT_TYPE(IREE_HAL_CACHING_ALLOCATOR_ID,
                           IREE_TRACING_PLOT_TYPE_MEMORY, /*step=*/true,
                           /*fill=*/true, /*color=*/0);
  IREE_TRACE_PLOT_VALUE_I64(IREE_HAL_CACHING_ALLOCATOR_ID,
                            out_pool->free_allocated_size);

  IREE_TRACE_ZONE_END(z0);
}

// Deinitializes |pool|; all allocated buffers must have been released.
static void iree_hal_caching_allocator_pool_deinitialize(
    iree_hal_caching_allocator_pool_t* pool) {
  IREE_TRACE_ZONE_BEGIN(z0);

  // Trim first to release all the buffers. There shouldn't be any live
  // allocations by the time we are deinitializing.
  iree_hal_caching_allocator_pool_trim(pool);
  IREE_ASSERT_EQ(pool->total_allocated_size, 0,
                 "must have released all allocations prior to deinit");
  IREE_ASSERT_EQ(pool->free_allocated_size, 0,
                 "must have released all allocations prior to deinit");
  IREE_ASSERT_EQ(pool->free_count, 0,
                 "must have released all allocations prior to deinit");

  iree_slim_mutex_deinitialize(&pool->mutex);

  IREE_TRACE_ZONE_END(z0);
}

// Pushes |buffer| on to the pool free list as the most recently used.
// The buffer will be retained in the list.
//
// Must be called with the pool mutex held.
static void iree_hal_caching_allocator_pool_push_buffer(
    iree_hal_caching_allocator_pool_t* pool, iree_hal_buffer_t* buffer) {
  // Retain the buffer; the caller must release it to complete the ownership
  // transfer.
  iree_hal_buffer_retain(buffer);

  IREE_ASSERT_LT(pool->free_count, pool->params.max_free_allocation_count);

  // Add to the end of the list (the most recent).
  iree_host_size_t i = pool->free_count++;
  pool->free_buffers[i] = buffer;

  // Track that we're now retaining unused memory.
  pool->free_allocated_size += buffer->allocation_size;
  IREE_TRACE_PLOT_VALUE_I64(IREE_HAL_CACHING_ALLOCATOR_ID,
                            pool->free_allocated_size);
}

// Takes the buffer in the |pool| free list at index |i| and returns ownership.
// If the list is large this is bad but in most cases it's just a few dozen
// elements.
//
// Must be called with the pool mutex held.
static iree_hal_buffer_t* iree_hal_caching_allocator_pool_take_buffer_at(
    iree_hal_caching_allocator_pool_t* pool, iree_host_size_t i) {
  iree_hal_buffer_t* buffer = pool->free_buffers[i];
  if (i < pool->free_count - 1) {
    // Shift the list down to keep it dense and in ascending recency order.
    memmove(&pool->free_buffers[i], &pool->free_buffers[i + 1],
            (pool->free_count - i - 1) * sizeof(pool->free_buffers[0]));
  }
  --pool->free_count;
  pool->free_allocated_size -= buffer->allocation_size;
  IREE_TRACE_PLOT_VALUE_I64(IREE_HAL_CACHING_ALLOCATOR_ID,
                            pool->free_allocated_size);
  return buffer;
}

// Scans the |pool| free list for a buffer matching the given requirements and
// returns ownership.
//
// Must be called with the pool mutex held.
static iree_hal_buffer_t* iree_hal_caching_allocator_pool_find_and_take_buffer(
    iree_hal_caching_allocator_pool_t* pool,
    const iree_hal_buffer_params_t* params,
    iree_device_size_t allocation_size) {
  // Walk backwards so that we check the most recently released buffers first.
  for (int i = (int)pool->free_count - 1; i >= 0; --i) {
    // NOTE: we are not currently checking alignment as we don't really have it.
    // We assume programs will use consistent alignments for a particular heap
    // (as the heap has a min alignment).
    iree_hal_buffer_t* buffer = pool->free_buffers[i];
    if (iree_all_bits_set(iree_hal_buffer_memory_type(buffer), params->type) &&
        iree_all_bits_set(iree_hal_buffer_allowed_usage(buffer),
                          params->usage) &&
        iree_hal_buffer_allocation_size(buffer) == allocation_size) {
      return iree_hal_caching_allocator_pool_take_buffer_at(pool, i);
    }
  }
  return NULL;  // nothing found
}

// Trims |pool| down to at most |target_size| of available allocations.
// The oldest allocations will be trimmed first.
//
// Thread-safe; multiple threads may concurrently access the |pool|.
static void iree_hal_caching_allocator_pool_trim_to_size(
    iree_hal_caching_allocator_pool_t* pool, iree_device_size_t target_size) {
  IREE_TRACE_ZONE_BEGIN(z0);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)target_size);

  iree_slim_mutex_lock(&pool->mutex);

  while (pool->free_count > 0 && pool->total_allocated_size > target_size) {
    // Take the oldest buffer in the list.
    iree_hal_buffer_t* dead_buffer =
        iree_hal_caching_allocator_pool_take_buffer_at(pool,
                                                       pool->free_count - 1);

    // NOTE: we've removed the buffer but have not subtracted the size from
    // the total yet - we want to do that only after releasing the buffer.
    // If we didn't it's possible for another thread to start an allocation
    // thinking that we've already released the buffer.
    iree_device_size_t allocation_size =
        iree_hal_buffer_allocation_size(dead_buffer);

    // Release the buffer without holding the lock as deallocation can be slow.
    iree_slim_mutex_unlock(&pool->mutex);
    iree_hal_allocator_deallocate_buffer(pool->device_allocator, dead_buffer);
    iree_slim_mutex_lock(&pool->mutex);

    // Update accounting to represent that we've released the buffer.
    IREE_ASSERT_GE(pool->total_allocated_size, allocation_size);
    pool->total_allocated_size -= allocation_size;
  }

  iree_slim_mutex_unlock(&pool->mutex);

  IREE_TRACE_ZONE_END(z0);
}

// Releases all unused buffers in |pool| to the underlying device allocator.
//
// The pool mutex must not be held by the caller.
static void iree_hal_caching_allocator_pool_trim(
    iree_hal_caching_allocator_pool_t* pool) {
  iree_hal_caching_allocator_pool_trim_to_size(pool, 0);
}

// Acquires a buffer of |allocation_size| from the |pool|.
// The buffer will have a memory type and usage compatible with the given types.
// Fails if the pool is empty and the underlying device fails the allocation.
//
// Thread-safe; multiple threads may concurrently access the |pool|.
static iree_status_t iree_hal_caching_allocator_pool_acquire(
    iree_hal_caching_allocator_pool_t* pool,
    const iree_hal_buffer_params_t* params, iree_device_size_t allocation_size,
    iree_hal_buffer_t** out_buffer) {
  IREE_TRACE_ZONE_BEGIN(z0);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)allocation_size);

  // Scan the free list to find an appropriate block.
  // If found we pop it off the list and return it without needing to allocate.
  iree_slim_mutex_lock(&pool->mutex);
  iree_hal_buffer_t* existing_buffer =
      iree_hal_caching_allocator_pool_find_and_take_buffer(pool, params,
                                                           allocation_size);
  if (!existing_buffer) {
    // We'll need to allocate so we add the size such that it'll be accounted
    // for by other threads allocating at the same time.
    pool->total_allocated_size += allocation_size;
  }
  iree_slim_mutex_unlock(&pool->mutex);
  if (existing_buffer) {
    // Found a buffer! Return it uninitialized.
    *out_buffer = existing_buffer;
    IREE_TRACE_ZONE_END(z0);
    return iree_ok_status();
  }

  // Trim first before allocating so that we don't go over peak.
  iree_hal_caching_allocator_pool_trim_to_size(
      pool, pool->params.max_allocation_capacity);

  // No existing buffer was found that could be used and we'll need to allocate
  // one. Note that we do this without holding the lock as the underlying
  // device allocator can be very slow. It's possible for buffers to be released
  // to the pool by another thread while we're allocating here but that's OK.
  iree_hal_buffer_t* buffer = NULL;
  iree_status_t status = iree_hal_allocator_allocate_buffer(
      pool->device_allocator, *params, allocation_size, &buffer);

  // If the allocation failed then remove the size from the total.
  if (iree_status_is_ok(status)) {
    *out_buffer = buffer;
  } else {
    if (buffer) iree_hal_buffer_release(buffer);
    iree_slim_mutex_lock(&pool->mutex);
    pool->total_allocated_size -= allocation_size;
    iree_slim_mutex_unlock(&pool->mutex);
  }

  IREE_TRACE_ZONE_END(z0);
  return status;
}

// Releases a |buffer| to the |pool| if there is capacity remaining.
//
// Thread-safe; multiple threads may concurrently access the |pool|.
static void iree_hal_caching_allocator_pool_release(
    iree_hal_caching_allocator_pool_t* pool, iree_hal_buffer_t* buffer) {
  IREE_TRACE_ZONE_BEGIN(z0);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(
      z0, (int64_t)iree_hal_buffer_allocation_size(buffer));

  // Try to add the buffer to the pool. If the pool is at capacity we'll just
  // release it back to the allocator.
  iree_slim_mutex_lock(&pool->mutex);

  const iree_device_size_t allocation_size =
      iree_hal_buffer_allocation_size(buffer);
  const bool under_capacity = pool->total_allocated_size - allocation_size <=
                              pool->params.max_allocation_capacity;
  const bool under_count =
      pool->free_count + 1 <= pool->params.max_free_allocation_count;
  if (under_capacity && under_count) {
    iree_hal_caching_allocator_pool_push_buffer(pool, buffer);
    buffer = NULL;
  }

  // If the buffer didn't fit in the pool we drop it here while we don't hold
  // the lock as deallocations can be very expensive.
  if (buffer) {
    iree_slim_mutex_unlock(&pool->mutex);
    iree_hal_allocator_deallocate_buffer(pool->device_allocator, buffer);
    iree_slim_mutex_lock(&pool->mutex);
    pool->total_allocated_size -= allocation_size;
  }

  iree_slim_mutex_unlock(&pool->mutex);

  IREE_TRACE_ZONE_END(z0);
}

//===----------------------------------------------------------------------===//
// iree_hal_caching_allocator_t
//===----------------------------------------------------------------------===//

struct iree_hal_caching_allocator_t {
  iree_hal_resource_t resource;
  iree_allocator_t host_allocator;

  // Underlying device allocator used to allocate storage blocks.
  // We also route down to it for things we don't support (import/export/etc).
  iree_hal_allocator_t* device_allocator;

  // Total number of pools.
  iree_host_size_t pool_count;

  // Pointers to pool storage.
  // The count and layout of pools is immutable while each pool has a mutex to
  // guard the pool state.
  iree_hal_caching_allocator_pool_t* pools[];
};

static const iree_hal_allocator_vtable_t iree_hal_caching_allocator_vtable;

iree_hal_caching_allocator_t* iree_hal_caching_allocator_cast(
    iree_hal_allocator_t* base_value) {
  IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_caching_allocator_vtable);
  return (iree_hal_caching_allocator_t*)base_value;
}

iree_status_t iree_hal_caching_allocator_create_unbounded(
    iree_hal_allocator_t* device_allocator, iree_allocator_t host_allocator,
    iree_hal_allocator_t** out_allocator) {
  IREE_ASSERT_ARGUMENT(device_allocator);
  IREE_ASSERT_ARGUMENT(out_allocator);
  IREE_TRACE_ZONE_BEGIN(z0);

  // Query heaps from the underlying allocator.
  iree_hal_allocator_memory_heap_t heaps[8];
  iree_host_size_t heap_count = 0;
  IREE_RETURN_IF_ERROR(iree_hal_allocator_query_memory_heaps(
      device_allocator, IREE_ARRAYSIZE(heaps), heaps, &heap_count));

  // Setup pool parameters for each heap.
  iree_hal_caching_allocator_pool_params_t* heap_pool_params =
      (iree_hal_caching_allocator_pool_params_t*)iree_alloca(
          sizeof(iree_hal_caching_allocator_pool_params_t) * heap_count);
  for (iree_host_size_t i = 0; i < heap_count; ++i) {
    iree_hal_caching_allocator_pool_params_initialize(heaps[i],
                                                      &heap_pool_params[i]);
  }

  iree_status_t status = iree_hal_caching_allocator_create_with_pools(
      heap_count, heap_pool_params, device_allocator, host_allocator,
      out_allocator);

  IREE_TRACE_ZONE_END(z0);
  return status;
}

iree_status_t iree_hal_caching_allocator_create_with_pools(
    iree_host_size_t pool_count,
    const iree_hal_caching_allocator_pool_params_t* pool_params,
    iree_hal_allocator_t* device_allocator, iree_allocator_t host_allocator,
    iree_hal_allocator_t** out_allocator) {
  IREE_ASSERT_ARGUMENT(!pool_count || pool_params);
  IREE_ASSERT_ARGUMENT(device_allocator);
  IREE_ASSERT_ARGUMENT(out_allocator);
  IREE_TRACE_ZONE_BEGIN(z0);

  // Allocate the allocator itself and then a trailing list of variable-length
  // pools based on their free list sizes.
  iree_hal_caching_allocator_t* allocator = NULL;
  iree_host_size_t pool_list_size = pool_count * sizeof(allocator->pools[0]);
  iree_host_size_t total_size = iree_host_align(
      iree_sizeof_struct(*allocator) + pool_list_size, iree_max_align_t);
  iree_host_size_t pool_offset = total_size;
  for (iree_host_size_t i = 0; i < pool_count; ++i) {
    iree_hal_caching_allocator_pool_t* pool = NULL;
    total_size += iree_host_align(
        sizeof(*pool) + sizeof(pool->free_buffers[0]) *
                            pool_params[i].max_free_allocation_count,
        iree_max_align_t);
  }
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0,
      iree_allocator_malloc(host_allocator, total_size, (void**)&allocator));

  // Initialize the allocator.
  iree_hal_resource_initialize(&iree_hal_caching_allocator_vtable,
                               &allocator->resource);
  allocator->host_allocator = host_allocator;
  allocator->device_allocator = device_allocator;
  iree_hal_allocator_retain(allocator->device_allocator);
  allocator->pool_count = pool_count;

  // Initialize each pool.
  uint8_t* pool_ptr = (uint8_t*)allocator + pool_offset;
  for (iree_host_size_t i = 0; i < pool_count; ++i) {
    iree_hal_caching_allocator_pool_t* pool =
        (iree_hal_caching_allocator_pool_t*)pool_ptr;
    pool_ptr += iree_host_align(
        sizeof(*pool) + sizeof(pool->free_buffers[0]) *
                            pool_params[i].max_free_allocation_count,
        iree_max_align_t);
    allocator->pools[i] = pool;
    iree_hal_caching_allocator_pool_initialize(pool_params[i], device_allocator,
                                               pool);
  }

  *out_allocator = (iree_hal_allocator_t*)allocator;
  IREE_TRACE_ZONE_END(z0);
  return iree_ok_status();
}

// Selects a heap from |heaps| matching the given |heap_key|.
// Fails if no heap matches the given key. Optionally a buffer usage bitfield
// can be provided. Wildcards can be used with either to match the first heap
// that meets either requirement.
//
// Examples:
//   *: first heap
//   *;transfer: first heap with transfer usage
//   device_local: first heap with the IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL bit
//   device_local|host_visible
//   device_local;transfer|dispatch_storage
static iree_status_t iree_hal_select_heap(
    iree_string_view_t heap_key, iree_host_size_t heap_count,
    const iree_hal_allocator_memory_heap_t* heaps,
    const iree_hal_allocator_memory_heap_t** out_heap) {
  iree_string_view_t memory_type_str = iree_string_view_empty();
  iree_string_view_t buffer_usage_str = iree_string_view_empty();
  iree_string_view_split(heap_key, ';', &memory_type_str, &buffer_usage_str);

  // Parse the provided filters, if any.
  iree_hal_memory_type_t memory_type = IREE_HAL_MEMORY_TYPE_NONE;
  iree_hal_buffer_usage_t buffer_usage = IREE_HAL_BUFFER_USAGE_NONE;
  if (!iree_string_view_is_empty(memory_type_str) &&
      !iree_string_view_equal(memory_type_str, IREE_SV("*"))) {
    IREE_RETURN_IF_ERROR(
        iree_hal_memory_type_parse(memory_type_str, &memory_type));
  }
  if (!iree_string_view_is_empty(buffer_usage_str) &&
      !iree_string_view_equal(buffer_usage_str, IREE_SV("*"))) {
    IREE_RETURN_IF_ERROR(
        iree_hal_buffer_usage_parse(buffer_usage_str, &buffer_usage));
  }

  // Return the first heap satisfying all filters.
  for (iree_host_size_t i = 0; i < heap_count; ++i) {
    if ((!memory_type || iree_all_bits_set(heaps[i].type, memory_type)) &&
        (!buffer_usage ||
         iree_all_bits_set(heaps[i].allowed_usage, buffer_usage))) {
      *out_heap = &heaps[i];
      return iree_ok_status();
    }
  }

  // No matching heap found; can happen if the device doesn't have the kind of
  // heaps the user was expecting with the configuration.
  return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                          "no heap matching requested config params "
                          "memory_type='%.*s', buffer_usage='%.*s'",
                          (int)memory_type_str.size, memory_type_str.data,
                          (int)buffer_usage_str.size, buffer_usage_str.data);
}

iree_status_t iree_hal_caching_allocator_create_from_spec(
    iree_string_view_t config_pairs, iree_hal_allocator_t* device_allocator,
    iree_allocator_t host_allocator, iree_hal_allocator_t** out_allocator) {
  if (iree_string_view_is_empty(config_pairs)) {
    // No parameters implies unbounded; we'll hang on to all memory forever.
    // This is only useful in very specific usage patterns such as statically
    // shaped and deterministic benchmarks that always allocate the same amounts
    // of memory per invocation.
    return iree_hal_caching_allocator_create_unbounded(
        device_allocator, host_allocator, out_allocator);
  }

  // Query all heaps from the base allocator. We'll use this list to match the
  // user-provided pool parameters to heaps. It's likely that not all heaps
  // will be selected by the user.
  iree_host_size_t heap_count = 0;
  iree_hal_allocator_memory_heap_t heaps[16];
  IREE_RETURN_IF_ERROR(iree_hal_allocator_query_memory_heaps(
      device_allocator, IREE_ARRAYSIZE(heaps), heaps, &heap_count));

  // Build a list of pools based on user specification.
  iree_host_size_t pool_count = 0;
  iree_hal_caching_allocator_pool_params_t pool_params_storage[16];
  do {
    if (pool_count + 1 > IREE_ARRAYSIZE(pool_params_storage)) {
      return iree_make_status(IREE_STATUS_OUT_OF_RANGE,
                              "too many pools specified");
    }

    // Pop the key=value config pair from the list.
    iree_string_view_t config_pair = iree_string_view_empty();
    iree_string_view_split(config_pairs, ',', &config_pair, &config_pairs);
    iree_string_view_t heap_key = iree_string_view_empty();
    iree_string_view_t pool_config = iree_string_view_empty();
    iree_string_view_split(config_pair, '=', &heap_key, &pool_config);
    heap_key = iree_string_view_trim(heap_key);
    if (iree_string_view_is_empty(heap_key)) {
      return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                              "heap key must specified in pool params");
    }

    // Select the heap based on the key.
    const iree_hal_allocator_memory_heap_t* heap = NULL;
    IREE_RETURN_IF_ERROR(
        iree_hal_select_heap(heap_key, heap_count, heaps, &heap));

    // Configure the pool based on the provided parameters.
    iree_hal_caching_allocator_pool_params_t* pool_params =
        &pool_params_storage[pool_count++];
    iree_hal_caching_allocator_pool_params_initialize(*heap, pool_params);

    iree_string_view_t max_allocation_size_str = iree_string_view_empty();
    iree_string_view_t max_allocation_capacity_str = iree_string_view_empty();
    iree_string_view_t max_free_allocation_count_str = iree_string_view_empty();
    iree_string_view_split(pool_config, ';', &max_allocation_size_str,
                           &pool_config);
    iree_string_view_split(pool_config, ';', &max_allocation_capacity_str,
                           &pool_config);
    iree_string_view_split(pool_config, ';', &max_free_allocation_count_str,
                           &pool_config);
    max_allocation_size_str = iree_string_view_trim(max_allocation_size_str);
    if (!iree_string_view_is_empty(max_allocation_size_str) &&
        !iree_string_view_equal(max_allocation_size_str, IREE_SV("*"))) {
      IREE_RETURN_IF_ERROR(
          iree_string_view_parse_device_size(max_allocation_size_str,
                                             &pool_params->max_allocation_size),
          "parsing max_allocation_size");
    }
    max_allocation_capacity_str =
        iree_string_view_trim(max_allocation_capacity_str);
    if (!iree_string_view_is_empty(max_allocation_capacity_str) &&
        !iree_string_view_equal(max_allocation_capacity_str, IREE_SV("*"))) {
      IREE_RETURN_IF_ERROR(iree_string_view_parse_device_size(
                               max_allocation_capacity_str,
                               &pool_params->max_allocation_capacity),
                           "parsing max_allocation_capacity");
    }
    max_free_allocation_count_str =
        iree_string_view_trim(max_free_allocation_count_str);
    if (!iree_string_view_is_empty(max_free_allocation_count_str) &&
        !iree_string_view_equal(max_free_allocation_count_str, IREE_SV("*"))) {
      uint32_t max_free_allocation_count = 0;
      if (!iree_string_view_atoi_uint32(max_free_allocation_count_str,
                                        &max_free_allocation_count)) {
        return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                                "invalid count '%.*s'",
                                (int)max_free_allocation_count_str.size,
                                max_free_allocation_count_str.data);
      }
      pool_params->max_free_allocation_count = max_free_allocation_count;
    }
  } while (!iree_string_view_is_empty(config_pairs));
  return iree_hal_caching_allocator_create_with_pools(
      pool_count, pool_params_storage, device_allocator, host_allocator,
      out_allocator);
}

static void iree_hal_caching_allocator_destroy(
    iree_hal_allocator_t* IREE_RESTRICT base_allocator) {
  iree_hal_caching_allocator_t* allocator =
      iree_hal_caching_allocator_cast(base_allocator);
  iree_allocator_t host_allocator = allocator->host_allocator;
  IREE_TRACE_ZONE_BEGIN(z0);

  // Deinitialize each pool, returning any available resources to the underlying
  // device allocator.
  for (iree_host_size_t i = 0; i < allocator->pool_count; ++i) {
    iree_hal_caching_allocator_pool_deinitialize(allocator->pools[i]);
  }

  iree_hal_allocator_release(allocator->device_allocator);
  iree_allocator_free(host_allocator, allocator);

  IREE_TRACE_ZONE_END(z0);
}

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

static iree_status_t iree_hal_caching_allocator_trim(
    iree_hal_allocator_t* IREE_RESTRICT base_allocator) {
  iree_hal_caching_allocator_t* allocator =
      iree_hal_caching_allocator_cast(base_allocator);
  for (iree_host_size_t i = 0; i < allocator->pool_count; ++i) {
    iree_hal_caching_allocator_pool_trim(allocator->pools[i]);
  }
  return iree_hal_allocator_trim(allocator->device_allocator);
}

static void iree_hal_caching_allocator_query_statistics(
    iree_hal_allocator_t* IREE_RESTRICT base_allocator,
    iree_hal_allocator_statistics_t* IREE_RESTRICT out_statistics) {
  iree_hal_caching_allocator_t* allocator =
      iree_hal_caching_allocator_cast(base_allocator);
  iree_hal_allocator_query_statistics(allocator->device_allocator,
                                      out_statistics);
}

static iree_status_t iree_hal_caching_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) {
  // We could expose just the heaps backing our pools but that would prevent the
  // use of any underlying heaps that we aren't pooling.
  iree_hal_caching_allocator_t* allocator =
      iree_hal_caching_allocator_cast(base_allocator);
  return iree_hal_allocator_query_memory_heaps(allocator->device_allocator,
                                               capacity, heaps, out_count);
}

static iree_hal_buffer_compatibility_t
iree_hal_caching_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) {
  // Defer to the base allocator.
  iree_hal_caching_allocator_t* allocator =
      iree_hal_caching_allocator_cast(base_allocator);
  return iree_hal_allocator_query_buffer_compatibility(
      allocator->device_allocator, *params, *allocation_size, params,
      allocation_size);
}

static iree_hal_caching_allocator_pool_t* iree_hal_caching_allocator_find_pool(
    iree_hal_caching_allocator_t* allocator, iree_hal_memory_type_t type,
    iree_hal_buffer_usage_t allowed_usage) {
  // Scan in order; the preferred pools are first.
  for (iree_host_size_t i = 0; i < allocator->pool_count; ++i) {
    iree_hal_allocator_memory_heap_t heap = allocator->pools[i]->params.heap;
    if (iree_all_bits_set(heap.type, type) &&
        iree_all_bits_set(heap.allowed_usage, allowed_usage)) {
      return allocator->pools[i];
    }
  }
  return NULL;
}

static iree_status_t iree_hal_caching_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_caching_allocator_t* allocator =
      iree_hal_caching_allocator_cast(base_allocator);

  bool can_pool = true;

  // If the buffer is constant/exported then we'll skip the pool.
  // Constant buffers are often mapped directly from files or host memory and
  // benefit from being routed down to the underlying allocator.
  if (iree_any_bit_set(params->usage,
                       IREE_HAL_BUFFER_USAGE_SHARING_EXPORT |
                           IREE_HAL_BUFFER_USAGE_SHARING_IMMUTABLE |
                           IREE_HAL_BUFFER_USAGE_SHARING_REPLICATE)) {
    can_pool = false;
  }

  // If we don't want to pool then early-exit to the backing allocator.
  if (!can_pool) {
    return iree_hal_allocator_allocate_buffer(
        allocator->device_allocator, *params, allocation_size, out_buffer);
  }

  // We need to ensure we have the same parameters the allocator will use so
  // that when we reuse memory we're using the broadest set to select from that
  // can service each request.
  iree_hal_buffer_params_t compat_params;
  if (!iree_all_bits_set(iree_hal_allocator_query_buffer_compatibility(
                             allocator->device_allocator, *params,
                             allocation_size, &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");
  }

  // Try to find a pool for the buffer parameters.
  iree_hal_caching_allocator_pool_t* pool =
      iree_hal_caching_allocator_find_pool(allocator, compat_params.type,
                                           compat_params.usage);
  if (!pool) {
    // Fallback to the underlying allocator.
    return iree_hal_allocator_allocate_buffer(allocator->device_allocator,
                                              compat_params, allocation_size,
                                              out_buffer);
  }

  // Acquire the buffer from the pool.
  IREE_RETURN_IF_ERROR(iree_hal_caching_allocator_pool_acquire(
      pool, &compat_params, allocation_size, out_buffer));

  // Point the buffer back to us for deallocation.
  //
  // TODO(#19159): remove iree_hal_allocator_deallocate_buffer when pooling no
  // longer requires the pooling_allocator on iree_hal_buffer_t. We should
  // instead be creating a new iree_hal_cached_buffer_t that we return as if it
  // were an allocated buffer and that can store a reference back to the pool.
  (*out_buffer)->pooling_allocator = base_allocator;

  return iree_ok_status();
}

static void iree_hal_caching_allocator_deallocate_buffer(
    iree_hal_allocator_t* IREE_RESTRICT base_allocator,
    iree_hal_buffer_t* IREE_RESTRICT buffer) {
  iree_hal_caching_allocator_t* allocator =
      iree_hal_caching_allocator_cast(base_allocator);

  // Try to find the pool we would want to release the buffer into.
  // Note that we are only going to get called if we had successfully placed the
  // buffer into a pool.
  iree_hal_caching_allocator_pool_t* pool =
      iree_hal_caching_allocator_find_pool(
          allocator, iree_hal_buffer_memory_type(buffer),
          iree_hal_buffer_allowed_usage(buffer));
  IREE_ASSERT(pool, "pool to return cached buffer to not found");
  if (!pool) return;

  // TODO(#19159): remove iree_hal_allocator_deallocate_buffer when pooling no
  // longer requires the pooling_allocator on iree_hal_buffer_t.

  // Release back to pool (which may deallocate).
  iree_hal_caching_allocator_pool_release(pool, buffer);
}

static iree_status_t iree_hal_caching_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) {
  // Bypass the caching allocator and directly ask the backing implementation.
  iree_hal_caching_allocator_t* allocator =
      iree_hal_caching_allocator_cast(base_allocator);
  return iree_hal_allocator_import_buffer(allocator->device_allocator, *params,
                                          external_buffer, release_callback,
                                          out_buffer);
}

static iree_status_t iree_hal_caching_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) {
  // Bypass the caching allocator and directly ask the backing implementation.
  //
  // TODO(benvanik): confirm we want to allow export? it can arbitrarily extend
  // lifetime and we want to ensure we aren't going to reuse buffers with
  // outstanding exports.
  iree_hal_caching_allocator_t* allocator =
      iree_hal_caching_allocator_cast(base_allocator);
  return iree_hal_allocator_export_buffer(allocator->device_allocator, buffer,
                                          requested_type, requested_flags,
                                          out_external_buffer);
}

static const iree_hal_allocator_vtable_t iree_hal_caching_allocator_vtable = {
    .destroy = iree_hal_caching_allocator_destroy,
    .host_allocator = iree_hal_caching_allocator_host_allocator,
    .trim = iree_hal_caching_allocator_trim,
    .query_statistics = iree_hal_caching_allocator_query_statistics,
    .query_memory_heaps = iree_hal_caching_allocator_query_memory_heaps,
    .query_buffer_compatibility =
        iree_hal_caching_allocator_query_buffer_compatibility,
    .allocate_buffer = iree_hal_caching_allocator_allocate_buffer,
    .deallocate_buffer = iree_hal_caching_allocator_deallocate_buffer,
    .import_buffer = iree_hal_caching_allocator_import_buffer,
    .export_buffer = iree_hal_caching_allocator_export_buffer,
};
