name: Setup Cuda
description: Setup Cuda on Hyperstack or GitHub instance

inputs:
  cuda-version:
    description: Version of Cuda to use
    required: true
  gcc-version:
    description: Version of GCC to use
    required: true
  github-instance:
    description: Instance is hosted on GitHub
    default: 'false'

runs:
  using: "composite"
  steps:
    # Mandatory on hyperstack since a bootable volume is not re-usable yet.
    - name: Install dependencies
      shell: bash
      run: |
        wget https://github.com/Kitware/CMake/releases/download/v"${CMAKE_VERSION}"/cmake-"${CMAKE_VERSION}"-linux-x86_64.sh
        echo "${CMAKE_SCRIPT_SHA} cmake-${CMAKE_VERSION}-linux-x86_64.sh" > checksum
        sha256sum -c checksum
        sudo bash cmake-"${CMAKE_VERSION}"-linux-x86_64.sh --skip-license --prefix=/usr/ --exclude-subdir
        sudo apt remove -y unattended-upgrades
        sudo apt update
        sudo apt install -y cmake-format libclang-dev
      env:
        CMAKE_VERSION: 3.29.6
        CMAKE_SCRIPT_SHA: "6e4fada5cba3472ae503a11232b6580786802f0879cead2741672bf65d97488a"

    - name: Install GCC
      if: inputs.github-instance == 'true'
      shell: bash
      env:
        GCC_VERSION: ${{ inputs.gcc-version }}
      run: |
        sudo apt-get install gcc-"{GCC_VERSION}" g++-"{GCC_VERSION}"
        sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-"{GCC_VERSION}" 20
        sudo update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-"{GCC_VERSION}" 20

    - name: Check GCC
      shell: bash
      env:
        GCC_VERSION: ${{ inputs.gcc-version }}
      run: |
        which gcc-"${GCC_VERSION}"

    - name: Install CUDA
      if: inputs.github-instance == 'true'
      shell: bash
      env:
        CUDA_VERSION: ${{ inputs.cuda-version }}
        CUDA_KEYRING_PACKAGE: cuda-keyring_1.1-1_all.deb
        CUDA_KEYRING_SHA: "d93190d50b98ad4699ff40f4f7af50f16a76dac3bb8da1eaaf366d47898ff8df"
      run: |
        # Use Sed to extract a value from a string, this cannot be done with the ${variable//search/replace} pattern.
        # shellcheck disable=SC2001
        TOOLKIT_VERSION="$(echo "${CUDA_VERSION}" | sed 's/\(.*\)\.\(.*\)/\1-\2/')"
        wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/${CUDA_KEYRING_PACKAGE}
        echo "${CUDA_KEYRING_SHA} ${CUDA_KEYRING_PACKAGE}" > checksum
        sha256sum -c checksum
        sudo dpkg -i "${CUDA_KEYRING_PACKAGE}"
        sudo apt update
        sudo apt -y install cuda-toolkit-"${TOOLKIT_VERSION}"

    - name: Export CUDA variables
      shell: bash
      run: |
        find /usr/local -executable -name "nvcc"
        CUDA_PATH=/usr/local/cuda-"${CUDA_VERSION}"
        {
          echo "CUDA_PATH=$CUDA_PATH";
          echo "LD_LIBRARY_PATH=$CUDA_PATH/lib64:$LD_LIBRARY_PATH";
          echo "CUDA_MODULE_LOADER=EAGER";
          echo "PATH=$PATH:$CUDA_PATH/bin"; 
        } >> "${GITHUB_ENV}"
        {
          echo "PATH=$PATH:$CUDA_PATH/bin"; 
        } >> "${GITHUB_PATH}"
      env:
        CUDA_VERSION: ${{ inputs.cuda-version }}

    # Specify the correct host compilers
    - name: Export gcc and g++ variables
      shell: bash
      run: |
        {
          echo "CC=/usr/bin/gcc-${GCC_VERSION}";
          echo "CXX=/usr/bin/g++-${GCC_VERSION}";
          echo "CUDAHOSTCXX=/usr/bin/g++-${GCC_VERSION}";
        } >> "${GITHUB_ENV}"
      env:
        GCC_VERSION: ${{ inputs.gcc-version }}

    - name: Check setup
      shell: bash
      run: |
        which nvcc

    - name: Check device is detected
      shell: bash
      run: nvidia-smi
