[ROCm] support for rocm5.3 wheel builds

common/install_miopen.sh

@@ -59,23 +59,26 @@ MIOPEN_CMAKE_COMMON_FLAGS="
 -DMIOPEN_BUILD_DRIVER=OFF
 "
 # Pull MIOpen repo and set DMIOPEN_EMBED_DB based on ROCm version
-if [[ $ROCM_INT -ge 50200 ]]; then
-    MIOPEN_CMAKE_DB_FLAGS="-DMIOPEN_EMBED_DB=gfx900_56;gfx906_60;gfx90878;gfx90a6e;gfx1030_36"
+if [[ $ROCM_INT -eq 50300 ]]; then
+    MIOPEN_CMAKE_DB_FLAGS="-DMIOPEN_EMBED_DB=gfx900_56;gfx906_60;gfx90878;gfx90a6e;gfx1030_36 -DMIOPEN_USE_MLIR=Off"
+    MIOPEN_BRANCH="release/rocm-rel-5.3-staging"
+elif [[ $ROCM_INT -eq 50200 ]]; then
+    MIOPEN_CMAKE_DB_FLAGS="-DMIOPEN_EMBED_DB=gfx900_56;gfx906_60;gfx90878;gfx90a6e;gfx1030_36 -DMIOPEN_USE_MLIR=Off"
     MIOPEN_BRANCH="release/rocm-rel-5.2-staging"
-elif [[ $ROCM_INT -ge 50100 ]]; then
+elif [[ $ROCM_INT -eq 50100 ]]; then
     MIOPEN_CMAKE_DB_FLAGS="-DMIOPEN_EMBED_DB=gfx900_56;gfx906_60;gfx90878;gfx90a6e;gfx1030_36"
     MIOPEN_BRANCH="release/rocm-rel-5.1-staging"
-elif [[ $ROCM_INT -ge 50000 ]]; then
+elif [[ $ROCM_INT -eq 50000 ]]; then
     MIOPEN_CMAKE_DB_FLAGS="-DMIOPEN_EMBED_DB=gfx900_56;gfx906_60;gfx90878;gfx90a6e;gfx1030_36"
     MIOPEN_BRANCH="release/rocm-rel-5.0-staging"
-elif [[ $ROCM_INT -ge 40500 ]]; then
+elif [[ $ROCM_INT -eq 40500 ]]; then
     MIOPEN_CMAKE_COMMON_FLAGS="${MIOPEN_CMAKE_COMMON_FLAGS} -DMIOPEN_USE_HIP_KERNELS=Off -DMIOPEN_DEFAULT_FIND_MODE=Normal"
     MIOPEN_CMAKE_DB_FLAGS="-DMIOPEN_EMBED_DB=gfx900_56;gfx906_60;gfx90878;gfx90a6e;gfx1030_36"
     MIOPEN_BRANCH="release/rocm-rel-4.5-staging"
-elif [[ $ROCM_INT -ge 40300 ]]; then
+elif [[ $ROCM_INT -eq 40300 ]]; then
     MIOPEN_CMAKE_DB_FLAGS="-DMIOPEN_EMBED_DB=gfx900_56;gfx900_64;gfx906_60;gfx906_64;gfx90878;gfx1030_36"
     MIOPEN_BRANCH="release/rocm-rel-4.3-staging"
-elif [[ $ROCM_INT -ge 40200 ]]; then
+elif [[ $ROCM_INT -eq 40200 ]]; then
     MIOPEN_CMAKE_DB_FLAGS="-DMIOPEN_EMBED_DB=gfx803_36;gfx803_64;gfx900_56;gfx900_64;gfx906_60;gfx906_64;gfx90878"
     MIOPEN_BRANCH="rocm-4.2.x-staging"
 else

libtorch/Dockerfile

@@ -71,9 +71,10 @@ ENV MKLROOT /opt/intel
 ADD ./common/install_rocm.sh install_rocm.sh
 ADD ./common/install_rocm_drm.sh install_rocm_drm.sh
 ADD ./common/install_rocm_magma.sh install_rocm_magma.sh
-# gfortran needed for building magma from source for ROCm
+# gfortran and python needed for building magma from source for ROCm
 RUN apt-get update -y && \
     apt-get install gfortran -y && \
+    apt-get install python -y && \
     apt-get clean
 
 FROM rocm as rocm5.1.1

manywheel/build_common.sh

@@ -162,11 +162,23 @@ else
     echo "BUILD_DEBUG_INFO was not set, skipping debug info"
 fi
 
+if [[ "$DISABLE_RCCL" = 1 ]]; then
+    echo "Disabling NCCL/RCCL in pyTorch"
+    USE_RCCL=0
+    USE_NCCL=0
+    USE_KINETO=0
+else
+    USE_RCCL=1
+    USE_NCCL=1
+    USE_KINETO=1
+fi
+
 echo "Calling setup.py bdist at $(date)"
 time CMAKE_ARGS=${CMAKE_ARGS[@]} \
-     EXTRA_CAFFE2_CMAKE_FLAGS=${EXTRA_CAFFE2_CMAKE_FLAGS[@]} \
-     BUILD_LIBTORCH_CPU_WITH_DEBUG=$BUILD_DEBUG_INFO \
-     python setup.py bdist_wheel -d /tmp/$WHEELHOUSE_DIR
+    EXTRA_CAFFE2_CMAKE_FLAGS=${EXTRA_CAFFE2_CMAKE_FLAGS[@]} \
+    BUILD_LIBTORCH_CPU_WITH_DEBUG=$BUILD_DEBUG_INFO \
+    USE_NCCL=${USE_NCCL} USE_RCCL=${USE_RCCL} USE_KINETO=${USE_KINETO} \
+    python setup.py bdist_wheel -d /tmp/$WHEELHOUSE_DIR
 echo "Finished setup.py bdist at $(date)"
 
 # Build libtorch packages

manywheel/build_rocm.sh

@@ -89,7 +89,80 @@ else
 fi
 ROCM_INT=$(($ROCM_VERSION_MAJOR * 10000 + $ROCM_VERSION_MINOR * 100 + $ROCM_VERSION_PATCH))
 
-if [[ $ROCM_INT -ge 50200 ]]; then
+if [[ $ROCM_INT -ge 50300 ]]; then
+DEPS_LIST=(
+    "/opt/rocm/lib/libMIOpen.so.1"
+    "/opt/rocm/lib/libamdhip64.so.5"
+    "/opt/rocm/lib/libhipblas.so.0"
+    "/opt/rocm/lib/libhipfft.so"
+    "/opt/rocm/lib/libhiprand.so.1"
+    "/opt/rocm/lib/libhipsparse.so.0"
+    "/opt/rocm/lib/libhsa-runtime64.so.1"
+    "/opt/rocm/lib/libamd_comgr.so.2"
+    "/opt/rocm/magma/lib/libmagma.so"
+    "/opt/rocm/lib/librccl.so.1"
+    "/opt/rocm/lib/librocblas.so.0"
+    "/opt/rocm/lib/librocfft-device-0.so.0"
+    "/opt/rocm/lib/librocfft-device-1.so.0"
+    "/opt/rocm/lib/librocfft-device-2.so.0"
+    "/opt/rocm/lib/librocfft-device-3.so.0"
+    "/opt/rocm/lib/librocfft.so.0"
+    "/opt/rocm/lib/librocm_smi64.so.5"
+    "/opt/rocm/lib/librocrand.so.1"
+    "/opt/rocm/lib/librocsolver.so.0"
+    "/opt/rocm/lib/librocsparse.so.0"
+    "/opt/rocm/lib/libroctracer64.so.4"
+    "/opt/rocm/lib/libroctx64.so.4"
+    "$LIBGOMP_PATH"
+    "$LIBNUMA_PATH"
+    "$LIBELF_PATH"
+    "$LIBTINFO_PATH"
+    "$LIBDRM_PATH"
+    "$LIBDRM_AMDGPU_PATH"
+)
+
+DEPS_SONAME=(
+    "libMIOpen.so.1"
+    "libamdhip64.so.5"
+    "libhipblas.so.0"
+    "libhipfft.so"
+    "libhiprand.so.1"
+    "libhipsparse.so.0"
+    "libhsa-runtime64.so.1"
+    "libamd_comgr.so.2"
+    "libmagma.so"
+    "librccl.so.1"
+    "librocblas.so.0"
+    "librocfft-device-0.so.0"
+    "librocfft-device-1.so.0"
+    "librocfft-device-2.so.0"
+    "librocfft-device-3.so.0"
+    "librocfft.so.0"
+    "librocm_smi64.so.5"
+    "librocrand.so.1"
+    "librocsolver.so.0"
+    "librocsparse.so.0"
+    "libroctracer64.so.4"
+    "libroctx64.so.4"
+    "libgomp.so.1"
+    "libnuma.so.1"
+    "libelf.so.1"
+    "libtinfo.so.5"
+    "libdrm.so.2"
+    "libdrm_amdgpu.so.1"
+)
+
+DEPS_AUX_SRCLIST=(
+    "/opt/rocm/lib/rocblas/library/*"
+    "/opt/amdgpu/share/libdrm/amdgpu.ids"
+)
+
+DEPS_AUX_DSTLIST=(
+    "lib/rocblas/library/."
+    "share/libdrm/amdgpu.ids"
+)
+
+elif [[ $ROCM_INT -ge 50200 ]]; then
 DEPS_LIST=(
     "/opt/rocm/lib/libMIOpen.so.1"
     "/opt/rocm/lib/libamdhip64.so.5"