57a91567b4
The script only hashed .cu/.cuh in crates/ml/. Missed: - build.rs changes (removing --use_fast_math → stale cubins cached) - crates/ml-dqn/src/*.cu (replay buffer, seg tree kernels) - crates/ml-core/src/cuda_autograd/*.cu - crates/ml-ppo/src/cuda_nn/*.cu Now hashes ALL kernel sources + ALL build.rs files. Any change to nvcc flags or kernel source invalidates the entire PTX cache. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
57 lines
2.2 KiB
Bash
Executable File
57 lines
2.2 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# ptx-cache-invalidate.sh — Invalidate PTX cache when CUDA kernels change.
|
|
#
|
|
# The Rust runtime caches compiled PTX in $CARGO_TARGET_DIR/.ptx_cache/ (or
|
|
# /tmp/.ptx_cache/). The cache key includes a SHA-256 of the kernel source,
|
|
# so changed .cu files naturally miss. However, if the CI PVC preserves the
|
|
# cache across runs and the build binary is stale, old PTX can survive.
|
|
#
|
|
# This script computes a content hash of all .cu/.cuh kernel source files and
|
|
# compares it against a stamp file in the cache directory. If the hash differs
|
|
# (any kernel source changed), the entire PTX cache is purged, forcing nvcc
|
|
# recompilation on the next run.
|
|
#
|
|
# Usage: scripts/ptx-cache-invalidate.sh [cache_dir]
|
|
# cache_dir defaults to $CARGO_TARGET_DIR/.ptx_cache, then /tmp/.ptx_cache
|
|
|
|
set -euo pipefail
|
|
|
|
KERNEL_DIRS="crates/ml/src/cuda_pipeline crates/ml-dqn/src crates/ml-core/src/cuda_autograd crates/ml-ppo/src/cuda_nn"
|
|
CACHE_DIR="${1:-${CARGO_TARGET_DIR:+${CARGO_TARGET_DIR}/.ptx_cache}}"
|
|
CACHE_DIR="${CACHE_DIR:-/tmp/.ptx_cache}"
|
|
|
|
# Compute content hash of CUDA sources AND build.rs (nvcc flags).
|
|
# build.rs changes (e.g. removing --use_fast_math) produce different cubins
|
|
# even when .cu source is unchanged — must invalidate cache.
|
|
HASH=$( {
|
|
find $KERNEL_DIRS -type f \( -name '*.cu' -o -name '*.cuh' \) -exec sha256sum {} + 2>/dev/null
|
|
find crates/ml crates/ml-dqn crates/ml-core crates/ml-ppo -maxdepth 1 -name 'build.rs' -exec sha256sum {} + 2>/dev/null
|
|
} | sort | sha256sum | cut -d' ' -f1)
|
|
|
|
STAMP_FILE="${CACHE_DIR}/.kernel_hash"
|
|
|
|
if [ -f "$STAMP_FILE" ]; then
|
|
OLD_HASH=$(cat "$STAMP_FILE")
|
|
if [ "$HASH" = "$OLD_HASH" ]; then
|
|
echo "PTX cache valid — kernel hash unchanged ($HASH)"
|
|
exit 0
|
|
fi
|
|
echo "PTX cache STALE — kernel hash changed"
|
|
echo " old: $OLD_HASH"
|
|
echo " new: $HASH"
|
|
else
|
|
echo "PTX cache stamp missing — first run or cache cleared"
|
|
fi
|
|
|
|
# Purge stale PTX cache
|
|
if [ -d "$CACHE_DIR" ]; then
|
|
PTX_COUNT=$(find "$CACHE_DIR" -name '*.ptx' -type f 2>/dev/null | wc -l)
|
|
echo "Purging ${PTX_COUNT} cached PTX files from ${CACHE_DIR}"
|
|
rm -f "${CACHE_DIR}"/*.ptx
|
|
fi
|
|
|
|
# Write new stamp
|
|
mkdir -p "$CACHE_DIR"
|
|
echo "$HASH" > "$STAMP_FILE"
|
|
echo "PTX cache stamp updated: $HASH"
|