[ROCm]: Dockerfile and build script updates

	Add hipblaslt in Dockerfile
	Update docker file to default to ROCm5.6
 	CI scripts update to handle multiple ROCm versions

build/rocm/Dockerfile.ms

@@ -2,10 +2,10 @@
 FROM rocm/dev-ubuntu-20.04:5.4-complete as rt_build
 MAINTAINER Rahul Batra<[email protected]>
 ################################################################################
-ARG ROCM_DEB_REPO=http://repo.radeon.com/rocm/apt/5.5/
+ARG ROCM_DEB_REPO=http://repo.radeon.com/rocm/apt/5.6/
 ARG ROCM_BUILD_NAME=ubuntu
 ARG ROCM_BUILD_NUM=main
-ARG ROCM_PATH=/opt/rocm-5.5.0
+ARG ROCM_PATH=/opt/rocm-5.6.0
 
 ARG DEBIAN_FRONTEND=noninteractive
 ARG PYTHON_VERSION=3.9.0
@@ -50,6 +50,7 @@ RUN apt-get update --allow-insecure-repositories && DEBIAN_FRONTEND=noninteracti
   libxmlsec1-dev \
   libffi-dev \
   liblzma-dev \
+  hipblaslt-dev \
   wget && \
   apt-get clean && \
   rm -rf /var/lib/apt/lists/*
@@ -101,6 +102,6 @@ ENV PATH $PYENV_ROOT/shims:$PYENV_ROOT/bin:$PATH
 
 RUN pyenv install $PYTHON_VERSION
 
-RUN eval "$(pyenv init -)" && pyenv local ${PYTHON_VERSION} && pip3 install --upgrade --force-reinstall setuptools pip==22.0 && pip install numpy==1.21.0 setuptools build wheel six auditwheel scipy pytest pytest-rerunfailures matplotlib absl-py
+RUN eval "$(pyenv init -)" && pyenv local ${PYTHON_VERSION} && pip3 install --upgrade --force-reinstall setuptools pip && pip install numpy setuptools build wheel six auditwheel scipy pytest pytest-rerunfailures matplotlib absl-py
 
 

build/rocm/ci_build.sh

@@ -26,6 +26,10 @@
 #
 # COMMAND: Command to be executed in the docker container
 #
+# ROCM_DEB_REPO_VERSION: ROCm debian repo version
+#
+# ROCM_PATH: ROCM path in the docker container
+#
 # Environment variables read by this script
 # WORKSPACE
 # XLA_REPO
@@ -44,6 +48,8 @@ DOCKERFILE_PATH="${SCRIPT_DIR}/Dockerfile.ms"
 DOCKER_CONTEXT_PATH="${SCRIPT_DIR}"
 KEEP_IMAGE="--rm"
 KEEP_CONTAINER="--rm"
+ROCM_DEB_REPO_VERSION="5.6" #default for now is 5.6
+ROCM_PATH="/opt/rocm-5.6.0"
 POSITIONAL_ARGS=()
 
 RUNTIME_FLAG=1
@@ -71,6 +77,14 @@ while [[ $# -gt 0 ]]; do
       KEEP_CONTAINER=""
       shift 1
       ;;
+    --rocm_deb_repo_version)
+      ROCM_DEB_REPO_VERSION="$2"
+      shift 2
+      ;;
+    --rocm_path)
+      ROCM_PATH="$2"
+      shift 2
+      ;;
 
     *)
       POSITIONAL_ARGS+=("$1")
@@ -119,7 +133,7 @@ echo "Python Version (${PYTHON_VERSION})"
 if [[ "${RUNTIME_FLAG}" -eq 1  ]]; then
   echo "Building (runtime) container (${DOCKER_IMG_NAME}) with Dockerfile($DOCKERFILE_PATH)..."
   docker build --target rt_build --tag ${DOCKER_IMG_NAME} \
-        --build-arg PYTHON_VERSION=$PYTHON_VERSION \
+        --build-arg PYTHON_VERSION=$PYTHON_VERSION --build-arg ROCM_DEB_REPO="http://repo.radeon.com/rocm/apt/"$ROCM_DEB_REPO_VERSION --build-arg ROCM_PATH=$ROCM_PATH\
       -f "${DOCKERFILE_PATH}" "${DOCKER_CONTEXT_PATH}"
 else
   echo "Building (CI) container (${DOCKER_IMG_NAME}) with Dockerfile($DOCKERFILE_PATH)..."