[Offload][OMPT] Add callbacks for (dis)associate_ptr (#99046)

This adds the OMPT callbacks for the API functions disassociate_ptr and
associate_ptr.

1 files changed, 29 insertions, 0 deletions

diff --git a/offload/include/OpenMP/OMPT/Interface.h b/offload/include/OpenMP/OMPT/Interface.h
index 327fadfcd4ac..0dc1bad8f7ec 100644
--- a/offload/include/OpenMP/OMPT/Interface.h
+++ b/offload/include/OpenMP/OMPT/Interface.h
@@ -109,6 +109,25 @@ public:
   /// 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 associate API
+  void beginTargetAssociatePointer(int64_t DeviceId, void *HstPtrBegin,
+                                   void *TgtPtrBegin, size_t Size, void *Code);
+
+  /// Top-level function for invoking callback after target associate API
+  void endTargetAssociatePointer(int64_t DeviceId, void *HstPtrBegin,
+                                 void *TgtPtrBegin, size_t Size, void *Code);
+
+  /// Top-level function for invoking callback before target disassociate API
+  void beginTargetDisassociatePointer(int64_t DeviceId, void *HstPtrBegin,
+                                      void *TgtPtrBegin, size_t Size,
+                                      void *Code);
+
+  /// Top-level function for invoking callback after target disassociate API
+  void endTargetDisassociatePointer(int64_t DeviceId, void *HstPtrBegin,
+                                    void *TgtPtrBegin, size_t Size, void *Code);
+
+  // Target kernel callbacks
+
   /// Top-level function for invoking callback before target construct
   void beginTarget(int64_t DeviceId, void *Code);
 
@@ -137,6 +156,16 @@ public:
       return std::make_pair(std::mem_fn(&Interface::beginTargetDataRetrieve),
                             std::mem_fn(&Interface::endTargetDataRetrieve));
 
+    if constexpr (OpType == ompt_target_data_associate)
+      return std::make_pair(
+          std::mem_fn(&Interface::beginTargetAssociatePointer),
+          std::mem_fn(&Interface::endTargetAssociatePointer));
+
+    if constexpr (OpType == ompt_target_data_disassociate)
+      return std::make_pair(
+          std::mem_fn(&Interface::beginTargetDisassociatePointer),
+          std::mem_fn(&Interface::endTargetDisassociatePointer));
+
     llvm_unreachable("Unhandled target data operation type!");
   }