#!/bin/bash
# Neuron port test suite. Run after tests/neuron/setup_trn2.sh.
#
#   bash tests/neuron/run_all.sh cpu              # host-only: kernel unit tests + end-to-end vs golden (~1 h)
#   CORES=0-7 bash tests/neuron/run_all.sh device # NeuronCore tests on the given logical cores
#   CORES=0-7 bash tests/neuron/run_all.sh all
#   bash tests/neuron/run_all.sh cpu deep_gemm flash_mla   # only selected stages
#
# Device stages use the first core of CORES for single-core runs and CORES for multi-rank deep_ep
# (world sizes 2/4/8 limited by the core count). They refuse to start if those cores are in use.
# Logs: $WORK/test_logs/<timestamp>/<stage>.log; a PASS/FAIL table is printed at the end.
set -uo pipefail

REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
WORK="${WORK:-$HOME/di}"
MODE="${1:-cpu}"; shift || true
ONLY=("$@")
CORES="${CORES:-0-7}"
LOGDIR="$WORK/test_logs/$(date +%Y%m%d-%H%M%S)"
mkdir -p "$LOGDIR"
# shellcheck disable=SC1091
source "$WORK/venv/bin/activate"
cd "$REPO" || exit 1
VENDOR="$REPO/deepinfer/neuron_ops/vendor"
T="$REPO/tests/neuron"
export DEEPINFER_ACCELERATOR=neuron
# Kernel unit tests compare bit-for-bit against references quantizing to the e4m3fn range (448); the end-to-end
# stages run the default trn2 e4m3 range (224, deepinfer/neuron_ops/fp8.py) and compare against the golden.
UNIT_FP8_MAX="${UNIT_FP8_MAX:-448}"
declare -a SUMMARY=()

first_core() { echo "${CORES%%[-,]*}"; }
core_count() { python3 - "$CORES" <<'EOF'
import sys
n = 0
for part in sys.argv[1].split(","):
    a, _, b = part.partition("-")
    n += (int(b) - int(a) + 1) if b else 1
print(n)
EOF
}
selected() { [ ${#ONLY[@]} -eq 0 ] && return 0; for s in "${ONLY[@]}"; do [ "$s" = "$1" ] && return 0; done; return 1; }

run() {  # run <stage> <command...>
  local stage="$1"; shift
  selected "$stage" || return 0
  echo "=== $stage"
  local t0=$SECONDS
  local fp8_max="$UNIT_FP8_MAX"
  case "$stage" in e2e_*|device_diff) fp8_max="${DEEPINFER_NEURON_FP8_MAX:-224}" ;; esac
  if DEEPINFER_NEURON_FP8_MAX="$fp8_max" "$@" > "$LOGDIR/$stage.log" 2>&1; then r=PASS; else r=FAIL; fi
  SUMMARY+=("$(printf '%-6s %-28s %5ss  %s' "$r" "$stage" $((SECONDS - t0)) "$LOGDIR/$stage.log")")
  echo "    $r ($((SECONDS - t0))s)"
  [ "$r" = FAIL ] && tail -5 "$LOGDIR/$stage.log" | sed 's/^/    | /'
}

cores_free() {  # are the NeuronCores in $CORES free? (other cores may be busy)
  python3 - "$CORES" <<'EOF'
import json, subprocess, sys
want = set()
for part in sys.argv[1].split(","):
    a, _, b = part.partition("-")
    want.update(range(int(a), int(b) + 1) if b else [int(a)])
busy = [d["neuron_device"] for d in json.loads(subprocess.check_output(["neuron-ls", "-j"]))
        if d.get("neuron_processes") and want & set(d["neuroncore_ids"])]
if busy:
    print(f"NeuronCores in {sys.argv[1]} are in use (devices {busy}); not starting device stages")
sys.exit(1 if busy else 0)
EOF
}

cpu_stages() {
  export DEEPINFER_NEURON_DEVICE=cpu
  run deep_gemm      env PYTHONPATH="$VENDOR" python "$T/deep_gemm/test_deep_gemm_neuron.py"
  run deep_gemm_tk   env PYTHONPATH="$VENDOR" python "$T/deep_gemm/test_interop_tile_kernels.py"
  run flash_mla      bash -c "cd $T/flash_mla && PYTHONPATH=$VENDOR python test_flash_mla.py"
  run deep_select    bash -c "cd $T/flash_mla && PYTHONPATH=$VENDOR python test_deep_select.py"
  run tilelang_stub  bash -c "cd $T/flash_mla && PYTHONPATH=$VENDOR python test_tilelang.py"
  run tile_kernels   python -m pytest -q -p no:cacheprovider "$T/tile_kernels/test_tile_kernels.py"
  for ws in 1 2 4; do
    run "deep_ep_ws$ws"     python "$T/deep_ep/test_ep.py" --num-processes $ws --port $((29531 + ws))
    run "engram_ws$ws"      python "$T/deep_ep/test_engram.py" --num-processes $ws --port $((29541 + ws))
    run "bucket_ws$ws"      python "$T/deep_ep/test_bucket.py" --num-processes $ws --port $((29551 + ws))
  done
  run sampling_kat   env PYTHONPATH="$REPO" python "$T/sampling/kat.py"
  run sampling_check env PYTHONPATH="$REPO" python "$T/sampling/check2.py"
  # Upstream interface tests of tilelang_ops (compare the Neuron backend against the *_ref implementations).
  for m in norm store_kv_cache load_paged_kv_cache fused_transform compress state_cache get_kv_index kv_index_cat \
           swa_index candidate_blocks flash_mla_utils gather prepare_v4_decode_metadata random_sampling fused_sample \
           sampler sum_smaller_probs structure_output; do
    run "tl_$m" env PYTHONPATH="$T" python -m pytest -q -p no:cacheprovider --noconftest -p nie_skip \
      "deepinfer/model_executor/layers/tilelang_ops/$m.py" -k "not benchmark and not 129280-4096"
  done
  run e2e_cpu python examples/neuron/run_deepinfer_tiny.py --device cpu --threads 32 --num-kv-blocks 32 \
    --out "$WORK/e2e_cpu" --golden-dir "$WORK/golden" --hf-ckpt "$WORK/tiny_ckpt" --compare
}

device_stages() {
  cores_free || return 0
  local c0 n
  c0="$(first_core)"; n="$(core_count)"
  export DEEPINFER_NEURON_DEVICE=neuron
  run device_primitives env NEURON_RT_VISIBLE_CORES="$c0" python "$T/device/test_device_primitives.py"
  for ws in 2 4 8 16; do
    [ "$ws" -le "$n" ] || continue
    for t in ep engram bucket; do
      run "${t}_neuron_ws$ws" env DEEP_EP_TEST_DEVICE=neuron DEEP_EP_TEST_CORE_BASE="$c0" \
        python "$T/deep_ep/test_$t.py" --num-processes $ws --port $((29600 + 10 * ws + ${#t}))
    done
  done
  run e2e_neuron env NEURON_RT_VISIBLE_CORES="$c0" CUDA_VISIBLE_DEVICES="$c0" MASTER_PORT=29733 \
    python examples/neuron/run_deepinfer_tiny.py --device neuron --num-kv-blocks 32 \
    --out "$WORK/e2e_neuron" --golden-dir "$WORK/golden" --hf-ckpt "$WORK/tiny_ckpt" --compare
  run device_diff env NEURON_RT_VISIBLE_CORES="$c0" CUDA_VISIBLE_DEVICES="$c0" MASTER_PORT=29734 \
    python examples/neuron/device_diff.py --prompts short --num-kv-blocks 32 \
    --out "$WORK/device_diff" --golden-dir "$WORK/golden" --hf-ckpt "$WORK/tiny_ckpt"
}

case "$MODE" in
  cpu) cpu_stages ;;
  device) device_stages ;;
  all) cpu_stages; device_stages ;;
  *) echo "usage: $0 cpu|device|all [stage...]"; exit 2 ;;
esac

echo; echo "===== summary ($LOGDIR)"
printf '%s\n' "${SUMMARY[@]}"
! printf '%s\n' "${SUMMARY[@]}" | grep -q '^FAIL'