#!/usr/bin/env bash
set -euo pipefail

root="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
out="${1:-$root/out}"
mkdir -p "$out"

if ! command -v nvidia-smi >/dev/null || ! command -v nvcc >/dev/null; then
  echo "nvidia-smi and nvcc are required" >&2
  exit 2
fi

cap="$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | sed -n '1p' | tr -d '.')"
case "$cap" in
  ''|*[!0-9]*) echo "unable to detect compute capability" >&2; exit 2 ;;
esac

nvcc -std=c++17 -arch="compute_${cap}" --ptx "$root/kernel_only.cu" -o "$out/kernel.ptx"
nvcc -std=c++17 -arch="sm_${cap}" --cubin "$root/kernel_only.cu" -o "$out/kernel.cubin"
nvcc -std=c++17 -arch="sm_${cap}" "$root/cuda_path.cu" -o "$out/cuda_path"
c++ -std=c++17 "$root/driver_loader.cpp" -I"${CUDA_HOME:-/usr/local/cuda}/include" \
  -L"${CUDA_HOME:-/usr/local/cuda}/lib64" -lcuda -o "$out/driver_loader"

cuobjdump --dump-elf "$out/kernel.cubin" > "$out/kernel.elf.txt"
nvdisasm "$out/kernel.cubin" > "$out/kernel.sass"

"$out/cuda_path"
"$out/driver_loader" "$out/kernel.cubin"
printf 'compute_capability=%s\nptx=%s\ncubin=%s\nsass=%s\n' \
  "$cap" "$out/kernel.ptx" "$out/kernel.cubin" "$out/kernel.sass"