Improve support for AMD ROCm GPU devices

This commit enhances the support for AMD ROCm devices in Open MPI. The core aspect of this commit is to enable using hipMemcpy* based operations for packing/unpacking derived datatypes.

At the moment, support for AMD GPUs is only available through ROCm enabled UCX. Note that applications using basic datatypes and contiguous buffers with ROCm enabled UCX already work even without this commit, and even applications using derived datatypes worked in most instances due to the fact that AMD GPUs require enabling PCIe large bar support.

The most notable limitations:
 - rocm mca variables are not listed with ompi_info, since there is no component (at the moment) that is using rocm.
 - the 'compiled but not enabled' has received very limited testing.

Compiling a hip/MPI code will require a few additional flags. A typical compile line is:
mpiCC -D__HIP_PLATFORM_AMD__ -I/opt/rocm/include/hip -I/opt/rocm/include hipMPItest.cc -L/opt/rocm/lib -lamdhip64

Signed-off-by: Edgar <edgar.gabriel@amd.com>

Author: edgargabriel

.gitignore

@@ -290,6 +290,8 @@ ompi/mpiext/cuda/c/MPIX_Query_cuda_support.3
 ompi/mpiext/cuda/c/mpiext_cuda_c.h
 ompi/mpiext/cuda/c/cuda_c.h
 
+ompi/mpiext/rocm/c/mpiext_rocm_c.h
+
 ompi/mpiext/pcollreq/c/MPIX_*.3
 ompi/mpiext/pcollreq/c/profile/pallgather_init.c
 ompi/mpiext/pcollreq/c/profile/pallgatherv_init.c

config/opal_check_rocm.m4

@@ -0,0 +1,82 @@
+dnl
+dnl Copyright (C) 2022      Advanced Micro Devices, Inc. All rights reserved.
+dnl $COPYRIGHT$
+dnl
+dnl Additional copyrights may follow
+dnl
+dnl $HEADER$
+dnl
+
+
+# OMPI_CHECK_ROCM(prefix, [action-if-found], [action-if-not-found])
+# --------------------------------------------------------
+# check if ROCM support can be found.  sets prefix_{CPPFLAGS,
+# LDFLAGS, LIBS} as needed and runs action-if-found if there is
+# support, otherwise executes action-if-not-found
+
+
+#
+# Check for ROCm  support
+#
+AC_DEFUN([OPAL_CHECK_ROCM],[
+
+     OPAL_VAR_SCOPE_PUSH([opal_check_rocm_happy rocm_save_CPPFLAGS rocm_save_LDFLAGS rocm_CPPFLAGS rocm_LDFLAGS])
+
+     rocm_save_CPPFLAGS="$CPPFLAGS"
+     rocm_save_LDFLAGS="$LDFLAGS"
+     
+     # Get some configuration information
+     AC_ARG_WITH([rocm],
+        [AS_HELP_STRING([--with-rocm(=DIR)],
+        [Build ROCm support, optionally adding DIR/include, DIR/lib, and DIR/lib64 to the search path for headers and libraries])])
+
+
+     AS_IF([ test -n "$with_rocm" && test "$with_rocm" = "yes" ],
+           [ with_rocm="/opt/rocm"] )
+
+     rocm_CPPFLAGS="-D__HIP_PLATFORM_AMD__"
+     rocm_LDFLAGS="-L${with_rocm}/lib/hip"
+
+     AS_IF([ test -n "$with_rocm" && test "$with_rocm" != "no" ],
+           [ OPAL_APPEND([CPPFLAGS], [$rocm_CPPFLAGS])
+	     OPAL_APPEND([LDFLAGS], [$rocm_LDFLAGS]) ])
+
+     OAC_CHECK_PACKAGE([rocm],
+                       [$1],
+                       [hip/hip_runtime.h],
+                       [amdhip64],
+		       [hipFree],
+                       [opal_check_rocm_happy="yes"],
+                       [opal_check_rocm_happy="no"])
+
+     LDFLAGS="$rocm_save_LDFLAGS"
+     OPAL_APPEND([CPPFLAGS], [${$1_CPPFLAGS}] )
+     
+     AS_IF([ test "$opal_check_rocm_happy" = "no" ],
+           [ CPPFLAGS="$rocm_save_CPPFLAGS"])
+
+     AS_IF([ test "$opal_check_rocm_happy" = "yes" ],
+           [ AC_DEFINE_UNQUOTED([OPAL_ROCM_SUPPORT], [1], [Enable ROCm support])
+             ROCM_SUPPORT=1 ],
+           [ AC_DEFINE_UNQUOTED([OPAL_ROCM_SUPPORT], [0], [Disable ROCm support])
+             ROCM_SUPPORT=0 ])
+
+     AS_IF([ test "$opal_check_rocm_happy" = "yes" ],
+            [$2],
+            [AS_IF([test -n "$with_rocm" && test "$with_rocm" != "no"],
+                   [AC_MSG_ERROR([ROCm support requested but not found.  Aborting])])
+            $3])
+
+     AM_CONDITIONAL([OPAL_rocm_support], [test "$opal_check_rocm_happy" = "yes"])
+     OPAL_VAR_SCOPE_POP
+])
+
+AC_DEFUN([OPAL_CHECK_ROCM_AFTER_OPAL_DL],[
+    # We cannot have ROCm support without OPAL DL support.  Error out
+    # if the user wants Rocm but we do not have OPAL DL support.
+     AS_IF([test $OPAL_HAVE_DL_SUPPORT -eq 0 && test "$opal_check_rocm_happy" = "yes"],
+          [AC_MSG_WARN([--with-rocm was specified, but dlopen support is disabled.])
+           AC_MSG_WARN([You must reconfigure Open MPI with dlopen ("dl") support.])
+           AC_MSG_ERROR([Cannot continue.])])
+
+])

config/opal_config_files.m4

@@ -18,6 +18,7 @@ AC_DEFUN([OPAL_CONFIG_FILES],[
     AC_CONFIG_FILES([
         opal/Makefile
         opal/cuda/Makefile
+        opal/rocm/Makefile
         opal/etc/Makefile
         opal/include/Makefile
         opal/datatype/Makefile

configure.ac

@@ -1003,8 +1003,24 @@ AC_CACHE_SAVE
 opal_show_title "System-specific tests"
 
 OPAL_CHECK_CUDA
+##################################
+# ROCm support
+##################################
+OPAL_CHECK_ROCM([opal_rocm],
+	        [opal_rocm_happy="yes"],
+	        [opal_rocm_happy="no"])
+OPAL_SUMMARY_ADD([Miscellaneous], [ROCm suport], [], [$opal_rocm_happy])
+
+AS_IF([test "$OPAL_CUDA_SUPPORT" = "1" && test "$OPAL_ROCM_SUPPORT" = "1"],
+          [AC_MSG_WARN([Cannot support both CUDA and ROCm.])
+           AC_MSG_WARN([You must reconfigure Open MPI choosing either CUDA or ROCm .])
+           AC_MSG_ERROR([Cannot continue.])])
+
+##################################
 OPAL_CHECK_OS_FLAVORS
 
+
+
 # Do we have _SC_NPROCESSORS_ONLN? (only going to pass if we also have
 # <unistd.h> and sysconf(), which is ok) OS X 10.4 has <unistd.h> and
 # sysconf(), but does not have _SC_NPROCESSORS_ONLN.  Doh!
@@ -1247,6 +1263,8 @@ AC_CACHE_SAVE
 
 OPAL_CHECK_CUDA_AFTER_OPAL_DL
 
+OPAL_CHECK_ROCM_AFTER_OPAL_DL
+
 ##################################
 # MPI Extended Interfaces
 ##################################

docs/man-openmpi/man3/MPIX_Query_rocm_support.3.rst

@@ -0,0 +1,72 @@
+.. _mpix_query_rocm_support:
+
+
+MPIX_Query_rocm_support
+=======================
+
+.. include_body
+
+**MPIX_Query_rocm_support** - Returns 1 if there is AMD ROCm aware support
+and 0 if there is not.
+
+
+SYNTAX
+------
+
+
+C Syntax
+^^^^^^^^
+
+.. code-block:: c
+
+   #include <mpi.h>
+   #include <mpi-ext.h>
+
+   int MPIX_Query_rocm_support(void)
+
+
+Fortran Syntax
+^^^^^^^^^^^^^^
+
+There is no Fortran binding for this function.
+
+
+C++ Syntax
+^^^^^^^^^^
+
+There is no C++ binding for this function.
+
+
+DESCRIPTION
+-----------
+
+This routine return 1 if MPI library is build with ROCm and runtime
+supports ROCm buffers. This routine must be called after MPI is
+initialized by a call to :ref:`MPI_Init` or :ref:`MPI_Init_thread`.
+
+
+Examples
+^^^^^^^^
+
+::
+
+
+   #include <stdio.h>
+   #include "mpi.h"
+
+   #include "mpi-ext.h" /* Needed for ROCm-aware check */
+
+   int main(int argc, char *argv[])
+   {
+
+       MPI_Init(&argc, &argv);
+
+       if (MPIX_Query_rocm_support()) {
+           printf("This MPI library has ROCm-aware support.);
+       } else {
+           printf("This MPI library does not have ROCm-aware support.);
+       }
+       MPI_Finalize();
+
+       return 0;
+   }

docs/man-openmpi/man3/index.rst

@@ -466,4 +466,5 @@ MPI API manual pages (section 3)
    MPI_Wtick.3.rst
    MPI_Wtime.3.rst
    MPIX_Query_cuda_support.3.rst
+   MPIX_Query_rocm_support.3.rst
    OMPI_Affinity_str.3.rst

ompi/mpiext/rocm/Makefile.am

@@ -0,0 +1,23 @@
+#
+# Copyright (c) 2004-2009 The Trustees of Indiana University and Indiana
+#                         University Research and Technology
+#                         Corporation.  All rights reserved.
+# Copyright (c) 2010-2012 Cisco Systems, Inc.  All rights reserved.
+# Copyright (c) 2015      NVIDIA, Inc. All rights reserved
+# Copyright (c) 2022      Advanced Micro Devices, Inc. All rights reserved
+# $COPYRIGHT$
+#
+# Additional copyrights may follow
+#
+# $HEADER$
+#
+
+# This Makefile is not traversed during a normal "make all" in an OMPI
+# build.  It *is* traversed during "make dist", however.  So you can
+# put EXTRA_DIST targets in here.
+#
+# You can also use this as a convenience for building this MPI
+# extension (i.e., "make all" in this directory to invoke "make all"
+# in all the subdirectories).
+
+SUBDIRS = c

ompi/mpiext/rocm/c/Makefile.am

@@ -0,0 +1,42 @@
+#
+# Copyright (c) 2004-2009 The Trustees of Indiana University and Indiana
+#                         University Research and Technology
+#                         Corporation.  All rights reserved.
+# Copyright (c) 2010-2014 Cisco Systems, Inc.  All rights reserved.
+# Copyright (c) 2015      NVIDIA, Inc. All rights reserved.
+# Copyright (c) 2018      Research Organization for Information Science
+#                         and Technology (RIST).  All rights reserved.
+# Copyright (c) 2022      Advanced Micro Devices, Inc. All rights reserved.
+#
+# $COPYRIGHT$
+#
+# Additional copyrights may follow
+#
+# $HEADER$
+#
+
+# This file builds the C bindings for MPI extensions.  It must be
+# present in all MPI extensions.
+
+# We must set these #defines so that the inner OMPI MPI prototype
+# header files do the Right Thing.
+AM_CPPFLAGS = -DOMPI_PROFILE_LAYER=0 -DOMPI_COMPILING_FORTRAN_WRAPPERS=1
+
+# Convenience libtool library that will be slurped up into libmpi.la.
+noinst_LTLIBRARIES = libmpiext_rocm_c.la
+
+# This is where the top-level header file (that is included in
+# <mpi-ext.h>) must be installed.
+ompidir = $(ompiincludedir)/mpiext
+
+# This is the header file that is installed.
+nodist_ompi_HEADERS = mpiext_rocm_c.h
+
+# Sources for the convenience libtool library.  Other than the one
+# header file, all source files in the extension have no file naming
+# conventions.
+libmpiext_rocm_c_la_SOURCES = \
+        $(ompi_HEADERS) \
+        mpiext_rocm.c
+libmpiext_rocm_c_la_LDFLAGS = -module -avoid-version
+

ompi/mpiext/rocm/c/mpiext_rocm.c

@@ -0,0 +1,53 @@
+/*
+ * Copyright (c) 2004-2009 The Trustees of Indiana University and Indiana
+ *                         University Research and Technology
+ *                         Corporation.  All rights reserved.
+ * Copyright (c) 2010-2012 Cisco Systems, Inc.  All rights reserved.
+ * Copyright (c) 2010      Oracle and/or its affiliates.  All rights reserved.
+ * Copyright (c) 2012      Los Alamos National Security, LLC.  All rights
+ *                         reserved.
+ * Copyright (c) 2015      NVIDIA, Inc. All rights reserved.
+ * Copyright (c) 2022      Advanced Micro Devices, Inc. All rights reserved.
+ * $COPYRIGHT$
+ *
+ * Additional copyrights may follow
+ *
+ * $HEADER$
+ *
+ */
+
+#include "ompi_config.h"
+
+#include <stdio.h>
+#include <string.h>
+
+#include "opal/constants.h"
+#include "opal/runtime/opal_params.h"
+#include "ompi/mpiext/rocm/c/mpiext_rocm_c.h"
+
+#if OPAL_ROCM_SUPPORT
+#include "opal/rocm/common_rocm_prototypes.h"
+#endif
+
+int MPIX_Query_rocm_support(void)
+{
+
+    if (!opal_built_with_rocm_support) {
+        return 0;
+    } else {
+        if ( opal_rocm_runtime_initialized ) {
+            return 1;
+        }
+#if OPAL_ROCM_SUPPORT
+        // There is a chance that the rocm runtime has simply not
+        // been initialized yet, since that is done during the first convertor creation
+        // Invoke a function that will trigger the rocm runtime initialized and
+        // check the value again after that.
+
+        int val1, val2;
+        mca_common_rocm_check_bufs((char *)&val1, (char *)&val2);
+#endif
+    }
+
+    return opal_rocm_runtime_initialized;
+}

ompi/mpiext/rocm/c/mpiext_rocm_c.h.in

@@ -0,0 +1,17 @@
+/*
+ * Copyright (c) 2004-2009 The Trustees of Indiana University.
+ *                         All rights reserved.
+ * Copyright (c) 2010-2012 Cisco Systems, Inc.  All rights reserved.
+ * Copyright (c) 2010      Oracle and/or its affiliates.  All rights reserved.
+ * Copyright (c) 2015      NVIDIA, Inc. All rights reserved.
+ * Copyright (c) 2022      Advanced Micro Devices, Inc. All rights reserved.
+ * $COPYRIGHT$
+ *
+ * Additional copyrights may follow
+ *
+ * $HEADER$
+ *
+ */
+
+#define MPIX_ROCM_AWARE_SUPPORT @MPIX_ROCM_AWARE_SUPPORT@
+OMPI_DECLSPEC int MPIX_Query_rocm_support(void);