[OMPT][Offload][OpenMP] Fixes for OMPT data used by libomptarget

offload/include/OpenMP/OMPT/Interface.h

@@ -25,12 +25,23 @@
 
 #define OMPT_IF_BUILT(stmt) stmt
 
+#define TargetTaskData                                                         \
+  ((OmptTaskInfoPtr == &OmptTaskInfo) ? nullptr                                \
+                                      : (&(OmptTaskInfoPtr->task_data)))
+#define TargetData (OmptTaskInfoPtr->target_data)
+
+/// Prefix of ompt_task_info_t struct from libomp
+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 +51,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 +227,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/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_tool_data.c

@@ -0,0 +1,152 @@
+// clang-format off
+// RUN: env LIBOMP_NUM_HIDDEN_HELPER_THREADS=1 %libomptarget-compile-run-and-check-generic
+// REQUIRES: ompt
+// clang-format on
+
+#include <inttypes.h>
+#include <omp-tools.h>
+#include <omp.h>
+#include <stdio.h>
+#include <string.h>
+
+#include "register_with_host.h"
+
+#define N 1000000
+#define M 1000
+
+int main() {
+  float *x = malloc(N * sizeof(float));
+  float *y = malloc(N * sizeof(float));
+
+  for (int i = 0; i < N; i++) {
+    x[i] = 1;
+    y[i] = 1;
+  }
+
+#pragma omp target enter data map(to : x[0 : N]) map(alloc : y[0 : N])
+#pragma omp target teams distribute parallel for
+  for (int i = 0; i < N; i++) {
+    for (int j = 0; j < M; j++) {
+      y[i] += 3 * x[i];
+    }
+  }
+
+#pragma omp target teams distribute parallel for
+  for (int i = 0; i < N; i++) {
+    for (int j = 0; j < M; j++) {
+      y[i] += 3 * x[i];
+    }
+  }
+
+#pragma omp target exit data map(release : x[0 : N]) map(from : y[0 : N])
+
+  printf("%f, %f\n", x[0], y[0]);
+
+  free(x);
+  free(y);
+  return 0;
+}
+
+// clang-format off
+/// CHECK: ompt_event_initial_task_begin
+/// CHECK-SAME: task_id=[[ENCOUNTERING_TASK:[0-f]+]]
+
+/// CHECK: Callback Target EMI: kind=ompt_target_enter_data endpoint=ompt_scope_begin
+/// CHECK-SAME: task_data=0x{{[0-f]+}} (0x[[ENCOUNTERING_TASK]])
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA:[0-f]+]])
+
+/// CHECK: Callback DataOp EMI: endpoint=ompt_scope_begin
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback DataOp EMI: endpoint=ompt_scope_end
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback DataOp EMI: endpoint=ompt_scope_begin
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback DataOp EMI: endpoint=ompt_scope_end
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback DataOp EMI: endpoint=ompt_scope_begin
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback DataOp EMI: endpoint=ompt_scope_end
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback Target EMI: kind=ompt_target_enter_data endpoint=ompt_scope_end
+/// CHECK-SAME: task_data=0x{{[0-f]+}} (0x[[ENCOUNTERING_TASK]])
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback Target EMI: kind=ompt_target endpoint=ompt_scope_begin
+/// CHECK-SAME: task_data=0x{{[0-f]+}} (0x[[ENCOUNTERING_TASK]])
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA:[0-f]+]])
+
+/// CHECK: Callback Submit EMI: endpoint=ompt_scope_begin
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback Submit EMI: endpoint=ompt_scope_end
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback Target EMI: kind=ompt_target endpoint=ompt_scope_end
+/// CHECK-SAME: task_data=0x{{[0-f]+}} (0x[[ENCOUNTERING_TASK]])
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback Target EMI: kind=ompt_target endpoint=ompt_scope_begin
+/// CHECK-SAME: task_data=0x{{[0-f]+}} (0x[[ENCOUNTERING_TASK]])
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA:[0-f]+]])
+
+/// CHECK: Callback Submit EMI: endpoint=ompt_scope_begin
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback Submit EMI: endpoint=ompt_scope_end
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback Target EMI: kind=ompt_target endpoint=ompt_scope_end
+/// CHECK-SAME: task_data=0x{{[0-f]+}} (0x[[ENCOUNTERING_TASK]])
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback Target EMI: kind=ompt_target_exit_data endpoint=ompt_scope_begin
+/// CHECK-SAME: task_data=0x{{[0-f]+}} (0x[[ENCOUNTERING_TASK]])
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA:[0-f]+]])
+
+/// CHECK: Callback DataOp EMI: endpoint=ompt_scope_begin
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback DataOp EMI: endpoint=ompt_scope_end
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback DataOp EMI: endpoint=ompt_scope_begin
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback DataOp EMI: endpoint=ompt_scope_end
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback DataOp EMI: endpoint=ompt_scope_begin
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback DataOp EMI: endpoint=ompt_scope_end
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])
+
+/// CHECK: Callback Target EMI: kind=ompt_target_exit_data endpoint=ompt_scope_end
+/// CHECK-SAME: task_data=0x{{[0-f]+}} (0x[[ENCOUNTERING_TASK]])
+/// CHECK-SAME: target_task_data=(nil) (0x0)
+/// CHECK-SAME: target_data=0x{{[0-f]+}} (0x[[TARGET_DATA]])