From caaf8099efa87a7ebca8920971b7d7f719808591 Mon Sep 17 00:00:00 2001 From: Jan Patrick Lehr Date: Wed, 17 Jul 2024 10:15:19 +0200 Subject: [PATCH] [Offload][OMPT] Add callbacks for (dis)associate_ptr (#99046) This adds the OMPT callbacks for the API functions disassociate_ptr and associate_ptr. --- offload/include/OpenMP/OMPT/Interface.h | 29 +++++++++++++ offload/src/OpenMP/API.cpp | 11 +++++ offload/src/OpenMP/OMPT/Callback.cpp | 57 +++++++++++++++++++++++++ offload/test/ompt/omp_api.c | 39 +++++++++++++++++ 4 files changed, 136 insertions(+) create mode 100644 offload/test/ompt/omp_api.c 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!"); } diff --git a/offload/src/OpenMP/API.cpp b/offload/src/OpenMP/API.cpp index 374c54163d6a..e59bdba8abf0 100644 --- a/offload/src/OpenMP/API.cpp +++ b/offload/src/OpenMP/API.cpp @@ -597,6 +597,12 @@ EXTERN int omp_target_associate_ptr(const void *HostPtr, const void *DevicePtr, FATAL_MESSAGE(DeviceNum, "%s", toString(DeviceOrErr.takeError()).c_str()); void *DeviceAddr = (void *)((uint64_t)DevicePtr + (uint64_t)DeviceOffset); + + OMPT_IF_BUILT(InterfaceRAII( + RegionInterface.getCallbacks(), DeviceNum, + const_cast(HostPtr), const_cast(DevicePtr), Size, + __builtin_return_address(0))); + int Rc = DeviceOrErr->getMappingInfo().associatePtr( const_cast(HostPtr), const_cast(DeviceAddr), Size); DP("omp_target_associate_ptr returns %d\n", Rc); @@ -625,6 +631,11 @@ EXTERN int omp_target_disassociate_ptr(const void *HostPtr, int DeviceNum) { if (!DeviceOrErr) FATAL_MESSAGE(DeviceNum, "%s", toString(DeviceOrErr.takeError()).c_str()); + OMPT_IF_BUILT(InterfaceRAII( + RegionInterface.getCallbacks(), DeviceNum, + const_cast(HostPtr), + /*DevicePtr=*/nullptr, /*Size=*/0, __builtin_return_address(0))); + int Rc = DeviceOrErr->getMappingInfo().disassociatePtr( const_cast(HostPtr)); DP("omp_target_disassociate_ptr returns %d\n", Rc); diff --git a/offload/src/OpenMP/OMPT/Callback.cpp b/offload/src/OpenMP/OMPT/Callback.cpp index f285843e39f3..f2964281eeb9 100644 --- a/offload/src/OpenMP/OMPT/Callback.cpp +++ b/offload/src/OpenMP/OMPT/Callback.cpp @@ -332,6 +332,63 @@ void Interface::endTargetUpdate(int64_t DeviceId, void *Code) { endTargetRegion(); } +void Interface::beginTargetAssociatePointer(int64_t DeviceId, void *HstPtrBegin, + void *TgtPtrBegin, size_t Size, + void *Code) { + beginTargetDataOperation(); + if (ompt_callback_target_data_op_emi_fn) { + ompt_callback_target_data_op_emi_fn( + ompt_scope_begin, TargetTaskData, &TargetData, &HostOpId, + ompt_target_data_associate, HstPtrBegin, omp_get_initial_device(), + TgtPtrBegin, DeviceId, Size, Code); + } else if (ompt_callback_target_data_op_fn) { + HostOpId = createOpId(); + ompt_callback_target_data_op_fn( + TargetData.value, HostOpId, ompt_target_data_associate, HstPtrBegin, + omp_get_initial_device(), TgtPtrBegin, DeviceId, Size, Code); + } +} + +void Interface::endTargetAssociatePointer(int64_t DeviceId, void *HstPtrBegin, + void *TgtPtrBegin, size_t Size, + void *Code) { + if (ompt_callback_target_data_op_emi_fn) { + ompt_callback_target_data_op_emi_fn( + ompt_scope_end, TargetTaskData, &TargetData, &HostOpId, + ompt_target_data_associate, HstPtrBegin, omp_get_initial_device(), + TgtPtrBegin, DeviceId, Size, Code); + } +} + +void Interface::beginTargetDisassociatePointer(int64_t DeviceId, + void *HstPtrBegin, + void *TgtPtrBegin, size_t Size, + void *Code) { + beginTargetDataOperation(); + if (ompt_callback_target_data_op_emi_fn) { + ompt_callback_target_data_op_emi_fn( + ompt_scope_begin, TargetTaskData, &TargetData, &HostOpId, + ompt_target_data_disassociate, HstPtrBegin, omp_get_initial_device(), + TgtPtrBegin, DeviceId, Size, Code); + } else if (ompt_callback_target_data_op_fn) { + HostOpId = createOpId(); + ompt_callback_target_data_op_fn( + TargetData.value, HostOpId, ompt_target_data_disassociate, HstPtrBegin, + omp_get_initial_device(), TgtPtrBegin, DeviceId, Size, Code); + } +} +void Interface::endTargetDisassociatePointer(int64_t DeviceId, + void *HstPtrBegin, + void *TgtPtrBegin, size_t Size, + void *Code) { + if (ompt_callback_target_data_op_emi_fn) { + ompt_callback_target_data_op_emi_fn( + ompt_scope_end, TargetTaskData, &TargetData, &HostOpId, + ompt_target_data_disassociate, HstPtrBegin, omp_get_initial_device(), + TgtPtrBegin, DeviceId, Size, Code); + } +} + void Interface::beginTarget(int64_t DeviceId, void *Code) { beginTargetRegion(); if (ompt_callback_target_emi_fn) { diff --git a/offload/test/ompt/omp_api.c b/offload/test/ompt/omp_api.c new file mode 100644 index 000000000000..a16ef7a64aa7 --- /dev/null +++ b/offload/test/ompt/omp_api.c @@ -0,0 +1,39 @@ +// RUN: %libomptarget-compile-run-and-check-generic +// REQUIRES: ompt +// REQUIRES: gpu + +#include "omp.h" +#include +#include + +#include "callbacks.h" +#include "register_non_emi.h" + +#define N 1024 + +int main(int argc, char **argv) { + int *h_a; + int *d_a; + + h_a = (int *)malloc(N * sizeof(int)); + memset(h_a, 0, N); + + d_a = (int *)omp_target_alloc(N * sizeof(int), omp_get_default_device()); + + omp_target_associate_ptr(h_a, d_a, N * sizeof(int), 0, + omp_get_default_device()); + omp_target_disassociate_ptr(h_a, omp_get_default_device()); + + omp_target_free(d_a, omp_get_default_device()); + free(h_a); + + return 0; +} + +// clang-format off +/// CHECK: Callback Init: +/// 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=5 +/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=6 +/// CHECK: Callback DataOp: target_id=[[TARGET_ID:[0-9]+]] host_op_id=[[HOST_OP_ID:[0-9]+]] optype=4 +/// CHECK: Callback Fini: