[OpenMP][OMPT] Fix reported target pointer for data alloc callback

[mhalk]

Summary

This patch fixes: https://github.com/llvm/llvm-project/issues/64671
DataOp EMI callbacks would not report the correct target pointer.
This is now alleviated by passing a void** into the function which
emits the actual callback, then evaluating that pointer.

Note: Since this is only done after the pointer has been properly
updated, only endpoint=2 callbacks will show a non-null value.

[mhalk]

If using void** as solution is despised, we could add an updateArguments method to InterfaceRAII.
That would replace the current tuple holding the arguments with a new, provided one.

void updateArguments(ArgsTy... NewArgs) {
  Arguments = std::make_tuple(NewArgs...);
}

We would then call updateArguments, after the actual data allocation.
So, the EMI callback where endpoint=2 will have the actual pointer value available.

void *DeviceTy::allocData(int64_t Size, void *HstPtr, int32_t Kind) {
  [...]
  TargetPtr = RTL->data_alloc(RTLDeviceID, Size, HstPtr, Kind);

  OMPT_IF_BUILT(TargetDataAllocRAII.updateArguments(
      RTLDeviceID, HstPtr, TargetPtr, Size,
      /* CodePtr */ OMPT_GET_RETURN_ADDRESS(0));)

  return TargetPtr;
}

[jdoerfert]

Why not just a reference? Anyway, this is fine too.

[mhalk]

In D157996#4589257, @jdoerfert wrote:

Why not just a reference? Anyway, this is fine too.

Thanks Johannes!

If you don't mind to elaborate w.r.t. the "reference", could you explain, please?
I'm curious -- sounds more robust to me.

[dhruvachak]

LGTM.

openmp/libomptarget/src/OmptCallback.cpp

 /// Create a new correlation id and update the target region id
 static uint64_t createRegionId() {
   uint64_t NewId = createId();
   RegionInterface.setTargetDataValue(NewId);
   return NewId;
 }
 
 void Interface::beginTargetDataAlloc(int64_t DeviceId, void *HstPtrBegin,
-                                     size_t Size, void *Code) {
+                                     void **TgtPtrBegin, size_t Size,
+                                     void *Code) {
   beginTargetDataOperation();
   if (ompt_callback_target_data_op_emi_fn) {
     // HostOpId will be set by the tool. Invoke the tool supplied data op EMI
     // callback
     ompt_callback_target_data_op_emi_fn(
         ompt_scope_begin, TargetTaskData, &TargetData, &TargetRegionOpId,
         ompt_target_data_alloc, HstPtrBegin,
-        /* SrcDeviceNum */ omp_get_initial_device(), /* TgtPtrBegin */ nullptr,
+        /* SrcDeviceNum */ omp_get_initial_device(), *TgtPtrBegin,
         /* TgtDeviceNum */ DeviceId, Size, Code);
   } else if (ompt_callback_target_data_op_fn) {
     // HostOpId is set by the runtime
     HostOpId = createOpId();
     // Invoke the tool supplied data op callback
     ompt_callback_target_data_op_fn(
         TargetData.value, HostOpId, ompt_target_data_alloc, HstPtrBegin,
-        /* SrcDeviceNum */ omp_get_initial_device(), /* TgtPtrBegin */ nullptr,
+        /* SrcDeviceNum */ omp_get_initial_device(), *TgtPtrBegin,
         /* TgtDeviceNum */ DeviceId, Size, Code);
   }
 }
 
 void Interface::endTargetDataAlloc(int64_t DeviceId, void *HstPtrBegin,
-                                   size_t Size, void *Code) {
+                                   void **TgtPtrBegin, size_t Size,
+                                   void *Code) {
   // Only EMI callback handles end scope
   if (ompt_callback_target_data_op_emi_fn) {
     // HostOpId will be set by the tool. Invoke the tool supplied data op EMI
     // callback
     ompt_callback_target_data_op_emi_fn(
         ompt_scope_end, TargetTaskData, &TargetData, &TargetRegionOpId,
         ompt_target_data_alloc, HstPtrBegin,
-        /* SrcDeviceNum */ omp_get_initial_device(), /* TgtPtrBegin */ nullptr,
+        /* SrcDeviceNum */ omp_get_initial_device(), *TgtPtrBegin,
         /* TgtDeviceNum */ DeviceId, Size, Code);
   }
   endTargetDataOperation();
 }
 
 void Interface::beginTargetDataSubmit(int64_t DeviceId, void *TgtPtrBegin,
                                       void *HstPtrBegin, size_t Size,
                                       void *Code) {

openmp/libomptarget/src/OmptInterface.h

 /// 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;
 
 /// Used to maintain execution state for this thread
 class Interface {
 public:
   /// Top-level function for invoking callback before device data allocation
-  void beginTargetDataAlloc(int64_t DeviceId, void *TgtPtrBegin, size_t Size,
-                            void *Code);
+  void beginTargetDataAlloc(int64_t DeviceId, void *HstPtrBegin,
+                            void **TgtPtrBegin, size_t Size, void *Code);
 
   /// Top-level function for invoking callback after device data allocation
-  void endTargetDataAlloc(int64_t DeviceId, void *TgtPtrBegin, size_t Size,
-                          void *Code);
+  void endTargetDataAlloc(int64_t DeviceId, void *HstPtrBegin,
+                          void **TgtPtrBegin, size_t Size, void *Code);
 
   /// Top-level function for invoking callback before data submit
   void beginTargetDataSubmit(int64_t DeviceId, void *HstPtrBegin,
                              void *TgtPtrBegin, size_t Size, void *Code);
 
   /// Top-level function for invoking callback after data submit
   void endTargetDataSubmit(int64_t DeviceId, void *HstPtrBegin,
                            void *TgtPtrBegin, size_t Size, void *Code);

openmp/libomptarget/src/device.cpp

 // Load binary to device.
 __tgt_target_table *DeviceTy::loadBinary(void *Img) {
   std::lock_guard<decltype(RTL->Mtx)> LG(RTL->Mtx);
   return RTL->load_binary(RTLDeviceID, Img);
 }
 
 void *DeviceTy::allocData(int64_t Size, void *HstPtr, int32_t Kind) {
   /// RAII to establish tool anchors before and after data allocation
+  void *TargetPtr = nullptr;
   OMPT_IF_BUILT(InterfaceRAII TargetDataAllocRAII(
                     RegionInterface.getCallbacks<ompt_target_data_alloc>(),
-                    RTLDeviceID, HstPtr, Size,
+                    RTLDeviceID, HstPtr, &TargetPtr, Size,
                     /* CodePtr */ OMPT_GET_RETURN_ADDRESS(0));)
 
-  return RTL->data_alloc(RTLDeviceID, Size, HstPtr, Kind);
+  TargetPtr = RTL->data_alloc(RTLDeviceID, Size, HstPtr, Kind);
+  return TargetPtr;
 }
 
 int32_t DeviceTy::deleteData(void *TgtAllocBegin, int32_t Kind) {
   /// RAII to establish tool anchors before and after data deletion
   OMPT_IF_BUILT(InterfaceRAII TargetDataDeleteRAII(
                     RegionInterface.getCallbacks<ompt_target_data_delete>(),
                     RTLDeviceID, TgtAllocBegin,
                     /* CodePtr */ OMPT_GET_RETURN_ADDRESS(0));)

openmp/libomptarget/test/ompt/veccopy_disallow_both.c

#pragma omp target teams distribute parallel for
   return rc;
 }
 
 /// CHECK: Callback Init:
 /// CHECK: Callback Load:
 /// CHECK: Callback Target EMI: kind=1 endpoint=1
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=1
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=1
+/// CHECK-NOT: dest=(nil)
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=1
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=1
+/// CHECK-NOT: dest=(nil)
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=2 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 EMI: endpoint=1 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=4
 /// CHECK: Callback Target EMI: kind=1 endpoint=2
 /// CHECK: Callback Target EMI: kind=1 endpoint=1
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=1
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=1
+/// CHECK-NOT: dest=(nil)
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=1
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=1
+/// CHECK-NOT: dest=(nil)
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=2
 /// CHECK: Callback Submit: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] req_num_teams=0
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=4
 /// CHECK: Callback Target EMI: kind=1 endpoint=2
 /// CHECK: Callback Fini:

openmp/libomptarget/test/ompt/veccopy_emi.c

#pragma omp target teams distribute parallel for
   return rc;
 }
 
 /// CHECK: Callback Init:
 /// CHECK: Callback Load:
 /// CHECK: Callback Target EMI: kind=1 endpoint=1
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=1
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=1
+/// CHECK-NOT: dest=(nil)
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=1
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=1
+/// CHECK-NOT: dest=(nil)
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=2
 /// CHECK: Callback Submit EMI: endpoint=1 req_num_teams=1
 /// CHECK: Callback Submit EMI: endpoint=2 req_num_teams=1
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=4
 /// CHECK: Callback Target EMI: kind=1 endpoint=2
 /// CHECK: Callback Target EMI: kind=1 endpoint=1
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=1
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=1
+/// CHECK-NOT: dest=(nil)
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=1
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=1
+/// CHECK-NOT: dest=(nil)
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=2
 /// CHECK: Callback Submit EMI: endpoint=1 req_num_teams=0
 /// CHECK: Callback Submit EMI: endpoint=2 req_num_teams=0
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=4
 /// CHECK: Callback Target EMI: kind=1 endpoint=2
 /// CHECK: Callback Fini:

openmp/libomptarget/test/ompt/veccopy_emi_map.c

 }
 
 /// CHECK: 0: Could not register callback 'ompt_callback_target_map_emi'
 /// CHECK: Callback Init:
 /// CHECK: Callback Load:
 /// CHECK: Callback Target EMI: kind=1 endpoint=1
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=1
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=1
+/// CHECK-NOT: dest=(nil)
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=1
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=1
+/// CHECK-NOT: dest=(nil)
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=2
 /// CHECK: Callback Submit EMI: endpoint=1 req_num_teams=1
 /// CHECK: Callback Submit EMI: endpoint=2 req_num_teams=1
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=4
 /// CHECK: Callback Target EMI: kind=1 endpoint=2
 /// CHECK: Callback Target EMI: kind=1 endpoint=1
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=1
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=1
+/// CHECK-NOT: dest=(nil)
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=1
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=1
+/// CHECK-NOT: dest=(nil)
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=2
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=2
 /// CHECK: Callback Submit EMI: endpoint=1 req_num_teams=0
 /// CHECK: Callback Submit EMI: endpoint=2 req_num_teams=0
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=3
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=1 optype=4
 /// CHECK: Callback DataOp EMI: endpoint=2 optype=4
 /// CHECK: Callback Target EMI: kind=1 endpoint=2
 /// CHECK: Callback Fini: