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

…backs

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

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()
 

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

openmp/libomptarget/include/ompt_connector.h → 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 @@ class OmptLibraryConnectorTy {
   std::string LibIdent;
 };
 
-#undef DEBUG_PREFIX
-
 #endif // OMPT_SUPPORT
 
-#endif // _OMPT_CONNECTOR_H
+#endif // _OMPTCONNECTOR_H

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(

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