#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
#
# ab.sh - CUDA A/B harness: run the same workload on the host (bare metal,
# real Turing GPU) and in the guest (Open-IOV emulated vGPU), then compare the
# machine-readable results. Mirrors the nvidia-smi host/guest comparison used
# earlier: the host run is the oracle.
#
# Usage:
#   tools/cuda-ab/ab.sh host     # build + run on the host, save run/host.txt
#   tools/cuda-ab/ab.sh guest    # push the binary + run in the running VM
#   tools/cuda-ab/ab.sh compare  # diff the stable result lines
#   tools/cuda-ab/ab.sh all      # host + guest + compare
#   tools/cuda-ab/ab.sh push     # push $WORKLOAD (default vector_add) to the VM
#   tools/cuda-ab/ab.sh victory  # push victory_mail to the VM (runner executes it)
#   tools/cuda-ab/ab.sh smiley   # push smiley (== WORKLOAD=smiley ./ab.sh push)
#
# WORKLOAD selects the program (<name>.cu -> run/<name>); guest execution uses
# AB_GUEST_TIMEOUT seconds (default 300) because the first emulated CUDA call
# can be much slower than a bare-metal launch.
#
# The binary links the CUDA runtime statically, so the guest needs no CUDA
# toolkit - only the guest driver's libcuda.so.
set -e

HERE="$(cd "$(dirname "$0")" && pwd)"
REPO="$(cd "$HERE/../.." && pwd)"
VMEXEC="$REPO/scripts/vgpu_vm.py"
RUN="$HERE/run"
WORKLOAD="${WORKLOAD:-vector_add}"
SRC="$HERE/$WORKLOAD.cu"
BIN="$RUN/$WORKLOAD"
AB_GUEST_TIMEOUT="${AB_GUEST_TIMEOUT:-300}"

mkdir -p "$RUN"

build_host() {
	command -v nvcc >/dev/null || { echo "nvcc not found (install nvidia-cuda-toolkit)"; exit 1; }
	nvcc -O2 --cudart=static -o "$BIN" "$SRC"
}

run_host() {
	build_host
	echo "[*] host run (oracle)"
	"$BIN" | tee "$RUN/host.txt"
}

push_bin() {
	[ -f "$SRC" ] || { echo "no source for WORKLOAD=$WORKLOAD ($SRC)"; exit 1; }
	# Rebuild when missing or older than the source, so a pushed binary is
	# never stale relative to its .cu.
	if [ ! -x "$BIN" ] || [ "$SRC" -nt "$BIN" ]; then
		build_host
	fi
	local b64 chunk off total size out
	b64=$(base64 -w0 "$BIN")
	total=${#b64}
	off=0
	size=70000
	echo "[*] pushing $(stat -c%s "$BIN") bytes to guest"
	"$VMEXEC" exec "rm -f /tmp/$WORKLOAD.b64" >/dev/null 2>&1
	while [ $off -lt $total ]; do
		chunk=${b64:$off:$size}
		"$VMEXEC" exec "echo -n '$chunk' >> /tmp/$WORKLOAD.b64" >/dev/null 2>&1
		off=$((off + size))
	done
	out=$("$VMEXEC" exec "base64 -d /tmp/$WORKLOAD.b64 > /tmp/$WORKLOAD && chmod +x /tmp/$WORKLOAD && rm -f /tmp/$WORKLOAD.b64 && echo PUSH_OK" 2>&1) || true
	printf '%s\n' "$out"
	case "$out" in
	*PUSH_OK*) ;;
	*) echo "push failed (no PUSH_OK)" >&2; exit 1 ;;
	esac
}

run_guest() {
	push_bin
	echo "[*] guest run (emulated vGPU)"
	"$VMEXEC" exec --timeout "$AB_GUEST_TIMEOUT" "/tmp/$WORKLOAD" 2>&1 | tee "$RUN/guest.txt"
}

compare() {
	local h g
	[ -f "$RUN/host.txt" ] || { echo "missing $RUN/host.txt (run host first)"; exit 1; }
	[ -f "$RUN/guest.txt" ] || { echo "missing $RUN/guest.txt (run guest first)"; exit 1; }
	h=$(grep -E '^(DEVICE|RESULT|STATUS)' "$RUN/host.txt")
	g=$(grep -E '^(DEVICE|RESULT|STATUS)' "$RUN/guest.txt")
	echo "--- host:";  echo "$h"
	echo "--- guest:"; echo "$g"
	# DEVICE lines carry the GPU name (differs by design); compare the rest.
	h_norm=$(echo "$h" | sed 's/^DEVICE .*/DEVICE <gpu>/')
	g_norm=$(echo "$g" | sed 's/^DEVICE .*/DEVICE <gpu>/')
	if [ "$h_norm" = "$g_norm" ]; then
		echo "A/B: MATCH"
	else
		echo "A/B: MISMATCH"
		diff <(echo "$h_norm") <(echo "$g_norm") || true
		exit 1
	fi
}

case "${1:-all}" in
host)    run_host ;;
guest)   run_guest ;;
compare) compare ;;
all)     run_host; run_guest; compare ;;
push)    push_bin ;;
victory)
	WORKLOAD=victory_mail
	SRC="$HERE/$WORKLOAD.cu"
	BIN="$RUN/$WORKLOAD"
	push_bin ;;
smiley)
	WORKLOAD=smiley
	SRC="$HERE/$WORKLOAD.cu"
	BIN="$RUN/$WORKLOAD"
	push_bin ;;
*)       echo "usage: $0 [host|guest|compare|all|push|victory|smiley]"; exit 1 ;;
esac