12 files changed, 245 insertions, 13 deletions
diff --git a/openmp/libomptarget/include/ompt_device_callbacks.h b/openmp/libomptarget/include/ompt_device_callbacks.h
index 3dbba3dedbee..127e18950249 100644
--- a/openmp/libomptarget/include/ompt_device_callbacks.h
+++ b/openmp/libomptarget/include/ompt_device_callbacks.h
@@ -49,6 +49,20 @@ public:
 #undef OmptBindCallback
   }
 
+  /// Used to find a callback given its name
+  ompt_interface_fn_t lookupCallback(const char *InterfaceFunctionName) {
+#define OmptLookup(Name, Type, Code)                                           \
+  if (strcmp(InterfaceFunctionName, #Name) == 0)                               \
+    return (ompt_interface_fn_t)Name##_fn;
+
+    FOREACH_OMPT_TARGET_CALLBACK(OmptLookup);
+#undef OmptLookup
+    return (ompt_interface_fn_t) nullptr;
+  }
+
+  /// Wrapper function to find a callback given its name
+  static ompt_interface_fn_t doLookup(const char *InterfaceFunctionName);
+
 private:
   /// Set to true if callbacks for this library have been initialized
   bool Enabled;
@@ -62,8 +76,6 @@ private:
 /// Device callbacks object for the library that performs the instantiation
 extern OmptDeviceCallbacksTy OmptDeviceCallbacks;
 
-#undef DEBUG_PREFIX
-
 #endif // OMPT_SUPPORT
 
 #endif // _OMPT_DEVICE_CALLBACKS_H
diff --git a/openmp/libomptarget/plugins-nextgen/CMakeLists.txt b/openmp/libomptarget/plugins-nextgen/CMakeLists.txt
index af02f050f467..a51da395680b 100644
--- a/openmp/libomptarget/plugins-nextgen/CMakeLists.txt
+++ b/openmp/libomptarget/plugins-nextgen/CMakeLists.txt
@@ -49,6 +49,7 @@ if(CMAKE_SYSTEM_PROCESSOR MATCHES "${tmachine}$")
         PRIVATE
         elf_common
         MemoryManager
+        OMPT
         PluginInterface
         ${LIBOMPTARGET_DEP_LIBFFI_LIBRARIES}
         ${OPENMP_PTHREAD_LIB}
diff --git a/openmp/libomptarget/plugins-nextgen/amdgpu/CMakeLists.txt b/openmp/libomptarget/plugins-nextgen/amdgpu/CMakeLists.txt
index b689ff5a38d5..0c807a402c68 100644
--- a/openmp/libomptarget/plugins-nextgen/amdgpu/CMakeLists.txt
+++ b/openmp/libomptarget/plugins-nextgen/amdgpu/CMakeLists.txt
@@ -82,6 +82,7 @@ add_llvm_library(omptarget.rtl.amdgpu.nextgen SHARED
   PRIVATE
   elf_common
   MemoryManager
+  OMPT
   PluginInterface
   ${LIBOMPTARGET_DEP_LIBRARIES}
   ${OPENMP_PTHREAD_LIB}
diff --git a/openmp/libomptarget/plugins-nextgen/common/CMakeLists.txt b/openmp/libomptarget/plugins-nextgen/common/CMakeLists.txt
index 1c5594eec5af..f0645d0d1753 100644
--- a/openmp/libomptarget/plugins-nextgen/common/CMakeLists.txt
+++ b/openmp/libomptarget/plugins-nextgen/common/CMakeLists.txt
@@ -10,4 +10,5 @@
 #
 ##===----------------------------------------------------------------------===##
 
+add_subdirectory(OMPT)
 add_subdirectory(PluginInterface)
diff --git a/openmp/libomptarget/plugins-nextgen/common/OMPT/CMakeLists.txt b/openmp/libomptarget/plugins-nextgen/common/OMPT/CMakeLists.txt
new file mode 100644
index 000000000000..c50dd18dce02
--- /dev/null
+++ b/openmp/libomptarget/plugins-nextgen/common/OMPT/CMakeLists.txt
@@ -0,0 +1,72 @@
+##===----------------------------------------------------------------------===##
+#
+# 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
+#
+##===----------------------------------------------------------------------===##
+#
+# Aggregation of parts which can be used by OpenMP tools
+#
+##===----------------------------------------------------------------------===##
+
+# NOTE: Don't try to build `OMPT` using `add_llvm_library` because we
+# don't want to export `OMPT` while `add_llvm_library` requires that.
+add_library(OMPT OBJECT
+  OmptCallback.cpp)
+
+# This is required when using LLVM libraries.
+llvm_update_compile_flags(OMPT)
+
+if (LLVM_LINK_LLVM_DYLIB)
+  set(llvm_libs LLVM)
+else()
+  llvm_map_components_to_libnames(llvm_libs
+    ${LLVM_TARGETS_TO_BUILD}
+    AggressiveInstCombine
+    Analysis
+    BinaryFormat
+    BitReader
+    BitWriter
+    CodeGen
+    Core
+    Extensions
+    InstCombine
+    Instrumentation
+    IPO
+    IRReader
+    Linker
+    MC
+    Object
+    Passes
+    Remarks
+    ScalarOpts
+    Support
+    Target
+    TargetParser
+    TransformUtils
+    Vectorize
+  )
+endif()
+
+target_link_libraries(OMPT
+  PUBLIC
+    ${llvm_libs}
+    elf_common
+    MemoryManager
+)
+
+# Define the TARGET_NAME and DEBUG_PREFIX.
+target_compile_definitions(OMPT PRIVATE
+  TARGET_NAME="OMPT"
+  DEBUG_PREFIX="OMPT"
+)
+
+target_include_directories(OMPT
+  INTERFACE ${CMAKE_CURRENT_SOURCE_DIR}
+  PRIVATE ${LIBOMPTARGET_INCLUDE_DIR}
+)
+
+set_target_properties(OMPT PROPERTIES
+  POSITION_INDEPENDENT_CODE ON
+  CXX_VISIBILITY_PRESET protected)
diff --git a/openmp/libomptarget/plugins-nextgen/common/OMPT/OmptCallback.cpp b/openmp/libomptarget/plugins-nextgen/common/OMPT/OmptCallback.cpp
new file mode 100644
index 000000000000..eb47f10a3f41
--- /dev/null
+++ b/openmp/libomptarget/plugins-nextgen/common/OMPT/OmptCallback.cpp
@@ -0,0 +1,83 @@
+//===---------- OmptCallback.cpp - Generic OMPT callbacks --------- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// OMPT support for PluginInterface
+//
+//===----------------------------------------------------------------------===//
+
+#ifdef OMPT_SUPPORT
+#include <atomic>
+#include <cstdio>
+#include <string.h>
+#include <vector>
+
+#include "Debug.h"
+#include "ompt_connector.h"
+#include "ompt_device_callbacks.h"
+
+/// Object maintaining all the callbacks in the plugin
+OmptDeviceCallbacksTy OmptDeviceCallbacks;
+
+/// Lookup function used for querying callback functions maintained
+/// by the plugin
+ompt_interface_fn_t
+OmptDeviceCallbacksTy::doLookup(const char *InterfaceFunctionName) {
+  // TODO This will be populated with device tracing functions
+  return (ompt_interface_fn_t) nullptr;
+}
+
+/// Used to indicate whether OMPT was enabled for this library
+static bool OmptEnabled = false;
+
+/// This function is passed to libomptarget as part of the OMPT connector
+/// object. It is called by libomptarget during initialization of OMPT in the
+/// plugin. \p lookup to be used to query callbacks registered with libomptarget
+/// \p initial_device_num Initial device num provided by libomptarget
+/// \p tool_data as provided by the tool
+static int OmptDeviceInit(ompt_function_lookup_t lookup, int initial_device_num,
+                          ompt_data_t *tool_data) {
+  DP("OMPT: Enter OmptDeviceInit\n");
+  OmptEnabled = true;
+  // The lookup parameter is provided by libomptarget which already has the tool
+  // callbacks registered at this point. The registration call below causes the
+  // same callback functions to be registered in the plugin as well.
+  OmptDeviceCallbacks.registerCallbacks(lookup);
+  DP("OMPT: Exit OmptDeviceInit\n");
+  return 0;
+}
+
+/// This function is passed to libomptarget as part of the OMPT connector
+/// object. It is called by libomptarget during finalization of OMPT in the
+/// plugin.
+static void OmptDeviceFini(ompt_data_t *tool_data) {
+  DP("OMPT: Executing OmptDeviceFini\n");
+}
+
+/// Used to initialize callbacks implemented by the tool. This interface will
+/// lookup the callbacks table in libomptarget and assign them to the callbacks
+/// table maintained in the calling plugin library.
+void OmptCallbackInit() {
+  DP("OMPT: Entering OmptCallbackInit\n");
+  /// Connect plugin instance with libomptarget
+  OmptLibraryConnectorTy LibomptargetConnector("libomptarget");
+  ompt_start_tool_result_t OmptResult;
+
+  // Initialize OmptResult with the init and fini functions that will be
+  // called by the connector
+  OmptResult.initialize = OmptDeviceInit;
+  OmptResult.finalize = OmptDeviceFini;
+  OmptResult.tool_data.value = 0;
+
+  // Initialize the device callbacks first
+  OmptDeviceCallbacks.init();
+
+  // Now call connect that causes the above init/fini functions to be called
+  LibomptargetConnector.connect(&OmptResult);
+  DP("OMPT: Exiting OmptCallbackInit\n");
+}
+#endif
diff --git a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/CMakeLists.txt b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/CMakeLists.txt
index 91d64f429823..c3950aff6187 100644
--- a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/CMakeLists.txt
+++ b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/CMakeLists.txt
@@ -60,6 +60,7 @@ target_link_libraries(PluginInterface
     ${llvm_libs}
     elf_common
     MemoryManager
+    OMPT
 )
 
 # Define the TARGET_NAME and DEBUG_PREFIX.
diff --git a/openmp/libomptarget/plugins-nextgen/cuda/CMakeLists.txt b/openmp/libomptarget/plugins-nextgen/cuda/CMakeLists.txt
index 397b06b67692..0ae1ddff5541 100644
--- a/openmp/libomptarget/plugins-nextgen/cuda/CMakeLists.txt
+++ b/openmp/libomptarget/plugins-nextgen/cuda/CMakeLists.txt
@@ -36,6 +36,7 @@ add_llvm_library(omptarget.rtl.cuda.nextgen SHARED
   LINK_LIBS PRIVATE
   elf_common
   MemoryManager
+  OMPT
   PluginInterface
   ${OPENMP_PTHREAD_LIB}
 
diff --git a/openmp/libomptarget/src/CMakeLists.txt b/openmp/libomptarget/src/CMakeLists.txt
index ef3a6270b78a..425121c9ef10 100644
--- a/openmp/libomptarget/src/CMakeLists.txt
+++ b/openmp/libomptarget/src/CMakeLists.txt
@@ -20,7 +20,7 @@ add_llvm_library(omptarget
   interface.cpp
   interop.cpp
   omptarget.cpp
-  ompt_callback.cpp
+  OmptCallback.cpp
   rtl.cpp
   LegacyAPI.cpp
 
diff --git a/openmp/libomptarget/src/ompt_callback.cpp b/openmp/libomptarget/src/OmptCallback.cpp
index 5715642ad3a9..2cfa487b123f 100644
--- a/openmp/libomptarget/src/ompt_callback.cpp
+++ b/openmp/libomptarget/src/OmptCallback.cpp
@@ -1,4 +1,4 @@
-//===-- ompt_callback.cpp - Target independent OpenMP target RTL -- C++ -*-===//
+//===-- OmptCallback.cpp - Target independent OpenMP target RTL --- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -19,6 +19,7 @@
 
 #include "omp-tools.h"
 
+#include "Debug.h"
 #include "ompt_connector.h"
 #include "ompt_device_callbacks.h"
 #include "private.h"
@@ -26,10 +27,39 @@
 #define fnptr_to_ptr(x) ((void *)(uint64_t)x)
 
 /// Used to indicate whether OMPT was enabled for this library
-bool ompt_enabled = false;
+bool OmptEnabled = false;
 /// Object maintaining all the callbacks for this library
 OmptDeviceCallbacksTy OmptDeviceCallbacks;
 
+/// Used to maintain the finalization function that is received
+/// from the plugin during connect
+class LibomptargetRtlFinalizer {
+public:
+  LibomptargetRtlFinalizer() : RtlFinalization(nullptr) {}
+  void registerRtl(ompt_finalize_t FinalizationFunction) {
+    assert((RtlFinalization == nullptr) &&
+           "RTL finalization may only be registered once");
+    RtlFinalization = FinalizationFunction;
+  }
+  void finalize() {
+    if (RtlFinalization)
+      RtlFinalization(nullptr /* tool_data */);
+    RtlFinalization = nullptr;
+  }
+
+private:
+  ompt_finalize_t RtlFinalization;
+};
+
+/// Object that will maintain the RTL finalizer from the plugin
+static LibomptargetRtlFinalizer LibraryFinalizer;
+
+/// Lookup function to be used by libomptarget library
+ompt_interface_fn_t
+OmptDeviceCallbacksTy::doLookup(const char *InterfaceFunctionName) {
+  return OmptDeviceCallbacks.lookupCallback(InterfaceFunctionName);
+}
+
 /// This is the function called by the higher layer (libomp) responsible
 /// for initializing OMPT in this library. This is passed to libomp
 /// as part of the OMPT connector object.
@@ -40,7 +70,7 @@ static int ompt_libomptarget_initialize(ompt_function_lookup_t lookup,
                                         int initial_device_num,
                                         ompt_data_t *tool_data) {
   DP("enter ompt_libomptarget_initialize!\n");
-  ompt_enabled = true;
+  OmptEnabled = true;
   // The lookup parameter is provided by libomp which already has the
   // tool callbacks registered at this point. The registration call
   // below causes the same callback functions to be registered in
@@ -50,9 +80,14 @@ static int ompt_libomptarget_initialize(ompt_function_lookup_t lookup,
   return 0;
 }
 
+/// This function is passed to libomp as part of the OMPT connector object.
+/// It is called by libomp during finalization of OMPT in libomptarget.
 static void ompt_libomptarget_finalize(ompt_data_t *data) {
   DP("enter ompt_libomptarget_finalize!\n");
-  ompt_enabled = false;
+  // Before disabling OMPT, call the finalizer (of the plugin) that was
+  // registered with this library
+  LibraryFinalizer.finalize();
+  OmptEnabled = false;
   DP("exit ompt_libomptarget_finalize!\n");
 }
 
@@ -61,11 +96,9 @@ static void ompt_libomptarget_finalize(ompt_data_t *data) {
  *****************************************************************************/
 /// Used to initialize callbacks implemented by the tool. This interface
 /// will lookup the callbacks table in libomp and assign them to the callbacks
-/// maintained in libomptarget. Using priority 102 to have this constructor
-/// run after the init target library constructor with priority 101 (see
-/// rtl.cpp).
-__attribute__((constructor(102))) static void ompt_init(void) {
-  DP("OMPT: Enter ompt_init\n");
+/// maintained in libomptarget.
+void InitOmptLibomp() {
+  DP("OMPT: Enter InitOmptLibomp\n");
   // Connect with libomp
   static OmptLibraryConnectorTy LibompConnector("libomp");
   static ompt_start_tool_result_t OmptResult;
@@ -81,7 +114,24 @@ __attribute__((constructor(102))) static void ompt_init(void) {
 
   // Now call connect that causes the above init/fini functions to be called
   LibompConnector.connect(&OmptResult);
-  DP("OMPT: Exit ompt_init\n");
+  DP("OMPT: Exit InitOmptLibomp\n");
 }
 
 #endif // OMPT_SUPPORT
+
+extern "C" {
+/// Used for connecting libomptarget with a plugin
+void ompt_libomptarget_connect(ompt_start_tool_result_t *result) {
+  DP("OMPT: Enter ompt_libomptarget_connect\n");
+  if (OmptEnabled && result) {
+    // Cache the fini function so that it can be invoked on exit
+    LibraryFinalizer.registerRtl(result->finalize);
+    // Invoke the provided init function with the lookup function maintained
+    // in this library so that callbacks maintained by this library are
+    // retrieved.
+    result->initialize(OmptDeviceCallbacksTy::doLookup,
+                       0 /* initial_device_num */, nullptr /* tool_data */);
+  }
+  DP("OMPT: Leave ompt_libomptarget_connect\n");
+}
+}
diff --git a/openmp/libomptarget/src/exports b/openmp/libomptarget/src/exports
index 6c3fdf0950ab..48591dd6c3fa 100644
--- a/openmp/libomptarget/src/exports
+++ b/openmp/libomptarget/src/exports
@@ -64,6 +64,7 @@ VERS1.0 {
     __tgt_interop_init;
     __tgt_interop_use;
     __tgt_interop_destroy;
+    ompt_libomptarget_connect;
   local:
     *;
 };
diff --git a/openmp/libomptarget/src/rtl.cpp b/openmp/libomptarget/src/rtl.cpp
index 9c7cc355d054..5360d9784cb6 100644
--- a/openmp/libomptarget/src/rtl.cpp
+++ b/openmp/libomptarget/src/rtl.cpp
@@ -43,6 +43,10 @@ PluginManager *PM;
 
 static char *ProfileTraceFile = nullptr;
 
+#ifdef OMPT_SUPPORT
+extern void InitOmptLibomp();
+#endif
+
 __attribute__((constructor(101))) void init() {
   DP("Init target library!\n");
 
@@ -65,6 +69,11 @@ __attribute__((constructor(101))) void init() {
   if (ProfileTraceFile)
     timeTraceProfilerInitialize(500 /* us */, "libomptarget");
 
+  #ifdef OMPT_SUPPORT
+    // Initialize OMPT first
+    InitOmptLibomp();
+  #endif
+
   PM->RTLs.loadRTLs();
   PM->registerDelayedLibraries();
 }