[Libomptarget] Add a utility function for checking existence of symbols

[jhuber6]

Summary:
There are now a few cases that check if a symbol is present before
continuing, effectively making them optional features if present in the
image. This was done in at least three locations and required an ugly
operation to consume the error. This patch makes a utility function to
handle that instead.

[llvmbot]

@llvm/pr-subscribers-backend-amdgpu

Author: Joseph Huber (jhuber6)

Changes

Summary:
There are now a few cases that check if a symbol is present before
continuing, effectively making them optional features if present in the
image. This was done in at least three locations and required an ugly
operation to consume the error. This patch makes a utility function to
handle that instead.

---

Full diff: https://github.com/llvm/llvm-project/pull/74550.diff

6 Files Affected:

• (modified) openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp (+1-4)
• (modified) openmp/libomptarget/plugins-nextgen/common/include/GlobalHandler.h (+5)
• (modified) openmp/libomptarget/plugins-nextgen/common/src/GlobalHandler.cpp (+20)
• (modified) openmp/libomptarget/plugins-nextgen/common/src/PluginInterface.cpp (+2-7)
• (modified) openmp/libomptarget/plugins-nextgen/common/src/RPC.cpp (+1-8)
• (modified) openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp (+1-4)

diff --git a/openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp b/openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp
index 69acfa54e6c96..0ffdabe5bcd42 100644
--- a/openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -2692,11 +2692,8 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
     // Perform a quick check for the named kernel in the image. The kernel
     // should be created by the 'amdgpu-lower-ctor-dtor' pass.
     GenericGlobalHandlerTy &Handler = Plugin.getGlobalHandler();
-    GlobalTy Global(Name, sizeof(void *));
-    if (auto Err = Handler.getGlobalMetadataFromImage(*this, Image, Global)) {
-      consumeError(std::move(Err));
+    if (!Handler.isSymbolInImage(*this, Image, Name))
       return Plugin::success();
-    }
 
     // Allocate and construct the AMDGPU kernel.
     AMDGPUKernelTy AMDGPUKernel(Name);
diff --git a/openmp/libomptarget/plugins-nextgen/common/include/GlobalHandler.h b/openmp/libomptarget/plugins-nextgen/common/include/GlobalHandler.h
index aff10b1d504fc..fa788c5e1d02b 100644
--- a/openmp/libomptarget/plugins-nextgen/common/include/GlobalHandler.h
+++ b/openmp/libomptarget/plugins-nextgen/common/include/GlobalHandler.h
@@ -120,6 +120,11 @@ class GenericGlobalHandlerTy {
   const ELF64LEObjectFile *
   getOrCreateELFObjectFile(const GenericDeviceTy &Device, DeviceImageTy &Image);
 
+  /// Returns whether the symbol named \p SymName is present in the given \p
+  /// Image.
+  bool isSymbolInImage(GenericDeviceTy &Device, DeviceImageTy &Image,
+                       StringRef SymName);
+
   /// Get the address and size of a global in the image. Address and size are
   /// return in \p ImageGlobal, the global name is passed in \p ImageGlobal.
   Error getGlobalMetadataFromImage(GenericDeviceTy &Device,
diff --git a/openmp/libomptarget/plugins-nextgen/common/src/GlobalHandler.cpp b/openmp/libomptarget/plugins-nextgen/common/src/GlobalHandler.cpp
index 6cb9a366f8b2c..a3d16d3a5bcff 100644
--- a/openmp/libomptarget/plugins-nextgen/common/src/GlobalHandler.cpp
+++ b/openmp/libomptarget/plugins-nextgen/common/src/GlobalHandler.cpp
@@ -105,6 +105,26 @@ Error GenericGlobalHandlerTy::moveGlobalBetweenDeviceAndHost(
   return Plugin::success();
 }
 
+bool GenericGlobalHandlerTy::isSymbolInImage(GenericDeviceTy &Device,
+                                             DeviceImageTy &Image,
+                                             StringRef SymName) {
+  // Get the ELF object file for the image. Notice the ELF object may already
+  // be created in previous calls, so we can reuse it. If this is unsuccessful
+  // just return false as we couldn't find it.
+  const ELF64LEObjectFile *ELFObj = getOrCreateELFObjectFile(Device, Image);
+  if (!ELFObj)
+    return false;
+
+  // Search the ELF symbol using the symbol name.
+  auto SymOrErr = utils::elf::getSymbol(*ELFObj, SymName);
+  if (!SymOrErr) {
+    consumeError(SymOrErr.takeError());
+    return false;
+  }
+
+  return *SymOrErr;
+}
+
 Error GenericGlobalHandlerTy::getGlobalMetadataFromImage(
     GenericDeviceTy &Device, DeviceImageTy &Image, GlobalTy &ImageGlobal) {
 
diff --git a/openmp/libomptarget/plugins-nextgen/common/src/PluginInterface.cpp b/openmp/libomptarget/plugins-nextgen/common/src/PluginInterface.cpp
index 7ca874bdfbf73..e9dc78dc0fe63 100644
--- a/openmp/libomptarget/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/openmp/libomptarget/plugins-nextgen/common/src/PluginInterface.cpp
@@ -785,14 +785,9 @@ Error GenericDeviceTy::deinit(GenericPluginTy &Plugin) {
     GenericGlobalHandlerTy &GHandler = Plugin.getGlobalHandler();
     for (auto *Image : LoadedImages) {
       DeviceMemoryPoolTrackingTy ImageDeviceMemoryPoolTracking = {0, 0, ~0U, 0};
-      GlobalTy TrackerGlobal("__omp_rtl_device_memory_pool_tracker",
-                             sizeof(DeviceMemoryPoolTrackingTy),
-                             &ImageDeviceMemoryPoolTracking);
-      if (auto Err =
-              GHandler.readGlobalFromDevice(*this, *Image, TrackerGlobal)) {
-        consumeError(std::move(Err));
+      if (!GHandler.isSymbolInImage(*this, *Image,
+                                    "__omp_rtl_device_memory_pool_tracker"))
         continue;
-      }
       DeviceMemoryPoolTracking.combine(ImageDeviceMemoryPoolTracking);
     }
 
diff --git a/openmp/libomptarget/plugins-nextgen/common/src/RPC.cpp b/openmp/libomptarget/plugins-nextgen/common/src/RPC.cpp
index a40bf147b2248..60e0540e96022 100644
--- a/openmp/libomptarget/plugins-nextgen/common/src/RPC.cpp
+++ b/openmp/libomptarget/plugins-nextgen/common/src/RPC.cpp
@@ -38,14 +38,7 @@ RPCServerTy::isDeviceUsingRPC(plugin::GenericDeviceTy &Device,
                               plugin::GenericGlobalHandlerTy &Handler,
                               plugin::DeviceImageTy &Image) {
 #ifdef LIBOMPTARGET_RPC_SUPPORT
-  void *ClientPtr;
-  plugin::GlobalTy Global(rpc_client_symbol_name, sizeof(void *), &ClientPtr);
-  if (auto Err = Handler.readGlobalFromImage(Device, Image, Global)) {
-    llvm::consumeError(std::move(Err));
-    return false;
-  }
-
-  return true;
+  return Handler.isSymbolInImage(Device, Image, rpc_client_symbol_name);
 #else
   return false;
 #endif
diff --git a/openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp b/openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp
index 698517179e4cd..7bad411b9d8e3 100644
--- a/openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp
+++ b/openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp
@@ -1055,11 +1055,8 @@ struct CUDADeviceTy : public GenericDeviceTy {
     // Perform a quick check for the named kernel in the image. The kernel
     // should be created by the 'nvptx-lower-ctor-dtor' pass.
     GenericGlobalHandlerTy &Handler = Plugin.getGlobalHandler();
-    GlobalTy Global(KernelName, sizeof(void *));
-    if (auto Err = Handler.getGlobalMetadataFromImage(*this, Image, Global)) {
-      consumeError(std::move(Err));
+    if (!Handler.isSymbolInImage(*this, Image, KernelName))
       return Plugin::success();
-    }
 
     // The Nvidia backend cannot handle creating the ctor / dtor array
     // automatically so we must create it ourselves. The backend will emit

[shiltian]

LG with one nit

[shiltian]

We always consume the error. It might be a good idea to print the error if debug is enabled, but of course it is out of the scope of this patch.

[jdoerfert]

This is the use of GHandler.readGlobalFromDevice for __omp_rtl_device_memory_pool_tracker that is actually needed.
The other use should be a "isSymbolInImage".