Add magma conda package install support via MAGMA_PACKAGE_SOURCE

common/install_rocm.sh

@@ -11,7 +11,7 @@ ver() {
 }
 
 # Map ROCm version to AMDGPU version
-declare -A AMDGPU_VERSIONS=( ["5.0"]="21.50" ["5.1.1"]="22.10.1" ["5.2"]="22.20" )
+declare -A AMDGPU_VERSIONS=( ["5.0"]="21.50" ["5.1.1"]="22.10.1" ["5.1.3"]="22.10.3" ["5.2"]="22.20" )
 
 install_ubuntu() {
     apt-get update

common/install_rocm_magma.sh

@@ -1,44 +1,86 @@
 #!/bin/bash
 
-# TODO upstream differences from this file is into the (eventual) one in pytorch
-# - (1) check for static lib mkl
-# - (2) MKLROOT as env var
-
 set -ex
 
-# TODO (2)
-MKLROOT=${MKLROOT:-/opt/intel}
-
-# "install" hipMAGMA into /opt/rocm/magma by copying after build
-git clone https://bitbucket.org/icl/magma.git
-pushd magma
-# fix for magma_queue memory leak issue
-git checkout c62d700d880c7283b33fb1d615d62fc9c7f7ca21
-cp make.inc-examples/make.inc.hip-gcc-mkl make.inc
-echo 'LIBDIR += -L$(MKLROOT)/lib' >> make.inc
-# TODO (1)
-if [[ -f "${MKLROOT}/lib/libmkl_core.a" ]]; then
-    echo 'LIB = -Wl,--start-group -lmkl_gf_lp64 -lmkl_gnu_thread -lmkl_core -Wl,--end-group -lpthread -lstdc++ -lm -lgomp -lhipblas -lhipsparse' >> make.inc
-fi
-echo 'LIB += -Wl,--enable-new-dtags -Wl,--rpath,/opt/rocm/lib -Wl,--rpath,$(MKLROOT)/lib -Wl,--rpath,/opt/rocm/magma/lib -ldl' >> make.inc
-echo 'DEVCCFLAGS += --gpu-max-threads-per-block=256' >> make.inc
-export PATH="${PATH}:/opt/rocm/bin"
-if [[ -n "$PYTORCH_ROCM_ARCH" ]]; then
-  amdgpu_targets=`echo $PYTORCH_ROCM_ARCH | sed 's/;/ /g'`
+function install_magma_source() {
+    # TODO upstream differences from this file is into the (eventual) one in pytorch
+    # - (1) check for static lib mkl
+    # - (2) MKLROOT as env var
+
+    # TODO (2)
+    MKLROOT=${MKLROOT:-/opt/intel}
+
+    # "install" hipMAGMA into /opt/rocm/magma by copying after build
+    git clone https://bitbucket.org/icl/magma.git
+    pushd magma
+    # fix for magma_queue memory leak issue
+    git checkout c62d700d880c7283b33fb1d615d62fc9c7f7ca21
+    cp make.inc-examples/make.inc.hip-gcc-mkl make.inc
+    echo 'LIBDIR += -L$(MKLROOT)/lib' >> make.inc
+    # TODO (1)
+    if [[ -f "${MKLROOT}/lib/libmkl_core.a" ]]; then
+        echo 'LIB = -Wl,--start-group -lmkl_gf_lp64 -lmkl_gnu_thread -lmkl_core -Wl,--end-group -lpthread -lstdc++ -lm -lgomp -lhipblas -lhipsparse' >> make.inc
+    fi
+    echo 'LIB += -Wl,--enable-new-dtags -Wl,--rpath,/opt/rocm/lib -Wl,--rpath,$(MKLROOT)/lib -Wl,--rpath,/opt/rocm/magma/lib -ldl' >> make.inc
+    echo 'DEVCCFLAGS += --gpu-max-threads-per-block=256' >> make.inc
+    export PATH="${PATH}:/opt/rocm/bin"
+    if [[ -n "$PYTORCH_ROCM_ARCH" ]]; then
+    amdgpu_targets=`echo $PYTORCH_ROCM_ARCH | sed 's/;/ /g'`
+    else
+    amdgpu_targets=`rocm_agent_enumerator | grep -v gfx000 | sort -u | xargs`
+    fi
+    for arch in $amdgpu_targets; do
+    echo "DEVCCFLAGS += --amdgpu-target=$arch" >> make.inc
+    done
+    # hipcc with openmp flag may cause isnan() on __device__ not to be found; depending on context, compiler may attempt to match with host definition
+    sed -i 's/^FOPENMP/#FOPENMP/g' make.inc
+    make -f make.gen.hipMAGMA -j $(nproc)
+    LANG=C.UTF-8 make lib/libmagma.so -j $(nproc) MKLROOT="${MKLROOT}"
+    make testing/testing_dgemm -j $(nproc) MKLROOT="${MKLROOT}"
+    popd
+    mkdir -p /opt/rocm/magma
+    mv magma/include /opt/rocm/magma
+    mv magma/lib /opt/rocm/magma
+    rm -rf magma
+}
+
+function install_magma_package() {
+
+    MAGMA_VERSION="2.6.2"
+    rocm_path="/opt/rocm"
+    tmp_dir=$(mktemp -d)
+    pushd ${tmp_dir}
+    wget --no-check-certificate -q $MAGMA_PACKAGE_SOURCE
+    tar -xf *.tar*
+    mkdir -p "${rocm_path}/magma"
+
+    if [ -e "$tmp_dir/magma/include" ]; then
+        mv "$tmp_dir/magma/include" "${rocm_path}/magma/include"
+        echo "Successfully installed MAGMA include files to ${rocm_path}/magma/include"
+    else
+        echo "Error: MAGMA include files not found in $tmp_dir/magma/include"
+    fi
+
+    if [ -e "$tmp_dir/magma/lib/" ]; then
+        mv "$tmp_dir/magma/lib/" "${rocm_path}/magma/lib"
+        echo "Successfully installed MAGMA library files to ${rocm_path}/magma/lib"
+    else
+        echo "Error: MAGMA library file not found in $tmp_dir/magma/lib"
+    fi
+
+    if [ ! -d "${rocm_path}/magma" ]; then
+        echo "Error: MAGMA installation failed"
+        exit 1
+    fi
+
+    popd
+    rm -rf $tmp_dir
+}
+
+if [ -z "$MAGMA_PACKAGE_SOURCE" ]; then
+    echo "MAGMA_PACKAGE_SOURCE is not set, building magma from source"
+    install_magma_source
 else
-  amdgpu_targets=`rocm_agent_enumerator | grep -v gfx000 | sort -u | xargs`
+    echo "MAGMA_PACKAGE_SOURCE is set, installing magma from source"
+    install_magma_package
 fi
-for arch in $amdgpu_targets; do
-  echo "DEVCCFLAGS += --amdgpu-target=$arch" >> make.inc
-done
-# hipcc with openmp flag may cause isnan() on __device__ not to be found; depending on context, compiler may attempt to match with host definition
-sed -i 's/^FOPENMP/#FOPENMP/g' make.inc
-make -f make.gen.hipMAGMA -j $(nproc)
-LANG=C.UTF-8 make lib/libmagma.so -j $(nproc) MKLROOT="${MKLROOT}"
-make testing/testing_dgemm -j $(nproc) MKLROOT="${MKLROOT}"
-popd
-mkdir -p /opt/rocm/magma
-mv magma/include /opt/rocm/magma
-mv magma/lib /opt/rocm/magma
-rm -rf magma
-

common/install_rocm_miopen_rpm.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+
+set -ex
+
+# Create the tmp dir to extract into
+EXTRACTDIR_ROOT=/extract_miopen_rpm
+mkdir -p ${EXTRACTDIR_ROOT}
+echo "Creating temporary directory for rpm download..."
+
+# Fail if rpm source is not available
+if ! wget -P ${EXTRACTDIR_ROOT} ${MIOPEN_RPM_SOURCE}; then
+  echo 'ERROR: Failed to download MIOpen package.'
+  exit 1
+fi
+echo "MIOpen package download complete..."
+
+# Extract rpm in EXTRACT_DIR
+cd ${EXTRACTDIR_ROOT}
+miopen_rpm=$(ls *.rpm)
+rpm2cpio ${miopen_rpm} | cpio -idmv 
+
+# Copy libMIOpen.so.1 over existing
+source_file=$(ls opt/rocm-*/lib/libMIOpen.so.1.0*)
+dest_file=$(ls /opt/rocm-${ROCM_VERSION}*/lib/libMIOpen.so.1.0*)
+if [ -e ${source_file} ] && [ -e ${dest_file} ]; then
+  echo "Source .so: ${source_file}"
+  echo "Dest .so: ${dest_file}"
+  cp $source_file $dest_file
+else
+  echo 'ERROR: either the source or destination path for libMIOpen.so.1.0 does not exist'
+  exit 1
+fi
+echo "libMIOpen so file from RPM copied to existing MIOpen install..."
+
+# Clean up extracted dir
+rm -rf ${EXTRACTDIR_ROOT}
+echo "Removed temporary directory..."

libtorch/Dockerfile

@@ -62,22 +62,22 @@ ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH}
 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
+ADD ./common/install_rocm_magma.sh install_rocm_magma.sh
 # 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.3
-RUN ROCM_VERSION=5.3 bash ./install_rocm.sh && rm install_rocm.sh
-RUN bash ./install_rocm_drm.sh && rm install_rocm_drm.sh
-#RUN bash ./install_rocm_magma.sh && rm install_rocm_magma.sh
-
 FROM rocm as rocm5.4.2
 RUN ROCM_VERSION=5.4.2 bash ./install_rocm.sh && rm install_rocm.sh
 RUN bash ./install_rocm_drm.sh && rm install_rocm_drm.sh
-#RUN bash ./install_rocm_magma.sh && rm install_rocm_magma.sh
+RUN bash ./install_rocm_magma.sh && rm install_rocm_magma.sh
+
+FROM rocm as rocm5.3
+RUN ROCM_VERSION=5.3 bash ./install_rocm.sh && rm install_rocm.sh
+RUN bash ./install_rocm_drm.sh && rm install_rocm_drm.sh
+RUN bash ./install_rocm_magma.sh && rm install_rocm_magma.sh
 
 FROM ${BASE_TARGET} as final
 # Install LLVM

libtorch/build_docker.sh

@@ -27,7 +27,7 @@ case ${GPU_ARCH_TYPE} in
     rocm)
         BASE_TARGET=rocm${GPU_ARCH_VERSION}
         DOCKER_TAG=rocm${GPU_ARCH_VERSION}
-        GPU_IMAGE=rocm/dev-ubuntu-20.04:${GPU_ARCH_VERSION}-magma
+        GPU_IMAGE=rocm/dev-ubuntu-20.04:${GPU_ARCH_VERSION}
         PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908"
         ROCM_REGEX="([0-9]+)\.([0-9]+)[\.]?([0-9]*)"
         if [[ $GPU_ARCH_VERSION =~ $ROCM_REGEX ]]; then

manywheel/Dockerfile

@@ -157,15 +157,14 @@ FROM cpu_final as rocm_final
 ARG ROCM_VERSION=3.7
 ARG PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH}
+ARG MAGMA_PACKAGE_SOURCE
+ENV MAGMA_PACKAGE_SOURCE ${MAGMA_PACKAGE_SOURCE}
 # Install ROCm
 ADD ./common/install_rocm.sh install_rocm.sh
 RUN ROCM_VERSION=${ROCM_VERSION} bash ./install_rocm.sh && rm install_rocm.sh
 ADD ./common/install_rocm_drm.sh install_rocm_drm.sh
 RUN bash ./install_rocm_drm.sh && rm install_rocm_drm.sh
+ADD ./common/install_rocm_magma.sh install_rocm_magma.sh
+RUN bash ./install_rocm_magma.sh && rm install_rocm_magma.sh
 # cmake3 is needed for the MIOpen build
 RUN ln -sf /usr/local/bin/cmake /usr/bin/cmake3
-### The following is now performed beforehand in a new GPU_IMAGE with magma and miopen preinstalled
-#ADD ./common/install_rocm_magma.sh install_rocm_magma.sh
-#RUN bash ./install_rocm_magma.sh && rm install_rocm_magma.sh
-#ADD ./common/install_miopen.sh install_miopen.sh
-#RUN bash ./install_miopen.sh ${ROCM_VERSION} && rm install_miopen.sh

manywheel/build_docker.sh

@@ -44,7 +44,7 @@ case ${GPU_ARCH_TYPE} in
         TARGET=rocm_final
         DOCKER_TAG=rocm${GPU_ARCH_VERSION}
         LEGACY_DOCKER_IMAGE=${DOCKER_REGISTRY}/pytorch/manylinux-rocm:${GPU_ARCH_VERSION}
-        GPU_IMAGE=rocm/dev-centos-7:${GPU_ARCH_VERSION}-magma-miopen-staging
+        GPU_IMAGE=rocm/dev-centos-7:${GPU_ARCH_VERSION}
         PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908"
         ROCM_REGEX="([0-9]+)\.([0-9]+)[\.]?([0-9]*)"
         if [[ $GPU_ARCH_VERSION =~ $ROCM_REGEX ]]; then

manywheel/build_rocm.sh

@@ -84,10 +84,6 @@ ROCM_SO_FILES=(
     "libmagma.so"
     "librccl.so"
     "librocblas.so"
-    "librocfft-device-0.so"
-    "librocfft-device-1.so"
-    "librocfft-device-2.so"
-    "librocfft-device-3.so" 
     "librocfft.so"
     "librocm_smi64.so"   
     "librocrand.so"
@@ -97,6 +93,13 @@ ROCM_SO_FILES=(
     "libroctx64.so"
 )
 
+if [[ $ROCM_INT -lt 50500 ]]; then
+    ROCM_SO_FILES+=("librocfft-device-0.so")
+    ROCM_SO_FILES+=("librocfft-device-1.so")
+    ROCM_SO_FILES+=("librocfft-device-2.so")
+    ROCM_SO_FILES+=("librocfft-device-3.so")
+fi
+
 if [[ $ROCM_INT -ge 50400 ]]; then
     ROCM_SO_FILES+=("libhiprtc.so")
 fi
@@ -146,6 +149,11 @@ ARCH_SPECIFIC_FILES=$(ls $ROCBLAS_LIB_SRC | grep -E $ARCH)
 OTHER_FILES=$(ls $ROCBLAS_LIB_SRC | grep -v gfx)
 ROCBLAS_LIB_FILES=($ARCH_SPECIFIC_FILES $OTHER_FILES)
 
+# MIOpen library files
+MIOPEN_SHARE_SRC=$ROCM_HOME/share/miopen/db
+MIOPEN_SHARE_DST=share/miopen/db
+MIOPEN_SHARE_FILES=($(ls $MIOPEN_SHARE_SRC | grep -E $ARCH))
+
 # ROCm library files
 ROCM_SO_PATHS=()
 for lib in "${ROCM_SO_FILES[@]}"
@@ -159,6 +167,10 @@ do
     if [[ -z $file_path ]]; then 
         file_path=($(find $ROCM_HOME/ -name "$lib")) # Then search in ROCM_HOME
     fi
+    if [[ -z $file_path ]]; then
+	echo "Error: Library file $lib is not found." >&2
+	exit 1
+    fi
     ROCM_SO_PATHS[${#ROCM_SO_PATHS[@]}]="$file_path" # Append lib to array
 done
 
@@ -174,11 +186,13 @@ DEPS_SONAME=(
 
 DEPS_AUX_SRCLIST=(
     "${ROCBLAS_LIB_FILES[@]/#/$ROCBLAS_LIB_SRC/}"
+    "${MIOPEN_SHARE_FILES[@]/#/$MIOPEN_SHARE_SRC/}"
     "/opt/amdgpu/share/libdrm/amdgpu.ids"
 )
 
 DEPS_AUX_DSTLIST=(
     "${ROCBLAS_LIB_FILES[@]/#/$ROCBLAS_LIB_DST/}"
+    "${MIOPEN_SHARE_FILES[@]/#/$MIOPEN_SHARE_DST/}"
     "share/libdrm/amdgpu.ids"
 )