BUILD: add support for specific gpu arch with ROCM (openucx#987)

MamziB · Jun 27, 2024 · 039403c · 039403c
1 parent afd7016
commit 039403c
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 8 deletions.
diff --git a/config/m4/rocm.m4 b/config/m4/rocm.m4
@@ -4,6 +4,17 @@
 # See file LICENSE for terms.
 #
 
+ROCM_ARCH_NATIVE="--offload-arch=native"
+ROCM_ARCH908="--offload-arch=gfx908"
+ROCM_ARCH90A="--offload-arch=gfx90a"
+ROCM_ARCH94="--offload-arch=gfx940 \
+--offload-arch=gfx941 \
+--offload-arch=gfx942"
+ROCM_ARCH10="--offload-arch=gfx1030"
+ROCM_ARCH11="--offload-arch=gfx1100 \
+--offload-arch=gfx1101 \
+--offload-arch=gfx1102"
+
 # ROCM_PARSE_FLAGS(ARG, VAR_LIBS, VAR_LDFLAGS, VAR_CPPFLAGS)
 # ----------------------------------------------------------
 # Parse whitespace-separated ARG into appropriate LIBS, LDFLAGS, and
@@ -73,7 +84,12 @@ AC_ARG_WITH([rocm],
  [Enable the use of ROCm (default is autodetect).])],
  [],
  [with_rocm=guess])
-
+AC_ARG_WITH([rocm-arch],
+ [AS_HELP_STRING([--with-rocm-arch=arch-code],
+ [Defines target GPU architecture,
+ see rocm documentation for valid --offload-arch options for details
+ 'all-arch-no-native' for all default architectures but not native])],
+ [], [with_rocm_arch=all])
 rocm_happy=no
 hip_happy=no
 AS_IF([test "x$with_rocm" != "xno"],
@@ -115,15 +131,26 @@ AS_IF([test "x$with_rocm" != "xno"],
  AC_SUBST([ROCM_ROOT])],
  [AC_MSG_WARN([ROCm not found])])
 
+
+ # Check whether we run on ROCm 6.0 or higher
+ CHECK_ROCM_VERSION(6, ROCM_VERSION_60_OR_GREATER)
+ AC_MSG_CHECKING([if ROCm version is 6.0 or above])
+
+ AS_IF([test "x$rocm_happy" = "xyes"],
+ [AS_IF([test "x$with_rocm_arch" = "xall"],
+ [ROCM_ARCH="${ROCM_ARCH908} ${ROCM_ARCH90A} ${ROCM_ARCH94} ${ROCM_ARCH10} ${ROCM_ARCH11} ${ROCM_ARCH_NATIVE}"],
+ [AS_IF([test "x$with_rocm_arch" = "xall-arch-no-native"],
+ [ROCM_ARCH="${ROCM_ARCH908} ${ROCM_ARCH90A} ${ROCM_ARCH94} ${ROCM_ARCH10} ${ROCM_ARCH11}"],
+ [ROCM_ARCH="$with_rocm_arch"])])
+ AS_IF([test "$ROCM_VERSION_60_OR_GREATER" = "1"],
+ AC_SUBST([ROCM_ARCH], ["$ROCM_ARCH"]),
+ AC_SUBST([ROCM_ARCH], [""]))])
  CPPFLAGS="$SAVE_CPPFLAGS"
  LDFLAGS="$SAVE_LDFLAGS"
  LIBS="$SAVE_LIBS"
 
  HIP_BUILD_FLAGS([$with_rocm], [HIP_LIBS], [HIP_LDFLAGS], [HIP_CPPFLAGS])
 
- # Check whether we run on ROCm 6.0 or higher
- CHECK_ROCM_VERSION(6, ROCM_VERSION_60_OR_GREATER)
- AC_MSG_CHECKING([if ROCm version is 6.0 or above])
  if test "$ROCM_VERSION_60_OR_GREATER" = "1" ; then
  AC_MSG_RESULT([yes])
  else

diff --git a/configure.ac b/configure.ac
@@ -265,6 +265,9 @@ AC_MSG_NOTICE([ C++ compiler: ${CXX} ${CXXFLAGS} ${BASE_CXXFLAGS}])
 AS_IF([test "x$cuda_happy" = "xyes"],[
 AC_MSG_NOTICE([ NVCC gencodes: ${NVCC_ARCH}])
 ])
+AS_IF([test "x$rocm_happy" = xyes],[
+AC_MSG_NOTICE([ROCM architectures: ${ROCM_ARCH}])
+])
 AC_MSG_NOTICE([ Perftest: ${mpi_enable}])
 AC_MSG_NOTICE([ Gtest: ${gtest_enable}])
 AC_MSG_NOTICE([ MC modules: <$(echo ${mc_modules}|tr ':' ' ') >])

diff --git a/cuda_lt.sh b/cuda_lt.sh
@@ -28,7 +28,7 @@ mkdir -p $pic_dir
 
 tmpcmd="${@:3}"
 if [[ "$tmpcmd" == *"amdclang"* ]]; then
- cmd="${@:3:2} -x hip -target x86_64-unknown-linux-gnu --offload-arch=gfx908 --offload-arch=gfx90a --offload-arch=gfx940 --offload-arch=gfx941 --offload-arch=gfx942 --offload-arch=gfx1030 --offload-arch=gfx1100 --offload-arch=gfx1101 --offload-arch=gfx1102 --offload-arch=native ${@:5} -fPIC -O3 -o ${pic_filepath}"
+ cmd="${@:3:2} -x hip -target x86_64-unknown-linux-gnu ${@:5} -fPIC -O3 -o ${pic_filepath}"
 elif [[ "$tmpcmd" == *"hipcc"* ]]; then
  cmd="${@:3} -fPIC -o ${pic_filepath}"
 else
@@ -38,7 +38,7 @@ echo $cmd
 $cmd
 
 if [[ "$tmpcmd" == *"amdclang"* ]]; then
- cmd="${@:3:2} -x hip -target x86_64-unknown-linux-gnu --offload-arch=gfx908 --offload-arch=gfx90a --offload-arch=gfx940 --offload-arch=gfx941 --offload-arch=gfx942 --offload-arch=gfx1030 --offload-arch=gfx1100 --offload-arch=gfx1101 --offload-arch=gfx1102 --offload-arch=native ${@:5} -O3 -o ${npic_filepath}"
+ cmd="${@:3:2} -x hip -target x86_64-unknown-linux-gnu ${@:5} -O3 -o ${npic_filepath}"
 else
  cmd="${@:3} -o ${npic_filepath}"
 fi

diff --git a/src/components/ec/rocm/kernel/Makefile.am b/src/components/ec/rocm/kernel/Makefile.am
@@ -17,10 +17,10 @@ HIPCCFLAGS = \
 LINK = $(LIBTOOL) --mode=link $(CC) -o $@
 
 .cu.o:
- $(HIPCC) -c $< -o $@ $(HIPCCFLAGS)
+ $(HIPCC) -c $< -o $@ $(ROCM_ARCH) $(HIPCCFLAGS)
 
 .cu.lo:
- /bin/bash $(top_srcdir)/cuda_lt.sh "$(LIBTOOL)" $@ $(HIPCC) -c $< $(HIPCCFLAGS)
+ /bin/bash $(top_srcdir)/cuda_lt.sh "$(LIBTOOL)" $@ $(HIPCC) -c $< $(ROCM_ARCH) $(HIPCCFLAGS)
 
 comp_noinst = libucc_ec_rocm_kernels.la