From 37d8bed53d88d6fff85b1fd86827f83e5573ca06 Mon Sep 17 00:00:00 2001
From: PeixuanZuo <94887879+PeixuanZuo@users.noreply.github.com>
Date: Tue, 14 Nov 2023 11:54:22 +0800
Subject: [PATCH] [ROCm] add migraphx into onnxruntime-training-rocm package
 (#18339)

---
 setup.py                                                      | 4 ++++
 .../github/azure-pipelines/orttraining-pai-ci-pipeline.yml    | 1 +
 tools/ci_build/github/azure-pipelines/templates/rocm.yml      | 1 +
 .../github/linux/docker/Dockerfile.manylinux2_28_rocm         | 2 +-
 tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile     | 4 +++-
 5 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index 9eca9845c9..1c04433c9a 100644
--- a/setup.py
+++ b/setup.py
@@ -219,6 +219,10 @@ try:
                     "librocm_smi64.so.5",
                     "libroctracer64.so.4",
                     "libtinfo.so.6",
+                    "libmigraphx_c.so.3",
+                    "libmigraphx.so.2",
+                    "libmigraphx_onnx.so.2",
+                    "libmigraphx_tf.so.2",
                 ]
 
                 tensorrt_dependencies = ["libnvinfer.so.8", "libnvinfer_plugin.so.8", "libnvonnxparser.so.8"]
diff --git a/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml
index ed84b514fb..8d02a5e580 100644
--- a/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml
+++ b/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml
@@ -108,6 +108,7 @@ jobs:
                 FETCHCONTENT_TRY_FIND_PACKAGE_MODE=NEVER \
               --use_cache \
               --use_rocm \
+              --use_migraphx \
               --rocm_version=$(RocmVersion) \
               --rocm_home ${ROCM_HOME} \
               --nccl_home ${ROCM_HOME}\
diff --git a/tools/ci_build/github/azure-pipelines/templates/rocm.yml b/tools/ci_build/github/azure-pipelines/templates/rocm.yml
index f41dc1c49f..2e9e6c6b35 100644
--- a/tools/ci_build/github/azure-pipelines/templates/rocm.yml
+++ b/tools/ci_build/github/azure-pipelines/templates/rocm.yml
@@ -79,6 +79,7 @@ jobs:
             /onnxruntime_src/tools/ci_build/build.py \
               --config ${{ parameters.BuildConfig }} \
               --use_rocm \
+              --use_migraphx \
               --rocm_version=${{ parameters.RocmVersion }} \
               --rocm_home=/opt/rocm \
               --nccl_home=/opt/rocm \
diff --git a/tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_rocm b/tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_rocm
index 19599c9f61..9e12fe8c75 100644
--- a/tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_rocm
+++ b/tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_rocm
@@ -31,7 +31,7 @@ RUN yum install -y hipify-clang
 RUN yum -y install wget
 
 # rocm lib
-RUN yum install -y miopen-hip-devel rocblas-devel rocrand-devel rccl-devel hipsparse-devel hipfft-devel hipcub-devel hipblas-devel rocthrust-devel
+RUN yum install -y miopen-hip-devel rocblas-devel rocrand-devel rccl-devel hipsparse-devel hipfft-devel hipcub-devel hipblas-devel rocthrust-devel migraphx-devel
 
 ENV AUDITWHEEL_POLICY=${POLICY} AUDITWHEEL_ARCH=${PLATFORM} AUDITWHEEL_PLAT=${POLICY}_${PLATFORM}
 ENV LC_ALL=en_US.UTF-8 LANG=en_US.UTF-8 LANGUAGE=en_US.UTF-8
diff --git a/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile b/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
index 412bc00d02..2ec826fc8f 100644
--- a/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
+++ b/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
@@ -80,7 +80,6 @@ RUN ln -sf /usr/lib/x86_64-linux-gnu/libstdc++.so.6 ${CONDA_ENVIRONMENT_PATH}/bi
 RUN pip install torch==2.0.1 torchvision==0.15.2 -f https://repo.radeon.com/rocm/manylinux/rocm-rel-${ROCM_VERSION}/ && \
     pip install torch-ort --no-dependencies
 
-
 ##### Install Cupy to decrease CPU utilization
 # Install non dev openmpi
 RUN wget https://download.open-mpi.org/release/open-mpi/v4.1/openmpi-4.1.5.tar.bz2 && \
@@ -130,6 +129,9 @@ RUN pip install \
     pytest-xdist \
     pytest-rerunfailures
 
+# Install migraphx
+RUN apt update && apt install -y migraphx
+
 ENV ORTMODULE_ONNX_OPSET_VERSION=15
 
 ARG BUILD_UID=1001