[OpenMP][OMPT] Change handling of target-related tool data in libomp

- store target_data in ompt_task_info_t to prevent data loss across
  scheduling of target region (both for deferred and undeferred target
  tasks)
- target_task_data is already in ompt_task_info_t, replace previous
  implementation which returned the wrong value
- combine queries for target_data and target_task_data and directly
  return ompt_task_info_t
- return correct task_data, OpenMP standard defines task_data in target
  callbacks as belonging to the generating (encountering) task

Author: kaloyan-ignatov

openmp/runtime/src/ompt-general.cpp

@@ -886,10 +886,12 @@ static ompt_interface_fn_t ompt_fn_lookup(const char *s) {
   return NULL;
 }
 
-static ompt_data_t *ompt_get_task_data() { return __ompt_get_task_data(); }
+static ompt_data_t *ompt_get_task_data() {
+  return __ompt_get_generating_task();
+}
 
-static ompt_data_t *ompt_get_target_task_data() {
-  return __ompt_get_target_task_data();
+static ompt_task_info_t *ompt_get_task_info_target() {
+  return __ompt_get_task_info_target();
 }
 
 /// Lookup function to query libomp callbacks registered by the tool
@@ -900,7 +902,7 @@ static ompt_interface_fn_t ompt_libomp_target_fn_lookup(const char *s) {
 
   provide_fn(ompt_get_callback);
   provide_fn(ompt_get_task_data);
-  provide_fn(ompt_get_target_task_data);
+  provide_fn(ompt_get_task_info_target);
 #undef provide_fn
 
 #define ompt_interface_fn(fn, type, code)                                      \

openmp/runtime/src/ompt-internal.h

@@ -57,8 +57,11 @@ typedef struct ompt_callbacks_active_s {
       (info->td_flags.merged_if0 ? ompt_task_mergeable : 0x0)
 
 typedef struct {
-  ompt_frame_t frame;
+  // liboffload only uses task_data and target_data. They must be the first
+  // elements!
   ompt_data_t task_data;
+  ompt_data_t target_data;
+  ompt_frame_t frame;
   struct kmp_taskdata *scheduling_parent;
   int thread_num;
   ompt_dispatch_chunk_t dispatch_chunk;

openmp/runtime/src/ompt-specific.cpp

@@ -346,14 +346,30 @@ void __ompt_lw_taskteam_unlink(kmp_info_t *thr) {
 // task support
 //----------------------------------------------------------
 
-ompt_data_t *__ompt_get_task_data() {
+ompt_data_t *__ompt_get_generating_task() {
   kmp_info_t *thr = ompt_get_thread();
-  ompt_data_t *task_data = thr ? OMPT_CUR_TASK_DATA(thr) : NULL;
-  return task_data;
+  if (thr) {
+    kmp_taskdata_t *taskdata = thr->th.th_current_task;
+    if (taskdata == NULL)
+      return NULL;
+    if (taskdata->td_flags.target)
+      return &(taskdata->td_parent->ompt_task_info.task_data);
+    else
+      return &(taskdata->ompt_task_info.task_data);
+  }
+  return NULL;
 }
 
-ompt_data_t *__ompt_get_target_task_data() {
-  return &__kmp_threads[__kmp_get_gtid()]->th.ompt_thread_info.target_task_data;
+ompt_task_info_t *__ompt_get_task_info_target() {
+  kmp_info_t *thr = ompt_get_thread();
+  if (thr) {
+    kmp_taskdata_t *taskdata = thr->th.th_current_task;
+    if (taskdata == NULL)
+      return NULL;
+    if (taskdata->td_flags.target)
+      return &taskdata->ompt_task_info;
+  }
+  return NULL;
 }
 
 int __ompt_get_task_info_internal(int ancestor_level, int *type,

openmp/runtime/src/ompt-specific.h

@@ -37,9 +37,9 @@ void __ompt_lw_taskteam_unlink(kmp_info_t *thr);
 
 ompt_team_info_t *__ompt_get_teaminfo(int depth, int *size);
 
-ompt_data_t *__ompt_get_task_data();
+ompt_data_t *__ompt_get_generating_task();
 
-ompt_data_t *__ompt_get_target_task_data();
+ompt_task_info_t *__ompt_get_task_info_target();
 
 ompt_task_info_t *__ompt_get_task_info_object(int depth);
 

openmp/runtime/test/ompt/callback.h

@@ -1018,6 +1018,8 @@ static void on_ompt_callback_error(ompt_severity_t severity,
          codeptr_ra);
 }
 
+#ifndef SKIP_CALLBACK_REGISTRATION
+
 int ompt_initialize(ompt_function_lookup_t lookup, int initial_device_num,
                     ompt_data_t *tool_data) {
   ompt_set_callback = (ompt_set_callback_t)lookup("ompt_set_callback");
@@ -1094,6 +1096,7 @@ ompt_start_tool_result_t *ompt_start_tool(unsigned int omp_version,
 #ifdef __cplusplus
 }
 #endif
+#endif // ifndef SKIP_CALLBACK_REGISTRATION
 #endif // ifndef USE_PRIVATE_TOOL
 #ifdef _OMPT_TESTS
 #undef _OMPT_TESTS