[DRAFT] Added magma conda build scripts for ROCm

common/install_conda.sh

@@ -5,7 +5,7 @@ set -ex
 # Anaconda
 wget -q https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh
 chmod +x  Miniconda3-latest-Linux-x86_64.sh
-./Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda
+bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda
 rm Miniconda3-latest-Linux-x86_64.sh
 export PATH=/opt/conda/bin:$PATH
 conda install -y conda-build anaconda-client git ninja

common/install_magma.sh

@@ -1,35 +1,21 @@
 #!/usr/bin/env bash
 
-set -eou pipefail
+#!/bin/bash
 
-MAGMA_VERSION="2.5.2"
+set -ex
 
 function do_install() {
-    cuda_version=$1
-    cuda_version_nodot=${1/./}
-
-    if [[ ${cuda_version_nodot} == 116 ]]; then
-        MAGMA_VERSION="2.6.1"
-        magma_archive="magma-cuda${cuda_version_nodot}-${MAGMA_VERSION}-0.tar.bz2"
-    elif [[ ${cuda_version_nodot} == 117 ]]; then
-	MAGMA_VERSION="2.6.1"
-	magma_archive="magma-cuda${cuda_version_nodot}-${MAGMA_VERSION}-0.tar.bz2"
-    else
-        magma_archive="magma-cuda${cuda_version_nodot}-${MAGMA_VERSION}-1.tar.bz2"
-    fi
-
-    cuda_dir="/usr/local/cuda-${cuda_version}"
-    (
-        set -x
-        tmp_dir=$(mktemp -d)
-        pushd ${tmp_dir}
-        wget -q https://anaconda.org/pytorch/magma-cuda${cuda_version_nodot}/${MAGMA_VERSION}/download/linux-64/${magma_archive}
-        tar -xvf "${magma_archive}"
-        mkdir -p "${cuda_dir}/magma"
-        mv include "${cuda_dir}/magma/include"
-        mv lib "${cuda_dir}/magma/lib"
-        popd
-    )
+    magma_archive_path=$1
+    magma_archive="$(basename "$magma_archive_path")"
+    rocm_path="/opt/rocm/"
+    tmp_dir=$(mktemp -d)
+    cp $magma_archive_path $tmp_dir
+    pushd ${tmp_dir}
+    tar -xvf "${magma_archive}" $tmp_dir
+    mkdir -p "${rocm_path}/magma"
+    mv $tmp_dir/include "${rocm_path}/magma/include"
+    mv $tmp_dir/lib "${rocm_path}/magma/lib"
+    popd
 }
 
 do_install $1

common/install_miopen.sh

@@ -59,23 +59,29 @@ 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 -ge 50400 ]] && [[ $ROCM_INT -lt 50500 ]]; 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.4-staging"
+elif [[ $ROCM_INT -ge 50300 ]] && [[ $ROCM_INT -lt 50400 ]]; 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 -ge 50200 ]] && [[ $ROCM_INT -lt 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.2-staging"
-elif [[ $ROCM_INT -ge 50100 ]]; then
+elif [[ $ROCM_INT -ge 50100 ]] && [[ $ROCM_INT -lt 50200 ]]; 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 -ge 50000 ]] && [[ $ROCM_INT -lt 50100 ]]; 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 -ge 40500 ]] && [[ $ROCM_INT -lt 50000 ]]; 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 -ge 40300 ]] && [[ $ROCM_INT -lt 40500 ]]; 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 -ge 40200 ]] && [[ $ROCM_INT -lt 40300 ]]; 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

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
@@ -33,7 +33,12 @@ install_ubuntu() {
     if [[ $(ver $ROCM_VERSION) -ge $(ver 4.5) ]]; then
         # Add amdgpu repository
         UBUNTU_VERSION_NAME=`cat /etc/os-release | grep UBUNTU_CODENAME | awk -F= '{print $2}'`
-        local amdgpu_baseurl="https://repo.radeon.com/amdgpu/${AMDGPU_VERSIONS[$ROCM_VERSION]}/ubuntu"
+        local amdgpu_baseurl
+        if [[ $(ver $ROCM_VERSION) -ge $(ver 5.3) ]]; then
+          amdgpu_baseurl="https://repo.radeon.com/amdgpu/${ROCM_VERSION}/ubuntu"
+        else
+          amdgpu_baseurl="https://repo.radeon.com/amdgpu/${AMDGPU_VERSIONS[$ROCM_VERSION]}/ubuntu"
+        fi
         echo "deb [arch=amd64] ${amdgpu_baseurl} ${UBUNTU_VERSION_NAME} main" > /etc/apt/sources.list.d/amdgpu.list
     fi
 
@@ -83,7 +88,12 @@ install_centos() {
 
   if [[ $(ver $ROCM_VERSION) -ge $(ver 4.5) ]]; then
       # Add amdgpu repository
-      local amdgpu_baseurl="https://repo.radeon.com/amdgpu/${AMDGPU_VERSIONS[$ROCM_VERSION]}/rhel/7.9/main/x86_64"
+      local amdgpu_baseurl
+      if [[ $(ver $ROCM_VERSION) -ge $(ver 5.3) ]]; then
+        amdgpu_baseurl="https://repo.radeon.com/amdgpu/${ROCM_VERSION}/rhel/7.9/main/x86_64"
+      else
+        amdgpu_baseurl="https://repo.radeon.com/amdgpu/${AMDGPU_VERSIONS[$ROCM_VERSION]}/rhel/7.9/main/x86_64"
+      fi
       echo "[AMDGPU]" > /etc/yum.repos.d/amdgpu.repo
       echo "name=AMDGPU" >> /etc/yum.repos.d/amdgpu.repo
       echo "baseurl=${amdgpu_baseurl}" >> /etc/yum.repos.d/amdgpu.repo

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

@@ -71,18 +71,29 @@ 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.2
+RUN ROCM_VERSION=5.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
+
+FROM rocm as rocm5.1.3
+RUN ROCM_VERSION=5.1.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.1.1
 RUN ROCM_VERSION=5.1.1 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.2
-RUN ROCM_VERSION=5.2 bash ./install_rocm.sh && rm install_rocm.sh
+FROM rocm as rocm5.1
+RUN ROCM_VERSION=5.1 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
 

magma/Makefile

@@ -1,60 +1,44 @@
-SHELL=/usr/bin/env bash
-
-DESIRED_CUDA ?= 11.3
-PACKAGE_NAME ?= magma-cuda113
-CUDA_ARCH_LIST ?= -gencode arch=compute_37,code=sm_37 -gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_70,code=sm_70
-
-DOCKER_RUN = set -eou pipefail; docker run --rm -i \
+ROCM_VERSION ?= 5.4
+PACKAGE_NAME ?= magma-rocm5.4
+PYTORCH_ROCM_ARCH ?= gfx906
+
+DOCKER_RUN = docker run -it \
+	--device=/dev/kfd \
+	--device=/dev/dri \
+	--group-add video \
+	--cap-add=SYS_PTRACE \
 	-v $(shell git rev-parse --show-toplevel):/builder \
 	-w /builder \
-	-e DESIRED_CUDA=${DESIRED_CUDA} \
+	-e ROCM_VERSION=${ROCM_VERSION} \
 	-e PACKAGE_NAME=${PACKAGE_NAME} \
-	-e CUDA_ARCH_LIST="${CUDA_ARCH_LIST}" \
-	"pytorch/conda-builder:cuda${DESIRED_CUDA}" \
-	magma/build_magma.sh
+	-e PACKAGE_VERSION=${PACKAGE_VERSION} \
+	-e PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH}" \
+	"${MANYLINUX_IMAGE}" \
+ 	/bin/bash magma/build_magma.sh	
 
 .PHONY: all
-all: magma-cuda117
-all: magma-cuda116
-all: magma-cuda115
-all: magma-cuda113
-all: magma-cuda102
+all: magma-rocm5.4
+all: magma-rocm5.3
 
 .PHONY:
 clean:
 	$(RM) -r magma-*
 	$(RM) -r output
 
-.PHONY: magma-cuda117
-magma-cuda117: DESIRED_CUDA := 11.7
-magma-cuda117: PACKAGE_NAME := magma-cuda117
-magma-cuda117: CUDA_ARCH_LIST += -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86
-magma-cuda117:
-	$(DOCKER_RUN)
-
-.PHONY: magma-cuda116
-magma-cuda116: DESIRED_CUDA := 11.6
-magma-cuda116: PACKAGE_NAME := magma-cuda116
-magma-cuda116: CUDA_ARCH_LIST += -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86
-magma-cuda116:
-	$(DOCKER_RUN)
-
-.PHONY: magma-cuda115
-magma-cuda115: DESIRED_CUDA := 11.5
-magma-cuda115: PACKAGE_NAME := magma-cuda115
-magma-cuda115: CUDA_ARCH_LIST += -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86
-magma-cuda115:
-	$(DOCKER_RUN)
-
-.PHONY: magma-cuda113
-magma-cuda113: DESIRED_CUDA := 11.3
-magma-cuda113: PACKAGE_NAME := magma-cuda113
-magma-cuda113: CUDA_ARCH_LIST += -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86
-magma-cuda113:
+.PHONY: magma-rocm5.4
+magma-rocm5.4: ROCM_VERSION := 5.4
+magma-rocm5.4: PACKAGE_NAME := magma-rocm
+magma-rocm5.4: PACKAGE_VERSION := rocm5.4_magma
+magma-rocm5.4: PYTORCH_ROCM_ARCH := gfx906
+magma-rocm5.4: MANYLINUX_IMAGE := pytorch/manylinux-builder:rocm5.4.2
+magma-rocm5.4:
 	$(DOCKER_RUN)
 
-.PHONY: magma-cuda102
-magma-cuda102: DESIRED_CUDA := 10.2
-magma-cuda102: PACKAGE_NAME := magma-cuda102
-magma-cuda102:
+.PHONY: magma-rocm5.3
+magma-rocm5.3: ROCM_VERSION := 5.3
+magma-rocm5.3: PACKAGE_NAME := magma-rocm
+magma-rocm5.3: PACKAGE_VERSION := rocm5.3-magma
+magma-rocm5.3: PYTORCH_ROCM_ARCH := gfx906
+magma-rocm5.3: MANYLINUX_IMAGE := pytorch/manylinux-builder:rocm5.3
+magma-rocm5.3:
 	$(DOCKER_RUN)

magma/README.md

@@ -2,13 +2,13 @@
 
 ## Building
 
-Look in the `Makefile` for available targets to build. To build any target, for example `magma-cuda111`, run
+Look in the `Makefile` for available targets to build. To build any target, for example `magma-rocm5.4`, run
 
 ```
-make magma-cuda111
+make magma-rocm5.4
 ```
 
-This should spawn a docker image with `devtoolset3` installed in order to keep binary size down. Within the docker image, it should run `build_magma.sh` with the correct environment variables set, which should package the necessary files with `conda build`.
+This should spawn a manylinux docker image. Within the docker image, it should run `build_magma.sh` with the correct environment variables set, which should package the necessary files with `conda build`.
 
 More specifically, `build_magma.sh` copies over the relevant files from the `package_files` directory depending on the CUDA version. More information on conda-build can be found [here](https://docs.conda.io/projects/conda-build/en/latest/concepts/recipe.html).
 
@@ -19,15 +19,15 @@ Outputted binaries should be in the `output` folder.
 Once you have built the binaries push them with:
 
 ```
-anaconda upload -u pytorch --force output/*/magma-cuda*.bz2
+anaconda upload -u pytorch --force output/*/magma-rocm*.bz2
 ```
 
-If you do not have upload permissions, please ping @seemethere or @soumith to gain access
+If you do not have upload permissions, please ping @jithunair or @jataylo to gain access
 
 ## New versions
 
-New CUDA versions can be added by creating a new make target with the next desired version. For CUDA version NN.n, the target should be named `magma-cudaNNn`.
+New ROCm versions can be added by creating a new make target with the next desired version. For ROCm version n.n.n, the target should be named `magma-rocmn.n`.
 
-Make sure to edit the appropriate environment variables (e.g., DESIRED_CUDA, CUDA_ARCH_LIST, PACKAGE_NAME) in the `Makefile` accordingly. Remember also to check `build_magma.sh` to ensure the logic for copying over the files remains correct.
+Make sure to edit the appropriate environment variables (e.g., ROCM_VERSION, PYTORCH_ROCM_ARCH) in the `Makefile` accordingly. Remember also to check `build_magma.sh` to ensure the logic for copying over the files remains correct.
 
-New patches can be added by editing `Makefile`, `build_magma.sh`, and `package_files/meta.yaml` the same way `cudaPointerAttributes.patch` is implemented.
+New patches can be added by editing `Makefile`, `build_magma.sh`, and `package_files/meta.yaml`.

magma/build_magma.sh

@@ -1,20 +1,22 @@
-#!/usr/bin/env bash
+#!/bin/bash
 
-set -eou pipefail
+set -e
 
-# Create a folder to be packaged
+# Install conda - remove once have conda manylinux images instead
+bash common/install_conda.sh
+export PATH=/opt/conda/bin:$PATH
+
+# Create the folder to be packaged
 PACKAGE_DIR=magma/${PACKAGE_NAME}
 PACKAGE_FILES=magma/package_files
 mkdir ${PACKAGE_DIR}
 cp ${PACKAGE_FILES}/build.sh ${PACKAGE_DIR}/build.sh
 cp ${PACKAGE_FILES}/meta.yaml ${PACKAGE_DIR}/meta.yaml
-cp ${PACKAGE_FILES}/thread_queue.patch ${PACKAGE_DIR}/thread_queue.patch
-cp ${PACKAGE_FILES}/cmakelists.patch ${PACKAGE_DIR}/cmakelists.patch
-cp ${PACKAGE_FILES}/getrf_shfl.patch ${PACKAGE_DIR}/getrf_shfl.patch
-cp ${PACKAGE_FILES}/getrf_nbparam.patch ${PACKAGE_DIR}/getrf_nbparam.patch
 
+# Conda prerequisites
 conda install -yq conda-build conda-verify
-. ./conda/switch_cuda_version.sh "${DESIRED_CUDA}"
+
+# Conda build
 (
     set -x
     conda build --output-folder magma/output "${PACKAGE_DIR}"