[OpenMP][Libomptarget] Introduce changes to support remote plugin

[atmnpatel]

Summary

In order to support remote execution, we need to be able to send the
target binary description to the remote host for registration (and
consequent deregistration). To support this, I added these two
optional new functions to the plugin API:

• __tgt_rtl_register_lib
• __tgt_rtl_unregister_lib

These functions will be called to properly manage the instance of
libomptarget running on the remote host.

[JonChesterfield]

Comments inline. I'd like to use the same hook to fix the memory manager lifetime error

openmp/libomptarget/src/interface.cpp

99

I suspect we should call loadrtls from the libomptarget init but that could be a different patch

116

Not sure whether the new calls should be before or after the existing one, but the order should be reversed for unregister relative to register to nest the lifetimes

openmp/libomptarget/src/rtl.cpp

26

Goes at the end, precedence given to first to ship

271

Moving this out of registerlib looks great

openmp/libomptarget/src/rtl.h

56

Functions should return non-void so they can indicate failure. Also, prefer function pointer to function type for clarity - you're missing an asterisk

[atmnpatel]

Addresses inline comments (partially, fixing).

[atmnpatel]

Added debug messages for when the new calls fail.

[atmnpatel]

openmp/libomptarget/src/interface.cpp

116

I kept it in this order because the remote plugin implementation receives __tgt_bin_desc(s) and uses that call to transfer the images/offload entries exactly once, and since RegisterLib asks if images are valid against RTLs, I needed to have libomptarget send the device images first, but if there's a compelling reason to swap the order, I'm happy to swap it.

118

I'm assuming we can get away with just outputting a debug message for when the calls fail? I don't think its right to abort...

[JonChesterfield]

openmp/libomptarget/src/interface.cpp

116

Pattern is usually:

• generic construction
• specific construction
• use the objects
• specific destruction
• generic destruction

Lines up with base/derived patterns and lexical scoping of lifetimes.

If swapping generic + specific makes more sense here (which it might) then cool. The order of destruction should definitely be the opposite of the order of construction though.

[JonChesterfield]

LGTM, thanks!

[atmnpatel]

Fixes.

[atmnpatel]

More fixes.

[atmnpatel]

Removed redundant third function register_requires, it may become necessary later but it isn't now.

[sylvestre.ledru]

@atmnpatel
Seems that this change caused:
https://bugs.llvm.org/show_bug.cgi?id=48973

[jdoerfert]

In D93293#2532699, @sylvestre.ledru wrote:

@atmnpatel
Seems that this change caused:
https://bugs.llvm.org/show_bug.cgi?id=48973

I don't understand when/why an RTLsTy object is copied at all. @atmnpatel What if we move LoadRTLs into a helper function that is not a member, pass a pointer into it instead. That might solve the issue. Or, pass &PM.RTLs instead, to match the former call_once call. Can you take a look?

openmp/libomptarget/src/interface.cpp

case tgt_mandatory:
     break;
   }
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 /// adds requires flags
 EXTERN void __tgt_register_requires(int64_t flags) {
   TIMESCOPE();
+  for (auto &RTL : PM->RTLs.UsedRTLs) {
+    if (RTL->register_requires) {
+      if ((*RTL->register_requires)(flags) != OFFLOAD_SUCCESS) {
+        DP("Could not register requires with %s", RTL->RTLName.c_str());
+      }
+    }
+  }
   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);
+  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);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 /// unloads a target shared library
 EXTERN void __tgt_unregister_lib(__tgt_bin_desc *desc) {
   TIMESCOPE();
   PM->RTLs.UnregisterLib(desc);
+  for (auto &RTL : PM->RTLs.UsedRTLs) {
+    if (RTL->unregister_lib) {
+      if ((*RTL->unregister_lib)(desc) != OFFLOAD_SUCCESS) {
+        DP("Could not register library with %s", RTL->RTLName.c_str());
+      }
+    }
+  }
 }
 
 /// creates host-to-target data mapping, stores it in the
 /// libomptarget.so internal structure (an entry in a stack of data maps)
 /// and passes the data to the device.
 EXTERN void __tgt_target_data_begin(int64_t device_id, int32_t arg_num,
     void **args_base, void **args, int64_t *arg_sizes, int64_t *arg_types) {
   TIMESCOPE();

openmp/libomptarget/src/rtl.h

struct RTLInfoTy {
   typedef int32_t(run_team_region_ty)(int32_t, void *, void **, ptrdiff_t *,
                                       int32_t, int32_t, int32_t, uint64_t);
   typedef int32_t(run_team_region_async_ty)(int32_t, void *, void **,
                                             ptrdiff_t *, int32_t, int32_t,
                                             int32_t, uint64_t,
                                             __tgt_async_info *);
   typedef int64_t(init_requires_ty)(int64_t);
   typedef int64_t(synchronize_ty)(int32_t, __tgt_async_info *);
+  typedef int32_t (*register_lib_ty)(__tgt_bin_desc *);
+  typedef int32_t (*register_requires_ty)(int64_t);
 
   int32_t Idx = -1;             // RTL index, index is the number of devices
                                 // of other RTLs that were registered before,
                                 // i.e. the OpenMP index of the first device
                                 // to be registered with this RTL.
   int32_t NumberOfDevices = -1; // Number of devices this RTL deals with.
 
   void *LibraryHandler = nullptr;
#endif
   data_exchange_async_ty *data_exchange_async = nullptr;
   data_delete_ty *data_delete = nullptr;
   run_region_ty *run_region = nullptr;
   run_region_async_ty *run_region_async = nullptr;
   run_team_region_ty *run_team_region = nullptr;
   run_team_region_async_ty *run_team_region_async = nullptr;
   init_requires_ty *init_requires = nullptr;
   synchronize_ty *synchronize = nullptr;
+  register_lib_ty *register_lib = nullptr;
+  register_lib_ty *unregister_lib = nullptr;
+  register_requires_ty *register_requires = nullptr;
 
   // Are there images associated with this RTL.
   bool isUsed = false;
 
   // Mutex for thread-safety when calling RTL interface functions.
   // It is easier to enforce thread-safety at the libomptarget level,
   // so that developers of new RTLs do not have to worry about it.
   std::mutex Mtx;
#endif
     data_delete = r.data_delete;
     run_region = r.run_region;
     run_region_async = r.run_region_async;
     run_team_region = r.run_team_region;
     run_team_region_async = r.run_team_region_async;
     init_requires = r.init_requires;
     isUsed = r.isUsed;
     synchronize = r.synchronize;
+    register_lib = r.register_lib;
+    unregister_lib = r.unregister_lib;
+    register_requires = r.register_requires;
   }
 };
 
 /// RTLs identified in the system.
-class RTLsTy {
-private:
-  // 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
-
-public:
+struct RTLsTy {
   // List of the detected runtime libraries.
   std::list<RTLInfoTy> AllRTLs;
 
   // Array of pointers to the detected runtime libraries that have compatible
   // binaries.
   std::vector<RTLInfoTy *> UsedRTLs;
 
   int64_t RequiresFlags = OMP_REQ_UNDEFINED;
 
   explicit RTLsTy() = default;
 
   // Register the clauses of the requires directive.
   void RegisterRequires(int64_t flags);
 
   // 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
+};
 
 /// 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
 /// are trying to (re)register an existing lib or really have a new one.
 struct TranslationTable {
   __tgt_target_table HostTable;
 

openmp/libomptarget/src/rtl.cpp

 #include <cstdlib>
 #include <cstring>
 #include <dlfcn.h>
 #include <mutex>
 #include <string>
 
 // List of all plugins that can support offloading.
 static const char *RTLNames[] = {
     /* PowerPC target       */ "libomptarget.rtl.ppc64.so",
     /* x86_64 target        */ "libomptarget.rtl.x86_64.so",
     /* CUDA target          */ "libomptarget.rtl.cuda.so",
     /* AArch64 target       */ "libomptarget.rtl.aarch64.so",
     /* SX-Aurora VE target  */ "libomptarget.rtl.ve.so",
     /* AMDGPU target        */ "libomptarget.rtl.amdgpu.so",
+    /* Remote target        */ "libomptarget.rtl.rpc.so",
 };
 
 PluginManager *PM;
 
 #if OMPTARGET_PROFILE_ENABLED
 static char *ProfileTraceFile = nullptr;
 #endif
 
#endif
     *((void **)&R.synchronize) = dlsym(dynlib_handle, "__tgt_rtl_synchronize");
     *((void **)&R.data_exchange) =
         dlsym(dynlib_handle, "__tgt_rtl_data_exchange");
     *((void **)&R.data_exchange_async) =
         dlsym(dynlib_handle, "__tgt_rtl_data_exchange_async");
     *((void **)&R.is_data_exchangable) =
         dlsym(dynlib_handle, "__tgt_rtl_is_data_exchangable");
 
+    *((void **)&R.register_lib) =
+        dlsym(dynlib_handle, "__tgt_rtl_register_lib");
+    *((void **)&R.unregister_lib) =
+        dlsym(dynlib_handle, "__tgt_rtl_unregister_lib");
+    *((void **)&R.register_requires) =
+        dlsym(dynlib_handle, "__tgt_rtl_register_requires");
+
     // No devices are supported by this RTL?
     if (!(R.NumberOfDevices = R.number_of_devices())) {
       DP("No devices supported in this RTL\n");
       continue;
     }
 
     DP("Registering RTL %s supporting %d devices!\n", R.RTLName.c_str(),
        R.NumberOfDevices);
void RTLsTy::RegisterRequires(int64_t flags) {
 
   // TODO: insert any other missing checks
 
   DP("New requires flags %" PRId64 " compatible with existing %" PRId64 "!\n",
      flags, RequiresFlags);
 }
 
 void RTLsTy::RegisterLib(__tgt_bin_desc *desc) {
-  // Attempt to load all plugins available in the system.
-  std::call_once(initFlag, &RTLsTy::LoadRTLs, this);
-
   PM->RTLsMtx.lock();
   // Register the images with the RTLs that understand them, if any.
   for (int32_t i = 0; i < desc->NumDeviceImages; ++i) {
     // Obtain the image.
     __tgt_device_image *img = &desc->DeviceImages[i];
 
     RTLInfoTy *FoundRTL = NULL;