#!/bin/bash
# Single-operator perf benchmark (manual debug).
# Edit OP/OP_FILE/INLINE_YAML below (or override via env: OP=... SHAPE_FILE=... ./run_pytest.sh).
# All artifacts go under runs/<op>_<timestamp>/: run.log, shapes.yaml,
# autotune_records/, ttgir/.
set -euo pipefail

OP="${OP:-fused_marlin_moe_mxfp4}"
OP_FILE="${OP_FILE:-fused_marlin_moe}"   # benchmark file is test_<OP_FILE>.py
SHAPE_FILE="${SHAPE_FILE:-}"

# Inline shape yaml (used when SHAPE_FILE is empty); top-level key must be the op name.
read -r -d '' INLINE_YAML <<'YAML' || true
fused_marlin_moe_mxfp4:
  # 4 MoE models x 4 token counts (M=1,16,64,256) = 16 shapes
  shapes:
    # Mixtral (E=8)
    - [1,   8, 4096, 14336,  2]
    - [16,  8, 4096, 14336,  2]
    - [64,  8, 4096, 14336,  2]
    - [256, 8, 4096, 14336,  2]
    # DeepSeek-V3 (E=256, H=7168)
    - [1,   256, 7168, 2048, 8]
    - [16,  256, 7168, 2048, 8]
    - [64,  256, 7168, 2048, 8]
    - [256, 256, 7168, 2048, 8]
    # Qwen3 (E=512)
    - [1,   512, 4096, 1024, 10]
    - [16,  512, 4096, 1024, 10]
    - [64,  512, 4096, 1024, 10]
    - [256, 512, 4096, 1024, 10]
    # DeepSeek-V4-Flash (E=256, H=4096)
    - [1,   256, 4096, 2048, 6]
    - [16,  256, 4096, 2048, 6]
    - [64,  256, 4096, 2048, 6]
    - [256, 256, 4096, 2048, 6]
  shape_desc: "num_tokens, num_experts, hidden_size, intermediate_size, topk"
YAML

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/FlagGems-dev}"
TEST_FILE="$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}"

OUT_DIR="$SCRIPT_DIR/runs/${OP}_$(date +%Y%m%d-%H%M%S)"
LOG_FILE="$OUT_DIR/run.log"
mkdir -p "$OUT_DIR"

export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
PLUGINS=(
  -p _seed_plugin
  -p _shape_inject_plugin
  -p _shape_iter_inject_plugin
  -p _bespoke_shape_plugin
  # -p _mm_cluster_fix_plugin
  -p _autotune_record_plugin
  -p _cudagraph_plugin
  -p _ir_meta_plugin
)
# Autotune record/replay (see _autotune_record_plugin.py):
# - default: record chosen configs to $OUT_DIR/autotune_records/<op>.json
# - REPLAY_FROM=<previous runs/<op>_<ts> dir>: replay that run's recorded
#   configs instead (for A/B runs where both sides must use the same config).
REPLAY_FROM="${REPLAY_FROM:-}"
if [[ -n "$REPLAY_FROM" ]]; then
  AUTOTUNE_ENV="FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR=$REPLAY_FROM/autotune_records"
  [[ -f "$REPLAY_FROM/autotune_records/$OP.json" ]] || \
    echo ">>> warning: $REPLAY_FROM/autotune_records/$OP.json not found; replay will fall back to autotune" >&2
else
  AUTOTUNE_ENV="FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR=$OUT_DIR/autotune_records"
  mkdir -p "$OUT_DIR/autotune_records"
fi
export FLAGGEMS_PERF_CURRENT_OP="$OP"

# Resolve the shape file: use SHAPE_FILE, or write the inline yaml to a temp file.
if [[ -z "$SHAPE_FILE" ]]; then
  SHAPE_FILE="$(mktemp --suffix=.yaml)"
  printf '%s\n' "$INLINE_YAML" > "$SHAPE_FILE"
  trap 'rm -f "$SHAPE_FILE"' EXIT
fi
cp -f "$SHAPE_FILE" "$OUT_DIR/shapes.yaml"  # archive the shape used

export PYTHONUNBUFFERED=1  # unbuffered live output

{
  # _ir_meta_plugin dumps organized ttgir (only variants actually launched
  # outside the autotune sweep) into TTGIR_DUMP_DIR at process exit, so the
  # dump happens even if a CUDA crash kills the run. Layout/legend: see
  # <dump>/naming.md and index.tsv (which also lists unused sweep losers).
  CACHE_DIR="$OUT_DIR/.triton_cache"
  rm -rf "$CACHE_DIR"; mkdir -p "$CACHE_DIR"
  status=0
  TRITON_CACHE_DIR="$CACHE_DIR" \
  FLAGGEMS_PERF_TTGIR_DUMP_DIR="$OUT_DIR/ttgir" \
  env "$AUTOTUNE_ENV" \
  USE_FLAGTUNE=1 python -u -m pytest -s "$TEST_FILE" \
    "${PLUGINS[@]}" \
    --shape_file "$SHAPE_FILE" \
    --level core --mode kernel || status=$?

  rm -rf "$CACHE_DIR"

  if (( status == 0 )); then
    echo ">>> done. outputs in $OUT_DIR"
  else
    echo ">>> FAILED (pytest exit $status). partial outputs in $OUT_DIR"
  fi
  exit "$status"
} 2>&1 | tee "$LOG_FILE"