Add support for ROCm5.6 for nightly wheels (#1442)

* Add msccl-algorithms directory to PyTorch wheel

* Bundle msccl-algorithms into wheel

* Use correct src path for msccl-algorithms

(cherry picked from commit 95b5af3)

* Add hipblaslt dependency for ROCm5.6 onwards

* Update build_all_docker.sh to ROCm5.6

Author: jithunnair-amd

libtorch/build_all_docker.sh

@@ -8,6 +8,6 @@ for cuda_version in 12.1 11.8; do
     GPU_ARCH_TYPE=cuda GPU_ARCH_VERSION="${cuda_version}" "${TOPDIR}/libtorch/build_docker.sh"
 done
 
-for rocm_version in 5.4.2 5.5; do
+for rocm_version in 5.5 5.6; do
     GPU_ARCH_TYPE=rocm GPU_ARCH_VERSION="${rocm_version}" "${TOPDIR}/libtorch/build_docker.sh"
 done

manywheel/build_all_docker.sh

@@ -14,7 +14,7 @@ for cuda_version in 12.1 11.8; do
     MANYLINUX_VERSION=2014 GPU_ARCH_TYPE=cuda GPU_ARCH_VERSION="${cuda_version}" "${TOPDIR}/manywheel/build_docker.sh"
 done
 
-for rocm_version in 5.4.2 5.5; do
+for rocm_version in 5.5 5.6; do
     GPU_ARCH_TYPE=rocm GPU_ARCH_VERSION="${rocm_version}" "${TOPDIR}/manywheel/build_docker.sh"
     MANYLINUX_VERSION=2014 GPU_ARCH_TYPE=rocm GPU_ARCH_VERSION="${rocm_version}" "${TOPDIR}/manywheel/build_docker.sh"
 done

manywheel/build_rocm.sh

@@ -94,6 +94,10 @@ ROCM_SO_FILES=(
     "libroctx64.so"
 )
 
+if [[ $ROCM_INT -ge 50600 ]]; then
+    ROCM_SO_FILES+=("libhipblaslt.so")
+fi
+
 if [[ $ROCM_INT -lt 50500 ]]; then
     ROCM_SO_FILES+=("librocfft-device-0.so")
     ROCM_SO_FILES+=("librocfft-device-1.so")
@@ -198,8 +202,17 @@ if [[ $ROCM_INT -ge 50500 ]]; then
 
     DEPS_AUX_SRCLIST+=(${MIOPEN_SHARE_FILES[@]/#/$MIOPEN_SHARE_SRC/})
     DEPS_AUX_DSTLIST+=(${MIOPEN_SHARE_FILES[@]/#/$MIOPEN_SHARE_DST/})
+elif [[ $ROCM_INT -ge 50600 ]]; then
+    # RCCL library files
+    RCCL_SHARE_SRC=$ROCM_HOME/lib/msccl-algorithms
+    RCCL_SHARE_DST=lib/msccl-algorithms
+    RCCL_SHARE_FILES=($(ls $RCCL_SHARE_SRC))
+
+    DEPS_AUX_SRCLIST+=(${RCCL_SHARE_FILES[@]/#/$RCCL_SHARE_SRC/})
+    DEPS_AUX_DSTLIST+=(${RCCL_SHARE_FILES[@]/#/$RCCL_SHARE_DST/})
 fi
 
+
 echo "PYTORCH_ROCM_ARCH: ${PYTORCH_ROCM_ARCH}"
 
 SCRIPTPATH="$( cd "$(dirname "$0")" ; pwd -P )"