#!/usr/bin/env bash
set -euo pipefail

root=/home/daniel/aeon-spark-test/h3/h3-blackwell-runtime
builds=/home/daniel/aeon-spark-test/h3/sageattention-p0-builds
output=/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks
extension=_qattn_sm89.cpython-312-aarch64-linux-gnu.so
variants=(baseline cap240 cap224 cap208 cap192 cap176 cap170 cap168 scopes pad-v-smem)

for variant in "${variants[@]}"; do
  report="$output/gb10-sage2-p0-$variant-counters.ncu-rep"
  if [[ -f "$report" ]]; then
    continue
  fi
  docker compose -f "$root/compose.spark.yml" run --rm --no-deps \
    --cap-add SYS_ADMIN --cap-add SYS_PTRACE \
    -v "$root/tools/profile_sage2_scheduler.py:/opt/h3-blackwell-runtime/tools/profile_sage2_scheduler.py:ro" \
    -v "$builds/$variant/lib/sageattention/$extension:/opt/venv/lib/python3.12/site-packages/sageattention/$extension:ro" \
    h3-blackwell-runtime \
    /opt/nvidia/nsight-compute/2025.3.1/ncu \
    --target-processes all \
    --profile-from-start off \
    --section LaunchStats \
    --section Occupancy \
    --section SchedulerStats \
    --section WarpStateStats \
    --section MemoryWorkloadAnalysis \
    --section SpeedOfLight \
    --kernel-name regex:qk_int_sv_f8_attn_kernel \
    --launch-count 1 \
    --force-overwrite \
    -o "/output/h3-blackwell-runtime/benchmarks/gb10-sage2-p0-$variant-counters" \
    python tools/profile_sage2_scheduler.py \
    --output "/output/h3-blackwell-runtime/benchmarks/gb10-sage2-p0-$variant-capture.json" \
    --cuda-profiler-capture \
    --skip-tail-study
done