xref: /aosp_15_r20/external/pytorch/.ci/docker/common/install_cusparselt.sh (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1#!/bin/bash
2
3set -ex
4
5# cuSPARSELt license: https://docs.nvidia.com/cuda/cusparselt/license.html
6mkdir tmp_cusparselt && cd tmp_cusparselt
7
8if [[ ${CUDA_VERSION:0:4} =~ ^12\.[2-4]$ ]]; then
9    arch_path='sbsa'
10    export TARGETARCH=${TARGETARCH:-$(uname -m)}
11    if [ ${TARGETARCH} = 'amd64' ] || [ "${TARGETARCH}" = 'x86_64' ]; then
12        arch_path='x86_64'
13    fi
14    CUSPARSELT_NAME="libcusparse_lt-linux-${arch_path}-0.6.2.3-archive"
15    curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-${arch_path}/${CUSPARSELT_NAME}.tar.xz
16elif [[ ${CUDA_VERSION:0:4} == "12.1" ]]; then
17    arch_path='sbsa'
18    export TARGETARCH=${TARGETARCH:-$(uname -m)}
19    if [ ${TARGETARCH} = 'amd64' ] || [ "${TARGETARCH}" = 'x86_64' ]; then
20        arch_path='x86_64'
21    fi
22    CUSPARSELT_NAME="libcusparse_lt-linux-${arch_path}-0.5.2.1-archive"
23    curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-${arch_path}/${CUSPARSELT_NAME}.tar.xz
24elif [[ ${CUDA_VERSION:0:4} == "11.8" ]]; then
25    CUSPARSELT_NAME="libcusparse_lt-linux-x86_64-0.4.0.7-archive"
26    curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-x86_64/${CUSPARSELT_NAME}.tar.xz
27fi
28
29tar xf ${CUSPARSELT_NAME}.tar.xz
30cp -a ${CUSPARSELT_NAME}/include/* /usr/local/cuda/include/
31cp -a ${CUSPARSELT_NAME}/lib/* /usr/local/cuda/lib64/
32cd ..
33rm -rf tmp_cusparselt
34ldconfig
35