Merge pull request #504 from AMD-Lightning-Internal/amd/dev/grodgers/fixrpc

[OPENMP] rpc fix from Joseph. Now consistent with trunk

Author: gregrodgers

offload/plugins-nextgen/common/include/PluginInterface.h

@@ -797,6 +797,12 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
   /// Setup the global device memory pool, if the plugin requires one.
   Error setupDeviceMemoryPool(GenericPluginTy &Plugin, DeviceImageTy &Image,
                               uint64_t PoolSize);
+
+  // Setup the RPC server for this device if needed. This may not run on some
+  // plugins like the CPU targets. By default, it will not be executed so it is
+  // up to the target to override this using the shouldSetupRPCServer function.
+  Error setupRPCServer(GenericPluginTy &Plugin, DeviceImageTy &Image);
+
   /// Synchronize the current thread with the pending operations on the
   /// __tgt_async_info structure.
   Error synchronize(__tgt_async_info *AsyncInfo);

offload/plugins-nextgen/common/src/PluginInterface.cpp

@@ -1133,6 +1133,10 @@ GenericDeviceTy::loadBinary(GenericPluginTy &Plugin,
     } else if (auto Err = setupDeviceMemoryPool(Plugin, *Image, HeapSize))
       return std::move(Err);
   }
+
+  if (auto Err = setupRPCServer(Plugin, *Image))
+    return std::move(Err);
+
 #ifdef OMPT_SUPPORT
   if (ompt::Initialized) {
     size_t Bytes =
@@ -1246,6 +1250,33 @@ Error GenericDeviceTy::setupDeviceMemoryPool(GenericPluginTy &Plugin,
   return GHandler.writeGlobalToDevice(*this, Image, DevEnvGlobal);
 }
 
+Error GenericDeviceTy::setupRPCServer(GenericPluginTy &Plugin,
+                                      DeviceImageTy &Image) {
+  // The plugin either does not need an RPC server or it is unavailable.
+  if (!shouldSetupRPCServer())
+    return Plugin::success();
+
+  // Check if this device needs to run an RPC server.
+  RPCServerTy &Server = Plugin.getRPCServer();
+  auto UsingOrErr =
+      Server.isDeviceUsingRPC(*this, Plugin.getGlobalHandler(), Image);
+  if (!UsingOrErr)
+    return UsingOrErr.takeError();
+
+  if (!UsingOrErr.get())
+    return Plugin::success();
+
+  if (auto Err = Server.initDevice(*this, Plugin.getGlobalHandler(), Image))
+    return Err;
+
+  if (auto Err = Server.startThread())
+    return Err;
+
+  RPCServer = &Server;
+  DP("Running an RPC server on device %d\n", getDeviceId());
+  return Plugin::success();
+}
+
 Error PinnedAllocationMapTy::insertEntry(void *HstPtr, void *DevAccessiblePtr,
                                          size_t Size, bool ExternallyLocked) {
   // Insert the new entry into the map.
@@ -1864,13 +1895,12 @@ Error GenericPluginTy::deinit() {
   if (GlobalHandler)
     delete GlobalHandler;
 
-#if RPC_FIXME
-  if (RPCServer) {
+  if (RPCServer && RPCServer->Thread->Running.load(std::memory_order_relaxed))
     if (Error Err = RPCServer->shutDown())
       return Err;
+
+  if (RPCServer)
     delete RPCServer;
-  }
-#endif
 
   if (RecordReplay)
     delete RecordReplay;