osc/rdma: Add checks for gpu support during dynamic attach

Previously, there were no checks for GPU support during the
dynamic attach process, only during window create. Added this
check and cleaned up some code.

Signed-off-by: William Zhang <wilzhang@amazon.com>

Author: wckzhang

ompi/mca/osc/rdma/osc_rdma.h

@@ -44,6 +44,7 @@
 #include "ompi/mca/osc/osc.h"
 #include "ompi/mca/osc/base/base.h"
 #include "opal/mca/btl/btl.h"
+#include "opal/mca/btl/base/base.h"
 #include "opal/mca/btl/base/btl_base_am_rdma.h"
 #include "ompi/memchecker.h"
 #include "ompi/op/op.h"
@@ -758,4 +759,36 @@ static inline int osc_rdma_is_accel(void *buf)
     uint64_t flags;
     return opal_accelerator.check_addr(buf, &dev_id, &flags);
 }
+
+/**
+ * @brief Checks whether any btl's are capable of accelerator support.
+ *
+ * @param[in] btls            The modules to check for btl accelerator support
+ *
+ * @returns true              If any btl module has accelerator support
+ * @returns false             Otherwise
+ */
+static inline bool osc_rdma_btl_accel_support(opal_list_t *btls)
+{
+    mca_btl_base_selected_module_t* selected_btl;
+    /* verify if we have any btls available.  Since we do not verify
+     * connectivity across all btls in the alternate case, this is as
+     * good a test as we are going to have for success. */
+    if (opal_list_is_empty(btls)) {
+        return false;
+    }
+
+    /* If we have any btls, we check again if any btl supports
+     * MCA_BTL_FLAGS_ACCELERATOR_RDMA */
+    OPAL_LIST_FOREACH(selected_btl, btls, mca_btl_base_selected_module_t) {
+        opal_output_verbose(MCA_BASE_VERBOSE_INFO, ompi_osc_base_framework.framework_output,
+                    "osc_rdma_component_query: check ACCELERATOR_RDMA flag: %s",
+                    selected_btl->btl_component->btl_version.mca_component_name);
+        mca_btl_base_module_t *btl = selected_btl->btl_module;
+        if (btl->btl_flags & MCA_BTL_FLAGS_ACCELERATOR_RDMA) {
+            return true;
+        }
+    }
+    return false;
+}
 #endif /* OMPI_OSC_RDMA_H */

ompi/mca/osc/rdma/osc_rdma_component.c

@@ -373,10 +373,6 @@ static int ompi_osc_rdma_component_query (struct ompi_win_t *win, void **base, s
                                           int flavor)
 {
 
-    opal_list_t *btls = NULL;
-    mca_btl_base_selected_module_t* selected_btl;
-    int gpu_check = 0;
-
     if (MPI_WIN_FLAVOR_SHARED == flavor) {
         return -1;
     }
@@ -385,7 +381,9 @@ static int ompi_osc_rdma_component_query (struct ompi_win_t *win, void **base, s
         uint64_t flags;
         int dev_id;
         if (opal_accelerator.check_addr(*base, &dev_id, &flags)) {
-            gpu_check = 1;
+            if (!osc_rdma_btl_accel_support(&mca_btl_base_modules_initialized)) {
+                return -1;
+            }
         }
     }
 
@@ -396,28 +394,6 @@ static int ompi_osc_rdma_component_query (struct ompi_win_t *win, void **base, s
         return -1;
     }
 
-    /* Not on GPU at all, skip the check */
-    if (!gpu_check) {
-        goto ok;
-    }
-
-    /* If we have any btls, we check again if any btl supports
-     * MCA_BTL_FLAGS_ACCELERATOR_RDMA */
-    btls = &mca_btl_base_modules_initialized;
-    OPAL_LIST_FOREACH(selected_btl, btls, mca_btl_base_selected_module_t) {
-        opal_output_verbose(MCA_BASE_VERBOSE_INFO, ompi_osc_base_framework.framework_output,
-                    "osc_rdma_component_query: check ACCELERATOR_RDMA flag: %s",
-                    selected_btl->btl_component->btl_version.mca_component_name);
-        mca_btl_base_module_t *btl = selected_btl->btl_module;
-        // Check flag: MCA_BTL_FLAGS_ACCELERATOR_RDMA
-        if (btl->btl_flags & MCA_BTL_FLAGS_ACCELERATOR_RDMA) {
-            goto ok;
-        }
-    }
-    /* No BTL supports the accelerator flag */
-    return -1;
-
-ok:
     return mca_osc_rdma_component.priority;
 }
 

ompi/mca/osc/rdma/osc_rdma_dynamic.c

@@ -172,6 +172,8 @@ int ompi_osc_rdma_attach (struct ompi_win_t *win, void *base, size_t len)
     intptr_t page_size = opal_getpagesize ();
     int region_index, ret;
     size_t aligned_len;
+    int dev_id;
+    uint64_t flags;
 
     if (module->flavor != MPI_WIN_FLAVOR_DYNAMIC) {
         return OMPI_ERR_RMA_FLAVOR;
@@ -182,6 +184,16 @@ int ompi_osc_rdma_attach (struct ompi_win_t *win, void *base, size_t len)
         return OMPI_SUCCESS;
     }
 
+    /* Component query check doesn't apply for dynamic window types.
+     * We have to check whether a GPU window is supported here. */
+    if (opal_accelerator.check_addr(base, &dev_id, &flags)) {
+        if (!osc_rdma_btl_accel_support(&mca_btl_base_modules_initialized)) {
+            OSC_RDMA_VERBOSE(MCA_BASE_VERBOSE_TRACE, "Failed to attach %p. No components capable of attaching accelerator buffers.",
+                             base);
+            return OMPI_ERR_NOT_SUPPORTED;
+        }
+    }
+
     OSC_RDMA_VERBOSE(MCA_BASE_VERBOSE_TRACE, "attach: %s, %p, %lu", win->w_name, base, (unsigned long) len);
 
     OPAL_THREAD_LOCK(&module->lock);