[OMPT][Offload] Refactor async-info data passing

This patch moves the OMPT-specific data type into its own header. At the
same time it changes the type of the pointer stores in the async info
object to a void* to not leak the OMPT specific data type into the
general interface.

The type change to void* was not part of the original WIP patch in
https://github.com/AMD-Lightning-Internal/llvm-project/pull/2282 but I
thought I can fold it into this refactoring. If this is unwanted, I'm
happy to revert that change.

Author: jplehr

offload/include/OpenMP/OMPT/Interface.h

@@ -16,6 +16,7 @@
 // Only provide functionality if target OMPT support is enabled
 #ifdef OMPT_SUPPORT
 #include "Callback.h"
+#include "OmptEventInfoTy.h"
 #include "Shared/APITypes.h"
 #include "Shared/Debug.h"
 #include "omp-tools.h"
@@ -462,14 +463,6 @@ template <typename FunctionPairTy, typename... ArgsTy>
 InterfaceRAII(FunctionPairTy Callbacks, ArgsTy... Args)
     -> InterfaceRAII<FunctionPairTy, ArgsTy...>;
 
-/// Holds info needed to fill asynchronous trace records
-struct OmptEventInfoTy {
-  /// The granted number of teams at runtime
-  uint64_t NumTeams;
-  /// Pointer to the actual buffer storage location
-  ompt_record_ompt_t *TraceRecord;
-};
-
 /// Similar to the original InterfaceRAII this class is used for tracing and
 /// extends the original with async capabilities. That is: It takes an
 /// additional AsyncInfo reference as argument to populate the relevant fields.
@@ -488,13 +481,15 @@ class TracerInterfaceRAII {
       auto Record = begin();
       // Gets freed in interface.cpp, functions
       // targetKernel and targetData once launching target operations returns.
-      if (!AI->OmptEventInfo)
-        AI->OmptEventInfo = new OmptEventInfoTy();
-      AI->OmptEventInfo->TraceRecord = Record;
-      AI->OmptEventInfo->NumTeams = 0;
+      if (!AI->ProfilerData)
+        AI->ProfilerData = new OmptEventInfoTy();
+      // TODO: Make sure this has the right type
+      auto OEI = reinterpret_cast<OmptEventInfoTy *>(AI->ProfilerData);
+      OEI->TraceRecord = Record;
+      OEI->NumTeams = 0;
     } else {
       // Actively prevent further tracing of this event
-      AI->OmptEventInfo = nullptr;
+      AI->ProfilerData = nullptr;
     }
   }
 

offload/include/OpenMP/OMPT/OmptEventInfoTy.h

@@ -0,0 +1,39 @@
+//===- OmptEventInfoTy.h - OMPT specific trace record data ------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Data structure used to communicate OMPT specific profiler data from the
+// high-level libomptarget into the vendor-specific plugins
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef OFFLOAD_INCLUDE_OPENMP_OMPT_OMPTEVENTINFOTY_H
+#define OFFLOAD_INCLUDE_OPENMP_OMPT_OMPTEVENTINFOTY_H
+
+#include "Shared/Debug.h"
+
+struct ompt_record_ompt_t;
+
+namespace llvm {
+namespace omp {
+namespace target {
+namespace ompt {
+
+/// Holds info needed to fill asynchronous trace records
+struct OmptEventInfoTy {
+  /// The granted number of teams at runtime
+  uint64_t NumTeams;
+  /// Pointer to the actual buffer storage location
+  ompt_record_ompt_t *TraceRecord;
+};
+
+} // namespace ompt
+} // namespace target
+} // namespace omp
+} // namespace llvm
+
+#endif // OFFLOAD_INCLUDE_OPENMP_OMPT_OMPTEVENTINFOTY_H

offload/include/Shared/APITypes.h

@@ -71,11 +71,6 @@ struct __tgt_device_binary {
 
 // clang-format on
 
-// OMPT: Forward declare for function pointer type
-namespace llvm::omp::target::ompt {
-struct OmptEventInfoTy;
-} // namespace llvm::omp::target::ompt
-
 /// This struct contains information exchanged between different asynchronous
 /// operations for device-dependent optimization and potential synchronization
 struct __tgt_async_info {
@@ -98,8 +93,11 @@ struct __tgt_async_info {
 
   /// Use for sync interface. When false => synchronous execution
   bool ExecAsync = true;
-  /// Maintain the actal data for the OMPT stuff
-  llvm::omp::target::ompt::OmptEventInfoTy *OmptEventInfo = nullptr;
+  /// Maintain the actal data for OMPT.
+  /// TODO: Moving forward, this should become a void* to not expose OMPT data
+  /// type in general types
+  // llvm::omp::target::ompt::OmptEventInfoTy *OmptEventInfo = nullptr;
+  void *ProfilerData = nullptr;
 };
 
 /// This struct contains all of the arguments to a target kernel region launch.

offload/libomptarget/interface.cpp

@@ -201,8 +201,8 @@ targetData(ident_t *Loc, int64_t DeviceId, int32_t ArgNum, void **ArgsBase,
   }
 
 #ifdef OMPT_SUPPORT
-  if (__tgt_async_info *AI = AsyncInfo; AI->OmptEventInfo)
-    delete AI->OmptEventInfo;
+  if (__tgt_async_info *AI = AsyncInfo; AI->ProfilerData)
+    delete AI->ProfilerData;
 #endif
 
   handleTargetOutcome(Rc == OFFLOAD_SUCCESS, Loc);
@@ -491,13 +491,13 @@ static inline int targetKernel(ident_t *Loc, int64_t DeviceId, int32_t NumTeams,
   }
 
 #ifdef OMPT_SUPPORT
-  if (__tgt_async_info *AI = AsyncInfo; AI->OmptEventInfo)
-    delete AI->OmptEventInfo;
+  if (__tgt_async_info *AI = AsyncInfo; AI->ProfilerData)
+    delete AI->ProfilerData;
 
   for (TargetAsyncInfoTy *LocalTAI : TargetAsyncInfos) {
     AsyncInfoTy &AsyncInfo = *LocalTAI;
-    if (__tgt_async_info *AI = AsyncInfo; AI->OmptEventInfo) {
-      delete AI->OmptEventInfo;
+    if (__tgt_async_info *AI = AsyncInfo; AI->ProfilerData) {
+      delete AI->ProfilerData;
     }
   }
 #endif

offload/plugins-nextgen/amdgpu/src/rtl.cpp

@@ -29,6 +29,7 @@
 #include "ErrorReporting.h"
 #include "OpenMP/OMPT/Interface.h"
 #include "OpenMP/OMPT/OmptCommonDefs.h"
+#include "OpenMP/OMPT/OmptEventInfoTy.h"
 #include "Shared/APITypes.h"
 #include "Shared/Debug.h"
 #include "Shared/Environment.h"
@@ -183,14 +184,14 @@ static void printOmptEventInfoTy(ompt::OmptEventInfoTy &OmptEventInfo) {
 static std::unique_ptr<ompt::OmptEventInfoTy>
 getOrNullOmptEventInfo(AsyncInfoWrapperTy &AsyncInfoWrapper) {
   __tgt_async_info *AI = AsyncInfoWrapper;
-  if (!AI || !AI->OmptEventInfo)
+  if (!AI || !AI->ProfilerData)
     return nullptr;
 
-  // We need to copy the content of the OmptEventInfo object to persist it
+  // We need to copy the content of the ProfilerData object to persist it
   // between multiple async operations.
-  auto LocalOmptEventInfo =
-      std::make_unique<ompt::OmptEventInfoTy>(*AI->OmptEventInfo);
-  printOmptEventInfoTy(*AI->OmptEventInfo);
+  auto LocalOmptEventInfo = std::make_unique<ompt::OmptEventInfoTy>(
+      *reinterpret_cast<ompt::OmptEventInfoTy *>(AI->ProfilerData));
+  // printOmptEventInfoTy(*AI->ProfilerData);
   printOmptEventInfoTy(*LocalOmptEventInfo);
   return LocalOmptEventInfo;
 }
@@ -4347,7 +4348,7 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
     // Implementation sanity checks: either unified_shared_memory or auto
     // zero-copy, not both
     if (isUnifiedSharedMemory && isAutoZeroCopy)
-      return Plugin::error(ErrorCode::UNKNOWN, 
+      return Plugin::error(ErrorCode::UNKNOWN,
                            "Internal runtime error: cannot be both "
                            "unified_shared_memory and auto zero-copy.");
 

offload/plugins-nextgen/common/src/PluginInterface.cpp

@@ -456,7 +456,7 @@ AsyncInfoWrapperTy::AsyncInfoWrapperTy(GenericDeviceTy &Device,
                                        __tgt_async_info *AsyncInfoPtr)
     : Device(Device),
       AsyncInfoPtr(AsyncInfoPtr ? AsyncInfoPtr : &LocalAsyncInfo) {
-  LocalAsyncInfo.OmptEventInfo = nullptr;
+  LocalAsyncInfo.ProfilerData = nullptr;
 }
 
 void AsyncInfoWrapperTy::finalize(Error &Err) {
@@ -525,7 +525,7 @@ Error GenericKernelTy::init(GenericDeviceTy &GenericDevice,
   } else {
     // Check that the retrieved execution mode is valid.
     if (!GenericKernelTy::isValidExecutionMode(ExecModeGlobal.getValue()))
-      return Plugin::error(ErrorCode::UNKNOWN, 
+      return Plugin::error(ErrorCode::UNKNOWN,
                            "Invalid execution mode %d for '%s'",
                            ExecModeGlobal.getValue(), getName());
     ExecutionMode = ExecModeGlobal.getValue();
@@ -613,22 +613,22 @@ GenericKernelTy::getKernelLaunchEnvironment(
        DPxPTR(&LocalKLE), DPxPTR(*AllocOrErr),
        sizeof(KernelLaunchEnvironmentTy));
 
-  // The OmptEventInfo at this point will have a callback for a kernel launch,
+  // The ProfilerData at this point will have a callback for a kernel launch,
   // not a data-op. This is due to the "external" operation being a kernel
   // launch and the data submit here being an implementation detail. We
-  // temporarily set the OmptEventInfo to nullptr, such that we disable the
+  // temporarily set the ProfilerData to nullptr, such that we disable the
   // timing etc further down to not trigger assertions or report implementation
   // detail.
   __tgt_async_info *AI = AsyncInfoWrapper;
-  if (AI && AI->OmptEventInfo) {
-    auto LocalOEI = AI->OmptEventInfo;
-    AI->OmptEventInfo = nullptr;
+  if (AI && AI->ProfilerData) {
+    auto LocalOEI = AI->ProfilerData;
+    AI->ProfilerData = nullptr;
     auto Err = GenericDevice.dataSubmit(*AllocOrErr, &LocalKLE,
                                         sizeof(KernelLaunchEnvironmentTy),
                                         AsyncInfoWrapper);
     if (Err)
       return Err;
-    AI->OmptEventInfo = LocalOEI;
+    AI->ProfilerData = LocalOEI;
     return static_cast<KernelLaunchEnvironmentTy *>(*AllocOrErr);
   }
 
@@ -783,10 +783,11 @@ Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, void **ArgPtrs,
   OMPT_IF_TRACING_ENABLED(if (llvm::omp::target::ompt::isTracedDevice(
                                   getDeviceId(&GenericDevice))) {
     __tgt_async_info *AI = AsyncInfoWrapper;
-    if (AI->OmptEventInfo != nullptr) {
+    if (AI->ProfilerData != nullptr) {
       // Set number of granted teams for OMPT
       setOmptGrantedNumTeams(NumBlocks[0]);
-      AI->OmptEventInfo->NumTeams = NumBlocks[0];
+      reinterpret_cast<OmptEventInfoTy *>(AI->ProfilerData)->NumTeams =
+          NumBlocks[0];
     }
   });