btl/ofi: Add FI_HMEM capability and memory registration

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

Author: wckzhang

opal/mca/btl/ofi/btl_ofi_component.c

@@ -106,7 +106,7 @@ static int validate_info(struct fi_info *info, uint64_t required_caps, char **in
     mr_mode = info->domain_attr->mr_mode;
 
     if (!(mr_mode == FI_MR_BASIC || mr_mode == FI_MR_SCALABLE
-          || (mr_mode & ~(FI_MR_VIRT_ADDR | FI_MR_ALLOCATED | FI_MR_PROV_KEY | FI_MR_ENDPOINT)) == 0)) {
+          || (mr_mode & ~(FI_MR_VIRT_ADDR | FI_MR_ALLOCATED | FI_MR_PROV_KEY | FI_MR_ENDPOINT | FI_MR_HMEM)) == 0)) {
         BTL_VERBOSE(("unsupported MR mode"));
         return OPAL_ERROR;
     }
@@ -256,8 +256,8 @@ static mca_btl_base_module_t **mca_btl_ofi_component_init(int *num_btl_modules,
     libfabric_api = fi_version();
 
     /* bail if OFI version is less than 1.5. */
-    if (libfabric_api < FI_VERSION(1, 5)) {
-        BTL_VERBOSE(("ofi btl disqualified because OFI version < 1.5."));
+    if (libfabric_api < FI_VERSION(1, 9)) {
+        BTL_VERBOSE(("ofi btl disqualified because OFI version < 1.9."));
         return NULL;
     }
 
@@ -339,16 +339,43 @@ static mca_btl_base_module_t **mca_btl_ofi_component_init(int *num_btl_modules,
 
     mca_btl_ofi_component.module_count = 0;
 
-    /* do the query. */
-    rc = fi_getinfo(FI_VERSION(1, 5), NULL, NULL, 0, &hints, &info_list);
+    /* Request device transfer capabilities, separate from required_caps */
+    hints.caps |= FI_HMEM;
+    hints.domain_attr->mr_mode |= FI_MR_HMEM;
+no_hmem:
+
+    /* Do the query. The earliest version that supports FI_HMEM hints is 1.9  */
+    rc = fi_getinfo(FI_VERSION(1, 9), NULL, NULL, 0, &hints, &info_list);
     if (0 != rc) {
+        if (hints.caps & FI_HMEM) {
+            /* Try again without FI_HMEM hints */
+            hints.caps &= ~FI_HMEM;
+            hints.domain_attr->mr_mode &= ~FI_MR_HMEM;
+            goto no_hmem;
+        }
         BTL_VERBOSE(("fi_getinfo failed with code %d: %s", rc, fi_strerror(-rc)));
         if (NULL != include_list) {
             opal_argv_free(include_list);
         }
         return NULL;
     }
 
+    /* If we get to this point with FI_HMEM hint set, we want it to be a
+     * required capability
+     */
+    if (hints.caps & FI_HMEM) {
+        /* The EFA provider has a bug where it incorrectly advertises FI_HMEM +
+         * FI_ATOMIC capability without being able to provide that support in
+         * versions before libfabric 1.18.0
+         */
+        if (libfabric_api < FI_VERSION(1, 18) && !strncasecmp(info_list->fabric_attr->prov_name, "efa", 3)) {
+            hints.caps &= ~FI_HMEM;
+            hints.domain_attr->mr_mode &= ~FI_MR_HMEM;
+            goto no_hmem;
+        }
+        required_caps |= FI_HMEM;
+    }
+
     /* count the number of resources/ */
     info = info_list;
     while (info) {

opal/mca/btl/ofi/btl_ofi_module.c

@@ -28,6 +28,8 @@
 #include "opal_config.h"
 #include "opal/class/opal_bitmap.h"
 #include "opal/datatype/opal_convertor.h"
+#include "opal/mca/accelerator/accelerator.h"
+#include "opal/mca/accelerator/base/base.h"
 #include "opal/mca/btl/btl.h"
 #include "opal/mca/mpool/base/base.h"
 #include "opal/mca/mpool/mpool.h"
@@ -234,14 +236,43 @@ static int mca_btl_ofi_deregister_mem(mca_btl_base_module_t *btl,
 int mca_btl_ofi_reg_mem(void *reg_data, void *base, size_t size,
                         mca_rcache_base_registration_t *reg)
 {
-    int rc;
+    int rc, dev_id;
+    uint64_t flags;
     static uint64_t access_flags = FI_REMOTE_WRITE | FI_REMOTE_READ | FI_READ | FI_WRITE;
+    struct fi_mr_attr attr = {0};
+    struct iovec iov = {0};
 
     mca_btl_ofi_module_t *btl = (mca_btl_ofi_module_t *) reg_data;
     mca_btl_ofi_reg_t *ur = (mca_btl_ofi_reg_t *) reg;
 
-    rc = fi_mr_reg(btl->domain, base, size, access_flags, 0, (uint64_t) reg, 0, &ur->ur_mr, NULL);
+    iov.iov_base = base;
+    iov.iov_len = size;
+    attr.mr_iov = &iov;
+    attr.iov_count = 1;
+    attr.access = access_flags;
+    attr.offset = 0;
+    attr.context = NULL;
+
+    if (OPAL_LIKELY(NULL != base)) {
+        rc = opal_accelerator.check_addr(base, &dev_id, &flags);
+        if (rc < 0) {
+            return rc;
+        } else if (rc > 0 ) {
+            if (0 == strcmp(opal_accelerator_base_selected_component.base_version.mca_component_name, "cuda")) {
+                attr.iface = FI_HMEM_CUDA;
+                opal_accelerator.get_device(&attr.device.cuda);
+	    } else if (0 == strcmp(opal_accelerator_base_selected_component.base_version.mca_component_name, "rocm")) {
+                attr.iface = FI_HMEM_ROCR;
+                opal_accelerator.get_device(&attr.device.cuda);
+            } else {
+                return OPAL_ERROR;
+            }
+        }
+    }
+
+    rc = fi_mr_regattr(btl->domain, &attr, 0, &ur->ur_mr);
     if (0 != rc) {
+        ur->ur_mr = NULL;
         return OPAL_ERR_OUT_OF_RESOURCE;
     }