[Offload][OMPT] Add callbacks for (dis)associate_ptr #99046
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This adds the OMPT callbacks for the API functions disassociate_ptr and associate_ptr.
|
@llvm/pr-subscribers-offload Author: Jan Patrick Lehr (jplehr) ChangesThis adds the OMPT callbacks for the API functions disassociate_ptr and associate_ptr. Full diff: https://github.com/llvm/llvm-project/pull/99046.diff 4 Files Affected:
diff --git a/offload/include/OpenMP/OMPT/Interface.h b/offload/include/OpenMP/OMPT/Interface.h
index 327fadfcd4acd..0dc1bad8f7ece 100644
--- a/offload/include/OpenMP/OMPT/Interface.h
+++ b/offload/include/OpenMP/OMPT/Interface.h
@@ -109,6 +109,25 @@ class Interface {
/// 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 @@ class Interface {
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 374c54163d6a4..e59bdba8abf0e 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<ompt_target_data_associate>(), DeviceNum,
+ const_cast<void *>(HostPtr), const_cast<void *>(DevicePtr), Size,
+ __builtin_return_address(0)));
+
int Rc = DeviceOrErr->getMappingInfo().associatePtr(
const_cast<void *>(HostPtr), const_cast<void *>(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<ompt_target_data_disassociate>(), DeviceNum,
+ const_cast<void *>(HostPtr),
+ /*DevicePtr=*/nullptr, /*Size=*/0, __builtin_return_address(0)));
+
int Rc = DeviceOrErr->getMappingInfo().disassociatePtr(
const_cast<void *>(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 f285843e39f38..f2964281eeb95 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 0000000000000..a16ef7a64aa7d
--- /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 <stdlib.h>
+#include <string.h>
+
+#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:
|
mhalk
approved these changes
Jul 16, 2024
There was a problem hiding this comment.
LGTM.
Note: Only had the chance to test this for amdgpu.
It seemed to me as if it would try to build for nvptx when running check-offload even when no corresponding targets were available, resulting in a fail. (Unsure. Maybe just a misconfiguration on my side.)
The changes should not touch any device specific pieces, so I assume this would work not only on amdgpu. |
yuxuanchen1997
pushed a commit
that referenced
this pull request
Jul 25, 2024
Summary: This adds the OMPT callbacks for the API functions disassociate_ptr and associate_ptr. Test Plan: Reviewers: Subscribers: Tasks: Tags: Differential Revision: https://phabricator.intern.facebook.com/D60250852
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This adds the OMPT callbacks for the API functions disassociate_ptr and associate_ptr.