[llvm] [Offload] Do not do not re-initialize RPC server (PR #218757)

[Robert Imschweiler via llvm-commits]

================
@@ -207,32 +207,40 @@ RPCServerTy::isDeviceUsingRPC(plugin::GenericDeviceTy &Device,
 Error RPCServerTy::initDevice(plugin::GenericDeviceTy &Device,
                               plugin::GenericGlobalHandlerTy &Handler,
                               plugin::DeviceImageTy &Image) {
+  std::lock_guard<decltype(BufferMutex)> Lock(BufferMutex);
   uint64_t NumPorts =
       std::min(Device.requestedRPCPortCount(), rpc::MAX_PORT_COUNT);
-  auto RPCBufferOrErr = Device.allocate(
-      rpc::Server::allocation_size(Device.getRPCNumLanes(), NumPorts), nullptr,
-      TARGET_ALLOC_HOST);
-  if (!RPCBufferOrErr)
-    return RPCBufferOrErr.takeError();
-
-  void *RPCBuffer = *RPCBufferOrErr;
-  if (!RPCBuffer)
-    return plugin::Plugin::error(
-        error::ErrorCode::UNKNOWN,
-        "failed to initialize RPC server for device %d", Device.getDeviceId());
-
-  // The doorbell is used by AMDGPU targets to let the server thread be
-  // descheduled. It is optional and will be ignored if the fields are null.
-  rpc::Doorbell Doorbell{};
-  if (auto Err = Device.Plugin.initRPCDoorbell(Doorbell.value, Doorbell.mailbox,
-                                               Doorbell.event_id))
-    return Err;
-
-  auto *DoorbellPtr = reinterpret_cast<rpc::Doorbell *>(
-      static_cast<uint8_t *>(RPCBuffer) + rpc::Server::doorbell_offset());
-  std::memcpy(DoorbellPtr, &Doorbell, sizeof(rpc::Doorbell));
+  void *RPCBuffer = Buffers[Device.getDeviceId()];
+  if (!RPCBuffer) {
+    auto RPCBufferOrErr = Device.allocate(
+        rpc::Server::allocation_size(Device.getRPCNumLanes(), NumPorts),
+        nullptr, TARGET_ALLOC_HOST);
+    if (!RPCBufferOrErr)
+      return RPCBufferOrErr.takeError();
+
+    RPCBuffer = *RPCBufferOrErr;
+    if (!RPCBuffer)
+      return plugin::Plugin::error(
+          error::ErrorCode::UNKNOWN,
+          "failed to initialize RPC server for device %d",
+          Device.getDeviceId());
+
+    // The doorbell is used by AMDGPU targets to let the server thread be
+    // descheduled. It is optional and will be ignored if the fields are null.
+    rpc::Doorbell Doorbell{};
+    if (auto Err = Device.Plugin.initRPCDoorbell(
+            Doorbell.value, Doorbell.mailbox, Doorbell.event_id))
+      return Err;
+
+    auto *DoorbellPtr = reinterpret_cast<rpc::Doorbell *>(
+        static_cast<uint8_t *>(RPCBuffer) + rpc::Server::doorbell_offset());
+    std::memcpy(DoorbellPtr, &Doorbell, sizeof(rpc::Doorbell));
+
+    Buffers[Device.getDeviceId()] = RPCBuffer;
+    Devices[Device.getDeviceId()] = &Device;
+  }
 
-  // Get the address of the RPC client from the device.
+  // Each image has its own client that must point at the shared buffer.
----------------
ro-i wrote:

ig from here on, it wouldn't be strictly necessary to still hold the lock? But it probably doesn't matter

https://github.com/llvm/llvm-project/pull/218757