[OpenMP] [OMPT] [2/8] Implemented a connector for communication of OMPT callbacks between libraries.

[dhruvachak]

Summary

This is part of a set of patches implementing OMPT target callback support and has been split out of the originally submitted https://reviews.llvm.org/D113728. The overall design can be found in https://rice.app.box.com/s/pf3gix2hs4d4o1aatwir1set05xmjljc

The purpose of this patch is to provide a way to register tool-provided callbacks into libomp when libomptarget is loaded.

Introduced a cmake variable LIBOMPTARGET_OMPT_SUPPORT that can be used to control OMPT target support. It follows host OMPT support, controlled by LIBOMP_HAVE_OMPT_SUPPORT.

Added a connector that can be used to communicate between OMPT implementations in libomp and libomptarget or libomptarget and a plugin.

Added a global constructor in libomptarget that uses the connector to force registration of tool-provided callbacks in libomp. A pair of init and fini functions are provided to libomp as part of the connect process which will be used to register the tool-provided callbacks in libomptarget.

Patch from John Mellor-Crummey <[e-mail]>
(With contributions from Dhruva Chakrabarti <[e-mail]>)

[dhruvachak]

Removed absolute path for location of omp-tools.h.

Set the location of OpenMP tools directory in runtime and use it in libomptarget.

[dhruvachak]

Rebased

[dreachem]

LGTM.

[dhruvachak]

Rebased

[dhruvachak]

This patch had to be reverted because of a link failure https://github.com/llvm/llvm-project/issues/58721

[dhruvachak]

Fixed issues reported in https://github.com/llvm/llvm-project/issues/58721

Handled LIBOMP_OMPT_SUPPORT and its interaction with LIBOMPTARGET_OMPT_SUPPORT.

Replaced dlsym with llvm library API.

[jhuber6]

Make sure the new files in libomptarget match the LLVM style for the variables and filenames. Otherwise looks file

[dhruvachak]

Changed new file name from ompt-connector.h to ompt_connector.h.

[dhruvachak]

In D123572#3921397, @jhuber6 wrote:

Make sure the new files in libomptarget match the LLVM style for the variables and filenames. Otherwise looks file

I changed the new include filename to ompt_connector.h. The filename formats appear mixed to me. Specifically, libomp seems to use "-" as the separator while libomptarget and plugins seem to use "_".

About variables, I tried to use CamelCase as much as possible for the new code. The problem is that the new code has to use OMPT APIs and they have a different naming format. So it is hard to avoid a mixed format. If you are referring to any specific use, let me know and I can fix it in a separate commit. Thanks.

[dhruvachak]

Rebased.

openmp/libomptarget/CMakeLists.txt

if(LIBOMPTARGET_CMAKE_BUILD_TYPE MATCHES debug)
   option(LIBOMPTARGET_ENABLE_DEBUG "Allow debug output with the environment variable LIBOMPTARGET_DEBUG=1" ON)
 else()
   option(LIBOMPTARGET_ENABLE_DEBUG "Allow debug output with the environment variable LIBOMPTARGET_DEBUG=1" OFF)
 endif()
 if(LIBOMPTARGET_ENABLE_DEBUG)
   add_definitions(-DOMPTARGET_DEBUG)
 endif()
 
+# OMPT support for libomptarget
+set(OMPT_TARGET_DEFAULT FALSE)
+if ((LIBOMP_HAVE_OMPT_SUPPORT) AND (NOT WIN32))
+  set (OMPT_TARGET_DEFAULT TRUE)
+endif()
+set(LIBOMPTARGET_OMPT_SUPPORT ${OMPT_TARGET_DEFAULT} CACHE BOOL "OMPT-target-support?")
+if (LIBOMPTARGET_OMPT_SUPPORT)
+  add_definitions(-DOMPT_SUPPORT=1)
+  message(STATUS "OMPT target enabled")
+else()
+  message(STATUS "OMPT target disabled")
+endif()
+
+pythonize_bool(LIBOMPTARGET_OMPT_SUPPORT)
+
 set(LIBOMPTARGET_INCLUDE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/include)
+message(STATUS "OpenMP tools dir in libomptarget: ${LIBOMP_OMP_TOOLS_INCLUDE_DIR}")
+include_directories(${LIBOMP_OMP_TOOLS_INCLUDE_DIR})
 
 # Build target agnostic offloading library.
 set(LIBOMPTARGET_SRC_DIR ${CMAKE_CURRENT_SOURCE_DIR}/src)
 add_subdirectory(${LIBOMPTARGET_SRC_DIR})
 
 # Definitions for testing, for reuse when testing libomptarget-nvptx.
 set(LIBOMPTARGET_OPENMP_HEADER_FOLDER "${LIBOMP_INCLUDE_DIR}" CACHE STRING
   "Path to folder containing omp.h")

openmp/libomptarget/include/ompt-connector.h

+//=== ompt-connector.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.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Support used by OMPT implementation to establish communication between
+// various OpenMP runtime libraries: host openmp library, target-independent
+// runtime library, and device-dependent runtime libraries.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _OMPT_CONNECTOR_H
+#define _OMPT_CONNECTOR_H
+
+#include <dlfcn.h>
+#include <string>
+
+#include "Debug.h"
+#include "omp-tools.h"
+#include "omptarget.h"
+
+#define LIBOMPTARGET_STRINGIFY(s) #s
+
+/// Type for the function to be invoked for connecting two libraries.
+typedef void (*OmptConnectRtnTy)(ompt_start_tool_result_t *result);
+
+/// Establish connection between openmp runtime libraries
+///
+/// This class is used to communicate between an OMPT implementation in
+/// libomptarget and libomp. It is also used to communicate between an
+/// OMPT implementation in a device-specific plugin and
+/// libomptarget. The decision whether OMPT is enabled or not needs to
+/// be made when the library is loaded before any functions in the
+/// library are invoked. For that reason, an instance of this class is
+/// intended to be defined in the constructor for libomptarget or a
+/// plugin so that the decision about whether OMPT is supposed to be
+/// enabled is known before any interface function in the library is
+/// invoked.
+class OmptLibraryConnectorTy {
+public:
+  /// Use \p LibName as the prefix of the global function used for connecting
+  /// two libraries, the source indicated by \p LibName and the destination
+  /// being the one that creates this object.
+  OmptLibraryConnectorTy(const char *LibName) {
+    LibConnRtn.append(LibName);
+    LibConnRtn.append("_connect");
+    IsInitialized = false;
+  }
+  OmptLibraryConnectorTy() = delete;
+  /// Use \p OmptResult init to connect the two libraries denoted by this
+  /// object. The init function of \p OmptResult will be used during connection
+  /// and the fini function of \p OmptResult will be used during teardown.
+  void connect(ompt_start_tool_result_t *OmptResult) {
+    initialize();
+    if (!LibConnHandle)
+      return;
+    // Call the function provided by the source library for connect
+    LibConnHandle(OmptResult);
+  }
+
+private:
+  void initialize() {
+    if (IsInitialized)
+      return;
+
+    DP("OMPT: Library connection routine = %s\n", LibConnRtn.c_str());
+
+    void *VPtr = dlsym(NULL, LibConnRtn.c_str());
+    // If dlsym fails, the handle will be null. connect() checks
+    // for this condition
+    LibConnHandle =
+        reinterpret_cast<OmptConnectRtnTy>(reinterpret_cast<long>(VPtr));
+    DP("OMPT: Library connection handle = %p\n", LibConnHandle);
+    IsInitialized = true;
+  }
+
+private:
+  /// Ensure initialization occurs only once
+  bool IsInitialized;
+  /// Handle of connect routine provided by source library
+  OmptConnectRtnTy LibConnHandle;
+  /// Name of connect routine provided by source library
+  std::string LibConnRtn;
+};
+
+#endif

openmp/libomptarget/src/CMakeLists.txt

 add_llvm_library(omptarget
   SHARED
 
   api.cpp
   device.cpp
   interface.cpp
   interop.cpp
   omptarget.cpp
+  ompt_callback.cpp
   rtl.cpp
   LegacyAPI.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LIBOMPTARGET_INCLUDE_DIR}
 
   LINK_COMPONENTS
   Support

openmp/libomptarget/src/ompt_callback.cpp

+//===-- ompt_callback.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.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Implementation of OMPT callback interfaces for target independent layer
+//
+//===----------------------------------------------------------------------===//
+
+#include <assert.h>
+#include <atomic>
+#include <cstdlib>
+#include <cstring>
+#include <dlfcn.h>
+
+#include "omp-tools.h"
+#include "ompt-connector.h"
+#include "private.h"
+
+#define fnptr_to_ptr(x) ((void *)(uint64_t)x)
+
+/// Used to indicate whether OMPT was enabled for this library
+bool ompt_enabled = false;
+
+#ifdef OMPT_SUPPORT
+/// 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");
+  ompt_enabled = true;
+  // TODO use the parameters to populate callbacks in libomptarget
+  DP("exit ompt_libomptarget_initialize!\n");
+  return 0;
+}
+
+static void ompt_libomptarget_finalize(ompt_data_t *data) {
+  DP("enter ompt_libomptarget_finalize!\n");
+  ompt_enabled = false;
+  DP("exit ompt_libomptarget_finalize!\n");
+}
+
+/*****************************************************************************
+ * 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. 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");
+  // Connect with libomp
+  static OmptLibraryConnectorTy LibompConnector("ompt_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.tool_data.value = 0;
+
+  // Now call connect that causes the above init/fini functions to be called
+  LibompConnector.connect(&OmptResult);
+  DP("OMPT: Exit ompt_init\n");
+}
+#endif

openmp/runtime/CMakeLists.txt

 endif()
 
 add_subdirectory(src)
 add_subdirectory(test)
 
 # make these variables available for tools:
 set(LIBOMP_LIBRARY_DIR ${LIBOMP_LIBRARY_DIR} PARENT_SCOPE)
 set(LIBOMP_INCLUDE_DIR ${LIBOMP_INCLUDE_DIR} PARENT_SCOPE)
+set(LIBOMP_OMP_TOOLS_INCLUDE_DIR ${LIBOMP_OMP_TOOLS_INCLUDE_DIR} PARENT_SCOPE)
 # make these variables available for tools/libompd:
 set(LIBOMP_SRC_DIR ${LIBOMP_SRC_DIR} PARENT_SCOPE)
 set(LIBOMP_OMPD_SUPPORT ${LIBOMP_OMPD_SUPPORT} PARENT_SCOPE)

openmp/runtime/cmake/config-ix.cmake

#      (LIBOMP_ARCH STREQUAL arm) OR
      AND # OS supported?
      ((WIN32 AND LIBOMP_HAVE_PSAPI) OR APPLE OR (NOT WIN32 AND LIBOMP_HAVE_WEAK_ATTRIBUTE)))
     set(LIBOMP_HAVE_OMPT_SUPPORT TRUE)
   else()
     set(LIBOMP_HAVE_OMPT_SUPPORT FALSE)
   endif()
 endif()
 
+set(LIBOMP_HAVE_OMPT_SUPPORT ${LIBOMP_HAVE_OMPT_SUPPORT} PARENT_SCOPE)
+
 # Check if HWLOC support is available
 if(${LIBOMP_USE_HWLOC})
   find_path(LIBOMP_HWLOC_INCLUDE_DIR NAMES hwloc.h HINTS ${LIBOMP_HWLOC_INSTALL_DIR} PATH_SUFFIXES include)
   set(CMAKE_REQUIRED_INCLUDES ${LIBOMP_HWLOC_INCLUDE_DIR})
   check_include_file(hwloc.h LIBOMP_HAVE_HWLOC_H)
   set(CMAKE_REQUIRED_INCLUDES)
   find_library(LIBOMP_HWLOC_LIBRARY
     NAMES hwloc libhwloc

openmp/runtime/src/CMakeLists.txt

install(
   FILES
   ${CMAKE_CURRENT_BINARY_DIR}/omp.h
   DESTINATION ${LIBOMP_HEADERS_INSTALL_PATH}
 )
 if(${LIBOMP_OMPT_SUPPORT})
   install(FILES ${CMAKE_CURRENT_BINARY_DIR}/omp-tools.h DESTINATION ${LIBOMP_HEADERS_INSTALL_PATH})
   # install under legacy name ompt.h
   install(FILES ${CMAKE_CURRENT_BINARY_DIR}/omp-tools.h DESTINATION ${LIBOMP_HEADERS_INSTALL_PATH} RENAME ompt.h)
+  set(LIBOMP_OMP_TOOLS_INCLUDE_DIR ${CMAKE_CURRENT_BINARY_DIR} PARENT_SCOPE)
 endif()
 if(${LIBOMP_FORTRAN_MODULES})
   install(FILES
     ${CMAKE_CURRENT_BINARY_DIR}/omp_lib.h
     ${CMAKE_CURRENT_BINARY_DIR}/omp_lib.mod
     ${CMAKE_CURRENT_BINARY_DIR}/omp_lib_kinds.mod
     DESTINATION ${LIBOMP_HEADERS_INSTALL_PATH}
   )
 endif()

openmp/runtime/src/exports_so.txt

global: # Exported symbols.
         #
         omp_*;     # Standard OpenMP functions.
         OMP_*;     # Standard OpenMP symbols.
 
         #
         # OMPT API
         #
         ompt_start_tool;     # OMPT start interface
+        ompt_libomp_connect; # OMPT libomptarget interface
 
         ompc_*;    # omp.h renames some standard functions to ompc_*.
         kmp_*;     # Intel extensions.
         kmpc_*;    # Intel extensions.
         __kmpc_*;  # Functions called by compiler-generated code.
         GOMP_*;    # GNU C compatibility functions.
 
         _You_must_link_with_*;     # Mutual detection/MS compatibility symbols.

openmp/runtime/src/ompt-general.cpp

 #if KMP_OS_WINDOWS
 static HMODULE ompt_tool_module = NULL;
 #define OMPT_DLCLOSE(Lib) FreeLibrary(Lib)
 #else
 static void *ompt_tool_module = NULL;
 #define OMPT_DLCLOSE(Lib) dlclose(Lib)
 #endif
 
+/// Used to track the initializer and the finalizer provided by libomptarget
+static ompt_start_tool_result_t *libomptarget_ompt_result = NULL;
+
 /*****************************************************************************
  * forward declarations
  ****************************************************************************/
 
 static ompt_interface_fn_t ompt_fn_lookup(const char *s);
 
 OMPT_API_ROUTINE ompt_data_t *ompt_get_thread_data(void);
 
fprintf(stderr,
             "  legal values are (NULL,\"\",\"disabled\","
             "\"enabled\").\n",
             ompt_env_var);
     break;
   }
   if (verbose_init && verbose_file != stderr && verbose_file != stdout)
     fclose(verbose_file);
 #if OMPT_DEBUG
-  printf("ompt_pre_init(): ompt_enabled = %d\n", ompt_enabled);
+  printf("ompt_pre_init(): ompt_enabled = %d\n", ompt_enabled.enabled);
 #endif
 }
 
 extern "C" int omp_get_initial_device(void);
 
 void ompt_post_init() {
   //--------------------------------------------------
   // Execute the post-initialization logic only once.
if (ompt_enabled.ompt_callback_implicit_task) {
           ompt_scope_begin, parallel_data, task_data, 1, 1, ompt_task_initial);
     }
 
     ompt_set_thread_state(root_thread, ompt_state_work_serial);
   }
 }
 
 void ompt_fini() {
-  if (ompt_enabled.enabled
-#if OMPD_SUPPORT
-      && ompt_start_tool_result && ompt_start_tool_result->finalize
-#endif
-  ) {
+  if (ompt_enabled.enabled) {
+    if (ompt_start_tool_result && ompt_start_tool_result->finalize) {
       ompt_start_tool_result->finalize(&(ompt_start_tool_result->tool_data));
     }
+    if (libomptarget_ompt_result && libomptarget_ompt_result->finalize) {
+      libomptarget_ompt_result->finalize(NULL);
+    }
+  }
 
   if (ompt_tool_module)
     OMPT_DLCLOSE(ompt_tool_module);
   memset(&ompt_enabled, 0, sizeof(ompt_enabled));
 }
 
 /*****************************************************************************
  * interface operations

 
 #define ompt_interface_fn(fn)                                                  \
   fn##_t fn##_f = fn;                                                          \
   if (strcmp(s, #fn) == 0)                                                     \
     return (ompt_interface_fn_t)fn##_f;
 
   FOREACH_OMPT_INQUIRY_FN(ompt_interface_fn)
 
+#undef ompt_interface_fn
+
   return NULL;
 }
+
+/// Lookup function to query libomp callbacks registered by the tool
+static ompt_interface_fn_t ompt_libomp_target_fn_lookup(const char *s) {
+#define ompt_interface_fn(fn, type, code)                                      \
+  if (strcmp(s, #fn) == 0)                                                     \
+    return (ompt_interface_fn_t)ompt_callbacks.ompt_callback(fn);
+
+  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;
+}
+
+/// 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");
+
+  // Ensure libomp callbacks have been added if not already
+  __ompt_force_initialization();
+
+  if (ompt_enabled.enabled &&
+      // Callbacks are initiated only if the device initialize callback
+      // has been registered by the tool
+      ompt_callbacks.ompt_callback(ompt_callback_device_initialize)) {
+    if (result) {
+      OMPT_VERBOSE_INIT_PRINT(
+          "libomp --> OMPT: Connecting with libomptarget\n");
+      // Pass in the libomp lookup function so that the already registered
+      // 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 */);
+      // 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");
+}

openmp/runtime/src/ompt-specific.h

 
 #include "kmp.h"
 
 #if OMPT_SUPPORT
 /*****************************************************************************
  * forward declarations
  ****************************************************************************/
 
+/// Entrypoint used by libomptarget to register callbacks in libomp, if not
+/// done already
+void __ompt_force_initialization();
+
 void __ompt_team_assign_id(kmp_team_t *team, ompt_data_t ompt_pid);
 void __ompt_thread_assign_wait_id(void *variable);
 
 void __ompt_lw_taskteam_init(ompt_lw_taskteam_t *lwt, kmp_info_t *thr, int gtid,
                              ompt_data_t *ompt_pid, void *codeptr);
 
 void __ompt_lw_taskteam_link(ompt_lw_taskteam_t *lwt, kmp_info_t *thr,
                              int on_heap, bool always = false);

openmp/runtime/src/ompt-specific.cpp

ompt_task_info_t *__ompt_get_scheduling_taskinfo(int depth) {
   }
 
   return info;
 }
 
 //******************************************************************************
 // interface operations
 //******************************************************************************
+//----------------------------------------------------------
+// initialization support
+//----------------------------------------------------------
+
+void __ompt_force_initialization() { __kmp_serial_initialize(); }
 
 //----------------------------------------------------------
 // thread support
 //----------------------------------------------------------
 
 ompt_data_t *__ompt_get_thread_data_internal() {
   if (__kmp_get_gtid() >= 0) {
     kmp_info_t *thread = ompt_get_thread();