[INFRA] Import NVIDIA/CCCL upstream as optimization reference library
CCCL (CUDA C++ Core Libraries) provides: - CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk) - Thrust: high-level parallel algorithms (transform_reduce, sort, scan) - libcudacxx: CUDA C++ standard library (atomics, barriers, memory) - cudax: experimental features (memory resources, allocators) - Tuning policies: per-SM hardware-specific algorithm parameters Competition optimization vectors mapped to CCCL: - Output TPS (83% weight): warp_reduce, block_reduce, device_topk - Input TPS (14% weight): device_scan, block_load, prefetch - Cache TPS (3% weight): prefix caching strategy patterns - Memory (0.9 util): pooled/cached/buddy allocators Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only) License: Apache-2.0
This commit is contained in:
EngineX CI
901
cccl_upstream/cub/cub/util_allocator.cuh
Normal file
901
cccl_upstream/cub/cub/util_allocator.cuh
Normal file
@@ -0,0 +1,901 @@
|
||||
// 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
|
||||
Reference in New Issue
Block a user