btl/uct: add helper to enable an active message handler on a context

This commit adds a new helper (mca_btl_uct_context_enable_am_handler) which will
enable an active message handler on the given context if it has not already been
enabled.

Signed-off-by: Nathan Hjelm <hjelmn@google.com>

Author: hjelmn

opal/mca/btl/uct/btl_uct_device_context.h

@@ -37,6 +37,15 @@ mca_btl_uct_device_context_t *mca_btl_uct_context_create(mca_btl_uct_module_t *m
  */
 void mca_btl_uct_context_destroy(mca_btl_uct_device_context_t *context);
 
+/**
+ * @brief Enable active messages on context if not already enabled
+ *
+ * @param[in] tl      TL this context belongs to
+ * @param[in] context Context to enable active messages on.
+ */
+void mca_btl_uct_context_enable_am_handler(mca_btl_uct_tl_t *tl,
+                                           mca_btl_uct_device_context_t *context);
+
 static inline bool mca_btl_uct_context_trylock(mca_btl_uct_device_context_t *context)
 {
     return OPAL_THREAD_TRYLOCK(&context->mutex);

opal/mca/btl/uct/btl_uct_tl.c

@@ -321,6 +321,20 @@ static int mca_btl_uct_populate_tl_attr(mca_btl_uct_module_t *module, mca_btl_uc
     return OPAL_SUCCESS;
 }
 
+void mca_btl_uct_context_enable_am_handler(mca_btl_uct_tl_t *tl,
+                                           mca_btl_uct_device_context_t *context)
+{
+    if (context->am_handler_installed) {
+        return;
+    }
+
+    BTL_VERBOSE(("installing AM handler for tl %s::%s context id %d",
+                 tl->uct_md->md_name, tl->uct_tl_name, context->context_id)); \
+    uct_iface_set_am_handler(context->uct_iface, MCA_BTL_UCT_FRAG, mca_btl_uct_am_handler,
+                             context, MCA_BTL_UCT_CB_FLAG_SYNC);
+    context->am_handler_installed = true;
+}
+
 mca_btl_uct_device_context_t *mca_btl_uct_context_create(mca_btl_uct_module_t *module,
                                                          mca_btl_uct_tl_t *tl, int context_id,
                                                          bool enable_progress)
@@ -383,9 +397,7 @@ mca_btl_uct_device_context_t *mca_btl_uct_context_create(mca_btl_uct_module_t *m
     }
 
     if (module != NULL && tl == module->am_tl) {
-        BTL_VERBOSE(("installing AM handler for tl %p context id %d", (void *) tl, context_id));
-        uct_iface_set_am_handler(context->uct_iface, MCA_BTL_UCT_FRAG, mca_btl_uct_am_handler,
-                                 context, MCA_BTL_UCT_CB_FLAG_SYNC);
+        mca_btl_uct_context_enable_am_handler(tl, context);
     }
 
     if (enable_progress) {

opal/mca/btl/uct/btl_uct_types.h

@@ -177,6 +177,9 @@ struct mca_btl_uct_device_context_t {
     /** progress is enabled on this context */
     bool progress_enabled;
 
+    /** communication AM handler is installed */
+    bool am_handler_installed;
+
     /** context is in AM callback */
     volatile bool in_am_callback;
 };