#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# CUDA Workload Validation
# Tests various CUDA workloads in the guest

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
OPEN_IOV_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"

echo "=== CUDA Workload Validation ==="
echo "Open-IOV root: $OPEN_IOV_ROOT"

# Check if we're root
if [ "$EUID" -ne 0 ]; then
    echo "ERROR: This script must be run as root"
    exit 1
fi

# This script assumes guest is already running with SSH on port 2222
# and NVIDIA driver is installed

GUEST_SSH="ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null root@localhost -p 2222"

echo "Testing CUDA samples..."
$GUEST_SSH "cd /usr/local/cuda/samples && find . -name 'Makefile' -execdir make \\; 2>&1 | tail -20"

echo "Running vectorAdd..."
$GUEST_SSH "cd /usr/local/cuda/samples/1_Utilities/vectorAdd && ./vectorAdd"

echo "Running matrixMul..."
$GUEST_SSH "cd /usr/local/cuda/samples/0_Simple/matrixMul && ./matrixMul"

echo "Running simpleCUBLAS..."
$GUEST_SSH "cd /usr/local/cuda/samples/7_CUDALibraries/simpleCUBLAS && ./simpleCUBLAS"

echo "Running simpleCUFFT..."
$GUEST_SSH "cd /usr/local/cuda/samples/7_CUDALibraries/simpleCUFFT && ./simpleCUFFT"

echo "Running deviceQuery..."
$GUEST_SSH "cd /usr/local/cuda/samples/1_Utilities/deviceQuery && ./deviceQuery"

echo "Running bandwidthTest..."
$GUEST_SSH "cd /usr/local/cuda/samples/1_Utilities/bandwidthTest && ./bandwidthTest"

echo "=== CUDA Workload Validation Complete ==="