[Offload] Repair and rename llvm-omp-device-info (to -offload-) (#100309)

The `llvm-omp-device-info` tool is very handy, but broke due to the lazy
evaluation of devices. This repairs the functionality and adds a test.
The tool is also renamed into `llvm-offload-device-info` as `-omp-` is
going away.

offload/cmake/OpenMPTesting.cmake

@@ -37,6 +37,17 @@ function(find_standalone_test_dependencies)
     return()
   endif()
 
+  find_program(OFFLOAD_DEVICE_INFO_EXECUTABLE
+    NAMES llvm-offload-device-info
+    PATHS ${OPENMP_LLVM_TOOLS_DIR})
+  if (NOT OFFLOAD_DEVICE_INFO_EXECUTABLE)
+    message(STATUS "Cannot find 'llvm-offload-device-info'.")
+    message(STATUS "Please put 'not' in your PATH, set OFFLOAD_DEVICE_INFO_EXECUTABLE to its full path, or point OPENMP_LLVM_TOOLS_DIR to its directory.")
+    message(WARNING "The check targets will not be available!")
+    set(ENABLE_CHECK_TARGETS FALSE PARENT_SCOPE)
+    return()
+  endif()
+
   find_program(OPENMP_NOT_EXECUTABLE
     NAMES not
     PATHS ${OPENMP_LLVM_TOOLS_DIR})
@@ -71,6 +82,7 @@ else()
     set(OPENMP_FILECHECK_EXECUTABLE ${LLVM_RUNTIME_OUTPUT_INTDIR}/FileCheck)
   endif()
   set(OPENMP_NOT_EXECUTABLE ${LLVM_RUNTIME_OUTPUT_INTDIR}/not)
+  set(OFFLOAD_DEVICE_INFO_EXECUTABLE ${LLVM_RUNTIME_OUTPUT_INTDIR}/llvm-offload-device-info)
 endif()
 
 # Macro to extract information about compiler from file. (no own scope)

offload/include/PluginManager.h

@@ -113,8 +113,14 @@ struct PluginManager {
     return Devices.getExclusiveAccessor();
   }
 
-  // Initialize all plugins.
-  void initAllPlugins();
+  /// Initialize \p Plugin. Returns true on success.
+  bool initializePlugin(GenericPluginTy &Plugin);
+
+  /// Initialize device \p DeviceNo of \p Plugin. Returns true on success.
+  bool initializeDevice(GenericPluginTy &Plugin, int32_t DeviceId);
+
+  /// Eagerly initialize all plugins and their devices.
+  void initializeAllDevices();
 
   /// Iterator range for all plugins (in use or not, but always valid).
   auto plugins() { return llvm::make_pointee_range(Plugins); }

offload/src/PluginManager.cpp

@@ -13,6 +13,7 @@
 #include "PluginManager.h"
 #include "Shared/Debug.h"
 #include "Shared/Profile.h"
+#include "device.h"
 
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -60,15 +61,61 @@ void PluginManager::deinit() {
   DP("RTLs unloaded!\n");
 }
 
-void PluginManager::initAllPlugins() {
-  for (auto &R : plugins()) {
-    if (auto Err = R.init()) {
-      [[maybe_unused]] std::string InfoMsg = toString(std::move(Err));
-      DP("Failed to init plugin: %s\n", InfoMsg.c_str());
+bool PluginManager::initializePlugin(GenericPluginTy &Plugin) {
+  if (Plugin.is_initialized())
+    return true;
+
+  if (auto Err = Plugin.init()) {
+    [[maybe_unused]] std::string InfoMsg = toString(std::move(Err));
+    DP("Failed to init plugin: %s\n", InfoMsg.c_str());
+    return false;
+  }
+
+  DP("Registered plugin %s with %d visible device(s)\n", Plugin.getName(),
+     Plugin.number_of_devices());
+  return true;
+}
+
+bool PluginManager::initializeDevice(GenericPluginTy &Plugin,
+                                     int32_t DeviceId) {
+  if (Plugin.is_device_initialized(DeviceId))
+    return true;
+
+  // Initialize the device information for the RTL we are about to use.
+  auto ExclusiveDevicesAccessor = getExclusiveDevicesAccessor();
+
+  int32_t UserId = ExclusiveDevicesAccessor->size();
+
+  // Set the device identifier offset in the plugin.
+#ifdef OMPT_SUPPORT
+  Plugin.set_device_identifier(UserId, DeviceId);
+#endif
+
+  auto Device = std::make_unique<DeviceTy>(&Plugin, UserId, DeviceId);
+  if (auto Err = Device->init()) {
+    [[maybe_unused]] std::string InfoMsg = toString(std::move(Err));
+    DP("Failed to init device %d: %s\n", DeviceId, InfoMsg.c_str());
+    return false;
+  }
+
+  ExclusiveDevicesAccessor->push_back(std::move(Device));
+
+  // We need to map between the plugin's device identifier and the one
+  // that OpenMP will use.
+  PM->DeviceIds[std::make_pair(&Plugin, DeviceId)] = UserId;
+
+  return true;
+}
+
+void PluginManager::initializeAllDevices() {
+  for (auto &Plugin : plugins()) {
+    if (!initializePlugin(Plugin))
       continue;
+
+    for (int32_t DeviceId = 0; DeviceId < Plugin.number_of_devices();
+         ++DeviceId) {
+      initializeDevice(Plugin, DeviceId);
     }
-    DP("Registered plugin %s with %d visible device(s)\n", R.getName(),
-       R.number_of_devices());
   }
 }
 
@@ -94,19 +141,12 @@ void PluginManager::registerLib(__tgt_bin_desc *Desc) {
 
     // Scan the RTLs that have associated images until we find one that supports
     // the current image.
-    for (auto &R : PM->plugins()) {
+    for (auto &R : plugins()) {
       if (!R.is_plugin_compatible(Img))
         continue;
 
-      if (!R.is_initialized()) {
-        if (auto Err = R.init()) {
-          [[maybe_unused]] std::string InfoMsg = toString(std::move(Err));
-          DP("Failed to init plugin: %s\n", InfoMsg.c_str());
-          continue;
-        }
-        DP("Registered plugin %s with %d visible device(s)\n", R.getName(),
-           R.number_of_devices());
-      }
+      if (!initializePlugin(R))
+        continue;
 
       if (!R.number_of_devices()) {
         DP("Skipping plugin %s with no visible devices\n", R.getName());
@@ -120,30 +160,8 @@ void PluginManager::registerLib(__tgt_bin_desc *Desc) {
         DP("Image " DPxMOD " is compatible with RTL %s device %d!\n",
            DPxPTR(Img->ImageStart), R.getName(), DeviceId);
 
-        if (!R.is_device_initialized(DeviceId)) {
-          // Initialize the device information for the RTL we are about to use.
-          auto ExclusiveDevicesAccessor = getExclusiveDevicesAccessor();
-
-          int32_t UserId = ExclusiveDevicesAccessor->size();
-
-          // Set the device identifier offset in the plugin.
-#ifdef OMPT_SUPPORT
-          R.set_device_identifier(UserId, DeviceId);
-#endif
-
-          auto Device = std::make_unique<DeviceTy>(&R, UserId, DeviceId);
-          if (auto Err = Device->init()) {
-            [[maybe_unused]] std::string InfoMsg = toString(std::move(Err));
-            DP("Failed to init device %d: %s\n", DeviceId, InfoMsg.c_str());
-            continue;
-          }
-
-          ExclusiveDevicesAccessor->push_back(std::move(Device));
-
-          // We need to map between the plugin's device identifier and the one
-          // that OpenMP will use.
-          PM->DeviceIds[std::make_pair(&R, DeviceId)] = UserId;
-        }
+        if (!initializeDevice(R, DeviceId))
+          continue;
 
         // Initialize (if necessary) translation table for this library.
         PM->TrlTblMtx.lock();
@@ -219,7 +237,7 @@ void PluginManager::unregisterLib(__tgt_bin_desc *Desc) {
 
     // Scan the RTLs that have associated images until we find one that supports
     // the current image. We only need to scan RTLs that are already being used.
-    for (auto &R : PM->plugins()) {
+    for (auto &R : plugins()) {
       if (R.is_initialized())
         continue;
 

offload/src/interface.cpp

@@ -57,7 +57,7 @@ EXTERN void __tgt_register_lib(__tgt_bin_desc *Desc) {
 /// Initialize all available devices without registering any image
 EXTERN void __tgt_init_all_rtls() {
   assert(PM && "Runtime not initialized");
-  PM->initAllPlugins();
+  PM->initializeAllDevices();
 }
 
 ////////////////////////////////////////////////////////////////////////////////

offload/test/lit.cfg

@@ -408,3 +408,5 @@ config.substitutions.append(("%flags_clang", config.test_flags_clang))
 config.substitutions.append(("%flags_flang", config.test_flags_flang))
 config.substitutions.append(("%flags", config.test_flags))
 config.substitutions.append(("%not", config.libomptarget_not))
+config.substitutions.append(("%offload-device-info",
+                             config.offload_device_info))

offload/test/lit.site.cfg.in

@@ -23,6 +23,7 @@ config.libomptarget_all_targets = "@LIBOMPTARGET_ALL_TARGETS@".split()
 config.libomptarget_current_target = "@CURRENT_TARGET@"
 config.libomptarget_filecheck = "@OPENMP_FILECHECK_EXECUTABLE@"
 config.libomptarget_not = "@OPENMP_NOT_EXECUTABLE@"
+config.offload_device_info = "@OFFLOAD_DEVICE_INFO_EXECUTABLE@"
 config.libomptarget_debug = @LIBOMPTARGET_DEBUG@
 config.has_libomptarget_ompt = @LIBOMPTARGET_OMPT_SUPPORT@
 config.libomptarget_has_libc = @LIBOMPTARGET_GPU_LIBC_SUPPORT@

offload/test/tools/llvm-omp-device-info.c

@@ -0,0 +1,6 @@
+// RUN: %offload-device-info | %fcheck-generic
+//
+// Just check any device was found and something is printed
+//
+// CHECK: Found {{[1-9].*}} devices:
+// CHECK: Device 0:

offload/tools/deviceinfo/CMakeLists.txt

@@ -1,13 +1,13 @@
-message(STATUS "Building the llvm-omp-device-info tool")
+message(STATUS "Building the llvm-offload-device-info tool")
 
-add_openmp_tool(llvm-omp-device-info llvm-omp-device-info.cpp)
+add_openmp_tool(llvm-offload-device-info llvm-offload-device-info.cpp)
 
-llvm_update_compile_flags(llvm-omp-device-info)
+llvm_update_compile_flags(llvm-offload-device-info)
 
-target_include_directories(llvm-omp-device-info PRIVATE
+target_include_directories(llvm-offload-device-info PRIVATE
   ${LIBOMPTARGET_INCLUDE_DIR}
 )
-target_link_libraries(llvm-omp-device-info PRIVATE
+target_link_libraries(llvm-offload-device-info PRIVATE
   omp
   omptarget
 )

offload/tools/deviceinfo/llvm-offload-device-info.cpp

@@ -1,4 +1,4 @@
-//===- llvm-omp-device-info.cpp - Obtain device info as seen from OpenMP --===//
+//===- llvm-offload-device-info.cpp - Device info as seen by LLVM/Offload -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -6,8 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This is a command line utility that, by using Libomptarget, and the device
-// plugins, list devices information as seen from the OpenMP Runtime.
+// This is a command line utility that, by using LLVM/Offload, and the device
+// plugins, list devices information as seen by the runtime.
 //
 //===----------------------------------------------------------------------===//
 
@@ -19,8 +19,9 @@ int main(int argc, char **argv) {
   __tgt_register_lib(&EmptyDesc);
   __tgt_init_all_rtls();
 
+  printf("Found %d devices:\n", omp_get_num_devices());
   for (int Dev = 0; Dev < omp_get_num_devices(); Dev++) {
-    printf("Device (%d):\n", Dev);
+    printf("  Device %d:\n", Dev);
     if (!__tgt_print_device_info(Dev))
       printf("    print_device_info not implemented\n");
     printf("\n");