[Offload][OMPT] Change handling of target_data and target_task_data

- store target_data and target_task_data in OMP runtime
- redefine ompt_task_info_t from libomp with only first two fields -
  target_task_data and target_data
- replace get_target_task_data with ompt_get_task_info_target
- define macro for transparent handling of task_data and
  target_task_data
- provide a local ompt_task_info_t struct for merged target regions
- provide tests to ensure correct values of target_data and
  target_task_data

Author: kaloyan-ignatov

offload/include/OpenMP/OMPT/Interface.h

@@ -25,12 +25,22 @@
 
 #define OMPT_IF_BUILT(stmt) stmt
 
+#define TargetTaskData                                                         \
+  ((OmptTaskInfoPtr == &OmptTaskInfo) ? nullptr                                \
+                                      : (&(OmptTaskInfoPtr->task_data)))
+#define TargetData (OmptTaskInfoPtr->target_data)
+
+typedef struct ompt_task_info_t {
+  ompt_data_t task_data;
+  ompt_data_t target_data;
+} ompt_task_info_t;
+
 /// Callbacks for target regions require task_data representing the
 /// encountering task.
 /// Callbacks for target regions and target data ops require
 /// target_task_data representing the target task region.
 typedef ompt_data_t *(*ompt_get_task_data_t)();
-typedef ompt_data_t *(*ompt_get_target_task_data_t)();
+typedef ompt_task_info_t *(*ompt_get_task_info_target_t)();
 
 namespace llvm {
 namespace omp {
@@ -40,7 +50,7 @@ namespace ompt {
 /// Function pointers that will be used to track task_data and
 /// target_task_data.
 static ompt_get_task_data_t ompt_get_task_data_fn;
-static ompt_get_target_task_data_t ompt_get_target_task_data_fn;
+static ompt_get_task_info_target_t ompt_get_task_info_target_fn;
 
 /// Used to maintain execution state for this thread
 class Interface {
@@ -216,16 +226,16 @@ class Interface {
 
 private:
   /// Target operations id
-  ompt_id_t HostOpId = 0;
-
-  /// Target region data
-  ompt_data_t TargetData = ompt_data_none;
+  ompt_id_t HostOpId{0};
 
   /// Task data representing the encountering task
-  ompt_data_t *TaskData = nullptr;
+  ompt_data_t *TaskData{nullptr};
+
+  /// TaskInfo contains target_data and task_data
+  ompt_task_info_t OmptTaskInfo{ompt_data_none, ompt_data_none};
 
-  /// Target task data representing the target task region
-  ompt_data_t *TargetTaskData = nullptr;
+  /// Ptr to TaskInfo in OpenMP runtime in case of deferred target tasks
+  ompt_task_info_t *OmptTaskInfoPtr{&OmptTaskInfo};
 
   /// Used for marking begin of a data operation
   void beginTargetDataOperation();

offload/libomptarget/OpenMP/OMPT/Callback.cpp

@@ -50,8 +50,8 @@ bool llvm::omp::target::ompt::Initialized = false;
 
 ompt_get_callback_t llvm::omp::target::ompt::lookupCallbackByCode = nullptr;
 ompt_function_lookup_t llvm::omp::target::ompt::lookupCallbackByName = nullptr;
-ompt_get_target_task_data_t ompt_get_target_task_data_fn = nullptr;
 ompt_get_task_data_t ompt_get_task_data_fn = nullptr;
+ompt_get_task_info_target_t ompt_get_task_info_target_fn = nullptr;
 
 /// Unique correlation id
 static std::atomic<uint64_t> IdCounter(1);
@@ -421,18 +421,17 @@ void Interface::beginTargetRegion() {
   // Set up task state
   assert(ompt_get_task_data_fn && "Calling a null task data function");
   TaskData = ompt_get_task_data_fn();
-  // Set up target task state
-  assert(ompt_get_target_task_data_fn &&
-         "Calling a null target task data function");
-  TargetTaskData = ompt_get_target_task_data_fn();
-  // Target state will be set later
-  TargetData = ompt_data_none;
+  // Set up target task and target state
+  assert(ompt_get_task_info_target_fn &&
+         "Calling a null target task info function");
+  if (ompt_task_info_t *TempTaskInfo = ompt_get_task_info_target_fn())
+    OmptTaskInfoPtr = TempTaskInfo;
 }
 
 void Interface::endTargetRegion() {
   TaskData = 0;
-  TargetTaskData = 0;
-  TargetData = ompt_data_none;
+  OmptTaskInfo = {ompt_data_none, ompt_data_none};
+  OmptTaskInfoPtr = &OmptTaskInfo;
 }
 
 /// Used to maintain the finalization functions that are received
@@ -471,7 +470,7 @@ int llvm::omp::target::ompt::initializeLibrary(ompt_function_lookup_t lookup,
 
   bindOmptFunctionName(ompt_get_callback, lookupCallbackByCode);
   bindOmptFunctionName(ompt_get_task_data, ompt_get_task_data_fn);
-  bindOmptFunctionName(ompt_get_target_task_data, ompt_get_target_task_data_fn);
+  bindOmptFunctionName(ompt_get_task_info_target, ompt_get_task_info_target_fn);
 #undef bindOmptFunctionName
 
   // Store pointer of 'ompt_libomp_target_fn_lookup' for use by libomptarget
@@ -480,8 +479,6 @@ int llvm::omp::target::ompt::initializeLibrary(ompt_function_lookup_t lookup,
   assert(lookupCallbackByCode && "lookupCallbackByCode should be non-null");
   assert(lookupCallbackByName && "lookupCallbackByName should be non-null");
   assert(ompt_get_task_data_fn && "ompt_get_task_data_fn should be non-null");
-  assert(ompt_get_target_task_data_fn &&
-         "ompt_get_target_task_data_fn should be non-null");
   assert(LibraryFinalizer == nullptr &&
          "LibraryFinalizer should not be initialized yet");
 

offload/test/ompt/callbacks.h

@@ -5,6 +5,37 @@
 // Tool related code below
 #include <omp-tools.h>
 
+static const char *ompt_target_data_op_t_values[] = {
+    "",
+    "ompt_target_data_alloc",
+    "ompt_target_data_transfer_to_device",
+    "ompt_target_data_transfer_from_device",
+    "ompt_target_data_delete",
+    "ompt_target_data_associate",
+    "ompt_target_data_disassociate",
+    "ompt_target_data_alloc_async",
+    "ompt_target_data_transfer_to_device_async",
+    "ompt_target_data_transfer_from_device_async",
+    "ompt_target_data_delete_async"};
+
+static const char *ompt_scope_endpoint_t_values[] = {
+    "", "ompt_scope_begin", "ompt_scope_end", "ompt_scope_beginend"};
+
+static const char *ompt_target_t_values[] = {"",
+                                             "ompt_target",
+                                             "ompt_target_enter_data",
+                                             "ompt_target_exit_data",
+                                             "ompt_target_update",
+                                             "",
+                                             "",
+                                             "",
+                                             "",
+                                             "",
+                                             "ompt_target_nowait",
+                                             "ompt_target_enter_data_nowait",
+                                             "ompt_target_exit_data_nowait",
+                                             "ompt_target_update_nowait"};
+
 // For EMI callbacks
 ompt_id_t next_op_id = 0x8000000000000001;
 
@@ -38,11 +69,11 @@ static void on_ompt_callback_target_data_op(
     void *src_addr, int src_device_num, void *dest_addr, int dest_device_num,
     size_t bytes, const void *codeptr_ra) {
   assert(codeptr_ra != 0 && "Unexpected null codeptr");
-  printf("  Callback DataOp: target_id=%lu host_op_id=%lu optype=%d src=%p "
+  printf("  Callback DataOp: target_id=%lu host_op_id=%lu optype=%s src=%p "
          "src_device_num=%d "
          "dest=%p dest_device_num=%d bytes=%lu code=%p\n",
-         target_id, host_op_id, optype, src_addr, src_device_num, dest_addr,
-         dest_device_num, bytes, codeptr_ra);
+         target_id, host_op_id, ompt_target_data_op_t_values[optype], src_addr,
+         src_device_num, dest_addr, dest_device_num, bytes, codeptr_ra);
 }
 
 static void on_ompt_callback_target(ompt_target_t kind,
@@ -51,9 +82,10 @@ static void on_ompt_callback_target(ompt_target_t kind,
                                     ompt_id_t target_id,
                                     const void *codeptr_ra) {
   assert(codeptr_ra != 0 && "Unexpected null codeptr");
-  printf("Callback Target: target_id=%lu kind=%d endpoint=%d device_num=%d "
+  printf("Callback Target: target_id=%lu kind=%s endpoint=%s device_num=%d "
          "code=%p\n",
-         target_id, kind, endpoint, device_num, codeptr_ra);
+         target_id, ompt_target_t_values[kind],
+         ompt_scope_endpoint_t_values[endpoint], device_num, codeptr_ra);
 }
 
 static void on_ompt_callback_target_submit(ompt_id_t target_id,
@@ -84,13 +116,15 @@ static void on_ompt_callback_target_data_op_emi(
   // target_task_data may be null, avoid dereferencing it
   uint64_t target_task_data_value =
       (target_task_data) ? target_task_data->value : 0;
-  printf("  Callback DataOp EMI: endpoint=%d optype=%d target_task_data=%p "
+  printf("  Callback DataOp EMI: endpoint=%s optype=%s target_task_data=%p "
          "(0x%lx) target_data=%p (0x%lx) host_op_id=%p (0x%lx) src=%p "
          "src_device_num=%d "
          "dest=%p dest_device_num=%d bytes=%lu code=%p\n",
-         endpoint, optype, target_task_data, target_task_data_value,
-         target_data, target_data->value, host_op_id, *host_op_id, src_addr,
-         src_device_num, dest_addr, dest_device_num, bytes, codeptr_ra);
+         ompt_scope_endpoint_t_values[endpoint],
+         ompt_target_data_op_t_values[optype], target_task_data,
+         target_task_data_value, target_data, target_data->value, host_op_id,
+         *host_op_id, src_addr, src_device_num, dest_addr, dest_device_num,
+         bytes, codeptr_ra);
 }
 
 static void on_ompt_callback_target_emi(ompt_target_t kind,
@@ -102,20 +136,21 @@ static void on_ompt_callback_target_emi(ompt_target_t kind,
   assert(codeptr_ra != 0 && "Unexpected null codeptr");
   if (endpoint == ompt_scope_begin)
     target_data->value = next_op_id++;
-  printf("Callback Target EMI: kind=%d endpoint=%d device_num=%d task_data=%p "
+  printf("Callback Target EMI: kind=%s endpoint=%s device_num=%d task_data=%p "
          "(0x%lx) target_task_data=%p (0x%lx) target_data=%p (0x%lx) code=%p\n",
-         kind, endpoint, device_num, task_data, task_data->value,
-         target_task_data, target_task_data->value, target_data,
-         target_data->value, codeptr_ra);
+         ompt_target_t_values[kind], ompt_scope_endpoint_t_values[endpoint],
+         device_num, task_data, task_data ? task_data->value : 0,
+         target_task_data, target_task_data ? target_task_data->value : 0,
+         target_data, target_data->value, codeptr_ra);
 }
 
 static void on_ompt_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) {
-  printf("  Callback Submit EMI: endpoint=%d  req_num_teams=%d target_data=%p "
+  printf("  Callback Submit EMI: endpoint=%s  req_num_teams=%d target_data=%p "
          "(0x%lx) host_op_id=%p (0x%lx)\n",
-         endpoint, requested_num_teams, target_data, target_data->value,
-         host_op_id, *host_op_id);
+         ompt_scope_endpoint_t_values[endpoint], requested_num_teams,
+         target_data, target_data->value, host_op_id, *host_op_id);
 }
 
 static void on_ompt_callback_target_map_emi(ompt_data_t *target_data,

offload/test/ompt/omp_api.c

@@ -1,6 +1,8 @@
+// clang-format off
 // RUN: %libomptarget-compile-run-and-check-generic
 // REQUIRES: ompt
 // REQUIRES: gpu
+// clang-format on
 
 #include "omp.h"
 #include <stdlib.h>
@@ -32,8 +34,8 @@ int main(int argc, char **argv) {
 
 // clang-format off
 /// CHECK: Callback Init:
-/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=1
-/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=5
-/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=6
-/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=4
+/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=ompt_target_data_alloc
+/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=ompt_target_data_associate
+/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=ompt_target_data_disassociate
+/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=ompt_target_data_delete
 /// CHECK: Callback Fini:

offload/test/ompt/register_with_host.h

@@ -0,0 +1,68 @@
+#define SKIP_CALLBACK_REGISTRATION 1
+
+#include "../../../openmp/runtime/test/ompt/callback.h"
+#include "callbacks.h"
+#include <omp-tools.h>
+
+// From openmp/runtime/test/ompt/callback.h
+#define register_ompt_callback_t(name, type)                                   \
+  do {                                                                         \
+    type f_##name = &on_##name;                                                \
+    if (ompt_set_callback(name, (ompt_callback_t)f_##name) == ompt_set_never)  \
+      printf("0: Could not register callback '" #name "'\n");                  \
+  } while (0)
+
+#define register_ompt_callback(name) register_ompt_callback_t(name, name##_t)
+
+// Init functions
+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");
+
+  if (!ompt_set_callback)
+    return 0; // failed
+
+  // host runtime functions
+  ompt_get_unique_id = (ompt_get_unique_id_t)lookup("ompt_get_unique_id");
+  ompt_get_thread_data = (ompt_get_thread_data_t)lookup("ompt_get_thread_data");
+  ompt_get_task_info = (ompt_get_task_info_t)lookup("ompt_get_task_info");
+
+  ompt_get_unique_id();
+
+  // host callbacks
+  register_ompt_callback(ompt_callback_sync_region);
+  register_ompt_callback_t(ompt_callback_sync_region_wait,
+                           ompt_callback_sync_region_t);
+  register_ompt_callback_t(ompt_callback_reduction,
+                           ompt_callback_sync_region_t);
+  register_ompt_callback(ompt_callback_implicit_task);
+  register_ompt_callback(ompt_callback_parallel_begin);
+  register_ompt_callback(ompt_callback_parallel_end);
+  register_ompt_callback(ompt_callback_task_create);
+  register_ompt_callback(ompt_callback_task_schedule);
+
+  // device callbacks
+  register_ompt_callback(ompt_callback_device_initialize);
+  register_ompt_callback(ompt_callback_device_finalize);
+  register_ompt_callback(ompt_callback_device_load);
+  register_ompt_callback(ompt_callback_target_data_op_emi);
+  register_ompt_callback(ompt_callback_target_emi);
+  register_ompt_callback(ompt_callback_target_submit_emi);
+
+  return 1; // success
+}
+
+void ompt_finalize(ompt_data_t *tool_data) {}
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+ompt_start_tool_result_t *ompt_start_tool(unsigned int omp_version,
+                                          const char *runtime_version) {
+  static ompt_start_tool_result_t ompt_start_tool_result = {&ompt_initialize,
+                                                            &ompt_finalize, 0};
+  return &ompt_start_tool_result;
+}
+#ifdef __cplusplus
+}
+#endif

offload/test/ompt/target_memcpy.c

@@ -1,6 +1,8 @@
+// clang-format off
 // RUN: %libomptarget-compile-run-and-check-generic
 // REQUIRES: ompt
 // REQUIRES: gpu
+// clang-format on
 
 /*
  * Verify that for the target OpenMP APIs, the return address is non-null and
@@ -46,26 +48,26 @@ int main() {
 }
 
 // clang-format off
-/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=1
+/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=ompt_target_data_alloc
 /// CHECK-SAME: src_device_num=[[HOST:[0-9]+]]
 /// CHECK-SAME: dest_device_num=[[DEVICE:[0-9]+]]
 /// CHECK-NOT: code=(nil)
 /// CHECK: code=[[CODE1:0x[0-f]+]]
-/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=2
+/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=ompt_target_data_transfer_to_device
 /// CHECK-SAME: src_device_num=[[HOST]] {{.+}} dest_device_num=[[DEVICE]]
 /// CHECK-NOT: code=(nil)
 /// CHECK-NOT: code=[[CODE1]]
 /// CHECK: code=[[CODE2:0x[0-f]+]]
-/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=3
+/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=ompt_target_data_transfer_from_device
 /// CHECK-SAME: src_device_num=[[DEVICE]] {{.+}} dest_device_num=[[DEVICE]]
 /// CHECK-NOT: code=(nil)
 /// CHECK-NOT: code=[[CODE2]]
 /// CHECK: code=[[CODE3:0x[0-f]+]]
-/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=3
+/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=ompt_target_data_transfer_from_device
 /// CHECK-SAME: src_device_num=[[DEVICE]] {{.+}} dest_device_num=[[HOST]]
 /// CHECK-NOT: code=(nil)
 /// CHECK-NOT: code=[[CODE3]]
 /// CHECK: code=[[CODE4:0x[0-f]+]]
-/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=4
+/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=ompt_target_data_delete
 /// CHECK-NOT: code=(nil)
 /// CHECK-NOT: code=[[CODE4]]