[OpenMP] [OMPT] [6/8] Added callback support for target data operations, target submit, and target regions.

[mhalk]

Summary

This patch adds support for invoking target callbacks but does not yet
invoke them. A new structure OmptInterface has been added that tracks
thread local states including correlation ids. This structure defines
methods that will be called from the device independent target library
with information related to a target entry point for which a callback
is invoked. (The callback signature is defined by the OMPT spec and can be
found in omp-tools.h.) These methods in turn use the callback functions maintained
by OmptDeviceCallbacksTy to invoke the tool supplied callbacks.

Depends on D124652

Patch from John Mellor-Crummey <[e-mail]>
(With contributions from Dhruva Chakrabarti <[e-mail]>)

[dhruvachak]

Rebased

[mhalk]

Rebased + spelling fixes

[mhalk]

Refactored patch for use with nextgen plugins.

[mhalk]

As discussed with Dhruva.

[jplehr]

Some initial comments

openmp/libomptarget/include/ompt_device_callbacks.h

101

The methods here should probably read in an active voice, e.g., invokeOmptTargetDataOpEmiCallback and alike.

219

No return needed.

openmp/libomptarget/src/OmptCallback.cpp

53 ↗

(On Diff #506968)

These functions should probably be updated to follow the LLVM naming convention, e.g., createId. Idk if this should be done right now, or in an NFC after the OMPT support has landed.

[mhalk]

openmp/libomptarget/src/OmptCallback.cpp

53 ↗

(On Diff #506968)

@NFC I was asking myself the same question.
But of course I could incorporate such NFCs in this set / stack of patches.

[mhalk]

openmp/libomptarget/include/ompt_device_callbacks.h

219

Done.

[mhalk]

Update w.r.t. the landing of D124652; implement current state of feedback.

[jplehr]

Added some nits.
Any more comments @jdoerfert @dreachem @dhruvachak before accepting this?

openmp/libomptarget/src/OmptCallback.cpp

63 ↗

(On Diff #539060)

This naming seems inconsistent, doesn't it?
Why is it called TargetDataValue in the OmptInterface but the function to access it getRegionId.

openmp/libomptarget/src/OmptInterface.h

43 ↗

(On Diff #539060)

This should probably be a class according to the coding standards (sorry if I'm only pointing to this now).
https://llvm.org/docs/CodingStandards.html#use-of-class-and-struct-keywords

112 ↗

(On Diff #539060)

Parameter names here and below spelling starting w/ uppercase?

[mhalk]

Thanks for pointing out these nits, I've prepped changes for the latter two (class & uppercase spelling) and will wait for further feedback.

openmp/libomptarget/src/OmptCallback.cpp

63 ↗

(On Diff #539060)

Yes, that is absolutely right.
Thanks for bringing this up.

Background:
According to the spec the underlying data type ompt_data_t is supposed to represent: "[...] data that is reserved for tool use and that is related to a thread or to a parallel or task region."

My observation so far is, that this is only used for (task / target) "regions".
From the use-case POV, maybe get[Target]RegionDataValue (get[Target]RegionDataPtr)?

But, "correctly", you'd call it getToolDataValue (getToolDataPtr) I guess -- while rather generic, that'd be my pick.

If anyone has a preference / other suggestion, I'm happy to oblige.

[jplehr]

openmp/libomptarget/src/OmptCallback.cpp

63 ↗

(On Diff #539060)

Ok, I re-read the spec. In the light of the spec the name regionId makes more sense.
Maybe a dumb question, but do we need this function, or is it better to return the ompt_data_t itself, so it is left to the caller what they want to do with it?
I mean, do we need / want / use this?

[mhalk]

openmp/libomptarget/src/OmptCallback.cpp

63 ↗

(On Diff #539060)

Now, at a closer look, after all the refactoring and moving we can simply remove getRegionId.
AFAIK and as discussed offline, getRegionId would have no other utility because of its visbility.

Its uses will be replaced with direct access to the corresponding TargetData.value within "non-EMI" callback functions.
The EMI callbacks already use the ompt_data_t TargetData field directly.

[mhalk]

Rebased + removed getRegionId as it could be replaced.

openmp/libomptarget/include/ompt_device_callbacks.h

 //===--------- ompt_device_callbacks.h - OMPT callbacks -- C++ ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 
 #define DEBUG_PREFIX "OMPT"
 
 #define FOREACH_OMPT_TARGET_CALLBACK(macro)                                    \
   FOREACH_OMPT_DEVICE_EVENT(macro)                                             \
   FOREACH_OMPT_NOEMI_EVENT(macro)                                              \
   FOREACH_OMPT_EMI_EVENT(macro)
 
+/// Function type def used for maintaining unique target region, target
+/// operations ids
+typedef uint64_t (*IdInterfaceTy)();
+
 /// Internal representation for OMPT device
 class OmptDeviceTy {
 public:
   OmptDeviceTy() { Enabled.store(false); }
   bool initialize() {
     bool Old = false;
     return Enabled.compare_exchange_strong(Old, true);
   }

 /// Internal representation for OMPT device callback functions.
 class OmptDeviceCallbacksTy {
 public:
   /// Invoked when a device is initialized
   void OmptCallbackDeviceInitialize(int DeviceNum, const char *Type) {
     if (ompt_callback_device_initialize_fn) {
       OmptDeviceTy *Device = lookupDevice(DeviceNum);
       if (Device && Device->initialize()) {
+        // Invoke the tool supplied device init callback
         ompt_callback_device_initialize_fn(
             DeviceNum, Type, reinterpret_cast<ompt_device_t *>(Device),
             doLookup, Documentation);
       }
     }
   }
 
   /// Invoked when a device is finalized
   void OmptCallbackDeviceFinalize(int DeviceNum) {
     if (ompt_callback_device_finalize_fn) {
       OmptDeviceTy *Device = lookupDevice(DeviceNum);
       if (Device && Device->finalize()) {
+        // Invoke the tool supplied device finalize callback
         ompt_callback_device_finalize_fn(DeviceNum);
       }
     }
   }
 
   /// Invoked when a device image is loaded
   void OmptCallbackDeviceLoad(int DeviceNum, const char *Filename,
                               int64_t OffsetInFile, void *VmaInFile,
                               size_t Bytes, void *HostAddr, void *DeviceAddr,
                               uint64_t ModuleId) {
     if (ompt_callback_device_load_fn) {
+      // Invoke the tool supplied device load callback
       ompt_callback_device_load_fn(DeviceNum, Filename, OffsetInFile, VmaInFile,
                                    Bytes, HostAddr, DeviceAddr, ModuleId);
     }
   }
 
   /// Invoked when a device image is unloaded
   void OmptCallbackDeviceUnload(int DeviceNum, uint64_t ModuleId) {
     if (ompt_callback_device_unload_fn) {
+      // Invoke the tool supplied device unload callback
       ompt_callback_device_unload_fn(DeviceNum, ModuleId);
     }
   }
 
+  /// Invoked when a data transfer is initiated regardless of whether the
+  /// external monitoring interface (EMI) callback is registered
+  void OmptCallbackTargetDataOpEmi(
+      ompt_scope_endpoint_t EndPoint, ompt_data_t *TargetTaskData,
+      ompt_data_t *TargetData, ompt_target_data_op_t OpType, void *SrcAddr,
+      int SrcDeviceNum, void *DestAddr, int DestDeviceNum, size_t Bytes,
+      const void *CodePtrRA, IdInterfaceTy IdInterface, ompt_id_t *HostOpId) {
+    // If the tool registered the EMI callback, invoke that. Note that the EMI
+    // function is invoked if both emi and non-emi are registered.
+    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(
+          EndPoint, TargetTaskData, TargetData, HostOpId, OpType, SrcAddr,
+          SrcDeviceNum, DestAddr, DestDeviceNum, Bytes, CodePtrRA);
+    } else if (EndPoint == ompt_scope_begin) {
+      // For non-EMI callback, proceed for begin scope only
+      OmptCallbackTargetDataOp(TargetData->value, OpType, SrcAddr, SrcDeviceNum,
+                               DestAddr, DestDeviceNum, Bytes, CodePtrRA,
+                               IdInterface, HostOpId);
+    }
+  }
+
+  /// Invoked when a data transfer is initiated and the non-EMI callback is
+  /// registered
+  void OmptCallbackTargetDataOp(
+      ompt_id_t TargetId, ompt_target_data_op_t OpType, void *SrcAddr,
+      int SrcDeviceNum, void *DestAddr, int DestDeviceNum, size_t Bytes,
+      const void *CodePtrRA, IdInterfaceTy IdInterface, ompt_id_t *HostOpId) {
+    if (ompt_callback_target_data_op_fn) {
+      // HostOpId is set by the runtime
+      *HostOpId = IdInterface();
+      // Invoke the tool supplied data op callback
+      ompt_callback_target_data_op_fn(TargetId, *HostOpId, OpType, SrcAddr,
+                                      SrcDeviceNum, DestAddr, DestDeviceNum,
+                                      Bytes, CodePtrRA);
+    }
+  }
+
+  /// Invoked when a target region is executed regardless of whether the
+  /// external monitoring interface (EMI) callback is registered
+  void OmptCallbackTargetEmi(ompt_target_t Kind, ompt_scope_endpoint_t EndPoint,
+                             int DeviceNum, ompt_data_t *TaskData,
+                             ompt_data_t *TargetTaskData,
+                             ompt_data_t *TargetData, const void *CodePtrRA,
+                             IdInterfaceTy IdInterface) {
+    if (ompt_callback_target_emi_fn) {
+      // Invoke the tool supplied target EMI callback
+      ompt_callback_target_emi_fn(Kind, EndPoint, DeviceNum, TaskData,
+                                  TargetTaskData, TargetData, CodePtrRA);
+    } else {
+      OmptCallbackTarget(Kind, EndPoint, DeviceNum, TaskData, CodePtrRA,
+                         TargetData, IdInterface);
+    }
+  }
+
+  /// Invoked when a target region is executed and the EMI callback is not
+  /// registered
+  void OmptCallbackTarget(ompt_target_t Kind, ompt_scope_endpoint_t EndPoint,
+                          int DeviceNum, ompt_data_t *TaskData,
+                          const void *CodePtrRA, ompt_data_t *TargetData,
+                          IdInterfaceTy IdInterface) {
+    // If we reach this point, ompt_callback_target_emi was not
+    // invoked so a tool didn't provide a target id. Thus, we must
+    // unconditionally get an id here, even if there is no
+    // OMPT callback target registered. We need to have an id for use by other
+    // callbacks.
+    // Note: On a scope_begin callback, IdInterface() will generate an id.
+    // On a scope_end callback, IdInterface() will return the existing
+    // id. It is safe to do the assignment again to TargetData->value.
+    TargetData->value = IdInterface();
+    if (ompt_callback_target_fn) {
+      // Invoke the tool supplied target callback
+      ompt_callback_target_fn(Kind, EndPoint, DeviceNum, TaskData,
+                              TargetData->value, CodePtrRA);
+    }
+  }
+
+  /// Invoked when a target map clause is executed regardless of whether the
+  /// external monitoring interface (EMI) callback is registered
+  void OmptCallbackTargetMapEmi(ompt_data_t *TargetData, unsigned int NItems,
+                                void **HostAddr, void **DeviceAddr,
+                                size_t *Bytes, unsigned int *MappingFlags,
+                                const void *CodePtrRA) {
+    if (ompt_callback_target_map_emi_fn) {
+      // Invoke the tool supplied map EMI callback
+      ompt_callback_target_map_emi_fn(TargetData, NItems, HostAddr, DeviceAddr,
+                                      Bytes, MappingFlags, CodePtrRA);
+    } else {
+      OmptCallbackTargetMap(TargetData->value, NItems, HostAddr, DeviceAddr,
+                            Bytes, MappingFlags, CodePtrRA);
+    }
+  }
+
+  /// Invoked when a target map clause is executed and the EMI callback is not
+  /// registered
+  void OmptCallbackTargetMap(ompt_id_t TargetId, unsigned int NItems,
+                             void **HostAddr, void **DeviceAddr, size_t *Bytes,
+                             unsigned int *MappingFlags,
+                             const void *CodePtrRA) {
+    if (ompt_callback_target_map_fn) {
+      // Invoke the tool supplied map callback
+      ompt_callback_target_map_fn(TargetId, NItems, HostAddr, DeviceAddr, Bytes,
+                                  MappingFlags, CodePtrRA);
+    }
+  }
+
+  /// Invoked when a target submit is executed regardless of whether the
+  /// external monitoring interface (EMI) callback is registered
+  void OmptCallbackTargetSubmitEmi(ompt_scope_endpoint_t EndPoint,
+                                   ompt_data_t *TargetData,
+                                   unsigned int RequestedNumTeams,
+                                   IdInterfaceTy IdInterface,
+                                   ompt_id_t *HostOpId) {
+    if (ompt_callback_target_submit_emi_fn) {
+      // HostOpId is set by the tool. Invoke the tool supplied target submit EMI
+      // callback
+      ompt_callback_target_submit_emi_fn(EndPoint, TargetData, HostOpId,
+                                         RequestedNumTeams);
+    } else if (EndPoint == ompt_scope_begin) {
+      return OmptCallbackTargetSubmit(TargetData->value, RequestedNumTeams,
+                                      IdInterface, HostOpId);
+    }
+  }
+
+  /// Invoked when a target submit is executed and the EMI callback is not
+  /// registered
+  void OmptCallbackTargetSubmit(ompt_id_t TargetId,
+                                unsigned int RequestedNumTeams,
+                                IdInterfaceTy IdInterface,
+                                ompt_id_t *HostOpId) {
+    if (ompt_callback_target_submit_fn) {
+      // HostOpId is set by the runtime.
+      *HostOpId = IdInterface();
+      // Invoke the tool supplied target submit callback
+      ompt_callback_target_submit_fn(TargetId, *HostOpId, RequestedNumTeams);
+    }
+  }
+
   /// Initialize the enabled flag and all the callbacks
   void init() {
     Enabled = false;
 #define initName(Name, Type, Code) Name##_fn = 0;
     FOREACH_OMPT_TARGET_CALLBACK(initName)
 #undef initName
   }
 
   bool isEnabled() { return Enabled; }
   /// Devices must be prepared before use
   void prepareDevices(int NumDevices) { resize(NumDevices); }
   /// Devices must be released after use
   void releaseDevices() { release(); }
 
   /// Used to register callbacks. \p Lookup is used to query a given callback
   /// by name and the result is assigned to the corresponding callback function.
   void registerCallbacks(ompt_function_lookup_t Lookup) {
     Enabled = true;
 #define OmptBindCallback(Name, Type, Code)                                     \
   Name##_fn = (Name##_t)Lookup(#Name);                                         \
-  DP("OMPT: class bound %s=%p\n", #Name, ((void *)(uint64_t)Name##_fn));
+  DP("OMPT: class bound %s=%p\n", #Name, (void *)(Name##_fn));
 
     FOREACH_OMPT_TARGET_CALLBACK(OmptBindCallback);
 #undef OmptBindCallback
   }
 
   /// Used to find a callback given its name
   ompt_interface_fn_t lookupCallback(const char *InterfaceFunctionName) {
 #define OmptLookup(Name, Type, Code)                                           \

openmp/libomptarget/src/ompt_callback.h

+//===------- ompt_callback.h - Target independent OpenMP target RTL ------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Declarations for OpenMP Tool callback dispatchers
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _OMPTARGET_OMPT_CALLBACK_H
+#define _OMPTARGET_OMPT_CALLBACK_H
+
+// If target OMPT support is compiled in
+#ifdef OMPT_SUPPORT
+#define OMPT_IF_BUILT(stmt) stmt
+#else
+#define OMPT_IF_BUILT(stmt)
+#endif
+
+#include "omp-tools.h"
+
+/// Used to maintain execution state for this thread
+struct OmptInterface {
+public:
+  /// Top-level function for invoking callback before device data allocation
+  void beginTargetDataAlloc(int64_t DeviceId, 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);
+
+  /// 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);
+
+  /// Top-level function for invoking callback before device data deallocation
+  void beginTargetDataDelete(int64_t DeviceId, void *TgtPtrBegin, void *Code);
+
+  /// Top-level function for invoking callback after device data deallocation
+  void endTargetDataDelete(int64_t DeviceId, void *TgtPtrBegin, void *Code);
+
+  /// Top-level function for invoking callback before data retrieve
+  void beginTargetDataRetrieve(int64_t DeviceId, void *HstPtrBegin,
+                               void *TgtPtrBegin, size_t Size, void *Code);
+
+  /// Top-level function for invoking callback after data retrieve
+  void endTargetDataRetrieve(int64_t DeviceId, void *HstPtrBegin,
+                             void *TgtPtrBegin, size_t Size, void *Code);
+
+  /// Top-level function for invoking callback before kernel dispatch
+  void beginTargetSubmit(unsigned int NumTeams = 1);
+
+  /// Top-level function for invoking callback after kernel dispatch
+  void endTargetSubmit(unsigned int NumTeams = 1);
+
+  // Target region callbacks
+
+  /// Top-level function for invoking callback before target enter data
+  /// construct
+  void beginTargetDataEnter(int64_t DeviceId, void *Code);
+
+  /// Top-level function for invoking callback after target enter data construct
+  void endTargetDataEnter(int64_t DeviceId, void *Code);
+
+  /// Top-level function for invoking callback before target exit data construct
+  void beginTargetDataExit(int64_t DeviceId, void *Code);
+
+  /// Top-level function for invoking callback after target exit data construct
+  void endTargetDataExit(int64_t DeviceId, void *Code);
+
+  /// Top-level function for invoking callback before target update construct
+  void beginTargetUpdate(int64_t DeviceId, void *Code);
+
+  /// Top-level function for invoking callback after target update construct
+  void endTargetUpdate(int64_t DeviceId, void *Code);
+
+  /// Top-level function for invoking callback before target construct
+  void beginTarget(int64_t DeviceId, void *Code);
+
+  /// Top-level function for invoking callback after target construct
+  void endTarget(int64_t DeviceId, void *Code);
+
+  /// Setters for target region and target operations correlation ids
+  void setHostOpId(ompt_id_t id) { HostOpId = id; }
+  void setTargetDataValue(uint64_t val) { TargetData.value = val; }
+  void setTargetDataPtr(void *ptr) { TargetData.ptr = ptr; }
+
+  /// Getters for target region and target operations correlation ids
+  uint64_t getTargetDataValue() { return TargetData.value; }
+  void *getTargetDataPtr() { return TargetData.ptr; }
+  ompt_id_t getHostOpId() { return HostOpId; }
+
+private:
+  /// Target operations id
+  ompt_id_t HostOpId = 0;
+
+  /// Target region data
+  ompt_data_t TargetData = ompt_data_none;
+
+  /// Task data representing the encountering task
+  ompt_data_t *TaskData = nullptr;
+
+  /// Target task data representing the target task region
+  ompt_data_t *TargetTaskData = nullptr;
+
+  /// Correlation id that is incremented with target operations
+  uint64_t TargetRegionOpId = 1;
+
+  /// Used for marking begin of a data operation
+  void beginTargetDataOperation();
+
+  /// Used for marking end of a data operation
+  void endTargetDataOperation();
+
+  /// Used for marking begin of a target region
+  void beginTargetRegion();
+
+  /// Used for marking end of a target region
+  void endTargetRegion();
+};
+
+extern thread_local OmptInterface ompt_interface;
+extern bool ompt_enabled;
+
+#endif // _OMPTARGET_OMPT_CALLBACK_H

openmp/libomptarget/src/ompt_callback.cpp

 //===-- ompt_callback.cpp - Target independent OpenMP target RTL -- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Implementation of OMPT callback interfaces for target independent layer
 //
 //===----------------------------------------------------------------------===//
 
 #ifdef OMPT_SUPPORT
 
 #include <assert.h>
 #include <atomic>
 #include <cstdlib>
 #include <cstring>
 
-#include "omp-tools.h"
-
 #include "Debug.h"
+#include "omp-tools.h"
 #include "ompt_connector.h"
+#include "ompt_callback.h"
 #include "ompt_device_callbacks.h"
 #include "private.h"
 
 #define fnptr_to_ptr(x) ((void *)(uint64_t)x)
 
 /// Used to indicate whether OMPT was enabled for this library
 bool ompt_enabled = false;
+
 /// Object maintaining all the callbacks for this library
 OmptDeviceCallbacksTy OmptDeviceCallbacks;
 
+/// Thread local state for target region and associated metadata
+thread_local OmptInterface ompt_interface;
+
+/// 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)();
+
+/// 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 = nullptr;
+static ompt_get_target_task_data_t ompt_get_target_task_data_fn = nullptr;
+
+/// Unique correlation id
+static std::atomic<uint64_t> unique_id_ticket(1);
+
+/// Used to create a new correlation id
+static uint64_t id_create() { return unique_id_ticket.fetch_add(1); }
+
+/// Create a new correlation id and update the operations id
+static uint64_t opid_create() {
+  uint64_t new_id = id_create();
+  ompt_interface.setHostOpId(new_id);
+  return new_id;
+}
+
+/// Get the current operations id
+static uint64_t opid_get() { return ompt_interface.getHostOpId(); }
+
+/// Create a new correlation id and update the target region id
+static uint64_t regionid_create() {
+  uint64_t new_id = id_create();
+  ompt_interface.setTargetDataValue(new_id);
+  return new_id;
+}
+
+/// Get the current target region id
+static uint64_t regionid_get() { return ompt_interface.getTargetDataValue(); }
+
 /// Used to maintain the finalization function that is received
 /// from the plugin during connect
 class LibomptargetRtlFinalizer {
 public:
   LibomptargetRtlFinalizer() : FiniFn(nullptr) {}
   void registerRtl(ompt_finalize_t fn) {
     assert(FiniFn == nullptr && "Finalizer already registered");
     FiniFn = fn;
   }
   void finalize() {
     if (FiniFn)
       FiniFn(nullptr /* tool_data */);
     FiniFn = nullptr;
   }
 
 private:
   ompt_finalize_t FiniFn;
 };
 
+void OmptInterface::beginTargetDataAlloc(int64_t DeviceId, void *HstPtrBegin,
+                                         size_t Size, void *Code) {
+  beginTargetDataOperation();
+  OmptDeviceCallbacks.OmptCallbackTargetDataOpEmi(
+      ompt_scope_begin, TargetTaskData, &TargetData, ompt_target_data_alloc,
+      HstPtrBegin, DeviceId, NULL, 0, Size, Code, opid_create,
+      &TargetRegionOpId);
+}
+
+void OmptInterface::endTargetDataAlloc(int64_t DeviceId, void *HstPtrBegin,
+                                       size_t Size, void *Code) {
+  OmptDeviceCallbacks.OmptCallbackTargetDataOpEmi(
+      ompt_scope_end, TargetTaskData, &TargetData, ompt_target_data_alloc,
+      HstPtrBegin, DeviceId, NULL, 0, Size, Code, opid_get, &TargetRegionOpId);
+  endTargetDataOperation();
+}
+
+void OmptInterface::beginTargetDataSubmit(int64_t DeviceId, void *TgtPtrBegin,
+                                          void *HstPtrBegin, size_t Size,
+                                          void *Code) {
+  beginTargetDataOperation();
+  OmptDeviceCallbacks.OmptCallbackTargetDataOpEmi(
+      ompt_scope_begin, TargetTaskData, &TargetData,
+      ompt_target_data_transfer_to_device, HstPtrBegin, 0, TgtPtrBegin,
+      DeviceId, Size, Code, opid_create, &TargetRegionOpId);
+}
+
+void OmptInterface::endTargetDataSubmit(int64_t DeviceId, void *TgtPtrBegin,
+                                        void *HstPtrBegin, size_t Size,
+                                        void *Code) {
+  OmptDeviceCallbacks.OmptCallbackTargetDataOpEmi(
+      ompt_scope_end, TargetTaskData, &TargetData,
+      ompt_target_data_transfer_to_device, HstPtrBegin, 0 /* DeviceId */,
+      TgtPtrBegin, DeviceId, Size, Code, opid_get, &TargetRegionOpId);
+  endTargetDataOperation();
+}
+
+void OmptInterface::beginTargetDataDelete(int64_t DeviceId, void *TgtPtrBegin,
+                                          void *Code) {
+  beginTargetDataOperation();
+  OmptDeviceCallbacks.OmptCallbackTargetDataOpEmi(
+      ompt_scope_begin, TargetTaskData, &TargetData, ompt_target_data_delete,
+      TgtPtrBegin, DeviceId, NULL /* DstAddr */, 0 /* DestDeviceNum */,
+      0 /* Bytes */, Code, opid_create, &TargetRegionOpId);
+}
+
+void OmptInterface::endTargetDataDelete(int64_t DeviceId, void *TgtPtrBegin,
+                                        void *Code) {
+  OmptDeviceCallbacks.OmptCallbackTargetDataOpEmi(
+      ompt_scope_end, TargetTaskData, &TargetData, ompt_target_data_delete,
+      TgtPtrBegin, DeviceId, NULL /* DstAddr */, 0 /* DestDeviceNum */,
+      0 /* Bytes */, Code, opid_get, &TargetRegionOpId);
+  endTargetDataOperation();
+}
+
+void OmptInterface::beginTargetDataRetrieve(int64_t DeviceId, void *HstPtrBegin,
+                                            void *TgtPtrBegin, size_t Size,
+                                            void *Code) {
+  beginTargetDataOperation();
+  OmptDeviceCallbacks.OmptCallbackTargetDataOpEmi(
+      ompt_scope_begin, TargetTaskData, &TargetData,
+      ompt_target_data_transfer_from_device, TgtPtrBegin, DeviceId, HstPtrBegin,
+      0 /* DeviceId */, Size, Code, opid_create, &TargetRegionOpId);
+}
+
+void OmptInterface::endTargetDataRetrieve(int64_t DeviceId, void *HstPtrBegin,
+                                          void *TgtPtrBegin, size_t Size,
+                                          void *Code) {
+  OmptDeviceCallbacks.OmptCallbackTargetDataOpEmi(
+      ompt_scope_end, TargetTaskData, &TargetData,
+      ompt_target_data_transfer_from_device, TgtPtrBegin, DeviceId, HstPtrBegin,
+      0 /* Device Id */, Size, Code, opid_get, &TargetRegionOpId);
+  endTargetDataOperation();
+}
+
+void OmptInterface::beginTargetSubmit(unsigned int numTeams) {
+  OmptDeviceCallbacks.OmptCallbackTargetSubmitEmi(
+      ompt_scope_begin, &TargetData, numTeams, opid_create, &TargetRegionOpId);
+}
+
+void OmptInterface::endTargetSubmit(unsigned int numTeams) {
+  OmptDeviceCallbacks.OmptCallbackTargetSubmitEmi(
+      ompt_scope_end, &TargetData, numTeams, opid_get, &TargetRegionOpId);
+}
+
+void OmptInterface::beginTargetDataEnter(int64_t DeviceId, void *Code) {
+  beginTargetRegion();
+  OmptDeviceCallbacks.OmptCallbackTargetEmi(
+      ompt_target_enter_data, ompt_scope_begin, DeviceId, TaskData,
+      TargetTaskData, &TargetData, Code, regionid_create);
+}
+
+void OmptInterface::endTargetDataEnter(int64_t DeviceId, void *Code) {
+  OmptDeviceCallbacks.OmptCallbackTargetEmi(
+      ompt_target_enter_data, ompt_scope_end, DeviceId, TaskData,
+      TargetTaskData, &TargetData, Code, regionid_get);
+  endTargetRegion();
+}
+
+void OmptInterface::beginTargetDataExit(int64_t DeviceId, void *Code) {
+  beginTargetRegion();
+  OmptDeviceCallbacks.OmptCallbackTargetEmi(
+      ompt_target_exit_data, ompt_scope_begin, DeviceId, TaskData,
+      TargetTaskData, &TargetData, Code, regionid_create);
+}
+
+void OmptInterface::endTargetDataExit(int64_t DeviceId, void *Code) {
+  OmptDeviceCallbacks.OmptCallbackTargetEmi(
+      ompt_target_exit_data, ompt_scope_end, DeviceId, TaskData, TargetTaskData,
+      &TargetData, Code, regionid_get);
+  endTargetRegion();
+}
+
+void OmptInterface::beginTargetUpdate(int64_t DeviceId, void *Code) {
+  beginTargetRegion();
+  OmptDeviceCallbacks.OmptCallbackTargetEmi(
+      ompt_target_update, ompt_scope_begin, DeviceId, TaskData, TargetTaskData,
+      &TargetData, Code, regionid_create);
+}
+
+void OmptInterface::endTargetUpdate(int64_t DeviceId, void *Code) {
+  OmptDeviceCallbacks.OmptCallbackTargetEmi(
+      ompt_target_update, ompt_scope_end, DeviceId, TaskData, TargetTaskData,
+      &TargetData, Code, regionid_get);
+  endTargetRegion();
+}
+
+void OmptInterface::beginTarget(int64_t DeviceId, void *Code) {
+  beginTargetRegion();
+  OmptDeviceCallbacks.OmptCallbackTargetEmi(
+      ompt_target, ompt_scope_begin, DeviceId, TaskData, TargetTaskData,
+      &TargetData, Code, regionid_create);
+}
+
+void OmptInterface::endTarget(int64_t DeviceId, void *Code) {
+  OmptDeviceCallbacks.OmptCallbackTargetEmi(
+      ompt_target, ompt_scope_end, DeviceId, TaskData, TargetTaskData,
+      &TargetData, Code, regionid_get);
+  endTargetRegion();
+}
+
+void OmptInterface::beginTargetDataOperation() {
+  DP("in ompt_target_region_begin (TargetRegionOpId = %lu)\n",
+     TargetData.value);
+}
+
+void OmptInterface::endTargetDataOperation() {
+  DP("in ompt_target_region_end (TargetRegionOpId = %lu)\n", TargetData.value);
+}
+
+void OmptInterface::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;
+}
+
+void OmptInterface::endTargetRegion() {
+  TaskData = 0;
+  TargetTaskData = 0;
+  TargetData = ompt_data_none;
+}
+
 /// Object that will maintain the finalizer of the plugin
 static LibomptargetRtlFinalizer LibOmptTargetRTLFinalizer;
 
 /// Lookup function to be used by libomptarget library
 ompt_interface_fn_t
 OmptDeviceCallbacksTy::doLookup(const char *InterfaceFunctionName) {
   return OmptDeviceCallbacks.lookupCallback(InterfaceFunctionName);
 }
 
 /// This is the function called by the higher layer (libomp) responsible
 /// for initializing OMPT in this library. This is passed to libomp
 /// as part of the OMPT connector object.
 /// \p lookup to be used to query callbacks registered with libomp
 /// \p initial_device_num Initial device num provided by libomp
 /// \p tool_data as provided by the tool
 static int ompt_libomptarget_initialize(ompt_function_lookup_t lookup,
                                         int initial_device_num,
                                         ompt_data_t *tool_data) {
   DP("enter ompt_libomptarget_initialize!\n");
   ompt_enabled = true;
+
+#define ompt_bind_name(fn)                                                     \
+  fn##_fn = (fn##_t)lookup(#fn);                                               \
+  DP("%s=%p\n", #fn, (void *)(fn##_fn));
+
+  // The lookup parameter is provided by libomp which already has the task
+  // and target_task functions registered at this point. Register those same
+  // functions in libomptarget as well.
+  ompt_bind_name(ompt_get_task_data);
+  ompt_bind_name(ompt_get_target_task_data);
+#undef ompt_bind_name
+
   // The lookup parameter is provided by libomp which already has the
   // tool callbacks registered at this point. The registration call
   // below causes the same callback functions to be registered in
   // libomptarget as well
   OmptDeviceCallbacks.registerCallbacks(lookup);
   DP("exit ompt_libomptarget_initialize!\n");
   return 0;
 }

openmp/runtime/src/ompt-general.cpp

 /*
  * ompt-general.cpp -- OMPT implementation of interface functions
  */
 
 //===----------------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
if (ompt_start_tool_result) {
     kmp_info_t *root_thread = ompt_get_thread();
 
     ompt_set_thread_state(root_thread, ompt_state_overhead);
 
     if (ompt_enabled.ompt_callback_thread_begin) {
       ompt_callbacks.ompt_callback(ompt_callback_thread_begin)(
           ompt_thread_initial, __ompt_get_thread_data_internal());
     }
-    ompt_data_t *task_data;
-    ompt_data_t *parallel_data;
+    ompt_data_t *task_data = nullptr;
+    ompt_data_t *parallel_data = nullptr;
     __ompt_get_task_info_internal(0, NULL, &task_data, NULL, &parallel_data,
                                   NULL);
     if (ompt_enabled.ompt_callback_implicit_task) {
       ompt_callbacks.ompt_callback(ompt_callback_implicit_task)(
           ompt_scope_begin, parallel_data, task_data, 1, 1, ompt_task_initial);
     }
 
     ompt_set_thread_state(root_thread, ompt_state_work_serial);
#define ompt_interface_fn(fn)                                                  \
 
   FOREACH_OMPT_INQUIRY_FN(ompt_interface_fn)
 
 #undef ompt_interface_fn
 
   return NULL;
 }
 
+static ompt_data_t *ompt_get_task_data() { return __ompt_get_task_data(); }
+
+static ompt_data_t *ompt_get_target_task_data() {
+  return __ompt_get_target_task_data();
+}
+
 /// Lookup function to query libomp callbacks registered by the tool
 static ompt_interface_fn_t ompt_libomp_target_fn_lookup(const char *s) {
+#define provide_fn(fn)                                                         \
+  if (strcmp(s, #fn) == 0)                                                     \
+    return (ompt_interface_fn_t)fn;
+
+  provide_fn(ompt_get_task_data);
+  provide_fn(ompt_get_target_task_data);
+
 #define ompt_interface_fn(fn, type, code)                                      \
   if (strcmp(s, #fn) == 0)                                                     \
     return (ompt_interface_fn_t)ompt_callbacks.ompt_callback(fn);
 
   FOREACH_OMPT_DEVICE_EVENT(ompt_interface_fn)
   FOREACH_OMPT_EMI_EVENT(ompt_interface_fn)
   FOREACH_OMPT_NOEMI_EVENT(ompt_interface_fn)
 

openmp/runtime/src/ompt-internal.h

 /*
  * ompt-internal.h - header of OMPT internal data structures
  */
 
 //===----------------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
typedef struct ompt_lw_taskteam_s {
   int heap;
   struct ompt_lw_taskteam_s *parent;
 } ompt_lw_taskteam_t;
 
 typedef struct {
   ompt_data_t thread_data;
   ompt_data_t task_data; /* stored here from implicit barrier-begin until
                             implicit-task-end */
+  ompt_data_t target_task_data; /* required by target support */
   void *return_address; /* stored here on entry of runtime */
   ompt_state_t state;
   ompt_wait_id_t wait_id;
   int ompt_task_yielded;
   int parallel_flags; // information for the last parallel region invoked
   void *idle_frame;
 } ompt_thread_info_t;
 

openmp/runtime/src/ompt-specific.h

 /*
  * ompt-specific.h - header of OMPT internal functions implementation
  */
 
 //===----------------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.

 
 void __ompt_lw_taskteam_link(ompt_lw_taskteam_t *lwt, kmp_info_t *thr,
                              int on_heap, bool always = false);
 
 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_target_task_data();
+
 ompt_task_info_t *__ompt_get_task_info_object(int depth);
 
 int __ompt_get_parallel_info_internal(int ancestor_level,
                                       ompt_data_t **parallel_data,
                                       int *team_size);
 
 int __ompt_get_task_info_internal(int ancestor_level, int *type,
                                   ompt_data_t **task_data,
                                   ompt_frame_t **task_frame,
                                   ompt_data_t **parallel_data, int *thread_num);
 
 ompt_data_t *__ompt_get_thread_data_internal();
 
 /*
  * Unused currently
 static uint64_t __ompt_get_get_unique_id_internal();
 */
 
 ompt_sync_region_t __ompt_get_barrier_kind(enum barrier_type, kmp_info_t *);
 
 /*****************************************************************************
  * macros
  ****************************************************************************/
 
-#define OMPT_CUR_TASK_INFO(thr) (&(thr->th.th_current_task->ompt_task_info))
+#define OMPT_CUR_TASK_INFO(thr) (&((thr)->th.th_current_task->ompt_task_info))
 #define OMPT_CUR_TASK_DATA(thr)                                                \
-  (&(thr->th.th_current_task->ompt_task_info.task_data))
-#define OMPT_CUR_TEAM_INFO(thr) (&(thr->th.th_team->t.ompt_team_info))
+  (&((thr)->th.th_current_task->ompt_task_info.task_data))
+#define OMPT_CUR_TEAM_INFO(thr) (&((thr)->th.th_team->t.ompt_team_info))
 #define OMPT_CUR_TEAM_DATA(thr)                                                \
-  (&(thr->th.th_team->t.ompt_team_info.parallel_data))
+  (&((thr)->th.th_team->t.ompt_team_info.parallel_data))
 
 #define OMPT_HAVE_WEAK_ATTRIBUTE KMP_HAVE_WEAK_ATTRIBUTE
 #define OMPT_HAVE_PSAPI KMP_HAVE_PSAPI
 #define OMPT_STR_MATCH(haystack, needle) __kmp_str_match(haystack, 0, needle)
 
 inline void *__ompt_load_return_address(int gtid) {
   kmp_info_t *thr = __kmp_threads[gtid];
   void *return_address = thr->th.ompt_thread_info.return_address;

openmp/runtime/src/ompt-specific.cpp

 /*
  * ompt-specific.cpp -- OMPT internal functions
  */
 
 //===----------------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
#endif
   }
   //    return lwtask;
 }
 
 //----------------------------------------------------------
 // task support
 //----------------------------------------------------------
 
+ompt_data_t *__ompt_get_task_data() {
+  kmp_info_t *thr = ompt_get_thread();
+  ompt_data_t *task_data = thr ? OMPT_CUR_TASK_DATA(thr) : NULL;
+  return task_data;
+}
+
+ompt_data_t *__ompt_get_target_task_data() {
+  return &__kmp_threads[__kmp_get_gtid()]->th.ompt_thread_info.target_task_data;
+}
+
 int __ompt_get_task_info_internal(int ancestor_level, int *type,
                                   ompt_data_t **task_data,
                                   ompt_frame_t **task_frame,
                                   ompt_data_t **parallel_data,
                                   int *thread_num) {
   if (__kmp_get_gtid() < 0)
     return 0;