[OpenMP] [OMPT] [amdgpu] [5/8] Implemented device init/fini/load callbacks

Added support in the generic plugin to invoke registered callbacks.

Depends on D124070

Patch from John Mellor-Crummey <johnmc@rice.edu>
(With contributions from Dhruva Chakrabarti <Dhruva.Chakrabarti@amd.com>)

Differential Revision: https://reviews.llvm.org/D124652

20 files changed, 319 insertions, 272 deletions

diff --git a/openmp/libomptarget/CMakeLists.txt b/openmp/libomptarget/CMakeLists.txt
index 50c59cb..844107b 100644
--- a/openmp/libomptarget/CMakeLists.txt
+++ b/openmp/libomptarget/CMakeLists.txt
@@ -83,6 +83,7 @@ if ((OMPT_TARGET_DEFAULT) AND (LIBOMPTARGET_OMPT_SUPPORT))
   add_definitions(-DOMPT_SUPPORT=1)
   message(STATUS "OMPT target enabled")
 else()
+  set(LIBOMPTARGET_OMPT_SUPPORT FALSE)
   message(STATUS "OMPT target disabled")
 endif()
 
diff --git a/openmp/libomptarget/include/OmptCallback.h b/openmp/libomptarget/include/OmptCallback.h
new file mode 100644
index 0000000..1393df9d
--- /dev/null
+++ b/openmp/libomptarget/include/OmptCallback.h
@@ -0,0 +1,85 @@
+//===---- OmptCallback.h - Target independent 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Interface used by target-independent runtimes to coordinate registration and
+// invocation of OMPT callbacks and initialization / finalization.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _OMPTCALLBACK_H
+#define _OMPTCALLBACK_H
+
+#ifdef OMPT_SUPPORT
+
+#include "omp-tools.h"
+
+#define DEBUG_PREFIX "OMPT"
+
+#define FOREACH_OMPT_TARGET_CALLBACK(macro)                                    \
+  FOREACH_OMPT_DEVICE_EVENT(macro)                                             \
+  FOREACH_OMPT_NOEMI_EVENT(macro)                                              \
+  FOREACH_OMPT_EMI_EVENT(macro)
+
+#define performOmptCallback(CallbackName, ...)                                 \
+  do {                                                                         \
+    if (ompt_callback_##CallbackName##_fn)                                     \
+      ompt_callback_##CallbackName##_fn(__VA_ARGS__);                          \
+  } while (0)
+
+namespace llvm {
+namespace omp {
+namespace target {
+namespace ompt {
+
+#define declareOmptCallback(Name, Type, Code) extern Name##_t Name##_fn;
+FOREACH_OMPT_NOEMI_EVENT(declareOmptCallback)
+FOREACH_OMPT_EMI_EVENT(declareOmptCallback)
+#undef declareOmptCallback
+
+/// This function will call an OpenMP API function. Which in turn will lookup a
+/// given enum value of type \p ompt_callbacks_t and copy the address of the
+/// corresponding callback funtion into the provided pointer.
+/// The pointer to the runtime function is passed during 'initializeLibrary'.
+/// \p which the enum value of the requested callback function
+/// \p callback the destination pointer where the address shall be copied
+extern ompt_get_callback_t lookupCallbackByCode;
+
+/// Lookup function to be used by the lower layer (e.g. the plugin). This
+/// function has to be provided when actually calling callback functions like
+/// 'ompt_callback_device_initialize_fn' (param: 'lookup').
+/// The pointer to the runtime function is passed during 'initializeLibrary'.
+/// \p InterfaceFunctionName the name of the OMPT callback function to look up
+extern ompt_function_lookup_t lookupCallbackByName;
+
+/// This is the function called by the higher layer (libomp / libomtarget)
+/// responsible for initializing OMPT in this library. This is passed to libomp
+/// as part of the OMPT connector object.
+/// \p lookup to be used to query callbacks registered with libomp
+/// \p initial_device_num initial device num (id) provided by libomp
+/// \p tool_data as provided by the tool
+int initializeLibrary(ompt_function_lookup_t lookup, int initial_device_num,
+                      ompt_data_t *tool_data);
+
+/// This function is passed to libomp / libomtarget as part of the OMPT
+/// connector object. It is called by libomp during finalization of OMPT in
+/// libomptarget -OR- by libomptarget during finalization of OMPT in the plugin.
+/// \p tool_data as provided by the tool
+void finalizeLibrary(ompt_data_t *tool_data);
+
+/// This function will connect the \p initializeLibrary and \p finalizeLibrary
+/// functions to their respective higher layer.
+void connectLibrary();
+
+} // namespace ompt
+} // namespace target
+} // namespace omp
+} // namespace llvm
+
+#endif // OMPT_SUPPORT
+
+#endif // _OMPTCALLBACK_H
diff --git a/openmp/libomptarget/include/ompt_connector.h b/openmp/libomptarget/include/OmptConnector.h
index 332c3a0..c0191a2 100644
--- a/openmp/libomptarget/include/ompt_connector.h
+++ b/openmp/libomptarget/include/OmptConnector.h
@@ -1,4 +1,4 @@
-//=== ompt_connector.h - Target independent OpenMP target RTL -- C++ ------===//
+//===- OmptConnector.h - 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.
@@ -12,8 +12,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef _OMPT_CONNECTOR_H
-#define _OMPT_CONNECTOR_H
+#ifndef _OMPTCONNECTOR_H
+#define _OMPTCONNECTOR_H
 
 #ifdef OMPT_SUPPORT
 
@@ -101,8 +101,6 @@ private:
   std::string LibIdent;
 };
 
-#undef DEBUG_PREFIX
-
 #endif // OMPT_SUPPORT
 
-#endif // _OMPT_CONNECTOR_H
+#endif // _OMPTCONNECTOR_H
diff --git a/openmp/libomptarget/include/ompt_device_callbacks.h b/openmp/libomptarget/include/ompt_device_callbacks.h
deleted file mode 100644
index 127e189..0000000
--- a/openmp/libomptarget/include/ompt_device_callbacks.h
+++ /dev/null
@@ -1,81 +0,0 @@
-//===--------- ompt_device_callbacks.h - 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
-//
-//===----------------------------------------------------------------------===//
-//
-// Interface used by both target-independent and device-dependent runtimes
-// to coordinate registration and invocation of OMPT callbacks
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef _OMPT_DEVICE_CALLBACKS_H
-#define _OMPT_DEVICE_CALLBACKS_H
-
-#ifdef OMPT_SUPPORT
-
-#include "Debug.h"
-#include <omp-tools.h>
-
-#define DEBUG_PREFIX "OMPT"
-
-#define FOREACH_OMPT_TARGET_CALLBACK(macro)                                    \
-  FOREACH_OMPT_DEVICE_EVENT(macro)                                             \
-  FOREACH_OMPT_NOEMI_EVENT(macro)                                              \
-  FOREACH_OMPT_EMI_EVENT(macro)
-
-/// Internal representation for OMPT device callback functions.
-class OmptDeviceCallbacksTy {
-public:
-  /// Initialize the enabled flag and all the callbacks
-  void init() {
-    Enabled = false;
-#define initName(Name, Type, Code) Name##_fn = 0;
-    FOREACH_OMPT_TARGET_CALLBACK(initName)
-#undef initName
-  }
-
-  /// Used to register callbacks. \p Lookup is used to query a given callback
-  /// by name and the result is assigned to the corresponding callback function.
-  void registerCallbacks(ompt_function_lookup_t Lookup) {
-    Enabled = true;
-#define OmptBindCallback(Name, Type, Code)                                     \
-  Name##_fn = (Name##_t)Lookup(#Name);                                         \
-  DP("OMPT: class bound %s=%p\n", #Name, ((void *)(uint64_t)Name##_fn));
-
-    FOREACH_OMPT_TARGET_CALLBACK(OmptBindCallback);
-#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;
-
-  /// Callback functions
-#define DeclareName(Name, Type, Code) Name##_t Name##_fn;
-  FOREACH_OMPT_TARGET_CALLBACK(DeclareName)
-#undef DeclareName
-};
-
-/// Device callbacks object for the library that performs the instantiation
-extern OmptDeviceCallbacksTy OmptDeviceCallbacks;
-
-#endif // OMPT_SUPPORT
-
-#endif // _OMPT_DEVICE_CALLBACKS_H
diff --git a/openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp b/openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp
index 45c542e..d56d3c9 100644
--- a/openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -23,6 +23,7 @@
 #include "Debug.h"
 #include "DeviceEnvironment.h"
 #include "GlobalHandler.h"
+#include "OmptCallback.h"
 #include "PluginInterface.h"
 #include "Utilities.h"
 #include "UtilitiesRTL.h"
@@ -2623,6 +2624,10 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
     // HSA functions from now on, e.g., hsa_shut_down.
     Initialized = true;
 
+#ifdef OMPT_SUPPORT
+    ompt::connectLibrary();
+#endif
+
     // Register event handler to detect memory errors on the devices.
     Status = hsa_amd_register_system_event_handler(eventHandler, nullptr);
     if (auto Err = Plugin::check(
diff --git a/openmp/libomptarget/plugins-nextgen/common/OMPT/OmptCallback.cpp b/openmp/libomptarget/plugins-nextgen/common/OMPT/OmptCallback.cpp
index eb47f10..30d07d3 100644
--- a/openmp/libomptarget/plugins-nextgen/common/OMPT/OmptCallback.cpp
+++ b/openmp/libomptarget/plugins-nextgen/common/OMPT/OmptCallback.cpp
@@ -11,73 +11,60 @@
 //===----------------------------------------------------------------------===//
 
 #ifdef OMPT_SUPPORT
-#include <atomic>
-#include <cstdio>
-#include <string.h>
-#include <vector>
+
+#include "llvm/Support/DynamicLibrary.h"
+
+#include <cstdlib>
+#include <cstring>
+#include <memory>
 
 #include "Debug.h"
-#include "ompt_connector.h"
-#include "ompt_device_callbacks.h"
+#include "OmptCallback.h"
+#include "OmptConnector.h"
 
-/// Object maintaining all the callbacks in the plugin
-OmptDeviceCallbacksTy OmptDeviceCallbacks;
+using namespace llvm::omp::target::ompt;
 
-/// 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;
-}
+ompt_get_callback_t llvm::omp::target::ompt::lookupCallbackByCode = nullptr;
+ompt_function_lookup_t llvm::omp::target::ompt::lookupCallbackByName = nullptr;
+
+int llvm::omp::target::ompt::initializeLibrary(ompt_function_lookup_t lookup,
+                                               int initial_device_num,
+                                               ompt_data_t *tool_data) {
+  DP("OMPT: Executing initializeLibrary (libomptarget)\n");
+#define bindOmptFunctionName(OmptFunction, DestinationFunction)                \
+  if (lookup)                                                                  \
+    DestinationFunction = (OmptFunction##_t)lookup(#OmptFunction);             \
+  DP("OMPT: initializeLibrary (libomptarget) bound %s=%p\n",                   \
+     #DestinationFunction, ((void *)(uint64_t)DestinationFunction));
 
-/// Used to indicate whether OMPT was enabled for this library
-static bool OmptEnabled = false;
+  bindOmptFunctionName(ompt_get_callback, lookupCallbackByCode);
+#undef bindOmptFunctionName
+
+  // Store pointer of 'ompt_libomp_target_fn_lookup' for use by the plugin
+  lookupCallbackByName = lookup;
 
-/// 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");
+void llvm::omp::target::ompt::finalizeLibrary(ompt_data_t *tool_data) {
+  DP("OMPT: Executing finalizeLibrary (libomptarget)\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");
+void llvm::omp::target::ompt::connectLibrary() {
+  DP("OMPT: Entering connectLibrary (libomptarget)\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.initialize = ompt::initializeLibrary;
+  OmptResult.finalize = ompt::finalizeLibrary;
   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");
+  DP("OMPT: Exiting connectLibrary (libomptarget)\n");
 }
+
 #endif
diff --git a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp
index 35c1159..79a968c 100644
--- a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp
+++ b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp
@@ -16,6 +16,11 @@
 #include "omptarget.h"
 #include "omptargetplugin.h"
 
+#ifdef OMPT_SUPPORT
+#include "OmptCallback.h"
+#include "omp-tools.h"
+#endif
+
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/JSON.h"
@@ -396,12 +401,37 @@ GenericDeviceTy::GenericDeviceTy(int32_t DeviceId, int32_t NumDevices,
       OMPX_InitialNumEvents("LIBOMPTARGET_NUM_INITIAL_EVENTS", 32),
       DeviceId(DeviceId), GridValues(OMPGridValues),
       PeerAccesses(NumDevices, PeerAccessState::PENDING), PeerAccessesLock(),
-      PinnedAllocs(*this), RPCHandle(nullptr) {}
+      PinnedAllocs(*this), RPCHandle(nullptr) {
+#ifdef OMPT_SUPPORT
+  OmptInitialized.store(false);
+  // Bind the callbacks to this device's member functions
+#define bindOmptCallback(Name, Type, Code)                                     \
+  if (ompt::lookupCallbackByCode)                                              \
+    ompt::lookupCallbackByCode((ompt_callbacks_t)(Code),                       \
+                               ((ompt_callback_t *)&(Name##_fn)));             \
+  DP("OMPT: class bound %s=%p\n", #Name, ((void *)(uint64_t)Name##_fn));
+
+  FOREACH_OMPT_DEVICE_EVENT(bindOmptCallback);
+#undef bindOmptCallback
+
+#endif
+}
 
 Error GenericDeviceTy::init(GenericPluginTy &Plugin) {
   if (auto Err = initImpl(Plugin))
     return Err;
 
+#ifdef OMPT_SUPPORT
+  bool ExpectedStatus = false;
+  if (OmptInitialized.compare_exchange_strong(ExpectedStatus, true))
+    performOmptCallback(device_initialize,
+                        /* device_num */ DeviceId,
+                        /* type */ getComputeUnitKind().c_str(),
+                        /* device */ reinterpret_cast<ompt_device_t *>(this),
+                        /* lookup */ ompt::lookupCallbackByName,
+                        /* documentation */ nullptr);
+#endif
+
   // Read and reinitialize the envars that depend on the device initialization.
   // Notice these two envars may change the stack size and heap size of the
   // device, so they need the device properly initialized.
@@ -444,7 +474,7 @@ Error GenericDeviceTy::init(GenericPluginTy &Plugin) {
 }
 
 Error GenericDeviceTy::deinit() {
-  // Delete the memory manager before deinitilizing the device. Otherwise,
+  // Delete the memory manager before deinitializing the device. Otherwise,
   // we may delete device allocations after the device is deinitialized.
   if (MemoryManager)
     delete MemoryManager;
@@ -457,9 +487,14 @@ Error GenericDeviceTy::deinit() {
     if (auto Err = RPCHandle->deinitDevice())
       return Err;
 
+#ifdef OMPT_SUPPORT
+  bool ExpectedStatus = true;
+  if (OmptInitialized.compare_exchange_strong(ExpectedStatus, false))
+    performOmptCallback(device_finalize, /* device_num */ DeviceId);
+#endif
+
   return deinitImpl();
 }
-
 Expected<__tgt_target_table *>
 GenericDeviceTy::loadBinary(GenericPluginTy &Plugin,
                             const __tgt_device_image *InputTgtImage) {
@@ -500,6 +535,19 @@ GenericDeviceTy::loadBinary(GenericPluginTy &Plugin,
   if (auto Err = setupRPCServer(Plugin, *Image))
     return std::move(Err);
 
+#ifdef OMPT_SUPPORT
+  size_t Bytes = getPtrDiff(InputTgtImage->ImageEnd, InputTgtImage->ImageStart);
+  performOmptCallback(device_load,
+                      /* device_num */ DeviceId,
+                      /* FileName */ nullptr,
+                      /* File Offset */ 0,
+                      /* VmaInFile */ nullptr,
+                      /* ImgSize */ Bytes,
+                      /* HostAddr */ InputTgtImage->ImageStart,
+                      /* DeviceAddr */ nullptr,
+                      /* FIXME: ModuleId */ 0);
+#endif
+
   // Return the pointer to the table of entries.
   return Image->getOffloadEntryTable();
 }
diff --git a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h
index f902d4e..ea49234 100644
--- a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h
+++ b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h
@@ -28,6 +28,10 @@
 #include "Utilities.h"
 #include "omptarget.h"
 
+#ifdef OMPT_SUPPORT
+#include "omp-tools.h"
+#endif
+
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include "llvm/Frontend/OpenMP/OMPGridValues.h"
@@ -854,6 +858,16 @@ protected:
   /// A pointer to an RPC server instance attached to this device if present.
   /// This is used to run the RPC server during task synchronization.
   RPCHandleTy *RPCHandle;
+
+#ifdef OMPT_SUPPORT
+  /// OMPT callback functions
+#define defineOmptCallback(Name, Type, Code) Name##_t Name##_fn = nullptr;
+  FOREACH_OMPT_DEVICE_EVENT(defineOmptCallback)
+#undef defineOmptCallback
+
+  /// Internal representation for OMPT device (initialize & finalize)
+  std::atomic<bool> OmptInitialized;
+#endif
 };
 
 /// Class implementing common functionalities of offload plugins. Each plugin
diff --git a/openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp b/openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp
index fbd89da..4ae9fae 100644
--- a/openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp
+++ b/openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp
@@ -19,6 +19,7 @@
 #include "Debug.h"
 #include "DeviceEnvironment.h"
 #include "GlobalHandler.h"
+#include "OmptCallback.h"
 #include "PluginInterface.h"
 
 #include "llvm/BinaryFormat/ELF.h"
@@ -946,6 +947,10 @@ struct CUDAPluginTy final : public GenericPluginTy {
       return 0;
     }
 
+#ifdef OMPT_SUPPORT
+    ompt::connectLibrary();
+#endif
+
     if (Res == CUDA_ERROR_NO_DEVICE) {
       // Do not initialize if there are no devices.
       DP("There are no devices supporting CUDA.\n");
diff --git a/openmp/libomptarget/src/OmptCallback.cpp b/openmp/libomptarget/src/OmptCallback.cpp
index 885757e..cde31023 100644
--- a/openmp/libomptarget/src/OmptCallback.cpp
+++ b/openmp/libomptarget/src/OmptCallback.cpp
@@ -12,123 +12,127 @@
 
 #ifdef OMPT_SUPPORT
 
-#include <assert.h>
-#include <atomic>
+#include "llvm/Support/DynamicLibrary.h"
+
 #include <cstdlib>
 #include <cstring>
-
-#include "omp-tools.h"
+#include <memory>
 
 #include "Debug.h"
-#include "ompt_connector.h"
-#include "ompt_device_callbacks.h"
-#include "private.h"
-
-#define fnptr_to_ptr(x) ((void *)(uint64_t)x)
-
-/// Used to indicate whether OMPT was enabled for this library
-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
+#include "OmptCallback.h"
+#include "OmptConnector.h"
+
+using namespace llvm::omp::target::ompt;
+
+// Define OMPT callback functions (bound to actual callbacks later on)
+#define defineOmptCallback(Name, Type, Code)                                   \
+  Name##_t llvm::omp::target::ompt::Name##_fn = nullptr;
+FOREACH_OMPT_NOEMI_EVENT(defineOmptCallback)
+FOREACH_OMPT_EMI_EVENT(defineOmptCallback)
+#undef defineOmptCallback
+
+/// Used to maintain the finalization functions that are received
+/// from the plugins during connect.
+/// Note: Currently, there are no plugin-specific finalizations, so each plugin
+/// will call the same (empty) function.
 class LibomptargetRtlFinalizer {
 public:
-  LibomptargetRtlFinalizer() : RtlFinalization(nullptr) {}
+  LibomptargetRtlFinalizer() {}
+
   void registerRtl(ompt_finalize_t FinalizationFunction) {
-    assert((RtlFinalization == nullptr) &&
-           "RTL finalization may only be registered once");
-    RtlFinalization = FinalizationFunction;
+    if (FinalizationFunction) {
+      RtlFinalizationFunctions.emplace_back(FinalizationFunction);
+    }
   }
+
   void finalize() {
-    if (RtlFinalization)
-      RtlFinalization(nullptr /* tool_data */);
-    RtlFinalization = nullptr;
+    for (auto FinalizationFunction : RtlFinalizationFunctions)
+      FinalizationFunction(/* tool_data */ nullptr);
+    RtlFinalizationFunctions.clear();
   }
 
 private:
-  ompt_finalize_t RtlFinalization;
+  llvm::SmallVector<ompt_finalize_t> RtlFinalizationFunctions;
 };
 
 /// Object that will maintain the RTL finalizer from the plugin
-static LibomptargetRtlFinalizer LibraryFinalizer;
+LibomptargetRtlFinalizer *LibraryFinalizer = nullptr;
 
-/// Lookup function to be used by libomptarget library
-ompt_interface_fn_t
-OmptDeviceCallbacksTy::doLookup(const char *InterfaceFunctionName) {
-  return OmptDeviceCallbacks.lookupCallback(InterfaceFunctionName);
-}
+ompt_get_callback_t llvm::omp::target::ompt::lookupCallbackByCode = nullptr;
+ompt_function_lookup_t llvm::omp::target::ompt::lookupCallbackByName = nullptr;
+
+int llvm::omp::target::ompt::initializeLibrary(ompt_function_lookup_t lookup,
+                                               int initial_device_num,
+                                               ompt_data_t *tool_data) {
+  DP("OMPT: Executing initializeLibrary (libomp)\n");
+#define bindOmptFunctionName(OmptFunction, DestinationFunction)                \
+  DestinationFunction = (OmptFunction##_t)lookup(#OmptFunction);               \
+  DP("OMPT: initializeLibrary (libomp) bound %s=%p\n", #DestinationFunction,   \
+     ((void *)(uint64_t)DestinationFunction));
+
+  bindOmptFunctionName(ompt_get_callback, lookupCallbackByCode);
+#undef bindOmptFunctionName
+
+  // Store pointer of 'ompt_libomp_target_fn_lookup' for use by libomptarget
+  lookupCallbackByName = lookup;
+
+  assert(lookupCallbackByCode && "lookupCallbackByCode should be non-null");
+  assert(lookupCallbackByName && "lookupCallbackByName should be non-null");
+  assert(LibraryFinalizer == nullptr &&
+         "LibraryFinalizer should not be initialized yet");
+
+  LibraryFinalizer = new LibomptargetRtlFinalizer();
 
-/// 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.
-/// \p lookup to be used to query callbacks registered with libomp
-/// \p initial_device_num Initial device num provided by libomp
-/// \p tool_data as provided by the tool
-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");
-  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
-  // libomptarget as well
-  OmptDeviceCallbacks.registerCallbacks(lookup);
-  DP("exit ompt_libomptarget_initialize!\n");
   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");
-  // 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");
+void llvm::omp::target::ompt::finalizeLibrary(ompt_data_t *data) {
+  DP("OMPT: Executing finalizeLibrary (libomp)\n");
+  // Before disabling OMPT, call the (plugin) finalizations that were registered
+  // with this library
+  LibraryFinalizer->finalize();
+  delete LibraryFinalizer;
 }
 
-/*****************************************************************************
- * constructor
- *****************************************************************************/
-/// 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.
-void InitOmptLibomp() {
-  DP("OMPT: Enter InitOmptLibomp\n");
+void llvm::omp::target::ompt::connectLibrary() {
+  DP("OMPT: Entering connectLibrary (libomp)\n");
   // Connect with libomp
   static OmptLibraryConnectorTy LibompConnector("libomp");
   static ompt_start_tool_result_t OmptResult;
 
   // Initialize OmptResult with the init and fini functions that will be
   // called by the connector
-  OmptResult.initialize = ompt_libomptarget_initialize;
-  OmptResult.finalize = ompt_libomptarget_finalize;
+  OmptResult.initialize = ompt::initializeLibrary;
+  OmptResult.finalize = ompt::finalizeLibrary;
   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
   LibompConnector.connect(&OmptResult);
-  DP("OMPT: Exit InitOmptLibomp\n");
+
+#define bindOmptCallback(Name, Type, Code)                                     \
+  if (lookupCallbackByCode)                                                    \
+    lookupCallbackByCode(                                                      \
+        (ompt_callbacks_t)(Code),                                              \
+        (ompt_callback_t *)&(llvm::omp::target::ompt::Name##_fn));
+  FOREACH_OMPT_NOEMI_EVENT(bindOmptCallback)
+  FOREACH_OMPT_EMI_EVENT(bindOmptCallback)
+#undef bindOmptCallback
+
+  DP("OMPT: Exiting connectLibrary (libomp)\n");
 }
 
 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);
+  if (result && LibraryFinalizer) {
+    // Cache each fini function, so that they 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 */);
+    result->initialize(lookupCallbackByName,
+                       /* initial_device_num */ 0, /* tool_data */ nullptr);
   }
   DP("OMPT: Leave ompt_libomptarget_connect\n");
 }
diff --git a/openmp/libomptarget/src/rtl.cpp b/openmp/libomptarget/src/rtl.cpp
index 6fe7b4e9..6a73f6e 100644
--- a/openmp/libomptarget/src/rtl.cpp
+++ b/openmp/libomptarget/src/rtl.cpp
@@ -12,6 +12,7 @@
 
 #include "llvm/Object/OffloadBinary.h"
 
+#include "OmptCallback.h"
 #include "device.h"
 #include "private.h"
 #include "rtl.h"
@@ -42,7 +43,7 @@ PluginManager *PM;
 static char *ProfileTraceFile = nullptr;
 
 #ifdef OMPT_SUPPORT
-extern void InitOmptLibomp();
+extern void ompt::connectLibrary();
 #endif
 
 __attribute__((constructor(101))) void init() {
@@ -67,10 +68,10 @@ __attribute__((constructor(101))) void init() {
   if (ProfileTraceFile)
     timeTraceProfilerInitialize(500 /* us */, "libomptarget");
 
-  #ifdef OMPT_SUPPORT
-    // Initialize OMPT first
-    InitOmptLibomp();
-  #endif
+#ifdef OMPT_SUPPORT
+  // Initialize OMPT first
+  ompt::connectLibrary();
+#endif
 
   PM->RTLs.loadRTLs();
   PM->registerDelayedLibraries();
diff --git a/openmp/libomptarget/test/ompt/veccopy.c b/openmp/libomptarget/test/ompt/veccopy.c
index cb8db1b..f7668b2 100644
--- a/openmp/libomptarget/test/ompt/veccopy.c
+++ b/openmp/libomptarget/test/ompt/veccopy.c
@@ -1,8 +1,5 @@
 // RUN: %libomptarget-compile-run-and-check-generic
 // REQUIRES: ompt
-// UNSUPPORTED: nvptx64-nvidia-cuda
-// UNSUPPORTED: nvptx64-nvidia-cuda-oldDriver
-// UNSUPPORTED: nvptx64-nvidia-cuda-LTO
 // UNSUPPORTED: x86_64-pc-linux-gnu
 // UNSUPPORTED: x86_64-pc-linux-gnu-oldDriver
 // UNSUPPORTED: x86_64-pc-linux-gnu-LTO
@@ -56,11 +53,10 @@ int main() {
   return rc;
 }
 
-/// CHECK: Could not register callback 'ompt_callback_device_initialize'
-/// CHECK: Could not register callback 'ompt_callback_device_finalize'
-/// CHECK: Could not register callback 'ompt_callback_device_load'
 /// CHECK: Could not register callback 'ompt_callback_target_data_op'
 /// CHECK: Could not register callback 'ompt_callback_target'
 /// CHECK: Could not register callback 'ompt_callback_target_submit'
 
-/// CHECK: Success
+/// CHECK: Callback Init:
+/// CHECK: Callback Load:
+/// CHECK: Callback Fini:
diff --git a/openmp/libomptarget/test/ompt/veccopy_disallow_both.c b/openmp/libomptarget/test/ompt/veccopy_disallow_both.c
index 85f95ef..69f6e90 100644
--- a/openmp/libomptarget/test/ompt/veccopy_disallow_both.c
+++ b/openmp/libomptarget/test/ompt/veccopy_disallow_both.c
@@ -1,8 +1,5 @@
 // RUN: %libomptarget-compile-run-and-check-generic
 // REQUIRES: ompt
-// UNSUPPORTED: nvptx64-nvidia-cuda
-// UNSUPPORTED: nvptx64-nvidia-cuda-oldDriver
-// UNSUPPORTED: nvptx64-nvidia-cuda-LTO
 // UNSUPPORTED: x86_64-pc-linux-gnu
 // UNSUPPORTED: x86_64-pc-linux-gnu-oldDriver
 // UNSUPPORTED: x86_64-pc-linux-gnu-LTO
@@ -59,11 +56,10 @@ int main() {
   return rc;
 }
 
-/// CHECK: Could not register callback 'ompt_callback_device_initialize'
-/// CHECK: Could not register callback 'ompt_callback_device_finalize'
-/// CHECK: Could not register callback 'ompt_callback_device_load'
 /// CHECK: Could not register callback 'ompt_callback_target_data_op'
 /// CHECK: Could not register callback 'ompt_callback_target'
 /// CHECK: Could not register callback 'ompt_callback_target_submit'
 
-/// CHECK: Success
+/// CHECK: Callback Init:
+/// CHECK: Callback Load:
+/// CHECK: Callback Fini:
diff --git a/openmp/libomptarget/test/ompt/veccopy_emi.c b/openmp/libomptarget/test/ompt/veccopy_emi.c
index 5e041a2..df9dd14 100644
--- a/openmp/libomptarget/test/ompt/veccopy_emi.c
+++ b/openmp/libomptarget/test/ompt/veccopy_emi.c
@@ -1,8 +1,5 @@
 // RUN: %libomptarget-compile-run-and-check-generic
 // REQUIRES: ompt
-// UNSUPPORTED: nvptx64-nvidia-cuda
-// UNSUPPORTED: nvptx64-nvidia-cuda-oldDriver
-// UNSUPPORTED: nvptx64-nvidia-cuda-LTO
 // UNSUPPORTED: x86_64-pc-linux-gnu
 // UNSUPPORTED: x86_64-pc-linux-gnu-oldDriver
 // UNSUPPORTED: x86_64-pc-linux-gnu-LTO
@@ -57,11 +54,10 @@ int main() {
   return rc;
 }
 
-/// CHECK: Could not register callback 'ompt_callback_device_initialize'
-/// CHECK: Could not register callback 'ompt_callback_device_finalize'
-/// CHECK: Could not register callback 'ompt_callback_device_load'
 /// CHECK: Could not register callback 'ompt_callback_target_data_op_emi'
 /// CHECK: Could not register callback 'ompt_callback_target_emi'
 /// CHECK: Could not register callback 'ompt_callback_target_submit_emi'
 
-/// CHECK: Success
+/// CHECK: Callback Init:
+/// CHECK: Callback Load:
+/// CHECK: Callback Fini:
diff --git a/openmp/libomptarget/test/ompt/veccopy_emi_map.c b/openmp/libomptarget/test/ompt/veccopy_emi_map.c
index e237a4d..ebc310d 100644
--- a/openmp/libomptarget/test/ompt/veccopy_emi_map.c
+++ b/openmp/libomptarget/test/ompt/veccopy_emi_map.c
@@ -1,8 +1,5 @@
 // RUN: %libomptarget-compile-run-and-check-generic
 // REQUIRES: ompt
-// UNSUPPORTED: nvptx64-nvidia-cuda
-// UNSUPPORTED: nvptx64-nvidia-cuda-oldDriver
-// UNSUPPORTED: nvptx64-nvidia-cuda-LTO
 // UNSUPPORTED: x86_64-pc-linux-gnu
 // UNSUPPORTED: x86_64-pc-linux-gnu-oldDriver
 // UNSUPPORTED: x86_64-pc-linux-gnu-LTO
@@ -57,12 +54,11 @@ int main() {
   return rc;
 }
 
-/// CHECK: Could not register callback 'ompt_callback_device_initialize'
-/// CHECK: Could not register callback 'ompt_callback_device_finalize'
-/// CHECK: Could not register callback 'ompt_callback_device_load'
 /// CHECK: Could not register callback 'ompt_callback_target_data_op_emi'
 /// CHECK: Could not register callback 'ompt_callback_target_emi'
 /// CHECK: Could not register callback 'ompt_callback_target_submit_emi'
 /// CHECK: Could not register callback 'ompt_callback_target_map_emi'
 
-/// CHECK: Success
+/// CHECK: Callback Init:
+/// CHECK: Callback Load:
+/// CHECK: Callback Fini:
diff --git a/openmp/libomptarget/test/ompt/veccopy_map.c b/openmp/libomptarget/test/ompt/veccopy_map.c
index 977562b..83aa539 100644
--- a/openmp/libomptarget/test/ompt/veccopy_map.c
+++ b/openmp/libomptarget/test/ompt/veccopy_map.c
@@ -1,8 +1,5 @@
 // RUN: %libomptarget-compile-run-and-check-generic
 // REQUIRES: ompt
-// UNSUPPORTED: nvptx64-nvidia-cuda
-// UNSUPPORTED: nvptx64-nvidia-cuda-oldDriver
-// UNSUPPORTED: nvptx64-nvidia-cuda-LTO
 // UNSUPPORTED: x86_64-pc-linux-gnu
 // UNSUPPORTED: x86_64-pc-linux-gnu-oldDriver
 // UNSUPPORTED: x86_64-pc-linux-gnu-LTO
@@ -56,11 +53,10 @@ int main() {
   return rc;
 }
 
-/// CHECK: Could not register callback 'ompt_callback_device_initialize'
-/// CHECK: Could not register callback 'ompt_callback_device_finalize'
-/// CHECK: Could not register callback 'ompt_callback_device_load'
 /// CHECK: Could not register callback 'ompt_callback_target_data_op'
 /// CHECK: Could not register callback 'ompt_callback_target'
 /// CHECK: Could not register callback 'ompt_callback_target_submit'
 
-/// CHECK: Success
+/// CHECK: Callback Init:
+/// CHECK: Callback Load:
+/// CHECK: Callback Fini:
diff --git a/openmp/libomptarget/test/ompt/veccopy_no_device_init.c b/openmp/libomptarget/test/ompt/veccopy_no_device_init.c
index 165ba11..0bb9c46 100644
--- a/openmp/libomptarget/test/ompt/veccopy_no_device_init.c
+++ b/openmp/libomptarget/test/ompt/veccopy_no_device_init.c
@@ -1,8 +1,5 @@
 // RUN: %libomptarget-compile-run-and-check-generic
 // REQUIRES: ompt
-// UNSUPPORTED: nvptx64-nvidia-cuda
-// UNSUPPORTED: nvptx64-nvidia-cuda-oldDriver
-// UNSUPPORTED: nvptx64-nvidia-cuda-LTO
 // UNSUPPORTED: x86_64-pc-linux-gnu
 // UNSUPPORTED: x86_64-pc-linux-gnu-oldDriver
 // UNSUPPORTED: x86_64-pc-linux-gnu-LTO
@@ -57,9 +54,10 @@ int main() {
   return rc;
 }
 
-/// CHECK: Could not register callback 'ompt_callback_device_load'
 /// CHECK: Could not register callback 'ompt_callback_target_data_op'
 /// CHECK: Could not register callback 'ompt_callback_target'
 /// CHECK: Could not register callback 'ompt_callback_target_submit'
 
-/// CHECK: Success
+/// CHECK-NOT: Callback Init:
+/// CHECK-NOT: Callback Load:
+/// CHECK-NOT: Callback Fini:
diff --git a/openmp/libomptarget/test/ompt/veccopy_wrong_return.c b/openmp/libomptarget/test/ompt/veccopy_wrong_return.c
index c840539..eef2094 100644
--- a/openmp/libomptarget/test/ompt/veccopy_wrong_return.c
+++ b/openmp/libomptarget/test/ompt/veccopy_wrong_return.c
@@ -1,8 +1,5 @@
 // RUN: %libomptarget-compile-run-and-check-generic
 // REQUIRES: ompt
-// UNSUPPORTED: nvptx64-nvidia-cuda
-// UNSUPPORTED: nvptx64-nvidia-cuda-oldDriver
-// UNSUPPORTED: nvptx64-nvidia-cuda-LTO
 // UNSUPPORTED: x86_64-pc-linux-gnu
 // UNSUPPORTED: x86_64-pc-linux-gnu-oldDriver
 // UNSUPPORTED: x86_64-pc-linux-gnu-LTO
@@ -57,11 +54,10 @@ int main() {
   return rc;
 }
 
-/// CHECK: Could not register callback 'ompt_callback_device_initialize'
-/// CHECK: Could not register callback 'ompt_callback_device_finalize'
-/// CHECK: Could not register callback 'ompt_callback_device_load'
 /// CHECK: Could not register callback 'ompt_callback_target_data_op'
 /// CHECK: Could not register callback 'ompt_callback_target'
 /// CHECK: Could not register callback 'ompt_callback_target_submit'
 
-/// CHECK: Success
+/// CHECK-NOT: Callback Init:
+/// CHECK-NOT: Callback Load:
+/// CHECK-NOT: Callback Fini:
diff --git a/openmp/runtime/src/ompt-event-specific.h b/openmp/runtime/src/ompt-event-specific.h
index f6c7022..5ac7f6d 100644
--- a/openmp/runtime/src/ompt-event-specific.h
+++ b/openmp/runtime/src/ompt-event-specific.h
@@ -64,10 +64,10 @@
 
 #define ompt_callback_control_tool_implemented ompt_event_MAY_ALWAYS
 
-#define ompt_callback_device_initialize_implemented ompt_event_UNIMPLEMENTED
-#define ompt_callback_device_finalize_implemented ompt_event_UNIMPLEMENTED
+#define ompt_callback_device_initialize_implemented ompt_event_MAY_ALWAYS
+#define ompt_callback_device_finalize_implemented ompt_event_MAY_ALWAYS
 
-#define ompt_callback_device_load_implemented ompt_event_UNIMPLEMENTED
+#define ompt_callback_device_load_implemented ompt_event_MAY_ALWAYS
 #define ompt_callback_device_unload_implemented ompt_event_UNIMPLEMENTED
 
 /*----------------------------------------------------------------------------
diff --git a/openmp/runtime/src/ompt-general.cpp b/openmp/runtime/src/ompt-general.cpp
index 55a9365..05a11d3 100644
--- a/openmp/runtime/src/ompt-general.cpp
+++ b/openmp/runtime/src/ompt-general.cpp
@@ -880,6 +880,13 @@ static ompt_interface_fn_t ompt_fn_lookup(const char *s) {
 
 /// Lookup function to query libomp callbacks registered by the tool
 static ompt_interface_fn_t ompt_libomp_target_fn_lookup(const char *s) {
+#define provide_fn(fn)                                                         \
+  if (strcmp(s, #fn) == 0)                                                     \
+    return (ompt_interface_fn_t)fn;
+
+  provide_fn(ompt_get_callback);
+#undef provide_fn
+
 #define ompt_interface_fn(fn, type, code)                                      \
   if (strcmp(s, #fn) == 0)                                                     \
     return (ompt_interface_fn_t)ompt_callbacks.ompt_callback(fn);
@@ -887,7 +894,6 @@ static ompt_interface_fn_t ompt_libomp_target_fn_lookup(const char *s) {
   FOREACH_OMPT_DEVICE_EVENT(ompt_interface_fn)
   FOREACH_OMPT_EMI_EVENT(ompt_interface_fn)
   FOREACH_OMPT_NOEMI_EVENT(ompt_interface_fn)
-
 #undef ompt_interface_fn
 
   return (ompt_interface_fn_t)0;
@@ -896,7 +902,7 @@ static ompt_interface_fn_t ompt_libomp_target_fn_lookup(const char *s) {
 /// This function is called by the libomptarget connector to assign
 /// callbacks already registered with libomp.
 _OMP_EXTERN void ompt_libomp_connect(ompt_start_tool_result_t *result) {
-  OMPT_VERBOSE_INIT_PRINT("libomp --> OMPT: Enter libomp_ompt_connect\n");
+  OMPT_VERBOSE_INIT_PRINT("libomp --> OMPT: Enter ompt_libomp_connect\n");
 
   // Ensure libomp callbacks have been added if not already
   __ompt_force_initialization();
@@ -912,11 +918,11 @@ _OMP_EXTERN void ompt_libomp_connect(ompt_start_tool_result_t *result) {
       // functions can be extracted and assigned to the callbacks in
       // libomptarget
       result->initialize(ompt_libomp_target_fn_lookup,
-                         0 /* initial_device_num */, nullptr /* tool_data */);
+                         /* initial_device_num */ 0, /* tool_data */ nullptr);
       // Track the object provided by libomptarget so that the finalizer can be
       // called during OMPT finalization
       libomptarget_ompt_result = result;
     }
   }
-  OMPT_VERBOSE_INIT_PRINT("libomp --> OMPT: Exit libomp_ompt_connect\n");
+  OMPT_VERBOSE_INIT_PRINT("libomp --> OMPT: Exit ompt_libomp_connect\n");
 }