3ce5bff10f
CUB_DISABLE_NAMESPACE_MAGIC requires CUB_WRAPPED_NAMESPACE. CUB_WRAPPED_NAMESPACE=cccl_preload wraps cub into cccl_preload::cub. Source must use cccl_preload::cub::CachingDeviceAllocator.
100 lines
2.6 KiB
Bash
Executable File
100 lines
2.6 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# 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]
|
|
|
|
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"
|
|
|
|
[[ -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 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 \
|
|
; do
|
|
if [[ -f "${candidate}/cuda_runtime_api.h" ]]; then
|
|
CUDA_INC="${candidate}"
|
|
break
|
|
fi
|
|
done
|
|
|
|
# Find CUDA libs
|
|
CUDA_LIB=""
|
|
for candidate in \
|
|
/usr/local/corex/lib64 \
|
|
/usr/local/cuda/lib64 \
|
|
; do
|
|
if [[ -f "${candidate}/libcudart.so" ]]; then
|
|
CUDA_LIB="${candidate}"
|
|
break
|
|
fi
|
|
done
|
|
|
|
echo "[build] CUDA include: ${CUDA_INC:-system}"
|
|
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 this for non-NVCC compilers
|
|
-D__CUDA_ARCH_LIST__=700
|
|
-DCUB_DISABLE_NAMESPACE_MAGIC
|
|
-DCUB_WRAPPED_NAMESPACE=cccl_preload
|
|
)
|
|
|
|
if [[ "${CXX}" == *clang++* ]]; then
|
|
"${CXX}" "${COMMON_FLAGS[@]}" -x c++ -o "${OUT}" "${SRC}" 2>&1
|
|
else
|
|
"${CXX}" "${COMMON_FLAGS[@]}" -x c++ -o "${OUT}" "${SRC}" 2>&1
|
|
fi
|
|
|
|
if [[ -f "${OUT}" ]]; then
|
|
SIZE=$(stat -c%s "${OUT}" 2>/dev/null || echo "?")
|
|
echo ""
|
|
echo "[build] SUCCESS: ${OUT} (${SIZE} bytes)"
|
|
echo ""
|
|
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
|
|
fi
|