#!/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
#
# 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"
SRC="$HERE/vector_add.cu"
BIN="$RUN/vector_add"

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() {
	[ -x "$BIN" ] || build_host
	local b64 chunk off total size
	b64=$(base64 -w0 "$BIN")
	total=${#b64}
	off=0
	size=70000
	echo "[*] pushing $(stat -c%s "$BIN") bytes to guest"
	"$VMEXEC" exec "rm -f /tmp/vadd.b64" >/dev/null 2>&1
	while [ $off -lt $total ]; do
		chunk=${b64:$off:$size}
		"$VMEXEC" exec "echo -n '$chunk' >> /tmp/vadd.b64" >/dev/null 2>&1
		off=$((off + size))
	done
	"$VMEXEC" exec "base64 -d /tmp/vadd.b64 > /tmp/vector_add && chmod +x /tmp/vector_add && rm -f /tmp/vadd.b64 && echo PUSH_OK"
}

run_guest() {
	push_bin
	echo "[*] guest run (emulated vGPU)"
	"$VMEXEC" exec "/tmp/vector_add" 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 ;;
*)       echo "usage: $0 [host|guest|compare|all]"; exit 1 ;;
esac