xref: /aosp_15_r20/external/pytorch/.ci/docker/common/install_ucc.sh (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1#!/bin/bash
2
3set -ex
4
5if [[ -d "/usr/local/cuda/" ]];  then
6  with_cuda=/usr/local/cuda/
7else
8  with_cuda=no
9fi
10
11function install_ucx() {
12  set -ex
13  git clone --recursive https://github.com/openucx/ucx.git
14  pushd ucx
15  git checkout ${UCX_COMMIT}
16  git submodule update --init --recursive
17
18  ./autogen.sh
19  ./configure --prefix=$UCX_HOME      \
20      --enable-mt                     \
21      --with-cuda=$with_cuda          \
22      --enable-profiling              \
23      --enable-stats
24  time make -j
25  sudo make install
26
27  popd
28  rm -rf ucx
29}
30
31function install_ucc() {
32  set -ex
33  git clone --recursive https://github.com/openucx/ucc.git
34  pushd ucc
35  git checkout ${UCC_COMMIT}
36  git submodule update --init --recursive
37
38  ./autogen.sh
39  # We only run distributed tests on Tesla M60 and A10G
40  NVCC_GENCODE="-gencode=arch=compute_52,code=sm_52 -gencode=arch=compute_86,code=compute_86"
41  ./configure --prefix=$UCC_HOME          \
42    --with-ucx=$UCX_HOME                  \
43    --with-cuda=$with_cuda                \
44    --with-nvcc-gencode="${NVCC_GENCODE}"
45  time make -j
46  sudo make install
47
48  popd
49  rm -rf ucc
50}
51
52install_ucx
53install_ucc
54