#!/bin/bash
# Helper for setting up vfio-pci for AMD GPU pass-through to a QEMU amd64 guest
#
# This has been tested on a Debian 12 ("bookworm") amd64 host, with a single
# RX 6800 XT card attached. However, it should work for any Debian derivative,
# and most of it should work for other Linux distros as well.
#
# Christian Kastner <ckk@kvr.at>
# License: MIT
set -eu
shopt -s nullglob

# Option defaults
regularuser=

function usage() {
	cat >&2 <<- EOF

	This utility will look for AMD GPUs on the system, check whether the system
	is already configured for PCI pass-through, and suggest the appropriate
	measures if not.

	If USER is given, then it will also check whether the given user has
	sufficient permissions to use PCI pass-through in non-privileged mode.

	Synopsis:
	  $0 -h

	  $0 [-u USER]

	Options:
	  -h      Show this help
	  -u USER Verify that USER has the necessary permissions and ulimits to use
	          PCI pass-through together with QEMU.

	Examples:

	  # Check whether root can use pass-through

	  \$ $0

	  # Check whether user 'someuser' can use pass-through

	  \$ $0 -u someuser

	EOF
	exit 0
}

while getopts "hu:" OPTNAME
do
	case $OPTNAME in
	h)  usage;;
	u)  regularuser="$OPTARG";;
	?)  usage;;
	esac
done
shift $((OPTIND - 1))

# Sanity checks
if [ ! -x /usr/bin/lspci ]
then
	echo "Utility 'lspci' is not available on this system." >&2
	exit 1
elif ! grep -q -E 'svm|vmx' /proc/cpuinfo || [ ! -c /dev/kvm ]
then
	echo "Host CPU doesn't support KVM -- make sure AMD-V/Intel VT are enabled in BIOS." >&2
	exit 1
fi

# Given a slot ID like 09:00.1, gets the "pretty" device name
device_name() {
	lspci -s "$1" -vmm | sed -nre 's/^Device:\s+(.*)/\1/p'
}

# Given a slot ID like 09:00.1, gets the PCI vendor_id:device_id
device_pci_id() {
	lspci -s "$1" -n | cut -d' ' -f3
}

# Array of PCI slot IDs
declare -a GpuDevices

# Associative array (indexed by device ID above) of subdevices, which we define
# as one of either
#  (a) another device on the same card, like an audio device
#  (b) a device in the same IOMMU group
# Both (a) and (b) also need to be passed through.
declare -A GpuSubDevices

# Find all AMD GPU devices
# 1002=AMD, 0300=VGA compatible controller, 0380=Display controller
mapfile -t GpuDevices < <( { lspci -D -d 1002::0300 ; lspci -D -d 1002::0380 ; } | cut -d' ' -f1)
if [ "${#GpuDevices[@]}" -eq 0 ]
then
	echo "No AMD GPU devices were found on this system." >&2
	exit 1
fi

# Find subdevices per GPU
for device in "${GpuDevices[@]}"
do
	GpuSubDevices[$device]=""

	# First, all consumer devices of this GPU (like the audio device on the card)
	cd /sys/bus/pci/devices/"$device"
	for consumer_raw in consumer:pci:*
	do
		consumer=$(echo "$consumer_raw" | cut -d: -f3,4,5)
		[[ ${GpuSubDevices[$device]} == *$consumer* ]] || GpuSubDevices[$device]+=" $consumer"
	done

	# IOMMU might be disabled on BIOS configuration
	[ -d /sys/bus/pci/devices/"$device"/iommu_group ] || continue

	# Then, all devices in the same IOMMU group
	cd /sys/bus/pci/devices/"$device"/iommu_group/devices
	for member in *
	do
		[ "$device" == "$member" ] && continue
		[[ ${GpuSubDevices[$device]} == *$member* ]] || GpuSubDevices[$device]+=" $member"
	done
done

# This is for amd64
packages_missing=
for pkgname in \
	autopkgtest \
	qemu-system-x86 \
	qemu-user-static \
	qemu-utils \
	ovmf \
	vmdb2 \
	zerofree
do
	dpkg -l $pkgname 2> /dev/null | grep -qE '^ii' || packages_missing+=" $pkgname"
done

# This is a strong SHOULD
has_cache=N
for pkgname in approx apt-cacher apt-cacher-ng
do
	dpkg -l $pkgname 2> /dev/null | grep -qE '^ii' && has_cache=Y
done

# Find modules that need loading
# TODO: Also check /etc/modules-load.d
modules_missing=""
for module in vfio vfio_iommu_type1 vfio_pci
do
	altmodule=$(echo $module | tr '_' '-')
	if ! grep -q "^$module$" /etc/modules && ! grep -q "^$altmodule$" /etc/modules
	then
		modules_missing+=" $module"
	fi
done

# Determine modprobe configuration
modprobe=""
modprobe_file=""
modprobe_missing=""
if [ ! -d /etc//modprobe.d ] || ! grep -Eqr '^options[[:space:]]+(vfio-pci|vfio_pci)' /etc/modprobe.d
then
	modprobe="options vfio_pci ids="

	for device in "${GpuDevices[@]}"
	do
		modprobe+=$(device_pci_id "$device"),

		for subdevice in ${GpuSubDevices[$device]}
		do
			modprobe+=$(device_pci_id "$subdevice"),
		done
	done
	# Remove the final trailing comma
	modprobe=${modprobe%,}
else
	modprobe_file="$(grep -Elr '^options[[:space:]]+(vfio-pci|vfio_pci)' /etc/modprobe.d)"

	for device in "${GpuDevices[@]}"
	do
		pci_id=$(device_pci_id "$device")
		grep -Eqr "$pci_id" /etc/modprobe.d || modprobe_missing+="$pci_id,"

		for subdevice in ${GpuSubDevices[$device]}
		do
			pci_id=$(device_pci_id "$subdevice")
			grep -Eqr "$pci_id" /etc/modprobe.d || modprobe_missing+="$pci_id,"
		done
	done
	# Remove the final trailing comma
	modprobe_missing=${modprobe_missing%,}
fi

# Check whether vfio devices are assigned to group "render"
udev_file=""
if [ -d /etc/udev/rules.d ] && grep -Eqr '^[^#]*SUBSYSTEM=="vfio"' /etc/udev/rules.d
then
	udev_file="$(grep -Elr '^[^#]*SUBSYSTEM==\"vfio\".*GROUP=\"render\"' /etc/udev/rules.d)"
fi

# If requested, validate user setup
user_exists=N
groups_missing=
limits_missing=
if [ -n "$regularuser" ]
then
	if getent passwd "$regularuser" &> /dev/null
	then
		user_exists=Y
		for groupname in kvm render
		do
			if ! groups "$regularuser" | grep -q "\b${groupname}\b"
			then
				groups_missing+="$groupname,"
			fi
		done
		groups_missing=${groups_missing%,}

		for limitname in hard soft
		do
			pattern="^(${regularuser}|\*)[[:space:]]+(${limitname}|-)[[:space:]]+memlock[[:space:]]"
			if ! grep -Eqr "$pattern" /etc/security/limits.*
			then
				limits_missing+="$limitname,"
			fi
		done
		limits_missing=${limits_missing%,}
	fi
fi


###############
# Output time #
###############
# bash really isn't the right tool for this...

echo
echo Devices
echo =======
echo "Device ID          PCI ID      Device Name"
for device in "${GpuDevices[@]}"
do
	echo "$device       $(device_pci_id "$device")   $(device_name "$device")"
	for subdevice in ${GpuSubDevices[$device]}
	do
		echo " └─ $subdevice   $(device_pci_id "$subdevice")   $(device_name "$subdevice")"
	done
done
echo

echo Checks
echo ======

echo -n "BIOS (kvm)  "
if ! grep -qE '(svm|vmx)' /proc/cpuinfo
then
	echo "[TODO] Enable CPU virtualization features (AMD SVM or Intel VT-x)"
else
	echo "  [OK] CPU virtualization enabled"
fi

echo -n "BIOS (IOMMU)"
if ! [ -d /sys/kernel/iommu_groups ] || [ -z "$(ls /sys/kernel/iommu_groups)" ]
then
	echo "[TODO] Enable AMD IOMMU or Intel VT-d support for device pass-through"
else
	echo "  [OK] Device pass-through support detected"
fi

echo -n "packages    "
if [ -n "$packages_missing" ]
then
	echo "[TODO] Key packages not installed:$packages_missing"
else
	echo "  [OK] Key packages are installed"
fi

echo -n "APT cache   "
if [ "$has_cache" = N ]
then
	cat <<-EOF
	 [OPT] No local APT cache detected. While not strictly necessary, it is
	                   strongly suggested that you install one of the approx, apt-cacher,
	                   or 'apt-cacher-ng' packages
	EOF
else
	echo "  [OK] Local APT cache detected, make sure to use it"
fi

echo -n "memory      "
if [ $(($(getconf _PHYS_PAGES) * $(getconf PAGE_SIZE) / (1024 * 1024))) -lt 31000 ]
then
	echo "[TODO] You really want a host with at *least* 32G of memory installed"
else
	echo "  [OK] Host has at least 32G of memory"
fi

echo -n "modules     "
if [ -n "$modules_missing" ]
then
	echo "[TODO] In /etc/modules, these modules are missing:$modules_missing"
else
	echo "  [OK] In /etc/modules, all necessary modules are loaded"
fi

echo -n "modprobe    "
if [ -n "$modprobe" ]
then
	echo "[TODO] Needs vfio-pci configuration to grab the devices listed above"
	echo "                   Create e.g. /etc/modprobe.d/vfio.conf with the following contents:"
	echo "                       $modprobe"
else
	if [ -n "$modprobe_missing" ]
	then
		echo "[TODO] In $modprobe_file, these devices are not yet assigned to vfio-pci: $modprobe_missing"
	else
	    echo "  [OK] In $modprobe_file, all devices assigned to vfio-pci"
	fi
fi

echo -n "udev        "
if [ -z "$udev_file" ]
then
	echo "[TODO] vfio devices need to be assigned to group 'render'"
	echo "                   Create e.g. /etc/udev/rules.d/99-vfio.rules with the following contents:"
	echo "                       SUBSYSTEM==\"vfio\", GROUP=\"render\", MODE=\"0660\""
else
	echo "  [OK] In $udev_file, all vfio devices are assigned to group 'render'"
fi

if [ -n "$regularuser" ]
then
	if [ $user_exists = N ]
	then
	    echo "user        [TODO] User '$regularuser' needs to be created first"
	else
		echo -n "user-groups "
		if [ -n "$groups_missing" ]
		then
			echo "[TODO] User '$regularuser' is missing from groups: $groups_missing"
			echo "                   You can fix this with usermod -a -G $groups_missing $regularuser"
		else
			echo "  [OK] User '$regularuser' is in all relevant groups"
		fi
		echo -n "user-limits "
		if [ -n "$limits_missing" ]
		then
			echo "[TODO] User '$regularuser' needs elevated memlock ulimit"
			echo "                   Create e.g. /etc/security/limits.d/$regularuser.conf with the following contents:"
			echo "                       $regularuser       soft    memlock       32000000000"
			echo "                       $regularuser       hard    memlock       32000000000"
			echo "                   Increase 32000000000 to 48000000000 if your system has the memory"
		else
			echo "  [OK] User '$regularuser' has elevated memlock limit"
		fi
	fi
fi

cat <<EOF

WARNING: If you implement the changes suggested above, then the host will no
longer have access to the GPU(s). That means NO VIDEO OUTPUT. Please ensure
that you have alternative means to access the host (like SSH), or you have a
backup configuration to boot from.

Don't forget to update your initramfs and to reboot after making changes.
EOF