Skip to content

Commit f26cd46

Browse files
committed
[OMPT] Fix for ROCm/ROCm#2611
Change-Id: I11c150cd34593e2ff24ad6d25f19a0f19ccfd501
1 parent 1567b09 commit f26cd46

File tree

1 file changed

+82
-28
lines changed

1 file changed

+82
-28
lines changed

openmp/libomptarget/src/interface.cpp

Lines changed: 82 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -47,19 +47,56 @@ class OMPTInvokeWrapper {
4747
: IsNullOpt(true), CodePtr(nullptr), ReturnFramePtr(nullptr),
4848
DeviceId(-1), Kind(ompt_target), ScopeKind(ompt_scope_begin) {}
4949
OMPTInvokeWrapper(void *CodePtr, void *ReturnFramePtr, int64_t DeviceId,
50-
ompt_target_t Kind, ompt_scope_endpoint_t ScopeKind)
50+
ompt_target_t Kind)
5151
: IsNullOpt(false), CodePtr(CodePtr), ReturnFramePtr(ReturnFramePtr),
52-
DeviceId(DeviceId), Kind(Kind), ScopeKind(ScopeKind) {}
52+
DeviceId(DeviceId), Kind(Kind), ScopeKind(ompt_scope_begin) {}
5353

5454
void setDeviceId(int64_t DevId) { DeviceId = DevId; }
5555

56-
void invoke() {
56+
void invokeBegin() {
5757
if (IsNullOpt)
5858
return;
5959

6060
ompt_interface.ompt_state_set(ReturnFramePtr, CodePtr);
61-
ompt_interface.target_data_enter_begin(DeviceId, CodePtr);
62-
ompt_interface.target_trace_record_gen(DeviceId, Kind, ScopeKind, CodePtr);
61+
switch (Kind) {
62+
case ompt_target_enter_data:
63+
case ompt_target_enter_data_nowait:
64+
ompt_interface.target_data_enter_begin(DeviceId, CodePtr);
65+
break;
66+
case ompt_target_exit_data:
67+
case ompt_target_exit_data_nowait:
68+
ompt_interface.target_data_exit_begin(DeviceId, CodePtr);
69+
break;
70+
case ompt_target_update:
71+
case ompt_target_update_nowait:
72+
ompt_interface.target_update_begin(DeviceId, CodePtr);
73+
break;
74+
}
75+
ompt_interface.target_trace_record_gen(DeviceId, Kind, ompt_scope_begin,
76+
CodePtr);
77+
}
78+
79+
void invokeEnd() {
80+
if (IsNullOpt)
81+
return;
82+
83+
ompt_interface.ompt_state_set(ReturnFramePtr, CodePtr);
84+
switch (Kind) {
85+
case ompt_target_enter_data:
86+
case ompt_target_enter_data_nowait:
87+
ompt_interface.target_data_enter_end(DeviceId, CodePtr);
88+
break;
89+
case ompt_target_exit_data:
90+
case ompt_target_exit_data_nowait:
91+
ompt_interface.target_data_exit_end(DeviceId, CodePtr);
92+
break;
93+
case ompt_target_update:
94+
case ompt_target_update_nowait:
95+
ompt_interface.target_update_end(DeviceId, CodePtr);
96+
break;
97+
}
98+
ompt_interface.target_trace_record_gen(DeviceId, Kind, ompt_scope_end,
99+
CodePtr);
63100
}
64101

65102
private:
@@ -113,21 +150,6 @@ EXTERN void __tgt_unregister_lib(__tgt_bin_desc *Desc) {
113150
}
114151
}
115152

116-
template <typename TargetAsyncInfoTy>
117-
static inline void
118-
targetDataMapper(ident_t *Loc, int64_t DeviceId, int32_t ArgNum,
119-
void **ArgsBase, void **Args, int64_t *ArgSizes,
120-
int64_t *ArgTypes, map_var_info_t *ArgNames, void **ArgMappers,
121-
TargetDataFuncPtrTy TargetDataFunction,
122-
const char *RegionTypeMsg, const char *RegionName) {
123-
124-
OMPTInvokeWrapper IWrapper;
125-
126-
targetDataMapper<TargetAsyncInfoTy>(
127-
Loc, DeviceId, ArgNum, ArgsBase, Args, ArgSizes, ArgTypes, ArgNames,
128-
ArgMappers, TargetDataFunction, RegionTypeMsg, RegionName, &IWrapper);
129-
}
130-
131153
template <typename TargetAsyncInfoTy>
132154
static inline void
133155
targetDataMapper(ident_t *Loc, int64_t DeviceId, int32_t ArgNum,
@@ -167,13 +189,16 @@ targetDataMapper(ident_t *Loc, int64_t DeviceId, int32_t ArgNum,
167189

168190
// DeviceId is only valid after the call to checkDeviceAndCtors, so we update
169191
// the DevicId in the Wrapper object before invoking OMPT
170-
OMPT_IF_ENABLED(OMPTInvoker->setDeviceId(DeviceId); OMPTInvoker->invoke(););
192+
OMPT_IF_ENABLED(OMPTInvoker->setDeviceId(DeviceId);
193+
OMPTInvoker->invokeBegin(););
171194

172195
int Rc = OFFLOAD_SUCCESS;
173196
Rc = TargetDataFunction(Loc, Device, ArgNum, ArgsBase, Args, ArgSizes,
174197
ArgTypes, ArgNames, ArgMappers, AsyncInfo,
175198
false /* FromMapper */);
176199

200+
OMPT_IF_ENABLED(OMPTInvoker->invokeEnd(););
201+
177202
if (Rc == OFFLOAD_SUCCESS)
178203
Rc = AsyncInfo.synchronize();
179204

@@ -194,7 +219,7 @@ EXTERN void __tgt_target_data_begin_mapper(ident_t *Loc, int64_t DeviceId,
194219
#ifdef OMPT_SUPPORT
195220
OMPTInvokeWrapper IWrapper(OMPT_GET_RETURN_ADDRESS(0),
196221
OMPT_GET_FRAME_ADDRESS(0), -1,
197-
ompt_target_enter_data, ompt_scope_begin);
222+
ompt_target_enter_data);
198223
#else
199224
OMPTInvokeWrapper IWrapper;
200225
#endif
@@ -211,9 +236,19 @@ EXTERN void __tgt_target_data_begin_nowait_mapper(
211236
void **ArgMappers, int32_t DepNum, void *DepList, int32_t NoAliasDepNum,
212237
void *NoAliasDepList) {
213238
TIMESCOPE_WITH_IDENT(Loc);
239+
240+
#ifdef OMPT_SUPPORT
241+
OMPTInvokeWrapper IWrapper(OMPT_GET_RETURN_ADDRESS(0),
242+
OMPT_GET_FRAME_ADDRESS(0), -1,
243+
ompt_target_enter_data_nowait);
244+
#else
245+
OMPTInvokeWrapper IWrapper;
246+
#endif
247+
214248
targetDataMapper<TaskAsyncInfoWrapperTy>(
215249
Loc, DeviceId, ArgNum, ArgsBase, Args, ArgSizes, ArgTypes, ArgNames,
216-
ArgMappers, targetDataBegin, "Entering OpenMP data region", "begin");
250+
ArgMappers, targetDataBegin, "Entering OpenMP data region", "begin",
251+
&IWrapper);
217252
}
218253

219254
/// passes data from the target, releases target memory and destroys
@@ -230,7 +265,7 @@ EXTERN void __tgt_target_data_end_mapper(ident_t *Loc, int64_t DeviceId,
230265
#ifdef OMPT_SUPPORT
231266
OMPTInvokeWrapper IWrapper(OMPT_GET_RETURN_ADDRESS(0),
232267
OMPT_GET_FRAME_ADDRESS(0), -1,
233-
ompt_target_exit_data, ompt_scope_begin);
268+
ompt_target_exit_data);
234269
#else
235270
OMPTInvokeWrapper IWrapper;
236271
#endif
@@ -246,9 +281,19 @@ EXTERN void __tgt_target_data_end_nowait_mapper(
246281
void **ArgMappers, int32_t DepNum, void *DepList, int32_t NoAliasDepNum,
247282
void *NoAliasDepList) {
248283
TIMESCOPE_WITH_IDENT(Loc);
284+
285+
#ifdef OMPT_SUPPORT
286+
OMPTInvokeWrapper IWrapper(OMPT_GET_RETURN_ADDRESS(0),
287+
OMPT_GET_FRAME_ADDRESS(0), -1,
288+
ompt_target_exit_data_nowait);
289+
#else
290+
OMPTInvokeWrapper IWrapper;
291+
#endif
292+
249293
targetDataMapper<TaskAsyncInfoWrapperTy>(
250294
Loc, DeviceId, ArgNum, ArgsBase, Args, ArgSizes, ArgTypes, ArgNames,
251-
ArgMappers, targetDataEnd, "Exiting OpenMP data region", "end");
295+
ArgMappers, targetDataEnd, "Exiting OpenMP data region", "end",
296+
&IWrapper);
252297
}
253298

254299
EXTERN void __tgt_target_data_update_mapper(ident_t *Loc, int64_t DeviceId,
@@ -262,8 +307,7 @@ EXTERN void __tgt_target_data_update_mapper(ident_t *Loc, int64_t DeviceId,
262307

263308
#ifdef OMPT_SUPPORT
264309
OMPTInvokeWrapper IWrapper(OMPT_GET_RETURN_ADDRESS(0),
265-
OMPT_GET_FRAME_ADDRESS(0), -1, ompt_target_update,
266-
ompt_scope_begin);
310+
OMPT_GET_FRAME_ADDRESS(0), -1, ompt_target_update);
267311
#else
268312
OMPTInvokeWrapper IWrapper;
269313
#endif
@@ -280,9 +324,19 @@ EXTERN void __tgt_target_data_update_nowait_mapper(
280324
void **ArgMappers, int32_t DepNum, void *DepList, int32_t NoAliasDepNum,
281325
void *NoAliasDepList) {
282326
TIMESCOPE_WITH_IDENT(Loc);
327+
328+
#ifdef OMPT_SUPPORT
329+
OMPTInvokeWrapper IWrapper(OMPT_GET_RETURN_ADDRESS(0),
330+
OMPT_GET_FRAME_ADDRESS(0), -1,
331+
ompt_target_update_nowait);
332+
#else
333+
OMPTInvokeWrapper IWrapper;
334+
#endif
335+
283336
targetDataMapper<TaskAsyncInfoWrapperTy>(
284337
Loc, DeviceId, ArgNum, ArgsBase, Args, ArgSizes, ArgTypes, ArgNames,
285-
ArgMappers, targetDataUpdate, "Updating OpenMP data", "update");
338+
ArgMappers, targetDataUpdate, "Updating OpenMP data", "update",
339+
&IWrapper);
286340
}
287341

288342
static KernelArgsTy *upgradeKernelArgs(KernelArgsTy *KernelArgs,

0 commit comments

Comments
 (0)