Add OFI BTL flags for OSC when using accelerator memory

Adds the capability for the OSC RDMA component to handle
GPU windows during the query API.

Also adds flags to the OFI BTL showing accelerator support.

Signed-off-by: Jingyin Tang <jytang@amazon.com>

Author: jtamzn

ompi/mca/osc/rdma/osc_rdma_component.c

@@ -372,16 +372,20 @@ static int ompi_osc_rdma_component_query (struct ompi_win_t *win, void **base, s
                                           struct ompi_communicator_t *comm, struct opal_info_t *info,
                                           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;
     }
 
-    /* GPU buffers are not supported by the rdma component */
     if (MPI_WIN_FLAVOR_CREATE == flavor) {
         uint64_t flags;
         int dev_id;
         if (opal_accelerator.check_addr(*base, &dev_id, &flags)) {
-            return -1;
+            gpu_check = 1;
         }
     }
 
@@ -392,6 +396,28 @@ 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;
 }
 

opal/mca/btl/ofi/btl_ofi_component.c

@@ -604,6 +604,10 @@ static int mca_btl_ofi_init_device(struct fi_info *info)
     module->use_fi_mr_bind = false;
     module->bypass_cache = false;
 
+    if (ofi_info->caps & FI_HMEM) {
+        module->super.btl_flags |= MCA_BTL_FLAGS_ACCELERATOR_RDMA;
+    }
+
     if (ofi_info->domain_attr->mr_mode == FI_MR_BASIC
         || ofi_info->domain_attr->mr_mode & FI_MR_VIRT_ADDR) {
         module->use_virt_addr = true;

opal/mca/btl/ofi/btl_ofi_module.c

@@ -252,6 +252,7 @@ int mca_btl_ofi_reg_mem(void *reg_data, void *base, size_t size,
     attr.access = access_flags;
     attr.offset = 0;
     attr.context = NULL;
+    attr.requested_key = (uint64_t) reg;
 
     if (OPAL_LIKELY(NULL != base)) {
         rc = opal_accelerator.check_addr(base, &dev_id, &flags);