[OpenMP][OMPT] Fix `target enter data` callback ordering & reported device num

[mhalk]

Summary

This patch fixes: https://github.com/llvm/llvm-project/issues/64738
We observed multiple issues, primarily that the DeviceId was reported as -1
in certain scenarios. The reason for this is simply that the device is not
initialized at that point. Hence, we need to move the RAII object creation just
after the checkDeviceAndCtors, closer to the actual call we want to observe.

This also solves an odering issue where one target enter data callback would
be executed before the Init callback.
Additionally, this change will also fix that the callbacks corresponding to
enter / exit data and update in conjunction with nowait would not result
in the emission of an OMPT callback.

Added a testcase to cover initialized device number and omp target constructs.

[mhalk]

Clarify the intent of the test and returning a nullptr in the helper function.

[dhruvachak]

openmp/libomptarget/src/interface.cpp

141

Can't you have the RAII here? That way, you don't need the new/delete.
You can add an assert that the type is one of the expected 3.
InterfaceRAII(

TargetDataFunction == targetDataBegin ? 
RegionInterface.getCallbacks<ompt_target_enter_data>() :
TargetDataFunction == targetDataEnd ?
RegionInterface.getCallbacks<ompt_target_exit_data>() :
RegionInterface.getCallbacks<ompt_target_update>(),
DeviceId,
OMPT_GET_RETURN_ADDRESS(0));

[mhalk]

Removed helper function in favor of assert + callback selection via ternary ops.
Thanks @dhruvachak, was not quite sure if that would be acceptable.
Either way, this seems less invasive to me.

[mhalk]

Rebase + edit commit message to link issue https://github.com/llvm/llvm-project/issues/64738

[mhalk]

Change testcase to use EMI callbacks.

Discovered a possible segfault which triggers only when using EMI callbacks.
Caused by target_task_data=(nil) and seems to be fixed by this patch, too.

[dhruvachak]

Looks good.

openmp/libomptarget/src/interface.cpp

 #include <cassert>
 #include <cstdio>
 #include <cstdlib>
 #include <mutex>
 #include <type_traits>
 
 #ifdef OMPT_SUPPORT
 using namespace llvm::omp::target::ompt;
+
+typedef InterfaceRAII<
+    std::pair<std::_Mem_fn<void (Interface::*)(long, void *)>,
+              std::_Mem_fn<void (Interface::*)(long, void *)>>,
+    long, void *> *TargetDataRAIIPtr;
+
+/// Helper function for OMPT callbacks related to target data operations.
+/// Ordering of callbacks as well as displayed data could be affected by the
+/// fact that device CTOR is called lazily. This function helps to create RAII
+/// objects at a later stage of the execution, when initialization is done.
+/// Note: This will allocate an object, which has to be manually deleted.
+inline TargetDataRAIIPtr
+createTargetDataRAII(int64_t DeviceId, TargetDataFuncPtrTy TargetDataFunction) {
+  if (TargetDataFunction == targetDataBegin)
+    return new InterfaceRAII(
+        RegionInterface.getCallbacks<ompt_target_enter_data>(), DeviceId,
+        OMPT_GET_RETURN_ADDRESS(0));
+
+  if (TargetDataFunction == targetDataEnd)
+    return new InterfaceRAII(
+        RegionInterface.getCallbacks<ompt_target_exit_data>(), DeviceId,
+        OMPT_GET_RETURN_ADDRESS(0));
+
+  if (TargetDataFunction == targetDataUpdate)
+    return new InterfaceRAII(RegionInterface.getCallbacks<ompt_target_update>(),
+                             DeviceId, OMPT_GET_RETURN_ADDRESS(0));
+
+  // Do not create an object / do not emit an OMPT callback
+  return static_cast<TargetDataRAIIPtr>(nullptr);
+}
 #endif
 
 ////////////////////////////////////////////////////////////////////////////////
 /// adds requires flags
 EXTERN void __tgt_register_requires(int64_t Flags) {
   TIMESCOPE();
   PM->RTLs.registerRequires(Flags);
 }
DP("Entry %2d: Base=" DPxMOD ", Begin=" DPxMOD ", Size=%" PRId64
        (ArgNames) ? getNameFromMapping(ArgNames[I]).c_str() : "unknown");
   }
 #endif
 
   DeviceTy &Device = *PM->Devices[DeviceId];
   TargetAsyncInfoTy TargetAsyncInfo(Device);
   AsyncInfoTy &AsyncInfo = TargetAsyncInfo;
 
+  /// RAII to establish tool anchors before and after data begin / end / update
+  OMPT_IF_BUILT(auto *TargetDataRAII =
+                    createTargetDataRAII(DeviceId, TargetDataFunction);)
+
   int Rc = OFFLOAD_SUCCESS;
   Rc = TargetDataFunction(Loc, Device, ArgNum, ArgsBase, Args, ArgSizes,
                           ArgTypes, ArgNames, ArgMappers, AsyncInfo,
                           false /* FromMapper */);
 
   if (Rc == OFFLOAD_SUCCESS)
     Rc = AsyncInfo.synchronize();
 
+  // Destroy allocated RAII object
+  OMPT_IF_BUILT(delete TargetDataRAII;)
+
   handleTargetOutcome(Rc == OFFLOAD_SUCCESS, Loc);
 }
 
 /// creates host-to-target data mapping, stores it in the
 /// libomptarget.so internal structure (an entry in a stack of data maps)
 /// and passes the data to the device.
 EXTERN void __tgt_target_data_begin_mapper(ident_t *Loc, int64_t DeviceId,
                                            int32_t ArgNum, void **ArgsBase,
                                            void **Args, int64_t *ArgSizes,
                                            int64_t *ArgTypes,
                                            map_var_info_t *ArgNames,
                                            void **ArgMappers) {
   TIMESCOPE_WITH_IDENT(Loc);
-  /// RAII to establish tool anchors before and after data begin
-  OMPT_IF_BUILT(InterfaceRAII TargetDataEnterRAII(
-                    RegionInterface.getCallbacks<ompt_target_enter_data>(),
-                    DeviceId,
-                    /* CodePtr */ OMPT_GET_RETURN_ADDRESS(0));)
-
   targetDataMapper<AsyncInfoTy>(Loc, DeviceId, ArgNum, ArgsBase, Args, ArgSizes,
                                 ArgTypes, ArgNames, ArgMappers, targetDataBegin,
                                 "Entering OpenMP data region", "begin");
 }
 
 EXTERN void __tgt_target_data_begin_nowait_mapper(
     ident_t *Loc, int64_t DeviceId, int32_t ArgNum, void **ArgsBase,
     void **Args, int64_t *ArgSizes, int64_t *ArgTypes, map_var_info_t *ArgNames,

 /// created by the last __tgt_target_data_begin.
 EXTERN void __tgt_target_data_end_mapper(ident_t *Loc, int64_t DeviceId,
                                          int32_t ArgNum, void **ArgsBase,
                                          void **Args, int64_t *ArgSizes,
                                          int64_t *ArgTypes,
                                          map_var_info_t *ArgNames,
                                          void **ArgMappers) {
   TIMESCOPE_WITH_IDENT(Loc);
-  /// RAII to establish tool anchors before and after data end
-  OMPT_IF_BUILT(InterfaceRAII TargetDataExitRAII(
-                    RegionInterface.getCallbacks<ompt_target_exit_data>(),
-                    DeviceId,
-                    /* CodePtr */ OMPT_GET_RETURN_ADDRESS(0));)
-
   targetDataMapper<AsyncInfoTy>(Loc, DeviceId, ArgNum, ArgsBase, Args, ArgSizes,
                                 ArgTypes, ArgNames, ArgMappers, targetDataEnd,
                                 "Exiting OpenMP data region", "end");
 }
 
 EXTERN void __tgt_target_data_end_nowait_mapper(
     ident_t *Loc, int64_t DeviceId, int32_t ArgNum, void **ArgsBase,
     void **Args, int64_t *ArgSizes, int64_t *ArgTypes, map_var_info_t *ArgNames,
     void **ArgMappers, int32_t DepNum, void *DepList, int32_t NoAliasDepNum,
     void *NoAliasDepList) {
   TIMESCOPE_WITH_IDENT(Loc);
   targetDataMapper<TaskAsyncInfoWrapperTy>(
       Loc, DeviceId, ArgNum, ArgsBase, Args, ArgSizes, ArgTypes, ArgNames,
       ArgMappers, targetDataEnd, "Exiting OpenMP data region", "end");
 }
 
 EXTERN void __tgt_target_data_update_mapper(ident_t *Loc, int64_t DeviceId,
                                             int32_t ArgNum, void **ArgsBase,
                                             void **Args, int64_t *ArgSizes,
                                             int64_t *ArgTypes,
                                             map_var_info_t *ArgNames,
                                             void **ArgMappers) {
   TIMESCOPE_WITH_IDENT(Loc);
-  /// RAII to establish tool anchors before and after data update
-  OMPT_IF_BUILT(InterfaceRAII TargetDataUpdateRAII(
-                    RegionInterface.getCallbacks<ompt_target_update>(),
-                    DeviceId,
-                    /* CodePtr */ OMPT_GET_RETURN_ADDRESS(0));)
-
   targetDataMapper<AsyncInfoTy>(
       Loc, DeviceId, ArgNum, ArgsBase, Args, ArgSizes, ArgTypes, ArgNames,
       ArgMappers, targetDataUpdate, "Updating OpenMP data", "update");
 }
 
 EXTERN void __tgt_target_data_update_nowait_mapper(
     ident_t *Loc, int64_t DeviceId, int32_t ArgNum, void **ArgsBase,
     void **Args, int64_t *ArgSizes, int64_t *ArgTypes, map_var_info_t *ArgNames,
DP("Entry %2d: Base=" DPxMOD ", Begin=" DPxMOD ", Size=%" PRId64
            ? getNameFromMapping(KernelArgs->ArgNames[I]).c_str()
            : "unknown");
   }
 #endif
 
   DeviceTy &Device = *PM->Devices[DeviceId];
   TargetAsyncInfoTy TargetAsyncInfo(Device);
   AsyncInfoTy &AsyncInfo = TargetAsyncInfo;
-  OMPT_IF_BUILT(InterfaceRAII TargetDataAllocRAII(
+  /// RAII to establish tool anchors before and after target region
+  OMPT_IF_BUILT(InterfaceRAII TargetRAII(
                     RegionInterface.getCallbacks<ompt_target>(), DeviceId,
                     /* CodePtr */ OMPT_GET_RETURN_ADDRESS(0));)
 
   int Rc = OFFLOAD_SUCCESS;
   Rc = target(Loc, Device, HostPtr, *KernelArgs, AsyncInfo);
 
   if (Rc == OFFLOAD_SUCCESS)
     Rc = AsyncInfo.synchronize();
EXTERN int __tgt_target_kernel_replay(ident_t *Loc, int64_t DeviceId,
                                       int32_t NumTeams, int32_t ThreadLimit,
                                       uint64_t LoopTripCount) {
 
   if (checkDeviceAndCtors(DeviceId, Loc)) {
     DP("Not offloading to device %" PRId64 "\n", DeviceId);
     return OMP_TGT_FAIL;
   }
   DeviceTy &Device = *PM->Devices[DeviceId];
-  OMPT_IF_BUILT(InterfaceRAII TargetDataAllocRAII(
+  /// RAII to establish tool anchors before and after target region
+  OMPT_IF_BUILT(InterfaceRAII TargetRAII(
                     RegionInterface.getCallbacks<ompt_target>(), DeviceId,
                     /* CodePtr */ OMPT_GET_RETURN_ADDRESS(0));)
 
   AsyncInfoTy AsyncInfo(Device);
   int Rc = target_replay(Loc, Device, HostPtr, DeviceMemory, DeviceMemorySize,
                          TgtArgs, TgtOffsets, NumArgs, NumTeams, ThreadLimit,
                          LoopTripCount, AsyncInfo);
   if (Rc == OFFLOAD_SUCCESS)

openmp/libomptarget/test/ompt/veccopy_data.c

+// RUN: %libomptarget-compile-run-and-check-generic
+// REQUIRES: ompt
+// UNSUPPORTED: aarch64-unknown-linux-gnu
+// UNSUPPORTED: aarch64-unknown-linux-gnu-LTO
+// UNSUPPORTED: x86_64-pc-linux-gnu
+// UNSUPPORTED: x86_64-pc-linux-gnu-oldDriver
+// UNSUPPORTED: x86_64-pc-linux-gnu-LTO
+
+/*
+ * Example OpenMP program that registers non-EMI callbacks.
+ * Explicitly testing for an initialized device num and
+ * #pragma omp target [data enter / data exit / update]
+ * The latter with the addition of a nowait clause.
+ */
+
+#include <omp.h>
+#include <stdio.h>
+
+#include "callbacks.h"
+#include "register_non_emi.h"
+
+#define N 100000
+
+#pragma omp declare target
+int c[N];
+#pragma omp end declare target
+
+int main() {
+  int a[N];
+  int b[N];
+
+  int i;
+
+  for (i = 0; i < N; i++)
+    a[i] = 0;
+
+  for (i = 0; i < N; i++)
+    b[i] = i;
+
+  for (i = 0; i < N; i++)
+    c[i] = 0;
+
+#pragma omp target enter data map(to : a)
+#pragma omp target parallel for
+  {
+    for (int j = 0; j < N; j++)
+      a[j] = b[j];
+  }
+#pragma omp target exit data map(from : a)
+
+#pragma omp target parallel for map(alloc : c)
+  {
+    for (int j = 0; j < N; j++)
+      c[j] = 2 * j + 1;
+  }
+#pragma omp target update from(c) nowait
+#pragma omp barrier
+
+  int rc = 0;
+  for (i = 0; i < N; i++) {
+    if (a[i] != i) {
+      rc++;
+      printf("Wrong value: a[%d]=%d\n", i, a[i]);
+    }
+  }
+
+  for (i = 0; i < N; i++) {
+    if (c[i] != 2 * i + 1) {
+      rc++;
+      printf("Wrong value: c[%d]=%d\n", i, c[i]);
+    }
+  }
+
+  if (!rc)
+    printf("Success\n");
+
+  return rc;
+}
+
+/// CHECK-NOT: Callback Target:
+/// CHECK-NOT: device_num=-1
+/// CHECK: Callback Init:
+/// CHECK: Callback Load:
+/// CHECK: Callback Target: target_id=[[TARGET_ID:[0-9]+]] kind=2 endpoint=1
+/// CHECK-NOT: device_num=-1
+/// 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=2
+/// CHECK: Callback Target: target_id=[[TARGET_ID:[0-9]+]] kind=2 endpoint=2
+/// CHECK-NOT: device_num=-1
+/// CHECK: Callback Target: target_id=[[TARGET_ID:[0-9]+]] kind=1 endpoint=1
+/// 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=2
+/// CHECK: Callback Submit: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] req_num_teams=1
+/// 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=4
+/// CHECK: Callback Target: target_id=[[TARGET_ID:[0-9]+]] kind=1 endpoint=2
+/// CHECK: Callback Target: target_id=[[TARGET_ID:[0-9]+]] kind=3 endpoint=1
+/// 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=4
+/// CHECK: Callback Target: target_id=[[TARGET_ID:[0-9]+]] kind=3 endpoint=2
+/// CHECK: Callback Target: target_id=[[TARGET_ID:[0-9]+]] kind=1 endpoint=1
+/// CHECK: Callback Submit: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] req_num_teams=1
+/// CHECK: Callback Target: target_id=[[TARGET_ID:[0-9]+]] kind=1 endpoint=2
+/// CHECK: Callback Target: target_id=[[TARGET_ID:[0-9]+]] kind=4 endpoint=1
+/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=3
+/// CHECK: Callback Target: target_id=[[TARGET_ID:[0-9]+]] kind=4 endpoint=2
+/// CHECK: Callback Fini: