#!/bin/bash
# Helper for running a command in a podman container with GPU support
#
# Author: Christian Kastner <ckk@kvr.at>
# License: MIT

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

	Run a command in a podman container with GPU support

	This is a thin wrapper around podman-run(1) that takes care of adding all
	of the arguments necessary for using the GPU in the container. All
	arguments to $0 are passed on straight to podman-run.

	Synopsis:
	  $0 -h

	  $0 [podman-run args]

	Options:
	  -h          Show this help

	Examples:

	  # Configure the system for GPU-in-container use

	  \$ rocm-podman-setup -u <user>

	  # Create an image first, if needed

	  \$ rocm-podman-create -m http://10.1.2.3:9999/debian rocm/debian:unstable

	  # Call just as if you would call 'podman run'

	  \$ $0 --rm -it rocm/debian:unstable

	EOF
	exit 0
}
[ -z "${1:-}" ] || [ "$1" = "-h" ] && usage

userNAME=$(whoami)
renderGID="$(getent group render | cut -d: -f3)"
# By policy
videoGID=44

# Sanity checks
if [ -z "$renderGID" ]
then
	cat >&2 <<-EOF
	Group 'render' does not exist on this system. Are you sure that you are on
	the right system? This group should have been autmatically created by the
	udev package."
	EOF
	exit 1
elif ! groups "$userNAME" | grep -q '\brender\b'
then
	echo "'$userNAME' is not in group 'render'." >&2
	exit 1
elif ! groups "$userNAME" | grep -q '\bvideo\b'
then
	echo "'$userNAME' is not in group 'video'." >&2
	exit 1
elif [ "$(cat /proc/sys/kernel/unprivileged_userns_clone)" != "1" ]
then
	echo "unprivileged_userns_clone not enabled.">&2
	exit 1
elif ! [ -c /dev/kfd ]
then
	echo "Device /dev/kfd does not exist - is the amdgpu module loaded?." >&2
	exit 1
elif ! [ -w /dev/kfd ]
then
	echo "No write permissions for /dev/kfd." >&2
	exit 1
elif ! grep -q "$userNAME:$renderGID:1" /etc/subgid
then
	echo "No subgid mapping for group 'render'. Run rocm-podman-setup" >&2
	exit 1
elif ! grep -q "$userNAME:$videoGID:1" /etc/subgid
then
	echo "No subgid mapping for group 'video'. Run rocm-podman-setup" >&2
	exit 1
elif ! grep -q -E "$userNAME:[0-9]{6,}:6553[4-6]" /etc/subgid
then
	echo "No large subgid mapping for '$(whoami)'. Run rocm-podman-setup" >&2
	exit 1
fi

exec podman run \
	--device=/dev/dri \
	--device=/dev/kfd \
	--gidmap=0:0:1 \
	--gidmap=44:1:1 \
	--gidmap="$renderGID":2:1 \
	--gidmap=1:3:43 \
	--gidmap=45:46:$(("$renderGID"-"$videoGID"-1)) \
	--gidmap=$(("$renderGID"+1)):$(("$renderGID"+2)):65429 \
	"$@"
