Merge pull request open-mpi#10250 from edgargabriel/topic/rocm-support

Improve support for AMD ROCm GPU devices

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);