diff --git a/qwen3_6_scripts/cccl_preload/build_cccl_preload.sh b/qwen3_6_scripts/cccl_preload/build_cccl_preload.sh index 0d7c53c7..b97fde7c 100755 --- a/qwen3_6_scripts/cccl_preload/build_cccl_preload.sh +++ b/qwen3_6_scripts/cccl_preload/build_cccl_preload.sh @@ -1,41 +1,43 @@ #!/usr/bin/env bash -# Build libcccl_allocator.so — LD_PRELOAD .so for CUB CachingDeviceAllocator +# Build libcccl_allocator.so +# +# Full CCCL dependency chain (288 headers) in ./include/ +# Source: cccl_upstream/cub/cub/util_allocator.cuh + transitive deps # # Usage: # bash build_cccl_preload.sh [output_dir] -# -# On BI-V100 with CoreX SDK: -# bash build_cccl_preload.sh /workspace/qwen3_6_scripts/cccl_preload -# -# The .so intercepts cudaMalloc/cudaFree and routes through CUB's -# caching allocator, bypassing CoreX's "expandable segment not supported" -# ASSERT in CUDACachingAllocator.cpp:545. set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" OUTPUT_DIR="${1:-${SCRIPT_DIR}}" SRC="${SCRIPT_DIR}/cccl_allocator_preload.cu" +INC="${SCRIPT_DIR}/include" OUT="${OUTPUT_DIR}/libcccl_allocator.so" -# Find CoreX clang++ (preferred) or system g++ -if [[ -x /usr/local/corex-3.2.3/bin/clang++ ]]; then - CXX=/usr/local/corex-3.2.3/bin/clang++ - echo "[build] Using CoreX clang++: ${CXX}" -elif [[ -x /usr/local/corex/bin/clang++ ]]; then - CXX=/usr/local/corex/bin/clang++ - echo "[build] Using CoreX clang++ (alt): ${CXX}" -else - CXX=g++ - echo "[build] CoreX clang++ not found, falling back to g++" -fi +[[ -d "${INC}/cub" ]] || { echo "CCCL include tree missing: ${INC}/cub"; exit 2; } +[[ -d "${INC}/cuda" ]] || { echo "CCCL include tree missing: ${INC}/cuda"; exit 2; } -# Find CUDA include path +# Find compiler +CXX="" +for candidate in \ + /usr/local/corex-3.2.3/bin/clang++ \ + /usr/local/corex/bin/clang++ \ + /usr/local/corex/lib64/clang/16/bin/clang++ \ + ; do + if [[ -x "${candidate}" ]]; then + CXX="${candidate}" + break + fi +done +[[ -n "${CXX}" ]] || { CXX=g++; echo "[build] no CoreX clang++, falling back to g++"; } +echo "[build] CXX=${CXX}" + +# Find CUDA headers (for cuda_runtime_api.h) CUDA_INC="" for candidate in \ /usr/local/corex/include \ /usr/local/cuda/include \ - /usr/local/corex/lib64/clang/16/include \ ; do if [[ -f "${candidate}/cuda_runtime_api.h" ]]; then CUDA_INC="${candidate}" @@ -43,7 +45,7 @@ for candidate in \ fi done -# Find CUDA lib path for linking +# Find CUDA libs CUDA_LIB="" for candidate in \ /usr/local/corex/lib64 \ @@ -55,57 +57,43 @@ for candidate in \ fi done -if [[ -z "${CUDA_INC}" ]]; then - echo "[WARN] cuda_runtime_api.h not found — trying compile anyway" -fi - echo "[build] CUDA include: ${CUDA_INC:-system}" -echo "[build] CUDA lib: ${CUDA_LIB:-system}" -echo "[build] Source: ${SRC}" -echo "[build] Output: ${OUT}" +echo "[build] CUDA lib: ${CUDA_LIB:-system}" +echo "[build] CCCL include: ${INC} ($(find "${INC}" -type f | wc -l) files)" +echo "[build] Source: ${SRC}" +echo "[build] Output: ${OUT}" + +COMMON_FLAGS=( + -shared -fPIC -O2 -std=c++17 + -I"${INC}" + ${CUDA_INC:+-I"${CUDA_INC}"} + ${CUDA_LIB:+-L"${CUDA_LIB}"} + -lcudart -ldl + # Suppress CCCL warnings that don't affect correctness + -Wno-unused-function + -Wno-unknown-pragmas + # CUB needs these for non-NVCC compilers + -D_CCCL_COMPILER_GCC=1 + -D__CUDA_ARCH_LIST__=700 + -DCUB_DISABLE_NAMESPACE_MAGIC + -DCUB_WRAPPED_NAMESPACE=cccl_preload +) -# Build as shared library -# -x cuda or -x c++ depending on compiler if [[ "${CXX}" == *clang++* ]]; then - # CoreX clang++ can compile .cu natively - ${CXX} \ - -shared -fPIC \ - -O2 \ - ${CUDA_INC:+-I"${CUDA_INC}"} \ - ${CUDA_LIB:+-L"${CUDA_LIB}"} \ - -lcudart \ - -ldl \ - -std=c++17 \ - -o "${OUT}" \ - "${SRC}" + "${CXX}" "${COMMON_FLAGS[@]}" -x c++ -o "${OUT}" "${SRC}" 2>&1 else - # g++ needs .cu renamed or treated as C++ - # cuda_runtime_api.h should still work with host compiler - ${CXX} \ - -shared -fPIC \ - -O2 \ - ${CUDA_INC:+-I"${CUDA_INC}"} \ - ${CUDA_LIB:+-L"${CUDA_LIB}"} \ - -lcudart \ - -ldl \ - -std=c++17 \ - -x c++ \ - -o "${OUT}" \ - "${SRC}" + "${CXX}" "${COMMON_FLAGS[@]}" -x c++ -o "${OUT}" "${SRC}" 2>&1 fi if [[ -f "${OUT}" ]]; then - SIZE=$(stat -c%s "${OUT}" 2>/dev/null || stat -f%z "${OUT}" 2>/dev/null || echo "?") + SIZE=$(stat -c%s "${OUT}" 2>/dev/null || echo "?") + echo "" echo "[build] SUCCESS: ${OUT} (${SIZE} bytes)" echo "" - echo "Usage:" - echo " LD_PRELOAD=${OUT} CCCL_ALLOC_DEBUG=1 python3 -c 'import torch; t=torch.zeros(1024, device=\"cuda\")'" - echo "" - echo "In computility-run.yaml, add to env:" - echo " - name: LD_PRELOAD" - echo " value: /workspace/qwen3_6_scripts/cccl_preload/libcccl_allocator.so" - echo " - name: PYTORCH_CUDA_ALLOC_CONF" - echo " value: expandable_segments:True" + echo "Test:" + echo " LD_PRELOAD=${OUT} CCCL_ALLOC_DEBUG=1 \\" + echo " PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \\" + echo " python3 verify_preload.py" else echo "[build] FAILED" exit 1 diff --git a/qwen3_6_scripts/cccl_preload/cccl_allocator_preload.cu b/qwen3_6_scripts/cccl_preload/cccl_allocator_preload.cu index 2812d2ae..50061632 100644 --- a/qwen3_6_scripts/cccl_preload/cccl_allocator_preload.cu +++ b/qwen3_6_scripts/cccl_preload/cccl_allocator_preload.cu @@ -1,409 +1,90 @@ /* * cccl_allocator_preload.cu * - * LD_PRELOAD .so that replaces PyTorch's CUDA memory allocator with - * CUB's CachingDeviceAllocator (extracted from CCCL upstream). + * LD_PRELOAD .so — CUB CachingDeviceAllocator from CCCL upstream. + * Full dependency chain (288 files) extracted into include/. * - * Purpose: CoreX's CUDACachingAllocator.cpp:545 asserts - * "expandable segment not supported". Instead of patching libtorch, - * we intercept cudaMalloc/cudaFree at the dynamic linker level and - * route them through CUB's battle-tested caching allocator. + * Intercepts cudaMalloc/cudaFree, routes through CUB's geometric-bin + * caching allocator. Strips expandable_segments from + * PYTORCH_CUDA_ALLOC_CONF before libtorch reads it. * - * Source: cccl_upstream/cub/cub/util_allocator.cuh - * License: BSD-3 (NVIDIA/CUB) - * - * Build (on BI-V100 with CoreX clang++): - * bash build_cccl_preload.sh - * - * Usage: - * LD_PRELOAD=/workspace/qwen3_6_scripts/cccl_preload/libcccl_allocator.so \ - * CCCL_ALLOC_DEBUG=0 \ - * PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512 \ - * python3 -m vllm.entrypoints.openai.api_server ... + * Source: CCCL cub/cub/util_allocator.cuh (BSD-3, NVIDIA) + * Build: bash build_cccl_preload.sh */ -#include +/* ---- CCCL include chain (288 files from cccl_upstream) ---- */ +#include + +/* ---- System ---- */ #include #include #include #include -#include -#include -#include +#include /* ======================================================================== - * CUB CachingDeviceAllocator — extracted from CCCL - * cccl_upstream/cub/cub/util_allocator.cuh + * Configuration for BI-V100 (32GB × 4 cards) * - * All CUB/CCCL macro dependencies replaced with plain C++. + * CUB CachingDeviceAllocator parameters: + * bin_growth = 2 (power-of-2 bins: 256B, 512B, 1KB, ... 4GB) + * min_bin = 8 (2^8 = 256B minimum allocation) + * max_bin = 32 (2^32 = 4GB maximum cached bin) + * max_cached = 8GB per device + * + * More granular bins (growth=2) than CUB default (growth=8) because + * PyTorch tensor sizes vary widely in inference. * ======================================================================== */ -static bool g_cccl_debug = false; +static constexpr unsigned int ALLOC_BIN_GROWTH = 2; +static constexpr unsigned int ALLOC_MIN_BIN = 8; /* 256 bytes */ +static constexpr unsigned int ALLOC_MAX_BIN = 32; /* 4 GB */ +static constexpr size_t ALLOC_MAX_CACHED = (size_t)8 * 1024 * 1024 * 1024; /* 8GB */ -#define CcclDebug(e) (e) -#define CcclLog(...) \ - do { \ - if (g_cccl_debug) { \ - fprintf(stderr, "[cccl_alloc] "); \ - fprintf(stderr, __VA_ARGS__); \ - } \ - } while (0) - -struct CachingDeviceAllocator -{ - static constexpr unsigned int INVALID_BIN = (unsigned int) -1; - static constexpr size_t INVALID_SIZE = (size_t) -1; - static constexpr int INVALID_DEVICE_ORDINAL = -1; - - struct BlockDescriptor - { - void* d_ptr; - size_t bytes; - unsigned int bin; - int device; - cudaStream_t associated_stream; - cudaEvent_t ready_event; - - BlockDescriptor(void* d_ptr_, int device_) - : d_ptr(d_ptr_), bytes(0), bin(INVALID_BIN), device(device_), - associated_stream(nullptr), ready_event(nullptr) {} - - BlockDescriptor(int device_) - : d_ptr(nullptr), bytes(0), bin(INVALID_BIN), device(device_), - associated_stream(nullptr), ready_event(nullptr) {} - - static bool PtrCompare(const BlockDescriptor& a, const BlockDescriptor& b) { - return (a.device == b.device) ? (a.d_ptr < b.d_ptr) : (a.device < b.device); - } - static bool SizeCompare(const BlockDescriptor& a, const BlockDescriptor& b) { - return (a.device == b.device) ? (a.bytes < b.bytes) : (a.device < b.device); - } - }; - - using Compare = bool (*)(const BlockDescriptor&, const BlockDescriptor&); - - struct TotalBytes { size_t free; size_t live; TotalBytes() : free(0), live(0) {} }; - - using CachedBlocks = std::multiset; - using BusyBlocks = std::multiset; - using GpuCachedBytes = std::map; - - static unsigned int IntPow(unsigned int base, unsigned int exp) { - unsigned int retval = 1; - while (exp > 0) { - if (exp & 1) retval *= base; - base *= base; - exp >>= 1; - } - return retval; - } - - void NearestPowerOf(unsigned int& power, size_t& rounded_bytes, - unsigned int base, size_t value) { - power = 0; - rounded_bytes = 1; - if (value * base < value) { - power = sizeof(size_t) * 8; - rounded_bytes = size_t(0) - 1; - return; - } - while (rounded_bytes < value) { - rounded_bytes *= base; - power++; - } - } - - std::mutex mutex; - unsigned int bin_growth; - unsigned int min_bin; - unsigned int max_bin; - size_t min_bin_bytes; - size_t max_bin_bytes; - size_t max_cached_bytes; - bool skip_cleanup; - GpuCachedBytes cached_bytes; - CachedBlocks cached_blocks; - BusyBlocks live_blocks; - - /* - * Constructor tuned for BI-V100 (32GB per card, 4 cards): - * bin_growth=8, min_bin=3 (512B), max_bin=13 (~550MB) - * max_cached_bytes = 4GB per device (reasonable for 32GB card) - * - * This replaces PyTorch's expandable_segments with a proven - * geometric-bin caching strategy from CUB/CCCL. - */ - CachingDeviceAllocator() - : bin_growth(8) - , min_bin(3) /* 8^3 = 512B minimum allocation */ - , max_bin(13) /* 8^13 = ~550MB maximum cached bin */ - , min_bin_bytes(IntPow(8, 3)) - , max_bin_bytes(IntPow(8, 13)) - , max_cached_bytes((size_t)4 * 1024 * 1024 * 1024) /* 4GB per device */ - , skip_cleanup(true) /* CoreX may tear down CUDA before our dtor */ - , cached_blocks(BlockDescriptor::SizeCompare) - , live_blocks(BlockDescriptor::PtrCompare) - { - CcclLog("CachingDeviceAllocator init: bin_growth=%u min_bin=%u " - "max_bin=%u max_cached=%.1fGB\n", - bin_growth, min_bin, max_bin, - (double)max_cached_bytes / (1024.0*1024.0*1024.0)); - } - - /* ---- Real cudaMalloc/cudaFree via dlsym(RTLD_NEXT) ---- */ - using RealMalloc_t = cudaError_t (*)(void**, size_t); - using RealFree_t = cudaError_t (*)(void*); - - static RealMalloc_t get_real_malloc() { - static RealMalloc_t fn = (RealMalloc_t)dlsym(RTLD_NEXT, "cudaMalloc"); - return fn; - } - static RealFree_t get_real_free() { - static RealFree_t fn = (RealFree_t)dlsym(RTLD_NEXT, "cudaFree"); - return fn; - } - - 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 = cudaGetDevice(&entrypoint_device); - if (error != cudaSuccess) return error; - device = entrypoint_device; - } - - 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) { - search_key.bin = INVALID_BIN; - search_key.bytes = bytes; - } else { - mutex.lock(); - if (search_key.bin < min_bin) { - search_key.bin = min_bin; - search_key.bytes = min_bin_bytes; - } - - 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)) - { - bool is_reusable = false; - if (active_stream == block_itr->associated_stream) { - is_reusable = true; - } else { - cudaError_t event_status = cudaEventQuery(block_itr->ready_event); - if (event_status != cudaErrorNotReady) { - is_reusable = true; - } - } - - if (is_reusable) { - found = true; - search_key = *block_itr; - search_key.associated_stream = active_stream; - live_blocks.insert(search_key); - cached_bytes[device].free -= search_key.bytes; - cached_bytes[device].live += search_key.bytes; - - CcclLog("reuse %p (%zu bytes) dev=%d\n", - search_key.d_ptr, search_key.bytes, device); - cached_blocks.erase(block_itr); - break; - } - block_itr++; - } - mutex.unlock(); - } - - if (!found) { - if (device != entrypoint_device) { - if (entrypoint_device == INVALID_DEVICE_ORDINAL) - cudaGetDevice(&entrypoint_device); - cudaSetDevice(device); - } - - /* Use real cudaMalloc, not ourselves */ - error = get_real_malloc()(&search_key.d_ptr, search_key.bytes); - - if (error == cudaErrorMemoryAllocation) { - CcclLog("OOM for %zu bytes on dev=%d, freeing cache...\n", - search_key.bytes, device); - cudaGetLastError(); /* reset */ - - mutex.lock(); - BlockDescriptor free_key(device); - CachedBlocks::iterator block_itr = cached_blocks.lower_bound(free_key); - while ((block_itr != cached_blocks.end()) && - (block_itr->device == device)) - { - error = get_real_free()(block_itr->d_ptr); - if (error != cudaSuccess) break; - cudaEventDestroy(block_itr->ready_event); - cached_bytes[device].free -= block_itr->bytes; - block_itr = cached_blocks.erase(block_itr); - } - mutex.unlock(); - - if (error != cudaSuccess) return error; - error = get_real_malloc()(&search_key.d_ptr, search_key.bytes); - if (error != cudaSuccess) return error; - } else if (error != cudaSuccess) { - return error; - } - - cudaEventCreateWithFlags(&search_key.ready_event, cudaEventDisableTiming); - - mutex.lock(); - live_blocks.insert(search_key); - cached_bytes[device].live += search_key.bytes; - mutex.unlock(); - - CcclLog("alloc %p (%zu bytes, bin=%u) dev=%d\n", - search_key.d_ptr, search_key.bytes, search_key.bin, device); - - if ((entrypoint_device != INVALID_DEVICE_ORDINAL) && - (entrypoint_device != device)) - cudaSetDevice(entrypoint_device); - } - - *d_ptr = search_key.d_ptr; - return cudaSuccess; - } - - cudaError_t DeviceAllocate(void** d_ptr, size_t bytes, - cudaStream_t active_stream = nullptr) { - return DeviceAllocate(INVALID_DEVICE_ORDINAL, d_ptr, bytes, active_stream); - } - - cudaError_t DeviceFree(int device, void* d_ptr) - { - int entrypoint_device = INVALID_DEVICE_ORDINAL; - cudaError_t error = cudaSuccess; - - if (d_ptr == nullptr) return cudaSuccess; - - if (device == INVALID_DEVICE_ORDINAL) { - error = cudaGetDevice(&entrypoint_device); - if (error != cudaSuccess) return error; - device = entrypoint_device; - } - - mutex.lock(); - bool recached = false; - BlockDescriptor search_key(d_ptr, device); - BusyBlocks::iterator block_itr = live_blocks.find(search_key); - - if (block_itr != live_blocks.end()) { - search_key = *block_itr; - live_blocks.erase(block_itr); - cached_bytes[device].live -= search_key.bytes; - - if ((search_key.bin != INVALID_BIN) && - (cached_bytes[device].free + search_key.bytes <= max_cached_bytes)) - { - recached = true; - cached_blocks.insert(search_key); - cached_bytes[device].free += search_key.bytes; - CcclLog("cache %p (%zu bytes) dev=%d\n", - d_ptr, search_key.bytes, device); - } - } - mutex.unlock(); - - if (device != entrypoint_device) { - if (entrypoint_device == INVALID_DEVICE_ORDINAL) - cudaGetDevice(&entrypoint_device); - cudaSetDevice(device); - } - - if (recached) { - cudaEventRecord(search_key.ready_event, search_key.associated_stream); - } else { - /* Not tracked or cache full — real free */ - CcclLog("free %p dev=%d (not cached)\n", d_ptr, device); - error = get_real_free()(d_ptr); - if (block_itr != live_blocks.end()) - cudaEventDestroy(search_key.ready_event); - } - - if ((entrypoint_device != INVALID_DEVICE_ORDINAL) && - (entrypoint_device != device)) - cudaSetDevice(entrypoint_device); - - return error; - } - - cudaError_t DeviceFree(void* d_ptr) { - return DeviceFree(INVALID_DEVICE_ORDINAL, d_ptr); - } - - 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()) { - CachedBlocks::iterator begin = cached_blocks.begin(); - if (entrypoint_device == INVALID_DEVICE_ORDINAL) - cudaGetDevice(&entrypoint_device); - if (begin->device != current_device) { - cudaSetDevice(begin->device); - current_device = begin->device; - } - get_real_free()(begin->d_ptr); - cudaEventDestroy(begin->ready_event); - cached_bytes[current_device].free -= begin->bytes; - cached_blocks.erase(begin); - } - mutex.unlock(); - - if (entrypoint_device != INVALID_DEVICE_ORDINAL) - cudaSetDevice(entrypoint_device); - return error; - } -}; - -/* ======================================================================== - * Global singleton + LD_PRELOAD intercepts - * ======================================================================== */ - -static CachingDeviceAllocator& get_allocator() { - static CachingDeviceAllocator instance; +/* ---- Global allocator singleton ---- */ +static cub::CachingDeviceAllocator& get_allocator() { + static cub::CachingDeviceAllocator instance( + ALLOC_BIN_GROWTH, + ALLOC_MIN_BIN, + ALLOC_MAX_BIN, + ALLOC_MAX_CACHED, + true /* skip_cleanup: CoreX may tear down CUDA before our dtor */ + ); return instance; } static bool g_preload_active = false; +static bool g_debug = false; -/* Called once at .so load time */ +/* ---- Real cudaMalloc/cudaFree via dlsym(RTLD_NEXT) ---- */ +using RealMalloc_t = cudaError_t (*)(void**, size_t); +using RealFree_t = cudaError_t (*)(void*); + +static RealMalloc_t get_real_malloc() { + static RealMalloc_t fn = (RealMalloc_t)dlsym(RTLD_NEXT, "cudaMalloc"); + return fn; +} +static RealFree_t get_real_free() { + static RealFree_t fn = (RealFree_t)dlsym(RTLD_NEXT, "cudaFree"); + return fn; +} + +/* ======================================================================== + * Constructor: runs at LD_PRELOAD load time + * ======================================================================== */ __attribute__((constructor)) static void cccl_preload_init() { const char* debug_env = getenv("CCCL_ALLOC_DEBUG"); - g_cccl_debug = (debug_env && atoi(debug_env) > 0); + g_debug = (debug_env && atoi(debug_env) > 0); const char* disable_env = getenv("CCCL_ALLOC_DISABLE"); if (disable_env && atoi(disable_env) > 0) { fprintf(stderr, "[cccl_alloc] DISABLED by CCCL_ALLOC_DISABLE=1\n"); - g_preload_active = false; return; } - /* Strip expandable_segments from PYTORCH_CUDA_ALLOC_CONF - * so CoreX's allocator doesn't hit the assert. - * We handle the caching ourselves. */ + /* Strip expandable_segments from PYTORCH_CUDA_ALLOC_CONF */ const char* alloc_conf = getenv("PYTORCH_CUDA_ALLOC_CONF"); if (alloc_conf) { - /* Build a new conf string without expandable_segments */ std::string conf(alloc_conf); std::string clean; size_t pos = 0; @@ -411,49 +92,51 @@ static void cccl_preload_init() { size_t comma = conf.find(',', pos); if (comma == std::string::npos) comma = conf.size(); std::string token = conf.substr(pos, comma - pos); - /* Skip expandable_segments:* */ if (token.find("expandable_segments") == std::string::npos) { if (!clean.empty()) clean += ","; clean += token; } pos = comma + 1; } - if (clean.empty()) { + if (clean.empty()) unsetenv("PYTORCH_CUDA_ALLOC_CONF"); - } else { + else setenv("PYTORCH_CUDA_ALLOC_CONF", clean.c_str(), 1); - } - fprintf(stderr, "[cccl_alloc] stripped expandable_segments from " - "PYTORCH_CUDA_ALLOC_CONF: \"%s\" -> \"%s\"\n", + + fprintf(stderr, "[cccl_alloc] PYTORCH_CUDA_ALLOC_CONF: \"%s\" -> \"%s\"\n", alloc_conf, clean.empty() ? "(unset)" : clean.c_str()); } - /* Force-initialize the allocator singleton */ - (void)get_allocator(); + /* Initialize allocator */ + auto& alloc = get_allocator(); + if (g_debug) { + alloc.debug = true; + } + g_preload_active = true; - fprintf(stderr, "[cccl_alloc] LD_PRELOAD active — CUB CachingDeviceAllocator " - "replacing cudaMalloc/cudaFree\n"); + fprintf(stderr, + "[cccl_alloc] LD_PRELOAD active — CUB CachingDeviceAllocator " + "(growth=%u, bins=[%u..%u], max_cached=%.1fGB)\n", + ALLOC_BIN_GROWTH, ALLOC_MIN_BIN, ALLOC_MAX_BIN, + (double)ALLOC_MAX_CACHED / (1024.0*1024.0*1024.0)); } -/* ---- cudaMalloc intercept ---- */ +/* ======================================================================== + * cudaMalloc / cudaFree intercepts + * ======================================================================== */ + extern "C" cudaError_t cudaMalloc(void** devPtr, size_t size) { if (!g_preload_active) { - /* Fallback to real cudaMalloc during init or if disabled */ - static auto real_fn = (CachingDeviceAllocator::RealMalloc_t) - dlsym(RTLD_NEXT, "cudaMalloc"); - return real_fn(devPtr, size); + return get_real_malloc()(devPtr, size); } return get_allocator().DeviceAllocate(devPtr, size); } -/* ---- cudaFree intercept ---- */ extern "C" cudaError_t cudaFree(void* devPtr) { if (!g_preload_active || devPtr == nullptr) { - static auto real_fn = (CachingDeviceAllocator::RealFree_t) - dlsym(RTLD_NEXT, "cudaFree"); - return real_fn(devPtr); + return get_real_free()(devPtr); } return get_allocator().DeviceFree(devPtr); }