[Openmp-commits] [llvm] [openmp] [OpenMP][OMPT] Map `ompt_get_num_devices` to runtime function (PR #200790)

[Jan André Reuter via Openmp-commits]

https://github.com/Thyre updated https://github.com/llvm/llvm-project/pull/200790

>From 1b62ca236791d9bda426cff4ad6ba4ba7228e57d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Andr=C3=A9=20Reuter?= <j.reuter at fz-juelich.de>
Date: Mon, 1 Jun 2026 13:52:07 +0200
Subject: [PATCH 1/2] [OpenMP][OMPT] Map `ompt_get_num_devices` to runtime func
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Since 82e94a593433f36734e2d34898d353a2ecb65b8b, `ompt_get_num_devices` was hard
coded to always provide a return value of `1`, regardless of the actual number
of devices. Tools relying on this return value were henced forced to use other
approaches to estimate the number of available devices.

To improve the situtation, map the returned value to the one provided by
`omp_get_num_devices()`. This does not resolve the issue completely though.
When OMPT is initialized through `libomptarget.so`, the PluginManager,
providing the number of devices, is not initialized until after OMPT has been
fully set up. This means that tools, trying to retrieve the number of devices
during the _initialize_ callback, will still receive the incorrect number of
devices.

Still, this is a significant improvement compared to the prior state, providing
an incorrect result for most of the program runtime.

Signed-off-by: Jan André Reuter <j.reuter at fz-juelich.de>
---
 openmp/runtime/src/ompt-general.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/openmp/runtime/src/ompt-general.cpp b/openmp/runtime/src/ompt-general.cpp
index 959457d380d03..a708a8d76fcf4 100644
--- a/openmp/runtime/src/ompt-general.cpp
+++ b/openmp/runtime/src/ompt-general.cpp
@@ -864,8 +864,10 @@ OMPT_API_ROUTINE int ompt_get_target_info(uint64_t *device_num,
   return 0; // thread is not in a target region
 }
 
+extern "C" int omp_get_num_devices(void);
+
 OMPT_API_ROUTINE int ompt_get_num_devices(void) {
-  return 1; // only one device (the current device) is available
+  return omp_get_num_devices();
 }
 
 /*****************************************************************************

>From 803b608bf8de25b897de35f04a50c65b33ee6f5e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Andr=C3=A9=20Reuter?= <j.reuter at fz-juelich.de>
Date: Wed, 10 Jun 2026 11:11:54 +0200
Subject: [PATCH 2/2] [Offload][NFC] Add test for ompt_get_num_devices
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Jan André Reuter <j.reuter at fz-juelich.de>
---
 offload/test/ompt/get_num_devices.c | 57 +++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)
 create mode 100644 offload/test/ompt/get_num_devices.c

diff --git a/offload/test/ompt/get_num_devices.c b/offload/test/ompt/get_num_devices.c
new file mode 100644
index 0000000000000..7fc8a8ee29ea0
--- /dev/null
+++ b/offload/test/ompt/get_num_devices.c
@@ -0,0 +1,57 @@
+// clang-format off
+// RUN: %libomptarget-compile-run-and-check-generic
+// REQUIRES: ompt
+// REQUIRES: gpu
+// clang-format on
+
+/*
+ * Test to ensure that a tool receives the same number of devices
+ * via the OMPT entry point that an application receives via the
+ * omp_get_num_devices runtime function.
+ *
+ * Note, that the get_num_devices() entry point will _not_ provide
+ * the correct amount of devices before the PluginManager has been
+ * fully initialized, i.e. until after `ompt_initialize`.
+ */
+
+#include <assert.h>
+#include <stdio.h>
+
+#include <omp-tools.h>
+#include <omp.h>
+
+// OMPT entry point handles
+static ompt_get_num_devices_t ompt_get_num_devices = 0;
+
+// Init functions
+int ompt_initialize(ompt_function_lookup_t lookup, int initial_device_num,
+                    ompt_data_t *tool_data) {
+  ompt_get_num_devices = (ompt_get_num_devices_t)lookup("ompt_get_num_devices");
+
+  if (!ompt_get_num_devices)
+    return 0; // failed
+
+  return 1; // success
+}
+
+void ompt_finalize(ompt_data_t *tool_data) {}
+
+ompt_start_tool_result_t *ompt_start_tool(unsigned int omp_version,
+                                          const char *runtime_version) {
+  static ompt_start_tool_result_t ompt_start_tool_result = {&ompt_initialize,
+                                                            &ompt_finalize, 0};
+  return &ompt_start_tool_result;
+}
+
+int main(void) {
+  const int NumDevices = omp_get_num_devices();
+  assert(ompt_get_num_devices != NULL &&
+         "ompt_get_num_devices() should not be NULL\n");
+  const int NumDevicesOmpt = ompt_get_num_devices();
+
+  printf("Num Devices: %d\n", NumDevices);
+  printf("Num Devices via OMPT: %d\n", NumDevicesOmpt);
+}
+
+/// CHECK: Num Devices: [[NUM:[0-9]+]]
+/// CHECK: Num Devices via OMPT: [[NUM]]