[Openmp-commits] [openmp] aa4c0f1 - [OpenMP] [OMPT] [3/8] Implemented callback registration in libomptarget

[Dhruva Chakrabarti via Openmp-commits]

Author: Dhruva Chakrabarti
Date: 2022-12-08T11:43:10-08:00
New Revision: aa4c0f116c93a91186291ca3255583fc5deae7f1

URL: https://github.com/llvm/llvm-project/commit/aa4c0f116c93a91186291ca3255583fc5deae7f1
DIFF: https://github.com/llvm/llvm-project/commit/aa4c0f116c93a91186291ca3255583fc5deae7f1.diff

LOG: [OpenMP] [OMPT] [3/8] Implemented callback registration in libomptarget

The purpose of this patch is to have tool-provided callbacks registered
in libomptarget. The overall design document is in
https://rice.app.box.com/s/pf3gix2hs4d4o1aatwir1set05xmjljc

Defined a class OmptDeviceCallbacksTy that will be used by libomptarget
and a plugin for callbacks registered by a tool. Once the callbacks are
registered in libomp, a lookup function is passed to libomptarget that is
used to retrieve the callbacks and register them in libomptarget.

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

Reviewed By: jplehr, tianshilei1992

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

Added: 
    openmp/libomptarget/include/ompt_device_callbacks.h

Modified: 
    openmp/libomptarget/include/ompt_connector.h
    openmp/libomptarget/src/ompt_callback.cpp

Removed: 
    


################################################################################
diff  --git a/openmp/libomptarget/include/ompt_connector.h b/openmp/libomptarget/include/ompt_connector.h
index 5d8648d694d3..332c3a0199ce 100644
--- a/openmp/libomptarget/include/ompt_connector.h
+++ b/openmp/libomptarget/include/ompt_connector.h
@@ -22,12 +22,12 @@
 #include <memory>
 #include <string>
 
-#include "Debug.h"
 #include "omp-tools.h"
+
+#include "Debug.h"
 #include "omptarget.h"
 
 #define DEBUG_PREFIX "OMPT"
-
 #define LIBOMPTARGET_STRINGIFY(s) #s
 
 /// Type for the function to be invoked for connecting two libraries.

diff  --git a/openmp/libomptarget/include/ompt_device_callbacks.h b/openmp/libomptarget/include/ompt_device_callbacks.h
new file mode 100644
index 000000000000..3dbba3dedbee
--- /dev/null
+++ b/openmp/libomptarget/include/ompt_device_callbacks.h
@@ -0,0 +1,69 @@
+//===--------- 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
+  }
+
+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;
+
+#undef DEBUG_PREFIX
+
+#endif // OMPT_SUPPORT
+
+#endif // _OMPT_DEVICE_CALLBACKS_H

diff  --git a/openmp/libomptarget/src/ompt_callback.cpp b/openmp/libomptarget/src/ompt_callback.cpp
index 4928951668ca..5715642ad3a9 100644
--- a/openmp/libomptarget/src/ompt_callback.cpp
+++ b/openmp/libomptarget/src/ompt_callback.cpp
@@ -18,13 +18,17 @@
 #include <cstring>
 
 #include "omp-tools.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 ompt_enabled = false;
+/// Object maintaining all the callbacks for this library
+OmptDeviceCallbacksTy OmptDeviceCallbacks;
 
 /// This is the function called by the higher layer (libomp) responsible
 /// for initializing OMPT in this library. This is passed to libomp
@@ -37,7 +41,11 @@ static int ompt_libomptarget_initialize(ompt_function_lookup_t lookup,
                                         ompt_data_t *tool_data) {
   DP("enter ompt_libomptarget_initialize!\n");
   ompt_enabled = true;
-  // TODO use the parameters to populate callbacks in libomptarget
+  // 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;
 }
@@ -68,6 +76,9 @@ __attribute__((constructor(102))) static void ompt_init(void) {
   OmptResult.finalize = ompt_libomptarget_finalize;
   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 ompt_init\n");