set MKLROOT env var in manywheel/libtorch Dockerfile; default to CI env value; use env var in all scripts

Author: jithunnair-amd

common/install_mkl.sh

@@ -4,12 +4,12 @@ set -ex
 # MKL
 MKL_VERSION=2024.2.0
 
-mkdir -p /opt/intel/
+mkdir -p ${MKLROOT}
 pushd /tmp
 
 python3 -mpip install wheel
 python3 -mpip download -d . mkl-static==${MKL_VERSION}
 python3 -m wheel unpack mkl_static-${MKL_VERSION}-py2.py3-none-manylinux1_x86_64.whl
 python3 -m wheel unpack mkl_include-${MKL_VERSION}-py2.py3-none-manylinux1_x86_64.whl
-mv mkl_static-${MKL_VERSION}/mkl_static-${MKL_VERSION}.data/data/lib /opt/intel/
-mv mkl_include-${MKL_VERSION}/mkl_include-${MKL_VERSION}.data/data/include /opt/intel/
+mv mkl_static-${MKL_VERSION}/mkl_static-${MKL_VERSION}.data/data/lib ${MKLROOT}
+mv mkl_include-${MKL_VERSION}/mkl_include-${MKL_VERSION}.data/data/include ${MKLROOT}

common/install_rocm_magma.sh

@@ -1,25 +1,17 @@
 #!/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}
-
+MKLROOT=${MKLROOT:-/opt/conda/envs/py_$ANACONDA_PYTHON_VERSION}
 # "install" hipMAGMA into /opt/rocm/magma by copying after build
 git clone https://bitbucket.org/icl/magma.git
 pushd magma
-if [[ $PYTORCH_BRANCH == "release/1.10.1" ]]; then
-    git checkout magma_ctrl_launch_bounds
-else
-    git checkout a1625ff4d9bc362906bd01f805dbbe12612953f6
-fi
+
+# Version 2.7.2 + ROCm related updates
+git checkout a1625ff4d9bc362906bd01f805dbbe12612953f6
+
 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

libtorch/Dockerfile

@@ -1,5 +1,6 @@
 ARG BASE_TARGET=base
 ARG GPU_IMAGE=ubuntu:20.04
+ARG MKLROOT /opt/intel
 FROM ${GPU_IMAGE} as base
 
 ENV DEBIAN_FRONTEND=noninteractive
@@ -69,8 +70,7 @@ RUN ln -sf /usr/local/cuda-12.4 /usr/local/cuda
 FROM cpu as rocm
 ARG PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH}
-ENV MKLROOT /opt/intel
-# Adding ROCM_PATH env var so that LoadHip.cmake (even with logic updated for ROCm6.0)
+# Adding ROCM_PATH env var so that LoadHip.cmake (even with logic updated for ROCm6.0) 
 # find HIP works for ROCm5.7. Not needed for ROCm6.0 and above.
 # Remove below when ROCm5.7 is not in support matrix anymore.
 ENV ROCM_PATH /opt/rocm

manywheel/Dockerfile

@@ -3,6 +3,7 @@ ARG ROCM_VERSION=3.7
 ARG BASE_CUDA_VERSION=11.8
 
 ARG GPU_IMAGE=centos:7
+ARG MKLROOT=/opt/intel
 FROM centos:7 as base
 
 ENV LC_ALL en_US.UTF-8
@@ -131,7 +132,7 @@ COPY --from=openssl            /opt/openssl                          /opt/openss
 COPY --from=python             /opt/python                           /opt/python
 COPY --from=python             /opt/_internal                        /opt/_internal
 COPY --from=python             /opt/python/cp39-cp39/bin/auditwheel /usr/local/bin/auditwheel
-COPY --from=intel              /opt/intel                            /opt/intel
+COPY --from=intel              ${MKLROOT}                            ${MKLROOT} 
 COPY --from=patchelf           /usr/local/bin/patchelf               /usr/local/bin/patchelf
 COPY --from=jni                /usr/local/include/jni.h              /usr/local/include/jni.h
 COPY --from=libpng             /usr/local/bin/png*                   /usr/local/bin/