[libomptarget] [OMPT] Fixed return address computation for OMPT events. (llvm#80498)

Currently, __builtin_return_address is used to generate the return
address when the callback invoker is created. However, this may result
in the return address pointing to an internal runtime function. This is
not what a tool would typically want. A tool would want to know the
corresponding user code from where the runtime entry point is invoked.

This change adds a thread local variable that is assigned the return
address at the OpenMP runtime entry points. An RAII is used to manage
the modifications to the thread local variable. Whenever the return
address is required for OMPT events, it is read from the thread local
variable.

Adapted to work on amd-staging.

Change-Id: I3c85bdf48f81edbc421b9030f26f344ed5b963b2

Author: dhruvachak

openmp/libomptarget/include/OmptCommonDefs.h

@@ -76,13 +76,6 @@
 #define OMPT_FRAME_POSITION_DEFAULT ompt_frame_framepointer
 #endif
 
-// Allow 0 as the only argument to avoid unpredictable effects.
-// The setter stores the return address to a thread local variable.
-#define OMPT_SET_RETURN_ADDRESS                                                \
-  llvm::omp::target::ompt::ReturnAddress = __builtin_return_address(0)
-// The getter returns the address stored in the thread local variable.
-#define OMPT_GET_RETURN_ADDRESS llvm::omp::target::ompt::ReturnAddress
-
 #define OMPT_PTR_UNKNOWN ((void *)0)
 
 #define performIfOmptInitialized(stmt)                                         \

openmp/libomptarget/include/OpenMP/OMPT/Interface.h

@@ -382,6 +382,8 @@ class Interface {
 extern thread_local Interface RegionInterface;
 
 /// Thread local variable holding the return address.
+/// When using __builtin_return_address to set the return address,
+/// allow 0 as the only argument to avoid unpredictable effects.
 extern thread_local void *ReturnAddress;
 
 template <typename FuncTy, typename ArgsTy, size_t... IndexSeq>
@@ -422,11 +424,42 @@ template <typename FunctionPairTy, typename... ArgsTy>
 InterfaceRAII(FunctionPairTy Callbacks, ArgsTy... Args)
     -> InterfaceRAII<FunctionPairTy, ArgsTy...>;
 
+/// Used to set and reset the thread-local return address. The RAII is expected
+/// to be created at a runtime entry point when the return address should be
+/// null. If so, the return address is set and \p IsSetter is set in the ctor.
+/// The dtor resets the return address only if the corresponding object set it.
+/// So if the RAII is called from a nested runtime function, the ctor/dtor will
+/// do nothing since the thread local return address is already set.
+class ReturnAddressSetterRAII {
+public:
+  ReturnAddressSetterRAII(void *RA) : IsSetter(false) {
+    // Handle nested calls. If already set, do not set again since it
+    // must be in a nested call.
+    if (ReturnAddress == nullptr) {
+      // Store the return address to a thread local variable.
+      ReturnAddress = RA;
+      IsSetter = true;
+    }
+  }
+  ~ReturnAddressSetterRAII() {
+    // Reset the return address if this object set it.
+    if (IsSetter)
+      ReturnAddress = nullptr;
+  }
+
+private:
+  // Did this object set the thread-local return address?
+  bool IsSetter;
+};
+
 } // namespace ompt
 } // namespace target
 } // namespace omp
 } // namespace llvm
 
+// The getter returns the address stored in the thread local variable.
+#define OMPT_GET_RETURN_ADDRESS llvm::omp::target::ompt::ReturnAddress
+
 #pragma pop_macro("DEBUG_PREFIX")
 
 #endif // OMPT_SUPPORT

openmp/libomptarget/src/LegacyAPI.cpp

@@ -10,15 +10,21 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "OpenMP/OMPT/Interface.h"
 #include "omptarget.h"
 #include "private.h"
 
 #include "Shared/Profile.h"
 
+#ifdef OMPT_SUPPORT
+using namespace llvm::omp::target::ompt;
+#endif
+
 EXTERN void __tgt_target_data_begin(int64_t DeviceId, int32_t ArgNum,
                                     void **ArgsBase, void **Args,
                                     int64_t *ArgSizes, int64_t *ArgTypes) {
   TIMESCOPE();
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   __tgt_target_data_begin_mapper(nullptr, DeviceId, ArgNum, ArgsBase, Args,
                                  ArgSizes, ArgTypes, nullptr, nullptr);
 }
@@ -30,7 +36,7 @@ EXTERN void __tgt_target_data_begin_nowait(int64_t DeviceId, int32_t ArgNum,
                                            int32_t NoAliasDepNum,
                                            void *NoAliasDepList) {
   TIMESCOPE();
-
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   __tgt_target_data_begin_mapper(nullptr, DeviceId, ArgNum, ArgsBase, Args,
                                  ArgSizes, ArgTypes, nullptr, nullptr);
 }
@@ -39,6 +45,7 @@ EXTERN void __tgt_target_data_end(int64_t DeviceId, int32_t ArgNum,
                                   void **ArgsBase, void **Args,
                                   int64_t *ArgSizes, int64_t *ArgTypes) {
   TIMESCOPE();
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   __tgt_target_data_end_mapper(nullptr, DeviceId, ArgNum, ArgsBase, Args,
                                ArgSizes, ArgTypes, nullptr, nullptr);
 }
@@ -47,6 +54,7 @@ EXTERN void __tgt_target_data_update(int64_t DeviceId, int32_t ArgNum,
                                      void **ArgsBase, void **Args,
                                      int64_t *ArgSizes, int64_t *ArgTypes) {
   TIMESCOPE();
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   __tgt_target_data_update_mapper(nullptr, DeviceId, ArgNum, ArgsBase, Args,
                                   ArgSizes, ArgTypes, nullptr, nullptr);
 }
@@ -56,7 +64,7 @@ EXTERN void __tgt_target_data_update_nowait(
     int64_t *ArgSizes, int64_t *ArgTypes, int32_t DepNum, void *DepList,
     int32_t NoAliasDepNum, void *NoAliasDepList) {
   TIMESCOPE();
-
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   __tgt_target_data_update_mapper(nullptr, DeviceId, ArgNum, ArgsBase, Args,
                                   ArgSizes, ArgTypes, nullptr, nullptr);
 }
@@ -68,7 +76,7 @@ EXTERN void __tgt_target_data_end_nowait(int64_t DeviceId, int32_t ArgNum,
                                          int32_t NoAliasDepNum,
                                          void *NoAliasDepList) {
   TIMESCOPE();
-
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   __tgt_target_data_end_mapper(nullptr, DeviceId, ArgNum, ArgsBase, Args,
                                ArgSizes, ArgTypes, nullptr, nullptr);
 }
@@ -78,6 +86,7 @@ EXTERN int __tgt_target_mapper(ident_t *Loc, int64_t DeviceId, void *HostPtr,
                                int64_t *ArgSizes, int64_t *ArgTypes,
                                map_var_info_t *ArgNames, void **ArgMappers) {
   TIMESCOPE_WITH_IDENT(Loc);
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   KernelArgsTy KernelArgs{1,        ArgNum,   ArgsBase,   Args, ArgSizes,
                           ArgTypes, ArgNames, ArgMappers, 0};
   return __tgt_target_kernel(Loc, DeviceId, -1, -1, HostPtr, &KernelArgs);
@@ -87,6 +96,7 @@ EXTERN int __tgt_target(int64_t DeviceId, void *HostPtr, int32_t ArgNum,
                         void **ArgsBase, void **Args, int64_t *ArgSizes,
                         int64_t *ArgTypes) {
   TIMESCOPE();
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   return __tgt_target_mapper(nullptr, DeviceId, HostPtr, ArgNum, ArgsBase, Args,
                              ArgSizes, ArgTypes, nullptr, nullptr);
 }
@@ -96,7 +106,7 @@ EXTERN int __tgt_target_nowait(int64_t DeviceId, void *HostPtr, int32_t ArgNum,
                                int64_t *ArgTypes, int32_t DepNum, void *DepList,
                                int32_t NoAliasDepNum, void *NoAliasDepList) {
   TIMESCOPE();
-
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   return __tgt_target_mapper(nullptr, DeviceId, HostPtr, ArgNum, ArgsBase, Args,
                              ArgSizes, ArgTypes, nullptr, nullptr);
 }
@@ -107,7 +117,7 @@ EXTERN int __tgt_target_nowait_mapper(
     map_var_info_t *ArgNames, void **ArgMappers, int32_t DepNum, void *DepList,
     int32_t NoAliasDepNum, void *NoAliasDepList) {
   TIMESCOPE_WITH_IDENT(Loc);
-
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   return __tgt_target_mapper(Loc, DeviceId, HostPtr, ArgNum, ArgsBase, Args,
                              ArgSizes, ArgTypes, ArgNames, ArgMappers);
 }
@@ -120,7 +130,7 @@ EXTERN int __tgt_target_teams_mapper(ident_t *Loc, int64_t DeviceId,
                                      void **ArgMappers, int32_t NumTeams,
                                      int32_t ThreadLimit) {
   TIMESCOPE_WITH_IDENT(Loc);
-
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   KernelArgsTy KernelArgs{1,        ArgNum,   ArgsBase,   Args, ArgSizes,
                           ArgTypes, ArgNames, ArgMappers, 0};
   return __tgt_target_kernel(Loc, DeviceId, NumTeams, ThreadLimit, HostPtr,
@@ -132,6 +142,7 @@ EXTERN int __tgt_target_teams(int64_t DeviceId, void *HostPtr, int32_t ArgNum,
                               int64_t *ArgTypes, int32_t NumTeams,
                               int32_t ThreadLimit) {
   TIMESCOPE();
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   return __tgt_target_teams_mapper(nullptr, DeviceId, HostPtr, ArgNum, ArgsBase,
                                    Args, ArgSizes, ArgTypes, nullptr, nullptr,
                                    NumTeams, ThreadLimit);
@@ -145,7 +156,7 @@ EXTERN int __tgt_target_teams_nowait(int64_t DeviceId, void *HostPtr,
                                      void *DepList, int32_t NoAliasDepNum,
                                      void *NoAliasDepList) {
   TIMESCOPE();
-
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   return __tgt_target_teams_mapper(nullptr, DeviceId, HostPtr, ArgNum, ArgsBase,
                                    Args, ArgSizes, ArgTypes, nullptr, nullptr,
                                    NumTeams, ThreadLimit);
@@ -158,7 +169,7 @@ EXTERN int __tgt_target_teams_nowait_mapper(
     int32_t ThreadLimit, int32_t DepNum, void *DepList, int32_t NoAliasDepNum,
     void *NoAliasDepList) {
   TIMESCOPE_WITH_IDENT(Loc);
-
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   return __tgt_target_teams_mapper(Loc, DeviceId, HostPtr, ArgNum, ArgsBase,
                                    Args, ArgSizes, ArgTypes, ArgNames,
                                    ArgMappers, NumTeams, ThreadLimit);
@@ -186,6 +197,7 @@ EXTERN int __tgt_target_kernel_nowait(ident_t *Loc, int64_t DeviceId,
                                       int32_t NoAliasDepNum,
                                       void *NoAliasDepList) {
   TIMESCOPE_WITH_IDENT(Loc);
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
   return __tgt_target_kernel(Loc, DeviceId, NumTeams, ThreadLimit, HostPtr,
                              KernelArgs);
 }