[openmp] Workaround for HSA in issue 60119

[JonChesterfield]

Summary

Move plugin initialization to libomptarget initialization.

Removes the call_once control, probably fractionally faster overall.
Fixes issue 60119 because the plugin initialization, which might
try to dlopen unrelated shared libraries, is no longer nested within
a call from application code.

Fixes #60119

[jhuber6]

Fails these tests. It seems we need to load the dynamic libraries multiple times because the plugin constructor is only called once.

libomptarget :: amdgcn-amd-amdhsa :: offloading/dynamic_module.c
libomptarget :: amdgcn-amd-amdhsa-LTO :: offloading/dynamic_module.c

[Maetveis]

This change avoids the deadlock, but now library registration happens before (during) the device runtimes are loaded essentially skipping the registration.

In D142249#4070041, @jhuber6 wrote:

Fails these tests. It seems we need to load the dynamic libraries multiple times because the plugin constructor is only called once.

libomptarget :: amdgcn-amd-amdhsa :: offloading/dynamic_module.c
libomptarget :: amdgcn-amd-amdhsa-LTO :: offloading/dynamic_module.c

They fail because the __target_register_lib calls from the shared library runs before PM->RTLs.AllRTLs is filled. The order of events should be this roughly:

1. init() starts executing
2. the amdgpu plugin initialization starts
3. the amdgpu plugin initializes HSA
4. HSA starts to dlopens objects. Importantly as far as I have seen __attribute__((constructor)) functions will only run for the libraries that did not start them yet, e.g. init() does not get called again.
5. The user library is opened: since this is its first load its constructors are ran, including the call to __tgt_register_lib
6. the user library is not registered because the device plugin (step 2) is not loaded yet (we're currently nested inside

[Maetveis]

openmp/libomptarget/src/rtl.cpp

47–50

This isn't necessary, glibc will not run the __attribute__((constructor(101))) calls again on a second dlopen.

[JonChesterfield]

I don't have the luxury of a working dev system, trying to bring that up in parallel with this.

I think you're describing the current order of execution accurately and I want to change said order to get rid of the recursion induced by HSA.

The intent of this patch is for libomptarget to be constructed entirely before any of the other shared libraries call into it, in which case the tgt_register call would be fine. If dlopen within HSA is picking up libraries that haven't been constructed yet (e.g. the user one which wants to call tgt_register) then this patch isn't going to be sufficient.

edit:

The amdgpu plugin doesn't use tgt_rtl_init_plugin, guess the patch that used that to replace the constructor was reverted. I don't think it would help here though.

The attemptLoadRTL call does indeed end up calling dlopen on the user library before libomptarget finishes construction. In this patch at least, the user library ctors are first invoked via this HSA mechanism. So this approach is indeed not sufficient, it just moves the crash around.

[JonChesterfield]

openmp/libomptarget/src/rtl.cpp

47–50

I'm going to stick with paranoia for the time being. Glibc has a per-library flag that it checks to see if it's already initialised, but it's not totally clear to me whether that flag is set at the start or at the end of initialisation, which affects behaviour on the dlopen of libomptarget. Empirically you're right, that branch is dead, but I'm going to leave it there so I have one fewer failure mode to think about.

[JonChesterfield]

• delay register calls

[JonChesterfield]

This is trending in the direction of D142008. It's essentially the same workaround but restructured aiming to minimise blast radius.

It doesn't worry about threads on the basis that whatever thread has started to initialise libomptarget is going to hold a lock in the loader until the constructor is done by which point the boolean is set.

It therefore doesn't attempt the general fix of D142008 for various threads doing this and possibly falling over. Once the first thread has managed to stand up the plugins (or given up), the hazard we're working around is no longer a concern.

[JonChesterfield]

• clean up

[JonChesterfield]

• clean up

[JonChesterfield]

• simplify

[jhuber6]

LGTM.

openmp/libomptarget/include/device.h

549–551

Nit.

openmp/libomptarget/src/interface.cpp

38–41

[Maetveis]

Looks good, also works for my original program, thank you! Just one nit and a question:

Do we care about de-registration happening during init? It does not happen in the reproducer case when the dynamic loader loads the user libraries, but I think it could happen if one of the libraries HSA loads does a dlopen and dlclose in its constructor on a third library that's using OpenMP. Typed out like that it sounds somewhat ridiculous and hypothetical, so I'm fine without.

openmp/libomptarget/include/device.h

557

Nit: Shouldn't need to be atomic, but I can accept the paranoia.

[Maetveis]

openmp/libomptarget/include/device.h

548–553

This will just add them to the delayed vector again.

[JonChesterfield]

• address review feedback

[JonChesterfield]

In D142249#4070864, @Maetveis wrote:

Do we care about de-registration happening during init? It does not happen in the reproducer case when the dynamic loader loads the user libraries, but I think it could happen if one of the libraries HSA loads does a dlopen and dlclose in its constructor on a third library that's using OpenMP. Typed out like that it sounds somewhat ridiculous and hypothetical, so I'm fine without.

Interesting question. We should have a test cases with a target region in a constructor as well, though I think the current registration system will make that work as long as the user heeds the priority < 100 is reserved contract. We're probably safe from arbitrary calls into libomptarget being written by the user as if they're doing that they have something specific in mind.

Your scenario would break I think. Likewise, a target region in a constructor in a dlopened library will (I think, the graph is getting large enough to make a whiteboard seem sensible) try to use libomptarget before it has managed to construct itself and also break.

I think the right response to such situations is to ignore those edges - noone is currently doing that as they would have started complaining around the 5.3 release if they were, HSA should stop doing this dlopen dance by the next release, and at that point when someone brings a repro along the lines of "I'm doing openmp offloading before main and it doesn't work with rocm 5.4" we'll be able to ask them to upgrade. So it's my hope that the failure mode you've described is latent and will remain so. Code to work around it would tend to stay in tree forever.

openmp/libomptarget/include/device.h

548–553

ha, good spot! Broke that in the refactor to methods

[Maetveis]

openmp/libomptarget/include/device.h

16–17

Nit: no longer needed.

[JonChesterfield]

• drop atomic

openmp/libomptarget/include/device.h

 //===----------- device.h - Target independent OpenMP target RTL ----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Declarations for managing devices that are handled by RTL plugins.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef _OMPTARGET_DEVICE_H
 #define _OMPTARGET_DEVICE_H
 
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <list>
 #include <map>
 #include <memory>
 #include <mutex>
 #include <set>
 #include <thread>
 
struct PluginManager {
 
   // Store target policy (disabled, mandatory, default)
   kmp_target_offload_kind_t TargetOffloadPolicy = tgt_default;
   std::mutex TargetOffloadMtx; ///< For TargetOffloadPolicy
 
   /// Flag to indicate if we use events to ensure the atomicity of
   /// map clauses or not. Can be modified with an environment variable.
   const bool UseEventsForAtomicTransfers;
+
+  // Work around for plugins that call dlopen on shared libraries that call
+  // tgt_register_lib during their initialisation. Stash the pointers in a
+  // vector until the plugins are all initialised and then register them.
+  bool maybeDelayRegisterLib(__tgt_bin_desc *Desc) {
+    if (!RTLsLoaded) {
+      // Only reachable from libomptarget constructor
+      DelayedBinDesc.push_back(Desc);
+      return true;
+    } else {
+      return false;
+    }
+  }
+
+  void registerDelayedLibraries() {
+    // Only called by libomptarget constructor
+    RTLsLoaded = true;
+    for (auto *Desc : DelayedBinDesc)
+      __tgt_register_lib(Desc);
+    DelayedBinDesc.clear();
+  }
+
+private:
+  bool RTLsLoaded = false;
+  llvm::SmallVector<__tgt_bin_desc *> DelayedBinDesc;
 };
 
 extern PluginManager *PM;
 
 #endif

openmp/libomptarget/include/rtl.h

struct RTLsTy {
   void initAllRTLs();
 
   // Register a shared library with all (compatible) RTLs.
   void registerLib(__tgt_bin_desc *Desc);
 
   // Unregister a shared library from all RTLs.
   void unregisterLib(__tgt_bin_desc *Desc);
 
-  // Mutex-like object to guarantee thread-safety and unique initialization
-  // (i.e. the library attempts to load the RTLs (plugins) only once).
-  std::once_flag InitFlag;
-  void loadRTLs(); // not thread-safe
+  // not thread-safe, called from global constructor (i.e. once)
+  void loadRTLs();
 
 private:
   static bool attemptLoadRTL(const std::string &RTLName, RTLInfoTy &RTL);
 };
 
 /// Map between the host entry begin and the translation table. Each
 /// registered library gets one TranslationTable. Use the map from
 /// __tgt_offload_entry so that we may quickly determine whether we

openmp/libomptarget/plugins/remote/server/Server.cpp

Status RemoteOffloadImpl::IsValidBinary(ServerContext *Context,
   SERVER_DBG("Checked if binary (%p) is valid",
              (void *)(DeviceImage->image_ptr()))
   return Status::OK;
 }
 
 Status RemoteOffloadImpl::GetNumberOfDevices(ServerContext *Context,
                                              const Null *Null,
                                              I32 *NumberOfDevices) {
-  std::call_once(PM->RTLs.initFlag, &RTLsTy::LoadRTLs, &PM->RTLs);
-
   int32_t Devices = 0;
   PM->RTLsMtx.lock();
   for (auto &RTL : PM->RTLs.AllRTLs)
     Devices += RTL.NumberOfDevices;
   PM->RTLsMtx.unlock();
 
   NumberOfDevices->set_number(Devices);
 

openmp/libomptarget/src/interface.cpp

EXTERN void __tgt_register_requires(int64_t Flags) {
   TIMESCOPE();
   PM->RTLs.registerRequires(Flags);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 /// adds a target shared library to the target execution image
 EXTERN void __tgt_register_lib(__tgt_bin_desc *Desc) {
   TIMESCOPE();
-  std::call_once(PM->RTLs.InitFlag, &RTLsTy::loadRTLs, &PM->RTLs);
+  if (PM->maybeDelayRegisterLib(Desc))
+    return;
+
   for (auto &RTL : PM->RTLs.AllRTLs) {
     if (RTL.register_lib) {
       if ((*RTL.register_lib)(Desc) != OFFLOAD_SUCCESS) {
         DP("Could not register library with %s", RTL.RTLName.c_str());
       }
     }
   }
   PM->RTLs.registerLib(Desc);
 }

openmp/libomptarget/src/rtl.cpp

static const char *RTLNames[] = {
     /* Remote target        */ "libomptarget.rtl.rpc",
 };
 
 PluginManager *PM;
 
 static char *ProfileTraceFile = nullptr;
 
 __attribute__((constructor(101))) void init() {
   DP("Init target library!\n");
 
   bool UseEventsForAtomicTransfers = true;
   if (const char *ForceAtomicMap = getenv("LIBOMPTARGET_MAP_FORCE_ATOMIC")) {
     std::string ForceAtomicMapStr(ForceAtomicMap);
     if (ForceAtomicMapStr == "false" || ForceAtomicMapStr == "FALSE")
       UseEventsForAtomicTransfers = false;
     else if (ForceAtomicMapStr != "true" && ForceAtomicMapStr != "TRUE")
       fprintf(stderr,
               "Warning: 'LIBOMPTARGET_MAP_FORCE_ATOMIC' accepts only "
               "'true'/'TRUE' or 'false'/'FALSE' as options, '%s' ignored\n",
               ForceAtomicMap);
   }
 
   PM = new PluginManager(UseEventsForAtomicTransfers);
 
   ProfileTraceFile = getenv("LIBOMPTARGET_PROFILE");
   // TODO: add a configuration option for time granularity
   if (ProfileTraceFile)
     timeTraceProfilerInitialize(500 /* us */, "libomptarget");
+
+  PM->RTLs.loadRTLs();
+  PM->registerDelayedLibraries();
 }
 
 __attribute__((destructor(101))) void deinit() {
   DP("Deinit target library!\n");
   delete PM;
 
   if (ProfileTraceFile) {
     // TODO: add env var for file output