From 52f2493f0e22a16c17c6e39c6b1f88e1a988c9ae Mon Sep 17 00:00:00 2001
From: Jeff Daily <jeff.daily@amd.com>
Date: Wed, 8 Feb 2023 14:03:07 -0800
Subject: [PATCH] rocm libtorch prebuild magma; fix manylinux cmake version
 (#1296)

---
 libtorch/Dockerfile      |  6 +++---
 libtorch/build_docker.sh |  2 +-
 manywheel/Dockerfile     | 13 ++++++-------
 3 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/libtorch/Dockerfile b/libtorch/Dockerfile
index 38f3da9d8..3a116f8b5 100644
--- a/libtorch/Dockerfile
+++ b/libtorch/Dockerfile
@@ -62,7 +62,7 @@ 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 && \
@@ -72,12 +72,12 @@ RUN apt-get update -y && \
 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
+#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 ${BASE_TARGET} as final
 # Install LLVM
diff --git a/libtorch/build_docker.sh b/libtorch/build_docker.sh
index bbf42b1d0..d5505471b 100755
--- a/libtorch/build_docker.sh
+++ b/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}
+        GPU_IMAGE=rocm/dev-ubuntu-20.04:${GPU_ARCH_VERSION}-magma
         PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908"
         ROCM_REGEX="([0-9]+)\.([0-9]+)[\.]?([0-9]*)"
         if [[ $GPU_ARCH_VERSION =~ $ROCM_REGEX ]]; then
diff --git a/manywheel/Dockerfile b/manywheel/Dockerfile
index 7dc20257b..0b08a3149 100644
--- a/manywheel/Dockerfile
+++ b/manywheel/Dockerfile
@@ -160,15 +160,14 @@ 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
+# cmake is already installed inside the rocm base image, but both 2 and 3 exist
+# cmake3 is needed for pytorch build, and formerly the MIOpen build
+RUN rpm -e cmake && \
+    pip3 install cmake==3.18.4 && \
+    ln -sf /usr/local/bin/cmake /usr/bin/cmake && \
+    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
-## cmake is already installed inside the rocm base image, but both 2 and 3 exist
-## cmake3 is needed for the later MIOpen custom build, so that step is last.
-#RUN rpm -e cmake && \
-#    pip3 install cmake==3.18.4 && \
-#    ln -sf /usr/local/bin/cmake /usr/bin/cmake && \
-#    ln -sf /usr/local/bin/cmake /usr/bin/cmake3
-#
 #ADD ./common/install_miopen.sh install_miopen.sh
 #RUN bash ./install_miopen.sh ${ROCM_VERSION} && rm install_miopen.sh