pml_ucx: add option to use opal memhooks instead of ucx internal hooks

Signed-off-by: Yossi Itigin <[email protected]>

(cherry picked from commit 564f80d)

Author: yosefe

ompi/mca/pml/ucx/pml_ucx.h

@@ -42,6 +42,7 @@ struct mca_pml_ucx_module {
     ompi_request_t            completed_send_req;
     size_t                    request_size;
     int                       num_disconnect;
+    bool                      opal_mem_hooks;
 
     /* Converters pool */
     mca_pml_ucx_freelist_t    convs;

ompi/mca/pml/ucx/pml_ucx_component.c

@@ -9,6 +9,11 @@
 
 #include "pml_ucx.h"
 
+#include "opal/memoryhooks/memory.h"
+#include "opal/mca/memory/base/base.h"
+
+#include <ucm/api/ucm.h>
+
 
 static int mca_pml_ucx_component_register(void);
 static int mca_pml_ucx_component_open(void);
@@ -70,13 +75,39 @@ static int mca_pml_ucx_component_register(void)
                                            OPAL_INFO_LVL_3,
                                            MCA_BASE_VAR_SCOPE_LOCAL,
                                            &ompi_pml_ucx.num_disconnect);
+
+    ompi_pml_ucx.opal_mem_hooks = 0;
+    (void) mca_base_component_var_register(&mca_pml_ucx_component.pmlm_version, "opal_mem_hooks",
+                                           "Use OPAL memory hooks, instead of UCX internal memory hooks",
+                                           MCA_BASE_VAR_TYPE_BOOL, NULL, 0, 0,
+                                           OPAL_INFO_LVL_3,
+                                           MCA_BASE_VAR_SCOPE_LOCAL,
+                                           &ompi_pml_ucx.opal_mem_hooks);
     return 0;
 }
 
+static void mca_pml_ucx_mem_release_cb(void *buf, size_t length,
+                                       void *cbdata, bool from_alloc)
+{
+    ucm_vm_munmap(buf, length);
+}
+
 static int mca_pml_ucx_component_open(void)
 {
     ompi_pml_ucx.output = opal_output_open(NULL);
     opal_output_set_verbosity(ompi_pml_ucx.output, ompi_pml_ucx.verbose);
+
+    /* Set memory hooks */
+    if (ompi_pml_ucx.opal_mem_hooks &&
+        (OPAL_MEMORY_FREE_SUPPORT | OPAL_MEMORY_MUNMAP_SUPPORT) ==
+        ((OPAL_MEMORY_FREE_SUPPORT | OPAL_MEMORY_MUNMAP_SUPPORT) &
+         opal_mem_hooks_support_level()))
+    {
+        PML_UCX_VERBOSE(1, "%s", "using OPAL memory hooks as external events");
+        ucm_set_external_event(UCM_EVENT_VM_UNMAPPED);
+        opal_mem_hooks_register_release(mca_pml_ucx_mem_release_cb, NULL);
+    }
+
     return mca_pml_ucx_open();
 }
 
@@ -89,6 +120,7 @@ static int mca_pml_ucx_component_close(void)
         return rc;
     }
 
+    opal_mem_hooks_unregister_release(mca_pml_ucx_mem_release_cb);
     opal_output_close(ompi_pml_ucx.output);
     return 0;
 }