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

GPU=0
TARGET_GIB=60
SOAK_MIN=10
CHUNK_MIB=256
PASSES=3
MODE="standard"

usage() {
  cat <<'USAGE'
DaShen GPU Lab · CMP 170HX 8G→64G 现场体质筛选（不改驱动、不改显存 geometry）

用法：
  sudo bash 170hx_40g_qualify.sh [选项]

选项：
  --gpu N          测试第 N 张 GPU，默认 0
  --quick          快速筛选：58GiB / 3min / 2 passes
  --full           深度筛选：61GiB / 20min / 5 passes
  --target-gib N   指定 GPU 显存测试目标 GiB
  --soak-min N     指定热态循环分钟数
  --passes N       指定最少完整 pattern passes
  --chunk-mib N    单块大小，默认 256MiB
  -h, --help       显示帮助

说明：
  1) 只测试当前驱动暴露出来的 64GB 区域，用于确认 8GB/20C2→64GB 当前解锁体质与热态稳定性。
  2) 测试全程不修改驱动、VBIOS、频率、电压、功耗上限。
  3) 显存校验全部在 GPU 上进行，不需要 64GB 主机内存。
USAGE
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --gpu) GPU="$2"; shift 2 ;;
    --quick) MODE="quick"; TARGET_GIB=58; SOAK_MIN=3; PASSES=2; shift ;;
    --full) MODE="full"; TARGET_GIB=61; SOAK_MIN=20; PASSES=5; shift ;;
    --target-gib) TARGET_GIB="$2"; shift 2 ;;
    --soak-min) SOAK_MIN="$2"; shift 2 ;;
    --passes) PASSES="$2"; shift 2 ;;
    --chunk-mib) CHUNK_MIB="$2"; shift 2 ;;
    -h|--help) usage; exit 0 ;;
    *) echo "未知参数: $1" >&2; usage; exit 2 ;;
  esac
done

for cmd in nvidia-smi lspci python3; do
  command -v "$cmd" >/dev/null 2>&1 || { echo "缺少命令: $cmd" >&2; exit 2; }
done

if ! python3 - <<'PY' >/dev/null 2>&1
import torch
assert torch.cuda.is_available()
PY
then
  echo "需要可用的 PyTorch + CUDA 环境（python3 import torch 且 torch.cuda.is_available() == True）。" >&2
  exit 2
fi

TS="$(date +%Y%m%d_%H%M%S)"
LOGDIR="${PWD}/170hx_qualify_${TS}_gpu${GPU}"
mkdir -p "$LOGDIR"

START_ISO="$(date --iso-8601=seconds 2>/dev/null || date '+%Y-%m-%d %H:%M:%S')"
START_EPOCH="$(date +%s)"

cleanup() {
  if [[ -n "${MON_PID:-}" ]] && kill -0 "$MON_PID" 2>/dev/null; then
    kill "$MON_PID" 2>/dev/null || true
    wait "$MON_PID" 2>/dev/null || true
  fi
}
trap cleanup EXIT

echo "=== DaShen GPU Lab · CMP 170HX qualification ===" | tee "$LOGDIR/summary.txt"
echo "mode=$MODE gpu=$GPU target=${TARGET_GIB}GiB soak=${SOAK_MIN}min passes=$PASSES chunk=${CHUNK_MIB}MiB" | tee -a "$LOGDIR/summary.txt"
echo "start=$START_ISO" | tee -a "$LOGDIR/summary.txt"

# --- 基础识别 ---
BUS_RAW="$(nvidia-smi -i "$GPU" --query-gpu=pci.bus_id --format=csv,noheader | xargs)"
BUS="$(echo "$BUS_RAW" | sed -E 's/^00000000:/0000:/')"
NAME="$(nvidia-smi -i "$GPU" --query-gpu=name --format=csv,noheader | xargs)"
MEM_TOTAL_MIB="$(nvidia-smi -i "$GPU" --query-gpu=memory.total --format=csv,noheader,nounits | xargs | cut -d. -f1)"
UUID="$(nvidia-smi -i "$GPU" --query-gpu=uuid --format=csv,noheader | xargs)"
VBIOS="$(nvidia-smi -i "$GPU" --query-gpu=vbios_version --format=csv,noheader | xargs || true)"
PCI_LINE="$(lspci -nn -s "$BUS" || true)"

{
  echo "name=$NAME"
  echo "uuid=$UUID"
  echo "vbios=$VBIOS"
  echo "pci_bus=$BUS"
  echo "memory_total_mib=$MEM_TOTAL_MIB"
  echo "lspci=$PCI_LINE"
} | tee -a "$LOGDIR/summary.txt"

DEVICE_OK=0
if grep -qi '10de:20c2' <<<"$PCI_LINE"; then
  DEVICE_OK=1
  echo "[PASS] PCI ID = 10de:20c2（8GB SKU）" | tee -a "$LOGDIR/summary.txt"
elif grep -qi '10de:2082' <<<"$PCI_LINE"; then
  echo "[INFO] PCI ID = 10de:2082（10GB SKU，不是本次要筛的8GB/20C2）" | tee -a "$LOGDIR/summary.txt"
else
  echo "[WARN] 未识别为 20c2/2082，请人工核对。" | tee -a "$LOGDIR/summary.txt"
fi

MEM_OK=0
if [[ "$MEM_TOTAL_MIB" -ge 63000 && "$MEM_TOTAL_MIB" -le 67000 ]]; then
  MEM_OK=1
  echo "[PASS] 当前显存约64GB，符合8G->64G profile" | tee -a "$LOGDIR/summary.txt"
else
  echo "[WARN] 当前显存不是约64GB：${MEM_TOTAL_MIB} MiB" | tee -a "$LOGDIR/summary.txt"
fi

lspci -nn -vv -s "$BUS" > "$LOGDIR/lspci_full.txt" 2>&1 || true
grep -E 'LnkCap:|LnkSta:' "$LOGDIR/lspci_full.txt" | tee "$LOGDIR/pcie_link.txt" || true
nvidia-smi -q -i "$GPU" > "$LOGDIR/nvidia_smi_before.txt" 2>&1 || true

# --- 内核错误基线 ---
DMESG_OK=1
if dmesg > "$LOGDIR/dmesg_before.txt" 2>/dev/null; then
  XID_BEFORE="$(grep -Eic 'NVRM: Xid|Xid \(' "$LOGDIR/dmesg_before.txt" || true)"
else
  DMESG_OK=0
  XID_BEFORE=0
  echo "[WARN] 无权读取 dmesg；建议 sudo 运行，否则无法可靠判断新增 Xid。" | tee -a "$LOGDIR/summary.txt"
fi

# --- 遥测 ---
FIELDS="timestamp,index,pstate,temperature.gpu,power.draw,clocks.sm,clocks.mem,utilization.gpu,memory.used,memory.total"
if nvidia-smi --help-query-gpu 2>/dev/null | grep -q 'temperature.memory'; then
  FIELDS="${FIELDS},temperature.memory"
fi
(
  nvidia-smi -i "$GPU" --query-gpu="$FIELDS" --format=csv -l 1
) > "$LOGDIR/telemetry.csv" 2>&1 &
MON_PID=$!

# --- GPU 显存 pattern + alias + 热态循环 ---
export QUAL_GPU="$GPU"
export QUAL_TARGET_GIB="$TARGET_GIB"
export QUAL_SOAK_MIN="$SOAK_MIN"
export QUAL_PASSES="$PASSES"
export QUAL_CHUNK_MIB="$CHUNK_MIB"

set +e
python3 - <<'PY' 2>&1 | tee "$LOGDIR/vram_test.log"
import os, sys, time, math
import torch

GPU = int(os.environ['QUAL_GPU'])
TARGET_GIB = float(os.environ['QUAL_TARGET_GIB'])
SOAK_MIN = float(os.environ['QUAL_SOAK_MIN'])
MIN_PASSES = int(os.environ['QUAL_PASSES'])
CHUNK_MIB = int(os.environ['QUAL_CHUNK_MIB'])

torch.cuda.set_device(GPU)
dev = torch.device(f'cuda:{GPU}')
props = torch.cuda.get_device_properties(GPU)
free_b, total_b = torch.cuda.mem_get_info(GPU)
GiB = 1024**3
MiB = 1024**2

# 为比较操作、CUDA context 和框架开销保留空间。
reserve_b = int(1.5 * GiB)
target_b = int(TARGET_GIB * GiB)
max_safe = max(0, free_b - reserve_b)
if target_b > max_safe:
    print(f'[FAIL] 目标 {TARGET_GIB:.2f} GiB，但当前可安全分配约 {max_safe/GiB:.2f} GiB（已预留1.5GiB）。')
    sys.exit(3)

chunk_b = CHUNK_MIB * MiB
n_int32 = chunk_b // 4
num_full = target_b // chunk_b
remainder = target_b % chunk_b
sizes = [n_int32] * int(num_full)
if remainder >= 16 * MiB:
    sizes.append(remainder // 4)

print(f'[INFO] GPU: {props.name}')
print(f'[INFO] CUDA visible total={total_b/GiB:.2f} GiB free={free_b/GiB:.2f} GiB')
print(f'[INFO] target={sum(s*4 for s in sizes)/GiB:.2f} GiB chunks={len(sizes)} chunk={CHUNK_MIB} MiB')
print('[INFO] allocating GPU-only buffers ...')

bufs=[]
try:
    for i,n in enumerate(sizes):
        bufs.append(torch.empty(n, dtype=torch.int32, device=dev))
        if (i+1) % 16 == 0 or i+1 == len(sizes):
            alloc = sum(x.numel()*x.element_size() for x in bufs)/GiB
            print(f'  allocated {alloc:.2f} GiB')
except Exception as e:
    print(f'[FAIL] cuda allocation failed: {e}')
    sys.exit(4)

torch.cuda.synchronize()

def pat(pass_id, idx):
    # 每块/每轮不同的32-bit pattern，能抓 bit flip，也能抓 chunk alias/fold。
    u = ((0x9E3779B9 * (idx + 1)) ^ (0x85EBCA6B * (pass_id + 1)) ^ 0xA5A5A5A5) & 0xffffffff
    return u - 0x100000000 if u >= 0x80000000 else u

def sweep(pass_id):
    t0=time.perf_counter()
    # 先全部写，再全部读；如果两个 allocation alias，后写会破坏先写，读回时会被抓到。
    patterns=[]
    for i,b in enumerate(bufs):
        p=pat(pass_id,i)
        patterns.append(p)
        b.fill_(p)
    torch.cuda.synchronize()

    errors=0
    first=[]
    order=range(len(bufs)) if pass_id % 2 == 0 else reversed(range(len(bufs)))
    for i in order:
        b=bufs[i]; p=patterns[i]
        bad=int(torch.count_nonzero(b != p).item())
        if bad:
            errors += bad
            if len(first) < 8:
                first.append((i,bad,p))
    torch.cuda.synchronize()
    dt=time.perf_counter()-t0
    touched=sum(b.numel()*b.element_size() for b in bufs)
    # 每轮至少一次完整写+一次完整读；这里只作为相对记录，不当成纯HBM带宽benchmark。
    approx=(2*touched/dt)/GiB if dt > 0 else 0
    print(f'[SWEEP] pass={pass_id:04d} errors={errors} elapsed={dt:.3f}s approx_rw={approx:.1f} GiB/s')
    if first:
        print('[ERROR] first_bad_chunks=', first)
    return errors

start=time.monotonic()
pass_id=0
total_errors=0
# 至少跑 MIN_PASSES，同时持续到热态 soak 时间满足。
while pass_id < MIN_PASSES or (time.monotonic()-start) < SOAK_MIN*60:
    e=sweep(pass_id)
    total_errors += e
    if e:
        print(f'[FAIL] silent/data corruption detected, cumulative_errors={total_errors}')
        sys.exit(10)
    pass_id += 1

print(f'[PASS] VRAM pattern/alias test clean: passes={pass_id}, tested={sum(s*4 for s in sizes)/GiB:.2f} GiB, errors=0')
sys.exit(0)
PY
VRAM_RC=${PIPESTATUS[0]}
set -e

# --- 收尾 ---
cleanup
MON_PID=""
nvidia-smi -q -i "$GPU" > "$LOGDIR/nvidia_smi_after.txt" 2>&1 || true

XID_NEW=-1
if [[ "$DMESG_OK" -eq 1 ]] && dmesg > "$LOGDIR/dmesg_after.txt" 2>/dev/null; then
  XID_AFTER="$(grep -Eic 'NVRM: Xid|Xid \(' "$LOGDIR/dmesg_after.txt" || true)"
  XID_NEW=$((XID_AFTER-XID_BEFORE))
  if [[ "$XID_NEW" -lt 0 ]]; then XID_NEW=0; fi
  grep -Ei 'NVRM|Xid|RmInitAdapter|GSP|MMU|illegal' "$LOGDIR/dmesg_after.txt" | tail -n 200 > "$LOGDIR/kernel_gpu_events_tail.txt" || true
fi

END_EPOCH="$(date +%s)"
ELAPSED=$((END_EPOCH-START_EPOCH))

{
  echo ""
  echo "=== VERDICT ==="
  echo "device_20c2=$DEVICE_OK"
  echo "memory_64g_profile=$MEM_OK"
  echo "vram_test_rc=$VRAM_RC"
  echo "new_xid=$XID_NEW"
  echo "elapsed_seconds=$ELAPSED"

  if [[ "$DEVICE_OK" -eq 1 && "$MEM_OK" -eq 1 && "$VRAM_RC" -eq 0 && ( "$XID_NEW" -eq 0 || "$XID_NEW" -eq -1 ) ]]; then
    echo "RESULT=PASS_64G_CANDIDATE"
    echo "解释：这是一张当前64GB profile下值得保留的8G/20C2候选卡。"
  else
    echo "RESULT=REVIEW_OR_FAIL"
    echo "解释：至少一项未通过；查看 vram_test.log、kernel_gpu_events_tail.txt、nvidia_smi_after.txt。"
  fi
} | tee -a "$LOGDIR/summary.txt"

echo ""
echo "报告目录：$LOGDIR"
echo "重点文件："
echo "  summary.txt"
echo "  vram_test.log"
echo "  telemetry.csv"
echo "  pcie_link.txt"
echo "  nvidia_smi_before.txt / nvidia_smi_after.txt"
echo "  kernel_gpu_events_tail.txt"

if [[ "$DEVICE_OK" -ne 1 || "$MEM_OK" -ne 1 ]]; then
  exit 20
fi
exit "$VRAM_RC"