#!/bin/bash
set -e
# Copyright (C) 2022 Mo Zhou <lumin@debian.org>
# MIT/Expat License.
#
# Nvidia CUDA Deep Neural Network Library installer script (Debian Specific)
# Borrowed bits from Archlinux:
# https://github.com/archlinux/svntogit-community/blob/packages/cudnn/trunk/PKGBUILD
# Useful References:
# https://developer.nvidia.com/cuDNN
# https://developer.nvidia.com/rdp/cudnn-archive
# https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/

CUDA_VERSION="11.4"
CUDNN_VERSION="8.2.4.15"
#TMPDIR=$(mktemp -d)
TMPDIR=/tmp/nvidia-cudnn
TGZ=${TMPDIR}/cudnn.tgz
MULTIARCH=$(dpkg-architecture -qDEB_HOST_MULTIARCH)
PREFIX=/usr

download_cudnn () {
	# Download cudnn tgz to TMPDIR
	test -n "${1}" || (echo "download_cudnn(): invalid argument"; exit 1)
	test -n "${2}" || (echo "download_cudnn(): invalid argument"; exit 1)
	local arch
	if test $(arch) = "x86_64"; then
		arch="x64"
	elif test $(arch) = "ppc64le"; then
		arch="ppc64le"
	else
		echo $0: Unsupported architecture ${arch}
	fi
	local cuda_ver="${1}"
	local cudnn_ver="${2}"
	local cudnn_trunc_ver="$(echo ${cudnn_ver} | sed -e 's/\.[0-9]*$//')"
	local url_dir="https://developer.download.nvidia.com/compute/redist/cudnn"
	local url_base="v${cudnn_trunc_ver}/cudnn-${cuda_ver}-linux-${arch}-v${cudnn_ver}.tgz"

	test -d ${TMPDIR} || mkdir ${TMPDIR}
	local cmd="wget --verbose --show-progress=on --progress=bar --hsts-file=/tmp/wget-hsts -c ${url_dir}/${url_base} -O ${TGZ}"
	echo ${cmd}
	bash -c "${cmd}" || bash -c "${cmd} --no-check-certificate"
	test -f ${TGZ} || (echo "Download failed."; exit 1)
}

install_cudnn () {
	test -n "${1}" || (echo "install_cudnn(): invalid argument"; exit 1)
	test -n "${2}" || (echo "install_cudnn(): invalid argument"; exit 1)
	# Install extracted cudnn from src to dst
	local src=${1}  # e.g. /tmp/nvidia-cudnn/
	local dst=${2}  # e.g. /usr/local/
	FILES=( $(find ${src} -type f,l) )
	for F in ${FILES[@]}; do
		(echo ${F} | grep -qo "cudnn.tgz") && continue
		if (echo ${F} | grep -qo ".*/libcudnn.*\.so.*"); then
			# shared object file
			if test -L ${F}; then
				mkdir -p ${dst}/lib/${MULTIARCH}/ || true
				cp -av ${F} ${dst}/lib/${MULTIARCH}/
			else
				install -vDm0644 -t ${dst}/lib/${MULTIARCH} ${F}
			fi
		elif $(echo ${F} | grep -qo ".*/libcudnn.*\.a"); then
			# static library file
			install -vDm0644 -t ${dst}/lib/${MULTIARCH} ${F}
		elif $(echo ${F} | grep -qo ".*/cudnn.*\.h"); then
			# header file
			install -vDm0644 -t ${dst}/include/${MULTIARCH} ${F}
		elif $(echo ${F} | grep -qo "NVIDIA_SLA_cuDNN_Support.txt"); then
			# copyright file
			install -vDm0644 -t ${dst}/share/doc/nvidia-cudnn/ ${F}
		else
			echo Skipped ${F}
		fi
	done
}

purge_cudnn () {
	test -n "${1}" || (echo "install_cudnn(): invalid argument"; exit 1)
	# Purge cudnn from the given path
	local dst="${1}"
	FILES=( $(find ${dst}/lib/${MULTIARCH} -type f,l -name "libcudnn*.so*") )
	FILES+=( $(find ${dst}/include/${MULTIARCH} -type f -name "cudnn*.h") )
	FILES+=( $(find ${dst}/lib/${MULTIARCH} -type f -name "libcudnn*.a") )
	FILES+=( ${dst}/share/doc/nvidia-cudnn/NVIDIA_SLA_cuDNN_Support.txt )
	for F in ${FILES[@]}; do
		(test -e ${F} || test -L ${F}) && rm -rv ${F}
	done
}

usage () {
	cat << EOF
Usage: $0 <-d|-u|-p|-h> [--cuda <cuda_ver>] [--cudnn <cudnn_ver>] 
Arguments:
 -d|--download           download only (will print file path) (default: 0)
 -u|--update             update cudnn installation (default: 0)
 -p|--purge              purge cudnn installation (default: 0)
 -h|--help               display this help message
 --keep                  keep the downloaded tarball after installation
 --cuda <cuda_ver>       override cuda version (default: ${CUDA_VERSION})
 --cudnn <cudnn_ver>     override cudnn version (default: ${CUDNN_VERSION})
 --prefix <path>         override install prefix (default: /usr)
EOF
}

# main ########################################################################
DOWNLOAD_ONLY=0
DO_UPDATE=0
DO_PURGE=0
NO_CLEANUP=0
OVERRIDE_CUDA=""
OVERRIDE_CUDNN=""
while [[ $# -gt 0 ]]; do
	case $1 in
		-d|--download)
			DOWNLOAD_ONLY=1; shift;;
		-u|--update)
			DO_UPDATE=1; shift;;
		-p|--purge)
			DO_PURGE=1; shift;;
		--cuda)
			OVERRIDE_CUDA="$2"; shift; shift;;
		--cudnn)
			OVERRIDE_CUDNN="$2"; shift; shift;;
		--prefix)
			if test -n "$2"; then
				PREFIX="$2"
			fi
			shift; shift;;
		-h|--help)
			usage; exit 0;;
		-*|--*)
			usage; exit 1;;
		*)
			usage; exit 1;;
	esac
done

# process arguments
test -n "${OVERRIDE_CUDA}" && CUDA_VERSION="${OVERRIDE_CUDA}"
test -n "${OVERRIDE_CUDNN}" && CUDNN_VERSION="${OVERRIDE_CUDNN}"

test ${DOWNLOAD_ONLY} -eq 0 && \
	test ${DO_UPDATE} -eq 0 && \
	test ${DO_PURGE} -eq 0 && \
	(usage; exit 0)

# trigger actions
if test "${DOWNLOAD_ONLY}" -ne 0; then
	echo Downloading cuDNN as ${TGZ}
	download_cudnn $CUDA_VERSION $CUDNN_VERSION
	exit 0
elif test "${DO_UPDATE}" -ne 0; then
	echo Will download and install cuDNN to ${PREFIX}
	download_cudnn $CUDA_VERSION $CUDNN_VERSION
	if ! test -e ${TGZ}.extracted; then
		echo Extracting ${TGZ}
		tar xf ${TGZ} -C ${TMPDIR}/
		touch ${TGZ}.extracted
	fi
	install_cudnn ${TMPDIR} ${PREFIX}
elif test "${DO_PURGE}" -ne 0; then
	echo Purging cuDNN installation from ${PREFIX}
	purge_cudnn ${PREFIX}
fi
