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

@@ -5,9 +5,9 @@ set -ex
 # MKL
 MKL_VERSION=2022.2.1
 MKL_BUILD=16993
-mkdir -p /opt/intel/lib
+mkdir -p ${MKLROOT}/lib
 pushd /tmp
 curl -fsSL https://anaconda.org/intel/mkl-static/${MKL_VERSION}/download/linux-64/mkl-static-${MKL_VERSION}-intel_${MKL_BUILD}.tar.bz2 | tar xjv
-mv lib/* /opt/intel/lib/
+mv lib/* ${MKLROOT}/lib/
 curl -fsSL https://anaconda.org/intel/mkl-include/${MKL_VERSION}/download/linux-64/mkl-include-${MKL_VERSION}-intel_${MKL_BUILD}.tar.bz2 | tar xjv
-mv include /opt/intel/
+mv 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
@@ -63,7 +64,6 @@ 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) 
 # find HIP works for ROCm5.7. Not needed for ROCm6.0 and above.
 # Remove below when ROCm5.7 is not in support matrix anymore.

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
@@ -118,7 +119,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/