#!/usr/bin/env bash
# scripts/pcie-probe.sh — is the decode bottleneck the LINK, or the engine's read path?
#
# Instruments live in scripts/probes/: hip_smoke.hip answers "is the host ROCm usable from a plain
# process", pcie_bw.hip measures host→device bandwidth at the engine's real expert-slab size, and
# pcie_probe.hip sweeps sizes so a stall can be attributed to the link rather than to the engine.
# Written during the RDNA4 expert-offload campaign, kept because the same cards are in Dankserver.
#
# LAW — never benchmark beside a live service. The guard below runs BEFORE anything is built or run,
# and it stops nothing: putting a production endpoint down is not this tool's decision to make. When
# you have deliberately decided that contention is acceptable, say so with --force.
#
# Trap carried in the sources: hipGetDeviceProperties SEGFAULTS on this ROCm build even though
# hipGetDeviceCount works fine and reports every device. Nothing here calls it. Keep it that way.
#
# usage: pcie-probe.sh [--device N]... [--build-only] [--force]
#
#   --device N    measure this device; repeat for each card. Without it the probes run once with the
#                 environment's own device selection — one shared run cannot attribute bandwidth.
#   --build-only  compile and stop; never touches the cards.
#   --force       measure even though the guard says the GPU is busy.
#
# Environment (defaults in parentheses) — the sysfs and kfd paths are overridable so the guard can
# be tested on a machine with no cards, which is the only honest way to test a refusal:
#   PCIE_PROBE_SRC     probe sources    (<repo>/scripts/probes)
#   PCIE_PROBE_BUILD   build output     ($HOME/.cache/pcie-probe)
#   PCIE_PROBE_SYSFS   drm sysfs root   (/sys/class/drm)
#   PCIE_PROBE_KFD     KFD device node  (/dev/kfd)
#   PCIE_PROBE_HIPCC   compiler         (hipcc)
#   PCIE_PROBE_MAX_BUSY  refusal threshold, percent (5)
#
# Exit status: 0 measured · 2 usage or missing sources · 3 REFUSED, GPU busy · 4 no toolchain or
# build failed.

set -uo pipefail

SRC="${PCIE_PROBE_SRC:-$(dirname "$(realpath "$0")")/probes}"
BUILD="${PCIE_PROBE_BUILD:-$HOME/.cache/pcie-probe}"
SYSFS="${PCIE_PROBE_SYSFS:-/sys/class/drm}"
KFD="${PCIE_PROBE_KFD:-/dev/kfd}"
HIPCC="${PCIE_PROBE_HIPCC:-hipcc}"
MAX_BUSY="${PCIE_PROBE_MAX_BUSY:-5}"

force=0
build_only=0
devices=()
while [[ $# -gt 0 ]]; do
  case "$1" in
    --device)     [[ $# -ge 2 ]] || { echo "pcie-probe: --device needs an argument" >&2; exit 2; }
                  devices+=("$2"); shift 2 ;;
    --build-only) build_only=1; shift ;;
    --force)      force=1; shift ;;
    *)            echo "usage: pcie-probe.sh [--device N]... [--build-only] [--force]" >&2; exit 2 ;;
  esac
done

mapfile -t sources < <(find "$SRC" -maxdepth 1 -name '*.hip' 2>/dev/null | sort)
if (( ${#sources[@]} == 0 )); then
  echo "pcie-probe: no .hip probe sources in $SRC" >&2
  exit 2
fi

# ---- the guard, before anything touches the card --------------------------------
if (( ! force )); then
  busy=""
  if [[ -e "$KFD" ]] && command -v fuser > /dev/null 2>&1; then
    holders="$(fuser "$KFD" 2>/dev/null | tr -s '[:space:]' ' ')"
    [[ -n "${holders// /}" ]] && busy="KFD in use by:${holders}"
  fi
  for card in "$SYSFS"/card*; do
    [[ -f "$card/device/gpu_busy_percent" ]] || continue
    pct="$(cat "$card/device/gpu_busy_percent" 2>/dev/null)"
    [[ "$pct" =~ ^[0-9]+$ ]] || continue
    (( pct > MAX_BUSY )) && busy="${busy:+$busy; }$(basename "$card") at ${pct}%"
  done
  if [[ -n "$busy" ]]; then
    printf 'pcie-probe: REFUSED — GPU busy (%s)\n' "$busy" >&2
    printf '           never benchmark beside a live service. Take the service down deliberately,\n' >&2
    printf '           or pass --force if you have decided the contention is acceptable.\n' >&2
    exit 3
  fi
fi

# ---- build ----------------------------------------------------------------------
if ! command -v "$HIPCC" > /dev/null 2>&1; then
  printf 'pcie-probe: no HIP toolchain (%s) — run this on the box with the cards, not here\n' "$HIPCC" >&2
  exit 4
fi
mkdir -p "$BUILD"
for src in "${sources[@]}"; do
  bin="$BUILD/$(basename "${src%.hip}")"
  if ! "$HIPCC" -O2 -o "$bin" "$src" 2>&1 | tail -3; then
    printf 'pcie-probe: build failed for %s\n' "$(basename "$src")" >&2
    exit 4
  fi
done
if (( build_only )); then
  printf 'pcie-probe: built %s probe(s) into %s\n' "${#sources[@]}" "$BUILD"
  exit 0
fi

# ---- topology: the link the numbers must be read against ------------------------
printf 'pcie-probe host=%s hipcc=%s probes=%s\n' "$(hostname)" "$HIPCC" "${#sources[@]}"
for card in "$SYSFS"/card*; do
  [[ -f "$card/device/current_link_width" ]] || continue
  printf '  link %-9s current x%s @ %s   max x%s @ %s\n' "$(basename "$card")" \
    "$(cat "$card/device/current_link_width")" \
    "$(cat "$card/device/current_link_speed" 2>/dev/null || echo '?')" \
    "$(cat "$card/device/max_link_width" 2>/dev/null || echo '?')" \
    "$(cat "$card/device/max_link_speed" 2>/dev/null || echo '?')"
done

# ---- measure, per device --------------------------------------------------------
(( ${#devices[@]} > 0 )) || devices=(default)
for d in "${devices[@]}"; do
  if [[ $d == default ]]; then
    printf '\n########## engine device selection ##########\n'
  else
    printf '\n########## device %s ##########\n' "$d"
  fi
  for name in hip_smoke pcie_bw pcie_probe; do
    bin="$BUILD/$name"
    [[ -x "$bin" ]] || continue
    if [[ $d == default ]]; then "$bin"; else HIP_VISIBLE_DEVICES="$d" "$bin" "$d"; fi
  done
done