[Libomptarget] Remove RPCHandleTy indirection

[jhuber6]

Summary

The 'RPCHandleTy' was intended to capture the intention that a specific
device owns its slot in the RPC server. However, this required creating
a temporary store to hold these pointers. This was causing really weird
spurious failure due to undefined behaviour in the order of library
teardown. For example, the x64 plugin would be torn down, set this to
some invalid memory, and then the CUDA plugin would crash. Rather than
spend the time to fully diagnose this problem I found it pertinent to
simply remove the failure mode.

This patch removes this indirection so now the usage of the RPC server
must always be done with the intended device. This just requires some
extra handling for the AMDGPU indirection where we need to store a
reference to the device.

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

struct AMDGPUSignalTy {
 
   /// Deinitialize the signal.
   Error deinit() {
     hsa_status_t Status = hsa_signal_destroy(Signal);
     return Plugin::check(Status, "Error in hsa_signal_destroy: %s");
   }
 
   /// Wait until the signal gets a zero value.
-  Error wait(const uint64_t ActiveTimeout = 0,
-             RPCHandleTy *RPCHandle = nullptr) const {
-    if (ActiveTimeout && !RPCHandle) {
+  Error wait(const uint64_t ActiveTimeout = 0, RPCServerTy *RPCServer = nullptr,
+             GenericDeviceTy *Device = nullptr) const {
+    if (ActiveTimeout && !RPCServer) {
       hsa_signal_value_t Got = 1;
       Got = hsa_signal_wait_scacquire(Signal, HSA_SIGNAL_CONDITION_EQ, 0,
                                       ActiveTimeout, HSA_WAIT_STATE_ACTIVE);
       if (Got == 0)
         return Plugin::success();
     }
 
     // If there is an RPC device attached to this stream we run it as a server.
-    uint64_t Timeout = RPCHandle ? 8192 : UINT64_MAX;
-    auto WaitState = RPCHandle ? HSA_WAIT_STATE_ACTIVE : HSA_WAIT_STATE_BLOCKED;
+    uint64_t Timeout = RPCServer ? 8192 : UINT64_MAX;
+    auto WaitState = RPCServer ? HSA_WAIT_STATE_ACTIVE : HSA_WAIT_STATE_BLOCKED;
     while (hsa_signal_wait_scacquire(Signal, HSA_SIGNAL_CONDITION_EQ, 0,
                                      Timeout, WaitState) != 0) {
-      if (RPCHandle)
-        if (auto Err = RPCHandle->runServer())
+      if (RPCServer && Device)
+        if (auto Err = RPCServer->runServer(*Device))
           return Err;
     }
     return Plugin::success();
   }
 
   /// Load the value on the signal.
   hsa_signal_value_t load() const { return hsa_signal_load_scacquire(Signal); }
 
private:
   hsa_agent_t Agent;
 
   /// The queue that the stream uses to launch kernels.
   AMDGPUQueueTy &Queue;
 
   /// The manager of signals to reuse signals.
   AMDGPUSignalManagerTy &SignalManager;
 
+  /// A reference to the associated device.
+  GenericDeviceTy &Device;
+
   /// Array of stream slots. Use std::deque because it can dynamically grow
   /// without invalidating the already inserted elements. For instance, the
   /// std::vector may invalidate the elements by reallocating the internal
   /// array if there is not enough space on new insertions.
   std::deque<StreamSlotTy> Slots;
 
   /// The next available slot on the queue. This is reset to zero each time the
   /// stream is synchronized. It also indicates the current number of consumed
   /// slots at a given time.
   uint32_t NextSlot;
 
   /// The synchronization id. This number is increased each time the stream is
   /// synchronized. It is useful to detect if an AMDGPUEventTy points to an
   /// operation that was already finalized in a previous stream sycnhronize.
   uint32_t SyncCycle;
 
   /// A pointer associated with an RPC server running on the given device. If
   /// RPC is not being used this will be a null pointer. Otherwise, this
   /// indicates that an RPC server is expected to be run on this stream.
-  RPCHandleTy *RPCHandle;
+  RPCServerTy *RPCServer;
 
   /// Mutex to protect stream's management.
   mutable std::mutex Mutex;
 
   /// Timeout hint for HSA actively waiting for signal value to change
   const uint64_t StreamBusyWaitMicroseconds;
 
   /// Return the current number of asychronous operations on the stream.
public:
   AMDGPUStreamTy(AMDGPUDeviceTy &Device);
 
   /// Intialize the stream's signals.
   Error init() { return Plugin::success(); }
 
   /// Deinitialize the stream's signals.
   Error deinit() { return Plugin::success(); }
 
-  /// Attach an RPC handle to this stream.
-  void setRPCHandle(RPCHandleTy *Handle) { RPCHandle = Handle; }
+  /// Attach an RPC server to this stream.
+  void setRPCServer(RPCServerTy *Server) { RPCServer = Server; }
 
   /// Push a asynchronous kernel to the stream. The kernel arguments must be
   /// placed in a special allocation for kernel args and must keep alive until
   /// the kernel finalizes. Once the kernel is finished, the stream will release
   /// the kernel args buffer to the specified memory manager.
   Error pushKernelLaunch(const AMDGPUKernelTy &Kernel, void *KernelArgs,
                          uint32_t NumThreads, uint64_t NumBlocks,
                          uint32_t GroupSize,
public:
   Error synchronize() {
     std::lock_guard<std::mutex> Lock(Mutex);
 
     // No need to synchronize anything.
     if (size() == 0)
       return Plugin::success();
 
     // Wait until all previous operations on the stream have completed.
-    if (auto Err =
-            Slots[last()].Signal->wait(StreamBusyWaitMicroseconds, RPCHandle))
+    if (auto Err = Slots[last()].Signal->wait(StreamBusyWaitMicroseconds,
+                                              RPCServer, &Device))
       return Err;
 
     // Reset the stream and perform all pending post actions.
     return complete();
   }
 
   /// Query the stream and complete pending post actions if operations finished.
   /// Return whether all the operations completed. This operation does not block
Error AMDGPUResourceRef<ResourceTy>::create(GenericDeviceTy &Device) {
 
   Resource = new ResourceTy(AMDGPUDevice);
 
   return Resource->init();
 }
 
 AMDGPUStreamTy::AMDGPUStreamTy(AMDGPUDeviceTy &Device)
     : Agent(Device.getAgent()), Queue(Device.getNextQueue()),
-      SignalManager(Device.getSignalManager()),
+      SignalManager(Device.getSignalManager()), Device(Device),
       // Initialize the std::deque with some empty positions.
-      Slots(32), NextSlot(0), SyncCycle(0), RPCHandle(nullptr),
+      Slots(32), NextSlot(0), SyncCycle(0), RPCServer(nullptr),
       StreamBusyWaitMicroseconds(Device.getStreamBusyWaitMicroseconds()) {}
 
 /// Class implementing the AMDGPU-specific functionalities of the global
 /// handler.
 struct AMDGPUGlobalHandlerTy final : public GenericGlobalHandlerTy {
   /// Get the metadata of a global from the device. The name and size of the
   /// global is read from DeviceGlobal and the address of the global is written
   /// to DeviceGlobal.
Error AMDGPUKernelTy::launchImpl(GenericDeviceTy &GenericDevice,
   if (KernelArgs.NumArgs)
     std::memcpy(AllArgs, *static_cast<void **>(Args),
                 sizeof(void *) * KernelArgs.NumArgs);
 
   AMDGPUDeviceTy &AMDGPUDevice = static_cast<AMDGPUDeviceTy &>(GenericDevice);
   AMDGPUStreamTy &Stream = AMDGPUDevice.getStream(AsyncInfoWrapper);
 
   // If this kernel requires an RPC server we attach its pointer to the stream.
-  if (GenericDevice.getRPCHandle())
-    Stream.setRPCHandle(GenericDevice.getRPCHandle());
+  if (GenericDevice.getRPCServer())
+    Stream.setRPCServer(GenericDevice.getRPCServer());
 
   // Push the kernel launch into the stream.
   return Stream.pushKernelLaunch(*this, AllArgs, NumThreads, NumBlocks,
                                  GroupSize, ArgsMemoryManager);
 }
 
 Error AMDGPUKernelTy::printLaunchInfoDetails(GenericDeviceTy &GenericDevice,
                                              KernelArgsTy &KernelArgs,

openmp/libomptarget/plugins-nextgen/common/PluginInterface/CMakeLists.txt

 # Include the RPC server from the `libc` project if availible.
 if(TARGET llvmlibc_rpc_server AND ${LIBOMPTARGET_GPU_LIBC_SUPPORT})
   target_link_libraries(PluginInterface PRIVATE llvmlibc_rpc_server)
   target_compile_definitions(PluginInterface PRIVATE LIBOMPTARGET_RPC_SUPPORT)
 elseif(${LIBOMPTARGET_GPU_LIBC_SUPPORT})
   find_library(llvmlibc_rpc_server NAMES llvmlibc_rpc_server
                PATHS ${LIBOMPTARGET_LLVM_LIBRARY_DIR} NO_DEFAULT_PATH)
   if(llvmlibc_rpc_server)
-    message(WARNING ${llvmlibc_rpc_server})
     target_link_libraries(PluginInterface PRIVATE llvmlibc_rpc_server)
     target_compile_definitions(PluginInterface PRIVATE LIBOMPTARGET_RPC_SUPPORT)
   endif()
 endif()
 
 if ((OMPT_TARGET_DEFAULT) AND (LIBOMPTARGET_OMPT_SUPPORT))
   target_link_libraries(PluginInterface PUBLIC OMPT)
 endif()

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

struct GenericDeviceTy : public DeviceAllocatorTy {
   /// Instead of using more threads we increase the outer (block/team)
   /// parallelism.
   /// @see OMPX_MinThreadsForLowTripCount
   virtual uint32_t getMinThreadsForLowTripCountLoop() {
     return OMPX_MinThreadsForLowTripCount;
   }
 
   /// Get the RPC server running on this device.
-  RPCHandleTy *getRPCHandle() const { return RPCHandle; }
+  RPCServerTy *getRPCServer() const { return RPCServer; }
 
 private:
   /// Register offload entry for global variable.
   Error registerGlobalOffloadEntry(DeviceImageTy &DeviceImage,
                                    const __tgt_offload_entry &GlobalEntry,
                                    __tgt_offload_entry &DeviceEntry);
 
   /// Register offload entry for kernel function.
protected:
   llvm::SmallVector<PeerAccessState> PeerAccesses;
   std::mutex PeerAccessesLock;
 
   /// Map of host pinned allocations used for optimize device transfers.
   PinnedAllocationMapTy PinnedAllocs;
 
   /// A pointer to an RPC server instance attached to this device if present.
   /// This is used to run the RPC server during task synchronization.
-  RPCHandleTy *RPCHandle;
+  RPCServerTy *RPCServer;
 
 #ifdef OMPT_SUPPORT
   /// OMPT callback functions
 #define defineOmptCallback(Name, Type, Code) Name##_t Name##_fn = nullptr;
   FOREACH_OMPT_DEVICE_EVENT(defineOmptCallback)
 #undef defineOmptCallback
 
   /// Internal representation for OMPT device (initialize & finalize)

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

: MemoryManager(nullptr), OMP_TeamLimit("OMP_TEAM_LIMIT"),
       // device initialization. These cannot be consulted until the device is
       // initialized correctly. We intialize them in GenericDeviceTy::init().
       OMPX_TargetStackSize(), OMPX_TargetHeapSize(),
       // By default, the initial number of streams and events are 32.
       OMPX_InitialNumStreams("LIBOMPTARGET_NUM_INITIAL_STREAMS", 32),
       OMPX_InitialNumEvents("LIBOMPTARGET_NUM_INITIAL_EVENTS", 32),
       DeviceId(DeviceId), GridValues(OMPGridValues),
       PeerAccesses(NumDevices, PeerAccessState::PENDING), PeerAccessesLock(),
-      PinnedAllocs(*this), RPCHandle(nullptr) {
+      PinnedAllocs(*this), RPCServer(nullptr) {
 #ifdef OMPT_SUPPORT
   OmptInitialized.store(false);
   // Bind the callbacks to this device's member functions
 #define bindOmptCallback(Name, Type, Code)                                     \
   if (ompt::lookupCallbackByCode)                                              \
     ompt::lookupCallbackByCode((ompt_callbacks_t)(Code),                       \
                                ((ompt_callback_t *)&(Name##_fn)));             \
   DP("OMPT: class bound %s=%p\n", #Name, ((void *)(uint64_t)Name##_fn));
Error GenericDeviceTy::deinit() {
   // we may delete device allocations after the device is deinitialized.
   if (MemoryManager)
     delete MemoryManager;
   MemoryManager = nullptr;
 
   if (RecordReplay.isRecordingOrReplaying())
     RecordReplay.deinit();
 
-  if (RPCHandle)
-    if (auto Err = RPCHandle->deinitDevice())
+  if (RPCServer)
+    if (auto Err = RPCServer->deinitDevice(*this))
       return Err;
 
 #ifdef OMPT_SUPPORT
   bool ExpectedStatus = true;
   if (OmptInitialized.compare_exchange_strong(ExpectedStatus, false))
     performOmptCallback(device_finalize, /* device_num */ DeviceId);
 #endif
 
if (!UsingOrErr)
     return UsingOrErr.takeError();
 
   if (!UsingOrErr.get())
     return Plugin::success();
 
   if (auto Err = Server.initDevice(*this, Plugin.getGlobalHandler(), Image))
     return Err;
 
-  auto DeviceOrErr = Server.getDevice(*this);
-  if (!DeviceOrErr)
-    return DeviceOrErr.takeError();
-  RPCHandle = *DeviceOrErr;
+  RPCServer = &Server;
   DP("Running an RPC server on device %d\n", getDeviceId());
   return Plugin::success();
 }
 
 Error GenericDeviceTy::registerOffloadEntries(DeviceImageTy &Image) {
   const __tgt_offload_entry *Begin = Image.getTgtImage()->EntriesBegin;
   const __tgt_offload_entry *End = Image.getTgtImage()->EntriesEnd;
   for (const __tgt_offload_entry *Entry = Begin; Entry != End; ++Entry) {

openmp/libomptarget/plugins-nextgen/common/PluginInterface/RPC.h

 class DeviceImageTy;
 } // namespace plugin
 
 /// A generic class implementing the interface between the RPC server provided
 /// by the 'libc' project and 'libomptarget'. If the RPC server is not availible
 /// these routines will perform no action.
 struct RPCServerTy {
 public:
-  /// A wrapper around a single instance of the RPC server for a given device.
-  /// This is provided to simplify ownership of the underlying device.
-  struct RPCHandleTy {
-    RPCHandleTy(RPCServerTy &Server, plugin::GenericDeviceTy &Device)
-        : Server(Server), Device(Device) {}
-
-    llvm::Error runServer() { return Server.runServer(Device); }
-
-    llvm::Error deinitDevice() { return Server.deinitDevice(Device); }
-
-  private:
-    RPCServerTy &Server;
-    plugin::GenericDeviceTy &Device;
-  };
-
   RPCServerTy(uint32_t NumDevices);
 
   /// Check if this device image is using an RPC server. This checks for the
   /// precense of an externally visible symbol in the device image that will
   /// be present whenever RPC code is called.
   llvm::Expected<bool> isDeviceUsingRPC(plugin::GenericDeviceTy &Device,
                                         plugin::GenericGlobalHandlerTy &Handler,
                                         plugin::DeviceImageTy &Image);
 
   /// Initialize the RPC server for the given device. This will allocate host
   /// memory for the internal server and copy the data to the client on the
   /// device. The device must be loaded before this is valid.
   llvm::Error initDevice(plugin::GenericDeviceTy &Device,
                          plugin::GenericGlobalHandlerTy &Handler,
                          plugin::DeviceImageTy &Image);
 
-  /// Gets a reference to this server for a specific device.
-  llvm::Expected<RPCHandleTy *> getDevice(plugin::GenericDeviceTy &Device);
-
   /// Runs the RPC server associated with the \p Device until the pending work
   /// is cleared.
   llvm::Error runServer(plugin::GenericDeviceTy &Device);
 
   /// Deinitialize the RPC server for the given device. This will free the
   /// memory associated with the k
   llvm::Error deinitDevice(plugin::GenericDeviceTy &Device);
 
   ~RPCServerTy();
-
-private:
-  llvm::SmallVector<std::unique_ptr<RPCHandleTy>> Handles;
 };
 
-using RPCHandleTy = RPCServerTy::RPCHandleTy;
-
 } // namespace llvm::omp::target
 
 #endif

openmp/libomptarget/plugins-nextgen/common/PluginInterface/RPC.cpp

 using namespace omp;
 using namespace target;
 
 RPCServerTy::RPCServerTy(uint32_t NumDevices) {
 #ifdef LIBOMPTARGET_RPC_SUPPORT
   // If this fails then something is catastrophically wrong, just exit.
   if (rpc_status_t Err = rpc_init(NumDevices))
     FATAL_MESSAGE(1, "Error initializing the RPC server: %d\n", Err);
-  Handles.resize(NumDevices);
 #endif
 }
 
 llvm::Expected<bool>
 RPCServerTy::isDeviceUsingRPC(plugin::GenericDeviceTy &Device,
                               plugin::GenericGlobalHandlerTy &Handler,
                               plugin::DeviceImageTy &Image) {
 #ifdef LIBOMPTARGET_RPC_SUPPORT
#ifdef LIBOMPTARGET_RPC_SUPPORT
   if (auto Err = Device.dataRetrieve(&ClientPtr, ClientGlobal.getPtr(),
                                      sizeof(void *), nullptr))
     return Err;
 
   const void *ClientBuffer = rpc_get_client_buffer(DeviceId);
   if (auto Err = Device.dataSubmit(ClientPtr, ClientBuffer,
                                    rpc_get_client_size(), nullptr))
     return Err;
-
-  Handles[DeviceId] = std::make_unique<RPCHandleTy>(*this, Device);
 #endif
   return Error::success();
 }
 
-llvm::Expected<RPCHandleTy *>
-RPCServerTy::getDevice(plugin::GenericDeviceTy &Device) {
-#ifdef LIBOMPTARGET_RPC_SUPPORT
-  uint32_t DeviceId = Device.getDeviceId();
-  if (!Handles[DeviceId] || !rpc_get_buffer(DeviceId) ||
-      !rpc_get_client_buffer(DeviceId))
-    return plugin::Plugin::error(
-        "Attempt to get an RPC device while not initialized");
-
-  return Handles[DeviceId].get();
-#else
-  return plugin::Plugin::error(
-      "Attempt to get an RPC device while not available");
-#endif
-}
-
 Error RPCServerTy::runServer(plugin::GenericDeviceTy &Device) {
 #ifdef LIBOMPTARGET_RPC_SUPPORT
   if (rpc_status_t Err = rpc_handle_server(Device.getDeviceId()))
     return plugin::Plugin::error(
         "Error while running RPC server on device %d: %d", Device.getDeviceId(),
         Err);
 #endif
   return Error::success();

openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp

struct CUDADeviceTy : public GenericDeviceTy {
   }
 
   /// Synchronize current thread with the pending operations on the async info.
   Error synchronizeImpl(__tgt_async_info &AsyncInfo) override {
     CUstream Stream = reinterpret_cast<CUstream>(AsyncInfo.Queue);
     CUresult Res;
     // If we have an RPC server running on this device we will continuously
     // query it for work rather than blocking.
-    if (!getRPCHandle()) {
+    if (!getRPCServer()) {
       Res = cuStreamSynchronize(Stream);
     } else {
       do {
         Res = cuStreamQuery(Stream);
-        if (auto Err = getRPCHandle()->runServer())
+        if (auto Err = getRPCServer()->runServer(*this))
           return Err;
       } while (Res == CUDA_ERROR_NOT_READY);
     }
 
     // Once the stream is synchronized, return it to stream pool and reset
     // AsyncInfo. This is to make sure the synchronization only works for its
     // own tasks.
     CUDAStreamManager.returnResource(Stream);