#!/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 \
    "$@"
