#!/bin/bash

# Init parameter
CUDA_VERSION=${1:-"cuda-11.1"}

# Install cuda lib
cd /usr/local
rm -f ${CUDA_VERSION}.tar.gz
wget https://tools.mindspore.cn/tools/driver/cuda/${CUDA_VERSION}/build-lib/${CUDA_VERSION}.tar.gz
rm -rf ${CUDA_VERSION}
tar zxf ${CUDA_VERSION}.tar.gz
chmod 755 -R /usr/local/${CUDA_VERSION}
rm -f /usr/local/${CUDA_VERSION}.tar.gz

# Add libcublas for cuda-10.1
if [ "${CUDA_VERSION}" = "cuda-10.1" ]; then
    cd /usr/lib/x86_64-linux-gnu
    rm -f libcublas.so.10.1.0.105
    wget https://tools.mindspore.cn/tools/driver/cuda/cuda-10.1/build-lib/libcublas.so.10.1.0.105
    chmod 755 libcublas.so.10.1.0.105
    ln -s libcublas.so.10.1.0.105 libcublas.so.10
    ln -s libcublas.so.10 libcublas.so
    ls -l /usr/lib/x86_64-linux-gnu/libcublas.so*
fi
