[OpenMP][OMPT] Transparently support emi and noemi callbacks in ompt-multiplex.h

[protze.joachim]

Summary

OpenMP 5.1 added extended versions of existing target related
event callbacks (emi). These callbacks provide all information of
the old callbacks (noemi) and add more information. Therefore,
noemi callbacks can be transparently wrapped into emi callbacks.
The OpenMP spec allows to only register the emi or noemi
version of each callback. Since a tool down in the multiplexer
recursion might register the emi callback, we delay the
transparent wrapping for each individual tool.

This patch adds the wrappers and the code to transparently
register emi or noemi callbacks.

openmp/tools/multiplex/ompt-multiplex.h

 #include <inttypes.h>
 #include <omp-tools.h>
 #include <omp.h>
 #include <stdio.h>
 #include <string.h>
 
 static ompt_set_callback_t ompt_multiplex_set_callback;
 static ompt_get_task_info_t ompt_multiplex_get_task_info;
+static ompt_get_unique_id_t ompt_multiplex_get_unique_id;
 static ompt_get_thread_data_t ompt_multiplex_get_thread_data;
 static ompt_get_parallel_info_t ompt_multiplex_get_parallel_info;
 
 // contains name of the environment var in which the tool path is specified
 #ifndef CLIENT_TOOL_LIBRARIES_VAR
 #error CLIENT_TOOL_LIBRARIES_VAR should be defined before including of ompt-multiplex.h
 #endif
 

 #define OMPT_LOAD_CLIENT_FOREACH_OMPT_EVENT(macro)                             \
   macro(callback_thread_begin, ompt_callback_thread_begin_t, 1);               \
   macro(callback_thread_end, ompt_callback_thread_end_t, 2);                   \
   macro(callback_parallel_begin, ompt_callback_parallel_begin_t, 3);           \
   macro(callback_parallel_end, ompt_callback_parallel_end_t, 4);               \
   macro(callback_task_create, ompt_callback_task_create_t, 5);                 \
   macro(callback_task_schedule, ompt_callback_task_schedule_t, 6);             \
   macro(callback_implicit_task, ompt_callback_implicit_task_t, 7);             \
-  macro(callback_target, ompt_callback_target_t, 8);                           \
-  macro(callback_target_data_op, ompt_callback_target_data_op_t, 9);           \
-  macro(callback_target_submit, ompt_callback_target_submit_t, 10);            \
   macro(callback_control_tool, ompt_callback_control_tool_t, 11);              \
   macro(callback_device_initialize, ompt_callback_device_initialize_t, 12);    \
   macro(callback_device_finalize, ompt_callback_device_finalize_t, 13);        \
   macro(callback_device_load, ompt_callback_device_load_t, 14);                \
   macro(callback_device_unload, ompt_callback_device_unload_t, 15);            \
   macro(callback_sync_region_wait, ompt_callback_sync_region_t, 16);           \
   macro(callback_mutex_released, ompt_callback_mutex_t, 17);                   \
   macro(callback_dependences, ompt_callback_dependences_t, 18);                \
   macro(callback_task_dependence, ompt_callback_task_dependence_t, 19);        \
   macro(callback_work, ompt_callback_work_t, 20);                              \
   macro(callback_masked, ompt_callback_masked_t, 21);                          \
-  macro(callback_target_map, ompt_callback_target_map_t, 22);                  \
   macro(callback_sync_region, ompt_callback_sync_region_t, 23);                \
   macro(callback_lock_init, ompt_callback_mutex_acquire_t, 24);                \
   macro(callback_lock_destroy, ompt_callback_mutex_t, 25);                     \
   macro(callback_mutex_acquire, ompt_callback_mutex_acquire_t, 26);            \
   macro(callback_mutex_acquired, ompt_callback_mutex_t, 27);                   \
   macro(callback_nest_lock, ompt_callback_nest_lock_t, 28);                    \
   macro(callback_flush, ompt_callback_flush_t, 29);                            \
   macro(callback_cancel, ompt_callback_cancel_t, 30);                          \
   macro(callback_reduction, ompt_callback_sync_region_t, 31);                  \
   macro(callback_dispatch, ompt_callback_dispatch_t, 32);                      \
   macro(callback_target_emi, ompt_callback_target_emi_t, 33);                  \
   macro(callback_target_data_op_emi, ompt_callback_target_data_op_emi_t, 34);  \
   macro(callback_target_submit_emi, ompt_callback_target_submit_emi_t, 35);    \
   macro(callback_target_map_emi, ompt_callback_target_map_emi_t, 36);          \
   macro(callback_error, ompt_callback_error_t, 37);
 
+#define OMPT_LOAD_CLIENT_FOREACH_OMPT_EVENT_NOEMI(macro)                       \
+  macro(callback_target, ompt_callback_target_t, 8);                           \
+  macro(callback_target_data_op, ompt_callback_target_data_op_t, 9);           \
+  macro(callback_target_submit, ompt_callback_target_submit_t, 10);            \
+  macro(callback_target_map, ompt_callback_target_map_t, 22);
 
 typedef struct ompt_multiplex_callbacks_s {
 #define ompt_event_macro(event, callback, eventid) callback ompt_##event
 
   OMPT_LOAD_CLIENT_FOREACH_OMPT_EVENT(ompt_event_macro)
+  OMPT_LOAD_CLIENT_FOREACH_OMPT_EVENT_NOEMI(ompt_event_macro)
 
 #undef ompt_event_macro
 } ompt_multiplex_callbacks_t;
 
 typedef struct ompt_multiplex_callback_implementation_status_s {
 #define ompt_event_macro(event, callback, eventid) int ompt_##event
 
   OMPT_LOAD_CLIENT_FOREACH_OMPT_EVENT(ompt_event_macro)
+  OMPT_LOAD_CLIENT_FOREACH_OMPT_EVENT_NOEMI(ompt_event_macro)
 
 #undef ompt_event_macro
 } ompt_multiplex_callback_implementation_status_t;
 
 ompt_start_tool_result_t *ompt_multiplex_own_fns;
 ompt_start_tool_result_t *ompt_multiplex_client_fns;
 ompt_function_lookup_t ompt_multiplex_lookup_function;
 ompt_multiplex_callbacks_t ompt_multiplex_own_callbacks,
static void ompt_multiplex_callback_dispatch(ompt_data_t *parallel_data,
   }
   if (ompt_multiplex_client_callbacks.ompt_callback_dispatch) {
     ompt_multiplex_client_callbacks.ompt_callback_dispatch(
         ompt_multiplex_get_client_parallel_data(parallel_data),
         ompt_multiplex_get_client_task_data(task_data), kind, instance);
   }
 }
 
+static void ompt_callback_target_emi_wrapper_own(
+    ompt_target_t kind, ompt_scope_endpoint_t endpoint, int device_num,
+    ompt_data_t *task_data, ompt_data_t *target_task_data,
+    ompt_data_t *target_data, const void *codeptr_ra) {
+  if (endpoint == ompt_scope_begin || endpoint == ompt_scope_beginend) {
+    target_data->value = ompt_multiplex_get_unique_id();
+  }
+  ompt_multiplex_own_callbacks.ompt_callback_target(
+      kind, endpoint, device_num, task_data, target_data->value, codeptr_ra);
+}
+
+static void ompt_callback_target_emi_wrapper_client(
+    ompt_target_t kind, ompt_scope_endpoint_t endpoint, int device_num,
+    ompt_data_t *task_data, ompt_data_t *target_task_data,
+    ompt_data_t *target_data, const void *codeptr_ra) {
+  if (endpoint == ompt_scope_begin || endpoint == ompt_scope_beginend) {
+    target_data->value = ompt_multiplex_get_unique_id();
+  }
+  ompt_multiplex_client_callbacks.ompt_callback_target(
+      kind, endpoint, device_num, task_data, target_data->value, codeptr_ra);
+}
 
 static void ompt_multiplex_callback_target_emi(
     ompt_target_t kind, ompt_scope_endpoint_t endpoint, int device_num,
     ompt_data_t *task_data, ompt_data_t *target_task_data,
     ompt_data_t *target_data, const void *codeptr_ra) {
   if (endpoint == ompt_scope_begin) {
 #ifndef OMPT_MULTIPLEX_CUSTOM_GET_CLIENT_TARGET_DATA
     ompt_multiplex_allocate_data_pair(target_data);

 #endif
 
 #if defined(OMPT_MULTIPLEX_CUSTOM_DELETE_TASK_DATA)
     OMPT_MULTIPLEX_CUSTOM_DELETE_TASK_DATA(target_task_data);
 #endif
   }
 }
 
+static void ompt_callback_target_data_op_emi_wrapper_own(
+    ompt_scope_endpoint_t endpoint, ompt_data_t *target_task_data,
+    ompt_data_t *target_data, ompt_id_t *host_op_id,
+    ompt_target_data_op_t optype, void *src_addr, int src_device_num,
+    void *dest_addr, int dest_device_num, size_t bytes,
+    const void *codeptr_ra) {
+  if (endpoint == ompt_scope_begin || endpoint == ompt_scope_beginend) {
+    if (*host_op_id == 0)
+      *host_op_id = ompt_multiplex_get_unique_id();
+  }
+  if (endpoint == ompt_scope_end || endpoint == ompt_scope_beginend) {
+    ompt_multiplex_own_callbacks.ompt_callback_target_data_op(
+        target_data->value, *host_op_id, optype, src_addr, src_device_num,
+        dest_addr, dest_device_num, bytes, codeptr_ra);
+  }
+}
+
+static void ompt_callback_target_data_op_emi_wrapper_client(
+    ompt_scope_endpoint_t endpoint, ompt_data_t *target_task_data,
+    ompt_data_t *target_data, ompt_id_t *host_op_id,
+    ompt_target_data_op_t optype, void *src_addr, int src_device_num,
+    void *dest_addr, int dest_device_num, size_t bytes,
+    const void *codeptr_ra) {
+  if (endpoint == ompt_scope_begin || endpoint == ompt_scope_beginend) {
+    if (*host_op_id == 0)
+      *host_op_id = ompt_multiplex_get_unique_id();
+  }
+  if (endpoint == ompt_scope_end || endpoint == ompt_scope_beginend) {
+    ompt_multiplex_client_callbacks.ompt_callback_target_data_op(
+        target_data->value, *host_op_id, optype, src_addr, src_device_num,
+        dest_addr, dest_device_num, bytes, codeptr_ra);
+  }
+}
+
 static void ompt_multiplex_callback_target_data_op_emi(
     ompt_scope_endpoint_t endpoint, ompt_data_t *target_task_data,
     ompt_data_t *target_data, ompt_id_t *host_op_id,
     ompt_target_data_op_t optype, void *src_addr, int src_device_num,
     void *dest_addr, int dest_device_num, size_t bytes,
     const void *codeptr_ra) {
   if (ompt_multiplex_own_callbacks.ompt_callback_target_data_op_emi) {
     ompt_multiplex_own_callbacks.ompt_callback_target_data_op_emi(
         endpoint, ompt_multiplex_get_own_task_data(target_task_data),
         ompt_multiplex_get_own_target_data(target_data), host_op_id, optype,
         src_addr, src_device_num, dest_addr, bytes, dest_device_num,
         codeptr_ra);
   }
   if (ompt_multiplex_client_callbacks.ompt_callback_target_data_op_emi) {
     ompt_multiplex_client_callbacks.ompt_callback_target_data_op_emi(
         endpoint, ompt_multiplex_get_client_task_data(target_task_data),
         ompt_multiplex_get_client_target_data(target_data), host_op_id, optype,
         src_addr, src_device_num, dest_addr, bytes, dest_device_num,
         codeptr_ra);
   }
 }
 
+static void ompt_callback_target_submit_emi_wrapper_own(
+    ompt_scope_endpoint_t endpoint, ompt_data_t *target_data,
+    ompt_id_t *host_op_id, unsigned int requested_num_teams) {
+  if (endpoint == ompt_scope_begin || endpoint == ompt_scope_beginend) {
+    if (*host_op_id == 0)
+      *host_op_id = ompt_multiplex_get_unique_id();
+    ompt_multiplex_own_callbacks.ompt_callback_target_submit(
+        target_data->value, *host_op_id, requested_num_teams);
+  }
+}
+
+static void ompt_callback_target_submit_emi_wrapper_client(
+    ompt_scope_endpoint_t endpoint, ompt_data_t *target_data,
+    ompt_id_t *host_op_id, unsigned int requested_num_teams) {
+  if (endpoint == ompt_scope_begin || endpoint == ompt_scope_beginend) {
+    if (*host_op_id == 0)
+      *host_op_id = ompt_multiplex_get_unique_id();
+    ompt_multiplex_client_callbacks.ompt_callback_target_submit(
+        target_data->value, *host_op_id, requested_num_teams);
+  }
+}
+
 static void ompt_multiplex_callback_target_submit_emi(
     ompt_scope_endpoint_t endpoint, ompt_data_t *target_data,
     ompt_id_t *host_op_id, unsigned int requested_num_teams) {
   if (ompt_multiplex_own_callbacks.ompt_callback_target_submit_emi) {
     ompt_multiplex_own_callbacks.ompt_callback_target_submit_emi(
         endpoint, ompt_multiplex_get_own_target_data(target_data), host_op_id,
         requested_num_teams);
   }
   if (ompt_multiplex_client_callbacks.ompt_callback_target_submit_emi) {
     ompt_multiplex_client_callbacks.ompt_callback_target_submit_emi(
         endpoint, ompt_multiplex_get_client_target_data(target_data),
         host_op_id, requested_num_teams);
   }
 }
 
+static void ompt_callback_target_map_emi_wrapper_own(
+    ompt_data_t *target_data, unsigned int nitems, void **host_addr,
+    void **device_addr, size_t *bytes, unsigned int *mapping_flags,
+    const void *codeptr_ra) {
+  ompt_multiplex_own_callbacks.ompt_callback_target_map(
+      target_data->value, nitems, host_addr, device_addr, bytes, mapping_flags,
+      codeptr_ra);
+}
+
+static void ompt_callback_target_map_emi_wrapper_client(
+    ompt_data_t *target_data, unsigned int nitems, void **host_addr,
+    void **device_addr, size_t *bytes, unsigned int *mapping_flags,
+    const void *codeptr_ra) {
+  ompt_multiplex_client_callbacks.ompt_callback_target_map(
+      target_data->value, nitems, host_addr, device_addr, bytes, mapping_flags,
+      codeptr_ra);
+}
+
 static void ompt_multiplex_callback_target_map_emi(
     ompt_data_t *target_data, unsigned int nitems, void **host_addr,
     void **device_addr, size_t *bytes, unsigned int *mapping_flags,
     const void *codeptr_ra) {
   if (ompt_multiplex_own_callbacks.ompt_callback_target_map_emi) {
     ompt_multiplex_own_callbacks.ompt_callback_target_map_emi(
         ompt_multiplex_get_own_target_data(target_data), nitems, host_addr,
         device_addr, bytes, mapping_flags, codeptr_ra);
ompt_multiplex_own_callbacks.ompt_callback_error(severity, message, length,
                                                      codeptr_ra);
   }
   if (ompt_multiplex_client_callbacks.ompt_callback_error) {
     ompt_multiplex_client_callbacks.ompt_callback_error(severity, message,
                                                         length, codeptr_ra);
   }
 }
 
-
 // runtime entry functions
 
 int ompt_multiplex_own_get_task_info(int ancestor_level, int *type,
                                      ompt_data_t **task_data,
                                      ompt_frame_t **task_frame,
                                      ompt_data_t **parallel_data,
                                      int *thread_num) {
   int ret = ompt_multiplex_get_task_info(ancestor_level, type, task_data,
if (ompt_multiplex_implementation_status.ompt_##event_name == -1)          \
                      (ompt_callback_t)&ompt_multiplex_##event_name);           \
     else                                                                       \
       return ompt_multiplex_implementation_status.ompt_##event_name
 
     OMPT_LOAD_CLIENT_FOREACH_OMPT_EVENT(ompt_event_macro)
 
 #undef ompt_event_macro
 
+#define ompt_event_macro(event_name, callback_type, event_id)                  \
+  case ompt_##event_name:                                                      \
+    ompt_multiplex_own_callbacks.ompt_##event_name = (callback_type)callback;  \
+    ompt_multiplex_own_callbacks.ompt_##event_name##_emi =                     \
+        ompt_##event_name##_emi_wrapper_own;                                   \
+    if (ompt_multiplex_implementation_status.ompt_##event_name##_emi == -1)    \
+      return ompt_multiplex_implementation_status.ompt_##event_name##_emi =    \
+                 ompt_multiplex_set_callback(                                  \
+                     ompt_##event_name##_emi,                                  \
+                     (ompt_callback_t)&ompt_multiplex_##event_name##_emi);     \
+    else                                                                       \
+      return ompt_multiplex_implementation_status.ompt_##event_name##_emi
+
+    OMPT_LOAD_CLIENT_FOREACH_OMPT_EVENT_NOEMI(ompt_event_macro)
+
+#undef ompt_event_macro
+
   default:
     return ompt_set_error;
   }
 }
 
 OMPT_API_ROUTINE int
 ompt_multiplex_client_set_callback(ompt_callbacks_t which,
                                    ompt_callback_t callback) {
if (ompt_multiplex_implementation_status.ompt_##event_name == -1)          \
                      (ompt_callback_t)&ompt_multiplex_##event_name);           \
     else                                                                       \
       return ompt_multiplex_implementation_status.ompt_##event_name
 
     OMPT_LOAD_CLIENT_FOREACH_OMPT_EVENT(ompt_event_macro)
 
 #undef ompt_event_macro
 
+#define ompt_event_macro(event_name, callback_type, event_id)                  \
+  case ompt_##event_name:                                                      \
+    ompt_multiplex_client_callbacks.ompt_##event_name =                        \
+        (callback_type)callback;                                               \
+    ompt_multiplex_client_callbacks.ompt_##event_name##_emi =                  \
+        ompt_##event_name##_emi_wrapper_client;                                \
+    if (ompt_multiplex_implementation_status.ompt_##event_name##_emi == -1)    \
+      return ompt_multiplex_implementation_status.ompt_##event_name##_emi =    \
+                 ompt_multiplex_set_callback(                                  \
+                     ompt_##event_name##_emi,                                  \
+                     (ompt_callback_t)&ompt_multiplex_##event_name##_emi);     \
+    else                                                                       \
+      return ompt_multiplex_implementation_status.ompt_##event_name##_emi
+
+    OMPT_LOAD_CLIENT_FOREACH_OMPT_EVENT_NOEMI(ompt_event_macro)
+
+#undef ompt_event_macro
+
   default:
     return ompt_set_error;
   }
 }
 
 ompt_interface_fn_t ompt_multiplex_own_lookup(const char *name) {
   if (!strcmp(name, "ompt_set_callback"))
     return (ompt_interface_fn_t)&ompt_multiplex_own_set_callback;

 
 int ompt_multiplex_initialize(ompt_function_lookup_t lookup,
                               int initial_device_num, ompt_data_t *data) {
   ompt_multiplex_lookup_function = lookup;
   ompt_multiplex_set_callback =
       (ompt_set_callback_t)lookup("ompt_set_callback");
   ompt_multiplex_get_task_info =
       (ompt_get_task_info_t)lookup("ompt_get_task_info");
+  ompt_multiplex_get_unique_id =
+      (ompt_get_unique_id_t)lookup("ompt_get_unique_id");
   ompt_multiplex_get_thread_data =
       (ompt_get_thread_data_t)lookup("ompt_get_thread_data");
   ompt_multiplex_get_parallel_info =
       (ompt_get_parallel_info_t)lookup("ompt_get_parallel_info");
 
   // initialize ompt_multiplex_implementation_status
 #define ompt_event_macro(event_name, callback_type, event_id)                  \
   ompt_multiplex_implementation_status.ompt_##event_name = -1