// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3

/******************************************************************************
 * Simple caching allocator for device memory allocations. The allocator is
 * thread-safe and capable of managing device allocations on multiple devices.
 ******************************************************************************/

#pragma once

#include <cub/config.cuh>

#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK
#  if _CCCL_COMPILER(NVRTC)
#    error \
      "Including <cub/util_allocator.cuh> is not supported when compiling with NVRTC, which supports device code only. You can define CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK to disable this warning."
#  endif // _CCCL_COMPILER(NVRTC)
#endif // CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK

#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
#  pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
#  pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
#  pragma system_header
#endif // no system header

#include <cub/util_debug.cuh>
#include <cub/util_namespace.cuh>

#include <cuda/std/__host_stdlib/math.h>

#include <map>
#include <mutex>
#include <set>

CUB_NAMESPACE_BEGIN

/******************************************************************************
 * CachingDeviceAllocator (host use)
 ******************************************************************************/

/**
 * @brief A simple caching allocator for device memory allocations.
 *
 * @par Overview
 * The allocator is thread-safe and stream-safe and is capable of managing cached
 * device allocations on multiple devices.  It behaves as follows:
 *
 * @par
 * - Allocations from the allocator are associated with an @p active_stream. Once freed,
 *   the allocation becomes available immediately for reuse within the @p active_stream
 *   with which it was associated with during allocation, and it becomes available for
 *   reuse within other streams when all prior work submitted to @p active_stream has completed.
 * - Allocations are categorized and cached by bin size. A new allocation request of
 *   a given size will only consider cached allocations within the corresponding bin.
 * - Bin limits progress geometrically in accordance with the growth factor
 *   @p bin_growth provided during construction. Unused device allocations within
 *   a larger bin cache are not reused for allocation requests that categorize to
 *   smaller bin sizes.
 * - Allocation requests below ( @p bin_growth ^ @p min_bin ) are rounded up to
 *   ( @p bin_growth ^ @p min_bin ).
 * - Allocations above ( @p bin_growth ^ @p max_bin ) are not rounded up to the nearest
 *   bin and are simply freed when they are deallocated instead of being returned
 *   to a bin-cache.
 * - If the total storage of cached allocations on a given device will exceed
 *   @p max_cached_bytes, allocations for that device are simply freed when they are
 *   deallocated instead of being returned to their bin-cache.
 *
 * @par
 * For example, the default-constructed CachingDeviceAllocator is configured with:
 * - @p bin_growth          = 8
 * - @p min_bin             = 3
 * - @p max_bin             = 7
 * - @p max_cached_bytes    = 6MB - 1B
 *
 * @par
 * which delineates five bin-sizes: 512B, 4KB, 32KB, 256KB, and 2MB
 * and sets a maximum of 6,291,455 cached bytes per device
 *
 */
struct CachingDeviceAllocator
{
  //---------------------------------------------------------------------
  // Constants
  //---------------------------------------------------------------------

  /// Out-of-bounds bin
  static constexpr unsigned int INVALID_BIN = (unsigned int) -1;

  /// Invalid size
  static constexpr size_t INVALID_SIZE = (size_t) -1;

#ifndef _CCCL_DOXYGEN_INVOKED // Do not document

  /// Invalid device ordinal
  static constexpr int INVALID_DEVICE_ORDINAL = -1;

  //---------------------------------------------------------------------
  // Type definitions and helper types
  //---------------------------------------------------------------------

  /**
   * Descriptor for device memory allocations
   */
  struct BlockDescriptor
  {
    // Device pointer
    void* d_ptr;

    // Size of allocation in bytes
    size_t bytes;

    // Bin enumeration
    unsigned int bin;

    // device ordinal
    int device;

    // Associated associated_stream
    cudaStream_t associated_stream;

    // Signal when associated stream has run to the point at which this block was freed
    cudaEvent_t ready_event;

    // Constructor (suitable for searching maps for a specific block, given its pointer and
    // device)
    BlockDescriptor(void* d_ptr, int device)
        : d_ptr(d_ptr)
        , bytes(0)
        , bin(INVALID_BIN)
        , device(device)
        , associated_stream(nullptr)
        , ready_event(nullptr)
    {}

    // Constructor (suitable for searching maps for a range of suitable blocks, given a device)
    BlockDescriptor(int device)
        : d_ptr(nullptr)
        , bytes(0)
        , bin(INVALID_BIN)
        , device(device)
        , associated_stream(nullptr)
        , ready_event(nullptr)
    {}

    // Comparison functor for comparing device pointers
    static bool PtrCompare(const BlockDescriptor& a, const BlockDescriptor& b)
    {
      if (a.device == b.device)
      {
        return (a.d_ptr < b.d_ptr);
      }
      else
      {
        return (a.device < b.device);
      }
    }

    // Comparison functor for comparing allocation sizes
    static bool SizeCompare(const BlockDescriptor& a, const BlockDescriptor& b)
    {
      if (a.device == b.device)
      {
        return (a.bytes < b.bytes);
      }
      else
      {
        return (a.device < b.device);
      }
    }
  };

  /// BlockDescriptor comparator function interface
  using Compare = bool (*)(const BlockDescriptor&, const BlockDescriptor&);

  class TotalBytes
  {
  public:
    size_t free;
    size_t live;
    TotalBytes()
    {
      free = live = 0;
    }
  };

  /// Set type for cached blocks (ordered by size)
  using CachedBlocks = std::multiset<BlockDescriptor, Compare>;

  /// Set type for live blocks (ordered by ptr)
  using BusyBlocks = std::multiset<BlockDescriptor, Compare>;

  /// Map type of device ordinals to the number of cached bytes cached by each device
  using GpuCachedBytes = std::map<int, TotalBytes>;

  //---------------------------------------------------------------------
  // Utility functions
  //---------------------------------------------------------------------

  /**
   * Integer pow function for unsigned base and exponent
   */
  static unsigned int IntPow(unsigned int base, unsigned int exp)
  {
    unsigned int retval = 1;
    while (exp > 0)
    {
      if (exp & 1)
      {
        retval = retval * base; // multiply the result by the current base
      }
      base = base * base; // square the base
      exp  = exp >> 1; // divide the exponent in half
    }
    return retval;
  }

  /**
   * Round up to the nearest power-of
   */
  void NearestPowerOf(unsigned int& power, size_t& rounded_bytes, unsigned int base, size_t value)
  {
    power         = 0;
    rounded_bytes = 1;

    if (value * base < value)
    {
      // Overflow
      power         = sizeof(size_t) * 8;
      rounded_bytes = size_t(0) - 1;
      return;
    }

    while (rounded_bytes < value)
    {
      rounded_bytes *= base;
      power++;
    }
  }

  //---------------------------------------------------------------------
  // Fields
  //---------------------------------------------------------------------

  /// Mutex for thread-safety
  std::mutex mutex;

  /// Geometric growth factor for bin-sizes
  unsigned int bin_growth;

  /// Minimum bin enumeration
  unsigned int min_bin;

  /// Maximum bin enumeration
  unsigned int max_bin;

  /// Minimum bin size
  size_t min_bin_bytes;

  /// Maximum bin size
  size_t max_bin_bytes;

  /// Maximum aggregate cached bytes per device
  size_t max_cached_bytes;

  /// Whether or not to skip a call to FreeAllCached() when destructor is called.
  /// (The CUDA runtime may have already shut down for statically declared allocators)
  const bool skip_cleanup;

  /// Whether or not to print (de)allocation events to stdout
  bool debug;

  /// Map of device ordinal to aggregate cached bytes on that device
  GpuCachedBytes cached_bytes;

  /// Set of cached device allocations available for reuse
  CachedBlocks cached_blocks;

  /// Set of live device allocations currently in use
  BusyBlocks live_blocks;

#endif // _CCCL_DOXYGEN_INVOKED

  //---------------------------------------------------------------------
  // Methods
  //---------------------------------------------------------------------

  /**
   * @brief Constructor.
   *
   * @param bin_growth
   *   Geometric growth factor for bin-sizes
   *
   * @param min_bin
   *   Minimum bin (default is bin_growth ^ 1)
   *
   * @param max_bin
   *   Maximum bin (default is no max bin)
   *
   * @param max_cached_bytes
   *   Maximum aggregate cached bytes per device (default is no limit)
   *
   * @param skip_cleanup
   *   Whether or not to skip a call to @p FreeAllCached() when the destructor is called (default
   *   is to deallocate)
   */
  CachingDeviceAllocator(
    unsigned int bin_growth,
    unsigned int min_bin    = 1,
    unsigned int max_bin    = INVALID_BIN,
    size_t max_cached_bytes = INVALID_SIZE,
    bool skip_cleanup       = false)
      : bin_growth(bin_growth)
      , min_bin(min_bin)
      , max_bin(max_bin)
      , min_bin_bytes(IntPow(bin_growth, min_bin))
      , max_bin_bytes(IntPow(bin_growth, max_bin))
      , max_cached_bytes(max_cached_bytes)
      , skip_cleanup(skip_cleanup)
      , debug(false)
      , cached_blocks(BlockDescriptor::SizeCompare)
      , live_blocks(BlockDescriptor::PtrCompare)
  {}

  /**
   * @brief Default constructor.
   *
   * Configured with:
   * @par
   * - @p bin_growth          = 8
   * - @p min_bin             = 3
   * - @p max_bin             = 7
   * - @p max_cached_bytes    = ( @p bin_growth ^ @p max_bin) * 3 ) - 1 = 6,291,455 bytes
   *
   * which delineates five bin-sizes: 512B, 4KB, 32KB, 256KB, and 2MB and
   * sets a maximum of 6,291,455 cached bytes per device
   */
  CachingDeviceAllocator(bool skip_cleanup = false, bool debug = false)
      : bin_growth(8)
      , min_bin(3)
      , max_bin(7)
      , min_bin_bytes(IntPow(bin_growth, min_bin))
      , max_bin_bytes(IntPow(bin_growth, max_bin))
      , max_cached_bytes((max_bin_bytes * 3) - 1)
      , skip_cleanup(skip_cleanup)
      , debug(debug)
      , cached_blocks(BlockDescriptor::SizeCompare)
      , live_blocks(BlockDescriptor::PtrCompare)
  {}

  /**
   * @brief Sets the limit on the number bytes this allocator is allowed to cache per device.
   *
   * Changing the ceiling of cached bytes does not cause any allocations (in-use or
   * cached-in-reserve) to be freed.  See \p FreeAllCached().
   */
  cudaError_t SetMaxCachedBytes(size_t max_cached_bytes_)
  {
    // Lock
    mutex.lock();

#ifdef CUB_DEBUG_LOG
    _CubLog(
      "Changing max_cached_bytes (%lld -> %lld)\n", (long long) this->max_cached_bytes, (long long) max_cached_bytes_);
#endif

    this->max_cached_bytes = max_cached_bytes_;

    // Unlock
    mutex.unlock();

    return cudaSuccess;
  }

  /**
   * @brief Provides a suitable allocation of device memory for the given size on the specified
   *        device.
   *
   * Once freed, the allocation becomes available immediately for reuse within the @p
   * active_stream with which it was associated with during allocation, and it becomes available
   * for reuse within other streams when all prior work submitted to @p active_stream has
   * completed.
   *
   * @param[in] device
   *   Device on which to place the allocation
   *
   * @param[out] d_ptr
   *   Reference to pointer to the allocation
   *
   * @param[in] bytes
   *   Minimum number of bytes for the allocation
   *
   * @param[in] active_stream
   *   The stream to be associated with this allocation
   */
  cudaError_t DeviceAllocate(int device, void** d_ptr, size_t bytes, cudaStream_t active_stream = nullptr)
  {
    *d_ptr                = nullptr;
    int entrypoint_device = INVALID_DEVICE_ORDINAL;
    cudaError_t error     = cudaSuccess;

    if (device == INVALID_DEVICE_ORDINAL)
    {
      error = CubDebug(cudaGetDevice(&entrypoint_device));
      if (cudaSuccess != error)
      {
        return error;
      }

      device = entrypoint_device;
    }

    // Create a block descriptor for the requested allocation
    bool found = false;
    BlockDescriptor search_key(device);
    search_key.associated_stream = active_stream;
    NearestPowerOf(search_key.bin, search_key.bytes, bin_growth, bytes);

    if (search_key.bin > max_bin)
    {
      // Bin is greater than our maximum bin: allocate the request
      // exactly and give out-of-bounds bin.  It will not be cached
      // for reuse when returned.
      search_key.bin   = INVALID_BIN;
      search_key.bytes = bytes;
    }
    else
    {
      // Search for a suitable cached allocation: lock
      mutex.lock();

      if (search_key.bin < min_bin)
      {
        // Bin is less than minimum bin: round up
        search_key.bin   = min_bin;
        search_key.bytes = min_bin_bytes;
      }

      // Iterate through the range of cached blocks on the same device in the same bin
      CachedBlocks::iterator block_itr = cached_blocks.lower_bound(search_key);
      while ((block_itr != cached_blocks.end()) && (block_itr->device == device) && (block_itr->bin == search_key.bin))
      {
        // To prevent races with reusing blocks returned by the host but still
        // in use by the device, only consider cached blocks that are
        // either (from the active stream) or (from an idle stream)
        bool is_reusable = false;
        if (active_stream == block_itr->associated_stream)
        {
          is_reusable = true;
        }
        else
        {
          const cudaError_t event_status = cudaEventQuery(block_itr->ready_event);
          if (event_status != cudaErrorNotReady)
          {
            CubDebug(event_status);
            is_reusable = true;
          }
        }

        if (is_reusable)
        {
          // Reuse existing cache block.  Insert into live blocks.
          found                        = true;
          search_key                   = *block_itr;
          search_key.associated_stream = active_stream;
          live_blocks.insert(search_key);

          // Remove from free blocks
          cached_bytes[device].free -= search_key.bytes;
          cached_bytes[device].live += search_key.bytes;

#ifdef CUB_DEBUG_LOG
          _CubLog("\tDevice %d reused cached block at %p (%lld bytes) for stream %lld (previously associated with "
                  "stream %lld).\n",
                  device,
                  search_key.d_ptr,
                  (long long) search_key.bytes,
                  (long long) search_key.associated_stream,
                  (long long) block_itr->associated_stream);
#endif

          cached_blocks.erase(block_itr);

          break;
        }
        block_itr++;
      }

      // Done searching: unlock
      mutex.unlock();
    }

    // Allocate the block if necessary
    if (!found)
    {
      // Set runtime's current device to specified device (entrypoint may not be set)
      if (device != entrypoint_device)
      {
        error = CubDebug(cudaGetDevice(&entrypoint_device));
        if (cudaSuccess != error)
        {
          return error;
        }

        error = CubDebug(cudaSetDevice(device));
        if (cudaSuccess != error)
        {
          return error;
        }
      }

      // Attempt to allocate
      error = CubDebug(cudaMalloc(&search_key.d_ptr, search_key.bytes));
      if (error == cudaErrorMemoryAllocation)
      {
        // The allocation attempt failed: free all cached blocks on device and retry
#ifdef CUB_DEBUG_LOG
        _CubLog("\tDevice %d failed to allocate %lld bytes for stream %lld, retrying after freeing cached allocations",
                device,
                (long long) search_key.bytes,
                (long long) search_key.associated_stream);
#endif

        error = cudaSuccess; // Reset the error we will return
        cudaGetLastError(); // Reset CUDART's error

        // Lock
        mutex.lock();

        // Iterate the range of free blocks on the same device
        BlockDescriptor free_key(device);
        CachedBlocks::iterator block_itr = cached_blocks.lower_bound(free_key);

        while ((block_itr != cached_blocks.end()) && (block_itr->device == device))
        {
          // No need to worry about synchronization with the device: cudaFree is
          // blocking and will synchronize across all kernels executing
          // on the current device

          // Free device memory and destroy stream event.
          error = CubDebug(cudaFree(block_itr->d_ptr));
          if (cudaSuccess != error)
          {
            break;
          }

          error = CubDebug(cudaEventDestroy(block_itr->ready_event));
          if (cudaSuccess != error)
          {
            break;
          }

          // Reduce balance and erase entry
          cached_bytes[device].free -= block_itr->bytes;

#ifdef CUB_DEBUG_LOG
          _CubLog("\tDevice %d freed %lld bytes.\n\t\t  %lld available blocks cached (%lld bytes), %lld live blocks "
                  "(%lld bytes) outstanding.\n",
                  device,
                  (long long) block_itr->bytes,
                  (long long) cached_blocks.size(),
                  (long long) cached_bytes[device].free,
                  (long long) live_blocks.size(),
                  (long long) cached_bytes[device].live);
#endif

          block_itr = cached_blocks.erase(block_itr);
        }

        // Unlock
        mutex.unlock();

        // Return under error
        if (error)
        {
          return error;
        }

        // Try to allocate again
        error = CubDebug(cudaMalloc(&search_key.d_ptr, search_key.bytes));
        if (cudaSuccess != error)
        {
          return error;
        }
      }

      // Create ready event
      error = CubDebug(cudaEventCreateWithFlags(&search_key.ready_event, cudaEventDisableTiming));

      if (cudaSuccess != error)
      {
        return error;
      }

      // Insert into live blocks
      mutex.lock();
      live_blocks.insert(search_key);
      cached_bytes[device].live += search_key.bytes;
      mutex.unlock();

#ifdef CUB_DEBUG_LOG
      _CubLog("\tDevice %d allocated new device block at %p (%lld bytes associated with stream %lld).\n",
              device,
              search_key.d_ptr,
              (long long) search_key.bytes,
              (long long) search_key.associated_stream);
#endif

      // Attempt to revert back to previous device if necessary
      if ((entrypoint_device != INVALID_DEVICE_ORDINAL) && (entrypoint_device != device))
      {
        error = CubDebug(cudaSetDevice(entrypoint_device));
        if (cudaSuccess != error)
        {
          return error;
        }
      }
    }

    // Copy device pointer to output parameter
    *d_ptr = search_key.d_ptr;

#ifdef CUB_DEBUG_LOG
    if (debug)
    {
      _CubLog("\t\t%lld available blocks cached (%lld bytes), %lld live blocks outstanding(%lld bytes).\n",
              (long long) cached_blocks.size(),
              (long long) cached_bytes[device].free,
              (long long) live_blocks.size(),
              (long long) cached_bytes[device].live);
    }
#endif

    return error;
  }

  /**
   * @brief Provides a suitable allocation of device memory for the given size on the current
   *        device.
   *
   * Once freed, the allocation becomes available immediately for reuse within the @p
   * active_stream with which it was associated with during allocation, and it becomes available
   * for reuse within other streams when all prior work submitted to @p active_stream has
   * completed.
   *
   * @param[out] d_ptr
   *   Reference to pointer to the allocation
   *
   * @param[in] bytes
   *   Minimum number of bytes for the allocation
   *
   * @param[in] active_stream
   *   The stream to be associated with this allocation
   */
  cudaError_t DeviceAllocate(void** d_ptr, size_t bytes, cudaStream_t active_stream = nullptr)
  {
    return DeviceAllocate(INVALID_DEVICE_ORDINAL, d_ptr, bytes, active_stream);
  }

  /**
   * @brief Frees a live allocation of device memory on the specified device, returning it to the
   *        allocator.
   *
   * Once freed, the allocation becomes available immediately for reuse within the
   * @p active_stream with which it was associated with during allocation, and it becomes
   * available for reuse within other streams when all prior work submitted to @p active_stream
   * has completed.
   */
  cudaError_t DeviceFree(int device, void* d_ptr)
  {
    int entrypoint_device = INVALID_DEVICE_ORDINAL;
    cudaError_t error     = cudaSuccess;

    if (device == INVALID_DEVICE_ORDINAL)
    {
      error = CubDebug(cudaGetDevice(&entrypoint_device));
      if (cudaSuccess != error)
      {
        return error;
      }
      device = entrypoint_device;
    }

    // Lock
    mutex.lock();

    // Find corresponding block descriptor
    bool recached = false;
    BlockDescriptor search_key(d_ptr, device);
    BusyBlocks::iterator block_itr = live_blocks.find(search_key);
    if (block_itr != live_blocks.end())
    {
      // Remove from live blocks
      search_key = *block_itr;
      live_blocks.erase(block_itr);
      cached_bytes[device].live -= search_key.bytes;

      // Keep the returned allocation if bin is valid and we won't exceed the max cached threshold
      if ((search_key.bin != INVALID_BIN) && (cached_bytes[device].free + search_key.bytes <= max_cached_bytes))
      {
        // Insert returned allocation into free blocks
        recached = true;
        cached_blocks.insert(search_key);
        cached_bytes[device].free += search_key.bytes;

#ifdef CUB_DEBUG_LOG
        _CubLog("\tDevice %d returned %lld bytes from associated stream %lld.\n\t\t %lld available blocks cached (%lld "
                "bytes), %lld live blocks outstanding. (%lld bytes)\n",
                device,
                (long long) search_key.bytes,
                (long long) search_key.associated_stream,
                (long long) cached_blocks.size(),
                (long long) cached_bytes[device].free,
                (long long) live_blocks.size(),
                (long long) cached_bytes[device].live);
#endif
      }
    }

    // Unlock
    mutex.unlock();

    // First set to specified device (entrypoint may not be set)
    if (device != entrypoint_device)
    {
      error = CubDebug(cudaGetDevice(&entrypoint_device));
      if (cudaSuccess != error)
      {
        return error;
      }

      error = CubDebug(cudaSetDevice(device));
      if (cudaSuccess != error)
      {
        return error;
      }
    }

    if (recached)
    {
      // Insert the ready event in the associated stream (must have current device set properly)
      error = CubDebug(cudaEventRecord(search_key.ready_event, search_key.associated_stream));
      if (cudaSuccess != error)
      {
        return error;
      }
    }

    if (!recached)
    {
      // Free the allocation from the runtime and cleanup the event.
      error = CubDebug(cudaFree(d_ptr));
      if (cudaSuccess != error)
      {
        return error;
      }

      error = CubDebug(cudaEventDestroy(search_key.ready_event));
      if (cudaSuccess != error)
      {
        return error;
      }

#ifdef CUB_DEBUG_LOG
      _CubLog("\tDevice %d freed %lld bytes from associated stream %lld.\n\t\t  %lld available blocks cached (%lld "
              "bytes), %lld live blocks (%lld bytes) outstanding.\n",
              device,
              (long long) search_key.bytes,
              (long long) search_key.associated_stream,
              (long long) cached_blocks.size(),
              (long long) cached_bytes[device].free,
              (long long) live_blocks.size(),
              (long long) cached_bytes[device].live);
#endif
    }

    // Reset device
    if ((entrypoint_device != INVALID_DEVICE_ORDINAL) && (entrypoint_device != device))
    {
      error = CubDebug(cudaSetDevice(entrypoint_device));
      if (cudaSuccess != error)
      {
        return error;
      }
    }

    return error;
  }

  /**
   * @brief Frees a live allocation of device memory on the current device, returning it to the
   *        allocator.
   *
   * Once freed, the allocation becomes available immediately for reuse within the @p
   * active_stream with which it was associated with during allocation, and it becomes available
   * for reuse within other streams when all prior work submitted to @p active_stream has
   * completed.
   */
  cudaError_t DeviceFree(void* d_ptr)
  {
    return DeviceFree(INVALID_DEVICE_ORDINAL, d_ptr);
  }

  /**
   * @brief Frees all cached device allocations on all devices
   */
  cudaError_t FreeAllCached()
  {
    cudaError_t error     = cudaSuccess;
    int entrypoint_device = INVALID_DEVICE_ORDINAL;
    int current_device    = INVALID_DEVICE_ORDINAL;

    mutex.lock();

    while (!cached_blocks.empty())
    {
      // Get first block
      CachedBlocks::iterator begin = cached_blocks.begin();

      // Get entry-point device ordinal if necessary
      if (entrypoint_device == INVALID_DEVICE_ORDINAL)
      {
        error = CubDebug(cudaGetDevice(&entrypoint_device));
        if (cudaSuccess != error)
        {
          break;
        }
      }

      // Set current device ordinal if necessary
      if (begin->device != current_device)
      {
        error = CubDebug(cudaSetDevice(begin->device));
        if (cudaSuccess != error)
        {
          break;
        }
        current_device = begin->device;
      }

      // Free device memory
      error = CubDebug(cudaFree(begin->d_ptr));
      if (cudaSuccess != error)
      {
        break;
      }

      error = CubDebug(cudaEventDestroy(begin->ready_event));
      if (cudaSuccess != error)
      {
        break;
      }

      // Reduce balance and erase entry
      const size_t block_bytes = begin->bytes;
      cached_bytes[current_device].free -= block_bytes;
      cached_blocks.erase(begin);

#ifdef CUB_DEBUG_LOG
      _CubLog("\tDevice %d freed %lld bytes.\n\t\t  %lld available blocks cached (%lld bytes), %lld live blocks (%lld "
              "bytes) outstanding.\n",
              current_device,
              (long long) block_bytes,
              (long long) cached_blocks.size(),
              (long long) cached_bytes[current_device].free,
              (long long) live_blocks.size(),
              (long long) cached_bytes[current_device].live);
#endif
    }

    mutex.unlock();

    // Attempt to revert back to entry-point device if necessary
    if (entrypoint_device != INVALID_DEVICE_ORDINAL)
    {
      error = CubDebug(cudaSetDevice(entrypoint_device));
      if (cudaSuccess != error)
      {
        return error;
      }
    }

    return error;
  }

  /**
   * @brief Destructor
   */
  virtual ~CachingDeviceAllocator()
  {
    if (!skip_cleanup)
    {
      FreeAllCached();
    }
  }
};

CUB_NAMESPACE_END