#!/bin/bash
set -e
set +x
# Copyright (C) 2022-2023 Mo Zhou <lumin@debian.org>
# MIT/Expat License.
#
# Nvidia CUDA Deep Neural Network Library installer script (Debian Specific)
#
# XXX: you can browse the following directory for updating the
#  URL_{amd64,arm64} shell variables below:
#
#   NEW: https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/
#   OLD: https://developer.download.nvidia.com/compute/redist/cudnn/
#
#  More references when Nvidia breaks the above link:
#
#   https://github.com/archlinux/svntogit-community/blob/packages/cudnn/trunk/PKGBUILD
#   https://gitlab.archlinux.org/archlinux/packaging/packages/cudnn/-/blob/main/PKGBUILD
#   https://github.com/pytorch/builder/blob/main/common/install_cuda.sh
#
# XXX: [maintainer/user notes]
#
#  In case you want to upgrade to a newer version of cuDNN, you can just
#  browse the link above, and find the binary tarballs you want.
#  Then copy the urls and update the corresponding URL_* variables
#  below. The rest of the shell code can remain unchanged as long
#  as the upstream does not alter the file paths in tarball.
#
#  To test whether the updated links work or not, you can just copy
#  and paste the commands at the end part of the usage() function.
#  You don't have to test all the three actions {-d, -u, -p}.
#  As long as the update (-u) action works without issue, the download
#  is good as well (-d).
#
#  For a more thorough testing of install/purge, use piuparts instead.

## configs ####################################################################
TMPDIR="$(mktemp -d)"
TMPDIR_IS_OVERRIDEN=0
ARCH="$(dpkg-architecture -qDEB_HOST_ARCH)"
MULTIARCH="$(dpkg-architecture -qDEB_HOST_MULTIARCH)"
PREFIX="/usr"
FILELIST=""
CUDA_VER="12"
CUDNN_VER="9.23.2.1"
URL_amd64="https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-9.23.2.1_cuda12-archive.tar.xz"
SHA512_amd64=0058cefe3e15acd7b312eeedced6cff482eb2cef7cb790b6ff81c7556bdf66a3839001962b348e4a672219bcf4c72de55d4effb53f572932d70ca886d1a21e8f
URL_arm64="https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-sbsa/cudnn-linux-sbsa-9.23.2.1_cuda12-archive.tar.xz"
SHA512_arm64=121c07c0fce01a930611db030ef28cf260ecf6311f481e309ead723f4a0e255647197c834e1dec1584c3af37c4e033a1e5cec1d345047fc8e389488ddb9f80e9

## usage ######################################################################
usage () {
    cat << EOF
Usage: $(basename "$0") <-d|-u|-p|-h> [--prefix <prefix>] [--tmpdir <tmpdir>] ...
Arguments:
 -d|--download           download only (default: 0)
 -u|--update             update cudnn installation (default: 0)
 -p|--purge              purge cudnn installation (default: 0)
 -h|--help               display this help message
 --arch <arch>           override architecture (default: $(dpkg-architecture -qDEB_HOST_ARCH))
 --multiarch <multiarch> override multiarch triplet (default: $(dpkg-architecture -qDEB_HOST_MULTIARCH))
 --prefix <path>         override install prefix (default: /usr)
 --tmpdir <dir>          override temporary directory (default: ${TMPDIR})
 --filelist <path>       write installed file list (default: ${FILELIST})
Testing this script cross-architecture:
 $ ./update-nvidia-cudnn --arch amd64   --multiarch x86_64-linux-gnu      --tmpdir . --prefix fake {-d,-u,-p}
 $ ./update-nvidia-cudnn --arch arm64   --multiarch aarch64-linux-gnu     --tmpdir . --prefix fake {-d,-u,-p}
Version: cuDNN ${CUDNN_VER} for CUDA ${CUDA_VER}
EOF
}

## argument parsing ###########################################################
DOWNLOAD_ONLY=0
DO_UPDATE=0
DO_PURGE=0
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;;
        --prefix)
            if test -n "$2"; then
                PREFIX="$2"
            fi
            shift; shift;;
        --arch)
            if test -n "$2"; then
                ARCH="$2"
            fi
            shift; shift;;
        --multiarch)
            if test -n "$2"; then
                MULTIARCH="$2"
            fi
            shift; shift;;
        --tmpdir)
            if test -n "$2"; then
                TMPDIR="$2"
                TMPDIR_IS_OVERRIDEN=1
            fi
            shift; shift;;
        --filelist)
            if test -n "$2"; then
                FILELIST="$2"
            fi
            shift; shift;;
        -h|--help)
            usage; exit 0;;
        -*|--*)
            usage; exit 1;;
        *)
            usage; exit 1;;
    esac
done

# post processing
if test "${ARCH}" = "amd64"; then
    URL="${URL_amd64}"; SHA512="${SHA512_amd64}"
elif test "${ARCH}" = "arm64"; then
    URL="${URL_arm64}"; SHA512="${SHA512_arm64}"
else
    echo "$0: Unsupported architecture ${ARCH}" >&2
    exit 1
fi

## functions ##################################################################
verify_checksum() {
    local dest="$1"
    if ! sha512sum -c - << EOF; then
${SHA512} ${dest}
EOF
        echo "$0: sha512 checksum mismatch. aborting." >&2
        exit 1
    fi
}

download_cudnn () {
    # Download cudnn tarball to ${TMPDIR}
    # args: $1: URL for upstream tarball
    # return: saved file destination
    test -n "${1}" || { echo "download_cudnn(): URL not specified" >&2; exit 1; }
    local url="${1}"
    test -d "${TMPDIR}" || mkdir "${TMPDIR}"
    local dest="${TMPDIR}/$(basename "${url}")"
    # detect downloader and download
    if command -v curl > /dev/null; then
        local cmd="curl --continue-at - -L \"${URL}\" --output \"${dest}\""
        local nocheck="--insecure"
    elif command -v wget > /dev/null; then
        local cmd="wget --continue --verbose --show-progress=off -c \"${URL}\" -O \"${dest}\""
        local nocheck="--no-check-certificate"
    else
        echo "$0: Error: no downloader available." >&2
        exit 255
    fi
    # already exists?
    if ! test -f "${dest}"; then
        echo "${cmd}" >&2
        bash -c "${cmd}" || bash -c "${cmd} ${nocheck}" >&2
    else
        echo "Skipping download as file already exists: ${dest}" >&2
    fi
    test -f "${dest}" || { echo "Download failed." >&2; exit 1; }
    # return string
    echo "${dest}"
}

install_cudnn () {
    test -n "${1}" || { echo "install_cudnn(): invalid argument" >&2; exit 1; }
    test -n "${2}" || { echo "install_cudnn(): invalid argument" >&2; exit 1; }
    # Install extracted cudnn from src to dst
    local src="${1}"  # e.g. /tmp/nvidia-cudnn/
    local dst="${2}"  # e.g. /usr/local/
    local installed=()
    
    while IFS= read -r -d '' F; do
        [[ "${F}" == *cudnn.txz* ]] && continue
        if [[ "${F}" =~ /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
            installed+=( "${dst}/lib/${MULTIARCH}/$(basename "${F}")" )
        elif [[ "${F}" =~ /libcudnn.*\.a$ ]]; then
            # static library file
            install -vDm0644 -t "${dst}/lib/${MULTIARCH}" "${F}"
            installed+=( "${dst}/lib/${MULTIARCH}/$(basename "${F}")" )
        elif [[ "${F}" =~ /cudnn.*\.h$ ]]; then
            # header file
            install -vDm0644 -t "${dst}/include/${MULTIARCH}" "${F}"
            installed+=( "${dst}/include/${MULTIARCH}/$(basename "${F}")" )
        elif [[ "${F}" == *NVIDIA_SLA_cuDNN_Support.txt ]]; then
            # copyright file
            install -vDm0644 -t "${dst}/share/doc/nvidia-cudnn/" "${F}"
            installed+=( "${dst}/share/doc/nvidia-cudnn/$(basename "${F}")" )
        else
            echo "Skipped ${F}"
        fi
    done < <(find "${src}" \( -type f -o -type l \) -print0)
    
    echo "$0: Number of installed files: ${#installed[@]}"
    if test -n "${FILELIST}"; then
        if ! test -d "$(dirname "${FILELIST}")"; then
            mkdir -p "$(dirname "${FILELIST}")"
        fi
        for i in "${installed[@]}"; do
            echo "${i}" >> "${FILELIST}"
        done
        echo "$0: The list of installed files is recorded at ${FILELIST}"
    fi
}

# this is a dispatcher. When a ${FILELIST} is given, we use it to safely
# delete the recorded installed files. If not, we fallback to the manual
# deletion based on manually written matching rules. As long as the user
# does not modify /usr/lib/ (irrelevant to /usr/local) on their own, both
# methods will lead to the same results.
purge_cudnn () {
    local prefix="${1}"
    if test -n "${FILELIST}"; then
        purge_cudnn_filelist "${FILELIST}"
    else
        purge_cudnn_fallback "${prefix}"
    fi
}

purge_cudnn_filelist () {
    test -e "${1}" || { echo "purge_cudnn_filelist(): invalid argument" >&2; exit 1; }
    # purge cudnn from the given file list
    local filelist="${1}"
    # first, validate the given file list
    local FL=()
    while IFS= read -r line; do
        FL+=( "$line" )
    done < "${filelist}"
    
    for i in "${FL[@]}"; do
        if ! test -e "${i}"; then
            echo "Error: the given file list ${filelist} is invalid: file ${i} does not exist" >&2
            exit 2
        fi
    done
    # then, remove the listed files
    for i in "${FL[@]}"; do
        rm -v "${i}"
    done
    rm -v "${filelist}"
}

purge_cudnn_fallback () {
    test -n "${1}" || { echo "purge_cudnn_fallback(): invalid argument" >&2; exit 1; }
    # Purge cudnn from the given path (prefix)
    local dst="${1}"
    local FILES=()
    
    while IFS= read -r -d '' F; do
        FILES+=( "${F}" )
    done < <(find "${dst}/lib/${MULTIARCH}" \( -type f -o -type l \) -name "libcudnn*.so*" -print0 2>/dev/null || true; \
             find "${dst}/include/${MULTIARCH}" -type f -name "cudnn*.h" -print0 2>/dev/null || true; \
             find "${dst}/lib/${MULTIARCH}" -type f -name "libcudnn*.a" -print0 2>/dev/null || true)
             
    if test -f "${dst}/share/doc/nvidia-cudnn/NVIDIA_SLA_cuDNN_Support.txt"; then
        FILES+=( "${dst}/share/doc/nvidia-cudnn/NVIDIA_SLA_cuDNN_Support.txt" )
    fi
    
    if test 0 -eq ${#FILES[@]}; then
        exit 0
    fi
    for F in "${FILES[@]}"; do
        if test -e "${F}" || test -L "${F}"; then
            rm -rv "${F}"
        fi
    done
}

# flag check: must select one valid action
if test "${DOWNLOAD_ONLY}" -eq 0 && \
   test "${DO_UPDATE}" -eq 0 && \
   test "${DO_PURGE}" -eq 0; then
    usage
    exit 0
fi

# trigger actions
if test "${DOWNLOAD_ONLY}" -ne 0; then
    path="$(download_cudnn "${URL}")"
    verify_checksum "${path}"
    echo "${path}"
    exit 0
elif test "${DO_UPDATE}" -ne 0; then
    path="$(download_cudnn "${URL}")"
    verify_checksum "${path}"
    tmpdir2="$(mktemp -d)"
    echo "Extracting files from the downloaded tarball..."
    tar xvf "${path}" -C "${tmpdir2}/"
    echo "Installing the files to system directories..."
    install_cudnn "${tmpdir2}" "${PREFIX}"
    rm -rf "${tmpdir2}"
    # cleanup
    if test 0 -eq "${TMPDIR_IS_OVERRIDEN}"; then
        rm -rv "${path}"
        rmdir -v "${TMPDIR}"
    fi
elif test "${DO_PURGE}" -ne 0; then
    echo "Purging cuDNN installation from ${PREFIX}"
    purge_cudnn "${PREFIX}" || true
fi
