Skip to content

Commit bb485fb

Browse files
dhruvachakdhruvachak
committed
[OMPT] Fix for ROCm/ROCm#2611
Change-Id: Ie61e550aa05e31f1adebc0302d22f0acb1c3360c
1 parent 3f30797 commit bb485fb

File tree

1 file changed

+85
-30
lines changed

1 file changed

+85
-30
lines changed

‎openmp/libomptarget/src/interface.cpp

Lines changed: 85 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -45,21 +45,60 @@ class OMPTInvokeWrapper {
4545
public:
4646
OMPTInvokeWrapper()
4747
: IsNullOpt(true), CodePtr(nullptr), ReturnFramePtr(nullptr),
48-
DeviceId(-1), Kind(ompt_target), ScopeKind(ompt_scope_begin) {}
48+
DeviceId(-1), Kind(ompt_target) {}
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) {}
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+
// The trace record must be written out before invoking the callback
85+
// interface. This is because the end-callback interface resets metadata.
86+
ompt_interface.target_trace_record_gen(DeviceId, Kind, ompt_scope_end,
87+
CodePtr);
88+
switch (Kind) {
89+
case ompt_target_enter_data:
90+
case ompt_target_enter_data_nowait:
91+
ompt_interface.target_data_enter_end(DeviceId, CodePtr);
92+
break;
93+
case ompt_target_exit_data:
94+
case ompt_target_exit_data_nowait:
95+
ompt_interface.target_data_exit_end(DeviceId, CodePtr);
96+
break;
97+
case ompt_target_update:
98+
case ompt_target_update_nowait:
99+
ompt_interface.target_update_end(DeviceId, CodePtr);
100+
break;
101+
}
63102
}
64103

65104
private:
@@ -68,7 +107,6 @@ class OMPTInvokeWrapper {
68107
void *ReturnFramePtr;
69108
int64_t DeviceId;
70109
ompt_target_t Kind;
71-
ompt_scope_endpoint_t ScopeKind;
72110
};
73111

74112
////////////////////////////////////////////////////////////////////////////////
@@ -113,21 +151,6 @@ EXTERN void __tgt_unregister_lib(__tgt_bin_desc *Desc) {
113151
}
114152
}
115153

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-
131154
template <typename TargetAsyncInfoTy>
132155
static inline void
133156
targetDataMapper(ident_t *Loc, int64_t DeviceId, int32_t ArgNum,
@@ -167,13 +190,16 @@ targetDataMapper(ident_t *Loc, int64_t DeviceId, int32_t ArgNum,
167190

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

172196
int Rc = OFFLOAD_SUCCESS;
173197
Rc = TargetDataFunction(Loc, Device, ArgNum, ArgsBase, Args, ArgSizes,
174198
ArgTypes, ArgNames, ArgMappers, AsyncInfo,
175199
false /* FromMapper */);
176200

201+
OMPT_IF_ENABLED(OMPTInvoker->invokeEnd(););
202+
177203
if (Rc == OFFLOAD_SUCCESS)
178204
Rc = AsyncInfo.synchronize();
179205

@@ -194,7 +220,7 @@ EXTERN void __tgt_target_data_begin_mapper(ident_t *Loc, int64_t DeviceId,
194220
#ifdef OMPT_SUPPORT
195221
OMPTInvokeWrapper IWrapper(OMPT_GET_RETURN_ADDRESS(0),
196222
OMPT_GET_FRAME_ADDRESS(0), -1,
197-
ompt_target_enter_data, ompt_scope_begin);
223+
ompt_target_enter_data);
198224
#else
199225
OMPTInvokeWrapper IWrapper;
200226
#endif
@@ -211,9 +237,19 @@ EXTERN void __tgt_target_data_begin_nowait_mapper(
211237
void **ArgMappers, int32_t DepNum, void *DepList, int32_t NoAliasDepNum,
212238
void *NoAliasDepList) {
213239
TIMESCOPE_WITH_IDENT(Loc);
240+
241+
#ifdef OMPT_SUPPORT
242+
OMPTInvokeWrapper IWrapper(OMPT_GET_RETURN_ADDRESS(0),
243+
OMPT_GET_FRAME_ADDRESS(0), -1,
244+
ompt_target_enter_data_nowait);
245+
#else
246+
OMPTInvokeWrapper IWrapper;
247+
#endif
248+
214249
targetDataMapper<TaskAsyncInfoWrapperTy>(
215250
Loc, DeviceId, ArgNum, ArgsBase, Args, ArgSizes, ArgTypes, ArgNames,
216-
ArgMappers, targetDataBegin, "Entering OpenMP data region", "begin");
251+
ArgMappers, targetDataBegin, "Entering OpenMP data region", "begin",
252+
&IWrapper);
217253
}
218254

219255
/// passes data from the target, releases target memory and destroys
@@ -230,7 +266,7 @@ EXTERN void __tgt_target_data_end_mapper(ident_t *Loc, int64_t DeviceId,
230266
#ifdef OMPT_SUPPORT
231267
OMPTInvokeWrapper IWrapper(OMPT_GET_RETURN_ADDRESS(0),
232268
OMPT_GET_FRAME_ADDRESS(0), -1,
233-
ompt_target_exit_data, ompt_scope_begin);
269+
ompt_target_exit_data);
234270
#else
235271
OMPTInvokeWrapper IWrapper;
236272
#endif
@@ -246,9 +282,19 @@ EXTERN void __tgt_target_data_end_nowait_mapper(
246282
void **ArgMappers, int32_t DepNum, void *DepList, int32_t NoAliasDepNum,
247283
void *NoAliasDepList) {
248284
TIMESCOPE_WITH_IDENT(Loc);
285+
286+
#ifdef OMPT_SUPPORT
287+
OMPTInvokeWrapper IWrapper(OMPT_GET_RETURN_ADDRESS(0),
288+
OMPT_GET_FRAME_ADDRESS(0), -1,
289+
ompt_target_exit_data_nowait);
290+
#else
291+
OMPTInvokeWrapper IWrapper;
292+
#endif
293+
249294
targetDataMapper<TaskAsyncInfoWrapperTy>(
250295
Loc, DeviceId, ArgNum, ArgsBase, Args, ArgSizes, ArgTypes, ArgNames,
251-
ArgMappers, targetDataEnd, "Exiting OpenMP data region", "end");
296+
ArgMappers, targetDataEnd, "Exiting OpenMP data region", "end",
297+
&IWrapper);
252298
}
253299

254300
EXTERN void __tgt_target_data_update_mapper(ident_t *Loc, int64_t DeviceId,
@@ -262,8 +308,7 @@ EXTERN void __tgt_target_data_update_mapper(ident_t *Loc, int64_t DeviceId,
262308

263309
#ifdef OMPT_SUPPORT
264310
OMPTInvokeWrapper IWrapper(OMPT_GET_RETURN_ADDRESS(0),
265-
OMPT_GET_FRAME_ADDRESS(0), -1, ompt_target_update,
266-
ompt_scope_begin);
311+
OMPT_GET_FRAME_ADDRESS(0), -1, ompt_target_update);
267312
#else
268313
OMPTInvokeWrapper IWrapper;
269314
#endif
@@ -280,9 +325,19 @@ EXTERN void __tgt_target_data_update_nowait_mapper(
280325
void **ArgMappers, int32_t DepNum, void *DepList, int32_t NoAliasDepNum,
281326
void *NoAliasDepList) {
282327
TIMESCOPE_WITH_IDENT(Loc);
328+
329+
#ifdef OMPT_SUPPORT
330+
OMPTInvokeWrapper IWrapper(OMPT_GET_RETURN_ADDRESS(0),
331+
OMPT_GET_FRAME_ADDRESS(0), -1,
332+
ompt_target_update_nowait);
333+
#else
334+
OMPTInvokeWrapper IWrapper;
335+
#endif
336+
283337
targetDataMapper<TaskAsyncInfoWrapperTy>(
284338
Loc, DeviceId, ArgNum, ArgsBase, Args, ArgSizes, ArgTypes, ArgNames,
285-
ArgMappers, targetDataUpdate, "Updating OpenMP data", "update");
339+
ArgMappers, targetDataUpdate, "Updating OpenMP data", "update",
340+
&IWrapper);
286341
}
287342

288343
static KernelArgsTy *upgradeKernelArgs(KernelArgsTy *KernelArgs,

0 commit comments

Comments
 (0)