From 8c41bb9c202bb9f5e4b7ea01df12eaab47ee320e Mon Sep 17 00:00:00 2001
From: Yuting Jiang <yutingjiang@microsoft.com>
Date: Wed, 6 Dec 2023 17:31:51 +0800
Subject: [PATCH] push

---
 dockerfile/rocm6.0.x.dockerfile                          | 2 ++
 setup.py                                                 | 6 +++++-
 superbench/benchmarks/micro_benchmarks/rocm_common.cmake | 2 +-
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/dockerfile/rocm6.0.x.dockerfile b/dockerfile/rocm6.0.x.dockerfile
index cf8bc824..1cf2ca52 100644
--- a/dockerfile/rocm6.0.x.dockerfile
+++ b/dockerfile/rocm6.0.x.dockerfile
@@ -173,6 +173,8 @@ ADD third_party third_party
 
 RUN make RCCL_HOME=/opt/rccl/build/ ROCBLAS_BRANCH=release-staging/rocm-rel-6.0 ROCM_VER=rocm-5.5.0 -C third_party rocm -o cpu_hpl -o cpu_stream
 
+ENV CXX="/opt/rocm/bin/hipcc"
+
 ADD . .
 RUN apt install rocm-cmake -y && \
     python3 -m pip install --upgrade wheel setuptools==65.7 && \
diff --git a/setup.py b/setup.py
index 23c79683..d00fdfb0 100644
--- a/setup.py
+++ b/setup.py
@@ -183,7 +183,7 @@ setup(
             **x,
             'develop': x['dev'] + x['test'],
             'cpuworker': x['torch'],
-            'amdworker': x['torch'] + x['ort'],
+            'amdworker': x['amd'] + x['ort'],
             'nvworker': x['torch'] + x['ort'] + x['nvidia'],
         }
     )(
@@ -216,6 +216,10 @@ setup(
                 'onnxruntime-gpu==1.10.0',
             ],
             'nvidia': ['py3nvml>=0.2.6'],
+            'amd' : [
+                'torchvision>=0.8.0a0',
+                'transformers>=4.3.3, <4.23.0',
+            ]
         }
     ),
     include_package_data=True,
diff --git a/superbench/benchmarks/micro_benchmarks/rocm_common.cmake b/superbench/benchmarks/micro_benchmarks/rocm_common.cmake
index 9758ae27..f884bd58 100644
--- a/superbench/benchmarks/micro_benchmarks/rocm_common.cmake
+++ b/superbench/benchmarks/micro_benchmarks/rocm_common.cmake
@@ -47,7 +47,7 @@ if(EXISTS ${HIP_PATH})
     # Search for hip in common locations
     list(APPEND CMAKE_PREFIX_PATH ${HIP_PATH} ${ROCM_PATH})
     set(CMAKE_PREFIX_PATH /opt/rocm ROCM_PATH)
-    set(CXX "${HIP_PATH}/hipcc")
+    set(CMAKE_CXX_COMPILER "${HIP_PATH}/hipcc")
     set(CMAKE_MODULE_PATH "${HIP_PATH}/cmake" ${CMAKE_MODULE_PATH})
     set(CMAKE_MODULE_PATH "${HIP_PATH}/lib/cmake/hip" ${CMAKE_MODULE_PATH})
 endif()