[Libomptarget] Begin implementing support for RPC services

[jhuber6]

Summary

This patch adds the intial support for running an RPC server in
libomptarget to handle host services. We interface with the library
provided by the libc project to stand up a basic server. We introduce
a new type that is controlled by the plugin and has each device
intialize its interface. We then run a basic server to check the RPC
buffer.

This patch does not fully implement the interface. In the future each
plugin will want to define special handlers via the interface to support
things like malloc or H2D copies coming from RPC. We will also want to
allow the plugin to specify t he number of ports. This is currently
capped in the implementation but will be adjusted soon.

Right now running the server is handled by whatever thread ends up doing
the waiting. This is probably not a completely sound solution but I am
not overly familiar with the behaviour of OpenMP tasks and what would be
required here. This works okay with synchrnous regions, and somewhat
fine with nowait regions, but I've observed some weird behavior when
one of those regions calls exit.

[jdoerfert]

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

533

This I don't get. We wait for 8k units, then we run the server once, then we wait again if the kernel is not done. Why isn't this part of the loop below with a timeout that adjusts based on the RPC needs?

1071

We can grab that information in the constructor.

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

558

Can you hoist Plugin.getRPCServer(), to make it more readable.

I'm not a fan of RPCDevice as a name, we have too many "devices"

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

854

docs

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

70–105

Can the GenericDevice register these handlers?
A vector of them somewhere, or just let it call a register callback.
I would not make these special and define them here.
Move all into the GenericDevice and it will call an Impl method for the subclasses to register more.

135

This seems like a bad idea. We should use the device ID to index the Devices array, and we really need to rename this.

[jhuber6]

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

533

I wasn't sure how to mesh it with the existing logic, since we will presumably have streams that didn't come from kernel launches and we have some active timeout exception here where I guess we query at a higher rate compared to the blocking one? I just worked around it but I could try to integrate it. AMDGPU was much more difficult here than CUDA because of the indirection.

1071

I tried that, turns out we initialize these streams before we load the binary image, which is requited to setup the RPC server since we first need to know where the symbol is to initialize it.

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

18

Forgot to remove this.

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

558

Sure, I'll try to think of a better name.

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

854

You mean explain this in more detail?

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

70–105

So this will probably only be necessary for things like malloc / memcpy. Figured those are already provided by the generic device so we should be able to hide that complexity in here right next to the implementation. Didn't feel like it warranted abstraction.

135

Yeah you can make multiple. Only reason I didn't want because you'd need to allocate them because of the references, but we can just std::make_unique.

[tianshilei1992]

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

259

looks like irrelevant changes

[jhuber6]

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

259

People somehow commit changes without clang-formatting and then it ends up here because my editor will auto format. Somewhat annoying because I've clang formatted all of these files at least twice by now. I'll trim it.

[jhuber6]

Addressing comments

[jdoerfert]

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

533

I'm not convinced the thread starting the thing should run the server, but there is merit to having this option. However, the timeout still seems wrong to me. The fact that we have two locations we wait on (finish and RPC) is also not great.

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

558

RPC adapter? RCP service, ...

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

259

Set your editor to only format the lines you changed. Or clang format and precommit.

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

70–105

But this takes away the extension point, and splits the callbacks. We might want to use different mallocs per kernel. And having all callbacks in "one place" seems like a good idea as well.

[jhuber6]

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

533

Yeah, the alternative is to run a dedicated server thread, but I asked Shilei and he said that this should be fine. When running the thread I ran into problems, e.g. a synchronous kernel can exit without waiting for any async RPC calls to finish (exit, free, (maybe an async print variant?). Secondly when exiting it was tough to tell when the thread was alive or who should clean things up.

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

70–105

I figured that this portion of the malloc would be pretty straightforward, since I think its usage should be more similar to setting the breakpoint, i.e. we only do the RPC path if we run out of some preallocated buffer space.

[jhuber6]

Merge loops, Instance -> Handle.

[jdoerfert]

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

913

Name, also elsewhere.

2857–2858

auto RPCHandle = ...
set(RPCHandle)

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

259

recommit

[jhuber6]

Address comments

[jhuber6]

Add a test. This creates some transitive dependencies on the libcgpu.a library
exported from libc. I had to change it to emit it in a location that we can
find in the tests at least for the build process. Secondly, this cannot use the
wrapper headers as those are create in-place by the installation logic in the
resource dir. We'll need to declare the things on the GPU manually. Third, this
cannot have a dependency on the libcgpu.a target as that could potentially be
in another CMake invocation. So be aware that changes to that will not be
triggered by a re-run of check-openmp if the libc project was changed.

[jdoerfert]

LG, minor nits

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

259

unrelated. precommit

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

51

I'd again swap the cases, also below.

137

Hoist getDeviceId, 4 times makes it hard to read.

140

Check spelling

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

45

private:

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

482

Swap the cases, simple first.

[jhuber6]

Address nits

[arsenm]

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

1145

Can you use unique_ptr?

[jhuber6]

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

1145

I was keeping it consistent with the GlobalHandler above. I think there was a specific reason for why we did that originally but I don't remember off the top of my head. Might've been something about ordering of the destructors.

[arsenm]

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

1145

Also don't need the null checks

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) const {
+  Error wait(const uint64_t ActiveTimeout = 0,
+             RPCInstanceTy *RPCDevice = nullptr) const {
+    // If there is an RPC device attached to this stream we run it as a server.
+    if (RPCDevice) {
+      while (hsa_signal_wait_scacquire(Signal, HSA_SIGNAL_CONDITION_EQ, 0,
+                                       /*timeout_hint=*/8192,
+                                       HSA_WAIT_STATE_ACTIVE) != 0) {
+        if (auto Err = RPCDevice->runServer())
+          return Err;
+      }
+      return Plugin::success();
+    }
+
     if (ActiveTimeout) {
       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();
     }
     while (hsa_signal_wait_scacquire(Signal, HSA_SIGNAL_CONDITION_EQ, 0,
private:
   /// 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.
+  RPCInstanceTy *RPCDevice;
+
   /// 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.
   uint32_t size() const { return NextSlot; }
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 device to this stream.
+  void setRPCDevice(RPCInstanceTy *Device) { RPCDevice = Device; }
+
   /// 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,
                          AMDGPUMemoryManagerTy &MemoryManager) {
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))
+    if (auto Err =
+            Slots[last()].Signal->wait(StreamBusyWaitMicroseconds, RPCDevice))
       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
constructKernelEntry(const __tgt_offload_entry &KernelEntry,
 
     return AMDKernel;
   }
 
   /// Set the current context to this device's context. Do nothing since the
   /// AMDGPU devices do not have the concept of contexts.
   Error setContext() override { return Plugin::success(); }
 
+  /// We want to set up the RPC server for host services to the GPU if it is
+  /// availible.
+  bool shouldSetupRPCServer() const override {
+    return libomptargetSupportsRPC();
+  }
+
   /// Get the stream of the asynchronous info sructure or get a new one.
   AMDGPUStreamTy &getStream(AsyncInfoWrapperTy &AsyncInfoWrapper) {
     AMDGPUStreamTy *&Stream = AsyncInfoWrapper.getQueueAs<AMDGPUStreamTy *>();
     if (!Stream)
       Stream = AMDGPUStreamManager.getResource();
     return *Stream;
   }
 
Error AMDGPUResourceRef<ResourceTy>::create(GenericDeviceTy &Device) {
 
   return Resource->init();
 }
 
 AMDGPUStreamTy::AMDGPUStreamTy(AMDGPUDeviceTy &Device)
     : Agent(Device.getAgent()), Queue(Device.getNextQueue()),
       SignalManager(Device.getSignalManager()),
       // Initialize the std::deque with some empty positions.
-      Slots(32), NextSlot(0), SyncCycle(0),
+      Slots(32), NextSlot(0), SyncCycle(0), RPCDevice(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,
   //     alternative to copying them twice.
   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.getRPCInstance())
+    Stream.setRPCDevice(GenericDevice.getRPCInstance());
+
   // Push the kernel launch into the stream.
   return Stream.pushKernelLaunch(*this, AllArgs, NumThreads, NumBlocks,
                                  GroupSize, ArgsMemoryManager);
 }
 
 Error AMDGPUKernelTy::printLaunchInfoDetails(GenericDeviceTy &GenericDevice,
                                              KernelArgsTy &KernelArgs,
                                              uint32_t NumThreads,

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

 ##===----------------------------------------------------------------------===##
 #
 # 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
 #
 ##===----------------------------------------------------------------------===##
 #
 # Common parts which can be used by all plugins
 #
 ##===----------------------------------------------------------------------===##
 
 # NOTE: Don't try to build `PluginInterface` using `add_llvm_library` because we
 # don't want to export `PluginInterface` while `add_llvm_library` requires that.
 add_library(PluginInterface OBJECT
-  PluginInterface.cpp GlobalHandler.cpp JIT.cpp)
+  PluginInterface.cpp GlobalHandler.cpp JIT.cpp RPC.cpp)
 
 # Only enable JIT for those targets that LLVM can support.
 string(TOUPPER "${LLVM_TARGETS_TO_BUILD}" TargetsSupported)
 foreach(Target ${TargetsSupported})
   target_compile_definitions(PluginInterface PRIVATE "LIBOMPTARGET_JIT_${Target}")
 endforeach()
 
 # This is required when using LLVM libraries.
 llvm_update_compile_flags(PluginInterface)
 

 
 target_link_libraries(PluginInterface
   PUBLIC
     ${llvm_libs}
     elf_common
     MemoryManager
 )
 
+# Include the RPC server from the `libc` project if availible.
+if(TARGET llvmlibc_rpc_server)
+  target_link_libraries(PluginInterface PRIVATE llvmlibc_rpc_server)
+  target_compile_definitions(PluginInterface PRIVATE LIBOMPTARGET_RPC_SUPPORT)
+else()
+  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()
 
 # Define the TARGET_NAME and DEBUG_PREFIX.
 target_compile_definitions(PluginInterface PRIVATE
   TARGET_NAME="PluginInterface"
   DEBUG_PREFIX="PluginInterface"

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

 #include <shared_mutex>
 #include <vector>
 
 #include "Debug.h"
 #include "DeviceEnvironment.h"
 #include "GlobalHandler.h"
 #include "JIT.h"
 #include "MemoryManager.h"
+#include "RPC.h"
 #include "Utilities.h"
 #include "omptarget.h"
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include "llvm/Frontend/OpenMP/OMPGridValues.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Error.h"
struct GenericKernelTy {
 
   /// Launch the kernel on the specific device. The device must be the same
   /// one used to initialize the kernel.
   Error launch(GenericDeviceTy &GenericDevice, void **ArgPtrs,
                ptrdiff_t *ArgOffsets, KernelArgsTy &KernelArgs,
                AsyncInfoWrapperTy &AsyncInfoWrapper) const;
   virtual Error launchImpl(GenericDeviceTy &GenericDevice, uint32_t NumThreads,
                            uint64_t NumBlocks,
                            KernelArgsTy &KernelArgs, void *Args,
                            AsyncInfoWrapperTy &AsyncInfoWrapper) const = 0;
 
   /// Get the kernel name.
   const char *getName() const { return Name; }
 
   /// Indicate whether an execution mode is valid.
   static bool isValidExecutionMode(OMPTgtExecModeFlags ExecutionMode) {
     switch (ExecutionMode) {
struct GenericDeviceTy : public DeviceAllocatorTy {
   virtual Expected<DeviceImageTy *>
   loadBinaryImpl(const __tgt_device_image *TgtImage, int32_t ImageId) = 0;
 
   /// Setup the device environment if needed. Notice this setup may not be run
   /// on some plugins. By default, it will be executed, but plugins can change
   /// this behavior by overriding the shouldSetupDeviceEnvironment function.
   Error setupDeviceEnvironment(GenericPluginTy &Plugin, DeviceImageTy &Image);
 
+  // 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);
+
   /// Register the offload entries for a specific image on the device.
   Error registerOffloadEntries(DeviceImageTy &Image);
 
   /// Synchronize the current thread with the pending operations on the
   /// __tgt_async_info structure.
   Error synchronize(__tgt_async_info *AsyncInfo);
   virtual Error synchronizeImpl(__tgt_async_info &AsyncInfo) = 0;
 
struct GenericDeviceTy : public DeviceAllocatorTy {
   /// The minimum number of threads we use for a low-trip count combined loop.
   /// 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.
+  RPCInstanceTy *getRPCInstance() const { return RPCInstance; }
+
 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.
   Error registerKernelOffloadEntry(DeviceImageTy &DeviceImage,
private:
   virtual Error getDeviceHeapSize(uint64_t &V) = 0;
   virtual Error setDeviceHeapSize(uint64_t V) = 0;
 
   /// Indicate whether the device should setup the device environment. Notice
   /// that returning false in this function will change the behavior of the
   /// setupDeviceEnvironment() function.
   virtual bool shouldSetupDeviceEnvironment() const { return true; }
 
+  /// Indicate whether or not the device should setup the RPC server. This is
+  /// only necessary for unhosted targets like the GPU.
+  virtual bool shouldSetupRPCServer() const { return false; }
+
   /// Pointer to the memory manager or nullptr if not available.
   MemoryManagerTy *MemoryManager;
 
   /// Environment variables defined by the OpenMP standard.
   Int32Envar OMP_TeamLimit;
   Int32Envar OMP_NumTeams;
   Int32Envar OMP_TeamsThreadLimit;
 
protected:
   /// the device I has a matrix PeerAccesses with PeerAccesses[J] == AVAILABLE,
   /// the device I can access device J's memory directly. However, notice this
   /// does not mean that device J can access device I's memory directly.
   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.
+  RPCInstanceTy *RPCInstance;
 };
 
 /// Class implementing common functionalities of offload plugins. Each plugin
 /// should define the specific plugin class, derive from this generic one, and
 /// implement the necessary virtual function members.
 struct GenericPluginTy {
 
   /// Construct a plugin instance.
GenericGlobalHandlerTy &getGlobalHandler() {
     assert(GlobalHandler && "Global handler not initialized");
     return *GlobalHandler;
   }
 
   /// Get the reference to the JIT used for all devices connected to this
   /// plugin.
   JITEngine &getJIT() { return JIT; }
 
+  /// Get a reference to the RPC server used to provide host services.
+  RPCServerTy &getRPCServer() {
+    assert(RPCServer && "RPC server not initialized");
+    return *RPCServer;
+  }
+
   /// Get the OpenMP requires flags set for this plugin.
   int64_t getRequiresFlags() const { return RequiresFlags; }
 
   /// Set the OpenMP requires flags for this plugin.
   void setRequiresFlag(int64_t Flags) { RequiresFlags = Flags; }
 
   /// Initialize a device within the plugin.
   Error initDevice(int32_t DeviceId);
private:
   /// Pointer to the global handler for this plugin.
   GenericGlobalHandlerTy *GlobalHandler;
 
   /// Internal allocator for different structures.
   BumpPtrAllocator Allocator;
 
   /// The JIT engine shared by all devices connected to this plugin.
   JITEngine JIT;
+
+  /// The interface between the plugin and the GPU for host services.
+  RPCServerTy *RPCServer;
 };
 
 /// Class for simplifying the getter operation of the plugin. Anywhere on the
 /// code, the current plugin can be retrieved by Plugin::get(). The class also
 /// declares functions to create plugin-specific object instances. The check(),
 /// createPlugin(), createDevice() and createGlobalHandler() functions should be
 /// defined by each plugin implementation.
 class Plugin {
private:
 
   /// The next available resource in the pool.
   uint32_t NextAvailable;
 
   /// The actual resource pool.
   std::deque<ResourceRef> ResourcePool;
 };
 
+/// A static check on whether or not we support RPC in libomptarget.
+const bool libomptargetSupportsRPC();
+
 } // namespace plugin
 } // namespace target
 } // namespace omp
 } // namespace llvm
 
 #endif // OPENMP_LIBOMPTARGET_PLUGINS_COMMON_PLUGININTERFACE_H

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) {}
+      PinnedAllocs(*this), RPCInstance(nullptr) {}
 
 Error GenericDeviceTy::init(GenericPluginTy &Plugin) {
   if (auto Err = initImpl(Plugin))
     return Err;
 
   // Read and reinitialize the envars that depend on the device initialization.
   // Notice these two envars may change the stack size and heap size of the
   // device, so they need the device properly initialized.
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 (RPCInstance)
+    if (auto Err = RPCInstance->deinitDevice())
+      return std::move(Err);
+
   return deinitImpl();
 }
 
 Expected<__tgt_target_table *>
 GenericDeviceTy::loadBinary(GenericPluginTy &Plugin,
                             const __tgt_device_image *InputTgtImage) {
   assert(InputTgtImage && "Expected non-null target image");
   DP("Load data from image " DPxMOD "\n", DPxPTR(InputTgtImage->ImageStart));
GenericDeviceTy::loadBinary(GenericPluginTy &Plugin,
   // Setup the device environment if needed.
   if (auto Err = setupDeviceEnvironment(Plugin, *Image))
     return std::move(Err);
 
   // Register all offload entries of the image.
   if (auto Err = registerOffloadEntries(*Image))
     return std::move(Err);
 
+  if (auto Err = setupRPCServer(Plugin, *Image))
+    return std::move(Err);
+
   // Return the pointer to the table of entries.
   return Image->getOffloadEntryTable();
 }
 
 Error GenericDeviceTy::setupDeviceEnvironment(GenericPluginTy &Plugin,
                                               DeviceImageTy &Image) {
   // There are some plugins that do not need this step.
   if (!shouldSetupDeviceEnvironment())
Error GenericDeviceTy::setupDeviceEnvironment(GenericPluginTy &Plugin,
   if (auto Err = GHandler.writeGlobalToDevice(*this, Image, DevEnvGlobal)) {
     DP("Missing symbol %s, continue execution anyway.\n",
        DevEnvGlobal.getName().data());
     consumeError(std::move(Err));
   }
   return Plugin::success();
 }
 
+Error GenericDeviceTy::setupRPCServer(GenericPluginTy &Plugin,
+                                      DeviceImageTy &Image) {
+  // The plugin either does not need an RPC server or it is unavailible.
+  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 std::move(Err);
+
+  auto DeviceOrErr = Server.getDevice(*this);
+  if (!DeviceOrErr)
+    return DeviceOrErr.takeError();
+  RPCInstance = *DeviceOrErr;
+  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) {
     // The host should have always something in the address to uniquely
     // identify the entry.
     if (!Entry->addr)
       return Plugin::error("Failure to register entry without address");
if (NumDevices == 0)
     return Plugin::success();
 
   assert(Devices.size() == 0 && "Plugin already initialized");
   Devices.resize(NumDevices, nullptr);
 
   GlobalHandler = Plugin::createGlobalHandler();
   assert(GlobalHandler && "Invalid global handler");
 
+  RPCServer = new RPCServerTy(NumDevices);
+  assert(RPCServer && "Invalid RPC server");
+
   return Plugin::success();
 }
 
 Error GenericPluginTy::deinit() {
   // There is no global handler if no device is available.
   if (GlobalHandler)
     delete GlobalHandler;
 
   // Deinitialize all active devices.
   for (int32_t DeviceId = 0; DeviceId < NumDevices; ++DeviceId) {
     if (Devices[DeviceId]) {
       if (auto Err = deinitDevice(DeviceId))
         return Err;
     }
     assert(!Devices[DeviceId] && "Device was not deinitialized");
   }
 
+  if (RPCServer)
+    delete RPCServer;
+
   // Perform last deinitializations on the plugin.
   return deinitImpl();
 }
 
 Error GenericPluginTy::initDevice(int32_t DeviceId) {
   assert(!Devices[DeviceId] && "Device already initialized");
 
   // Create the device and save the reference.
Error GenericPluginTy::deinitDevice(int32_t DeviceId) {
 
   // Delete the device and invalidate its reference.
   delete Devices[DeviceId];
   Devices[DeviceId] = nullptr;
 
   return Plugin::success();
 }
 
+const bool llvm::omp::target::plugin::libomptargetSupportsRPC() {
+#ifdef LIBOMPTARGET_RPC_SUPPORT
+  return true;
+#else
+  return false;
+#endif
+}
+
 /// Exposed library API function, basically wrappers around the GenericDeviceTy
 /// functionality with the same name. All non-async functions are redirected
 /// to the async versions right away with a NULL AsyncInfoPtr.
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 int32_t __tgt_rtl_init_plugin() {

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

+//===- RPC.h - Interface for remote procedure calls from the GPU ----------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file provides the interface to support remote procedure calls (RPC) from
+// the GPU. This is required to implement host services like printf or malloc.
+// The interface to the RPC server is provided by the 'libc' project in LLVM.
+// For more information visit https://libc.llvm.org/gpu/.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef OPENMP_LIBOMPTARGET_PLUGINS_NEXTGEN_COMMON_RPC_H
+#define OPENMP_LIBOMPTARGET_PLUGINS_NEXTGEN_COMMON_RPC_H
+
+#include "llvm/Support/Error.h"
+
+#include <stdint.h>
+
+namespace llvm::omp::target {
+namespace plugin {
+struct GenericDeviceTy;
+struct GenericGlobalHandlerTy;
+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 RPCInstanceTy {
+    RPCInstanceTy(RPCServerTy &Server, plugin::GenericDeviceTy &Device)
+        : Server(Server), Device(Device) {}
+
+    llvm::Error runServer() { return Server.runServer(Device); }
+
+    llvm::Error deinitDevice() { return Server.deinitDevice(Device); }
+
+    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<RPCInstanceTy *> 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<RPCInstanceTy>> Instances;
+};
+
+using RPCInstanceTy = RPCServerTy::RPCInstanceTy;
+
+} // namespace llvm::omp::target
+
+#endif

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

+//===- RPC.h - Interface for remote procedure calls from the GPU ----------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "RPC.h"
+#include "Debug.h"
+#include "PluginInterface.h"
+
+// This header file may be present in-tree or from an LLVM installation. The
+// installed version lives alongside the GPU headers so we do not want to
+// include it directly.
+#if __has_include(<gpu-none-llvm/rpc_server.h>)
+#include <gpu-none-llvm/rpc_server.h>
+#elif defined(LIBOMPTARGET_RPC_SUPPORT)
+#include <rpc_server.h>
+#endif
+
+using namespace llvm;
+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);
+  Instances.resize(NumDevices);
+#endif
+}
+
+llvm::Expected<bool>
+RPCServerTy::isDeviceUsingRPC(plugin::GenericDeviceTy &Device,
+                              plugin::GenericGlobalHandlerTy &Handler,
+                              plugin::DeviceImageTy &Image) {
+#ifdef LIBOMPTARGET_RPC_SUPPORT
+  void *ClientPtr;
+  plugin::GlobalTy Global(rpc_client_symbol_name, sizeof(void *), &ClientPtr);
+  if (auto Err = Handler.readGlobalFromImage(Device, Image, Global)) {
+    llvm::consumeError(std::move(Err));
+    return false;
+  }
+
+  return true;
+#else
+  return false;
+#endif
+}
+
+Error RPCServerTy::initDevice(plugin::GenericDeviceTy &Device,
+                              plugin::GenericGlobalHandlerTy &Handler,
+                              plugin::DeviceImageTy &Image) {
+#ifdef LIBOMPTARGET_RPC_SUPPORT
+  auto Alloc = [](uint64_t Size, void *Data) {
+    plugin::GenericDeviceTy &Device =
+        *reinterpret_cast<plugin::GenericDeviceTy *>(Data);
+    return Device.allocate(Size, nullptr, TARGET_ALLOC_HOST);
+  };
+  // TODO: Allow the device to declare its requested port count.
+  if (rpc_status_t Err =
+          rpc_server_init(Device.getDeviceId(), RPC_MAXIMUM_PORT_COUNT,
+                          Device.getWarpSize(), Alloc, &Device))
+    return plugin::Plugin::error(
+        "Failed to initialize RPC server for device %d: %d",
+        Device.getDeviceId(), Err);
+
+  // Register a custom opcode handler to perform plugin specific allocation.
+  // FIXME: We need to make sure this uses asynchronous allocations on CUDA.
+  auto MallocHandler = [](rpc_port_t Port, void *Data) {
+    rpc_recv_and_send(
+        Port,
+        [](rpc_buffer_t *Buffer, void *Data) {
+          plugin::GenericDeviceTy &Device =
+              *reinterpret_cast<plugin::GenericDeviceTy *>(Data);
+          Buffer->data[0] = reinterpret_cast<uintptr_t>(
+              Device.allocate(Buffer->data[0], nullptr, TARGET_ALLOC_DEVICE));
+        },
+        Data);
+  };
+  if (rpc_status_t Err = rpc_register_callback(Device.getDeviceId(), RPC_MALLOC,
+                                               MallocHandler, &Device))
+    return plugin::Plugin::error(
+        "Failed to register RPC malloc handler for device %d: %d\n",
+        Device.getDeviceId(), Err);
+
+  // Register a custom opcode handler to perform plugin specific deallocation.
+  auto FreeHandler = [](rpc_port_t Port, void *Data) {
+    rpc_recv(
+        Port,
+        [](rpc_buffer_t *Buffer, void *Data) {
+          plugin::GenericDeviceTy &Device =
+              *reinterpret_cast<plugin::GenericDeviceTy *>(Data);
+          Device.free(reinterpret_cast<void *>(Buffer->data[0]),
+                      TARGET_ALLOC_DEVICE);
+        },
+        Data);
+  };
+  if (rpc_status_t Err = rpc_register_callback(Device.getDeviceId(), RPC_FREE,
+                                               FreeHandler, &Device))
+    return plugin::Plugin::error(
+        "Failed to register RPC free handler for device %d: %d\n",
+        Device.getDeviceId(), Err);
+
+  // Get the address of the RPC client from the device.
+  void *ClientPtr;
+  plugin::GlobalTy ClientGlobal(rpc_client_symbol_name, sizeof(void *));
+  if (auto Err =
+          Handler.getGlobalMetadataFromDevice(Device, Image, ClientGlobal))
+    return Err;
+
+  if (auto Err = Device.dataRetrieve(&ClientPtr, ClientGlobal.getPtr(),
+                                     sizeof(void *), nullptr))
+    return Err;
+
+  const void *ClientBuffer = rpc_get_client_buffer(Device.getDeviceId());
+  if (auto Err = Device.dataSubmit(ClientPtr, ClientBuffer,
+                                   rpc_get_client_size(), nullptr))
+    return Err;
+
+  Instances[Device.getDeviceId()] =
+      std::make_unique<RPCInstanceTy>(*this, Device);
+#endif
+  return Error::success();
+}
+
+llvm::Expected<RPCInstanceTy *>
+RPCServerTy::getDevice(plugin::GenericDeviceTy &Device) {
+#ifdef LIBOMPTARGET_RPC_SUPPORT
+  if (!Instances[Device.getDeviceId()] ||
+      !rpc_get_buffer(Device.getDeviceId()) ||
+      !rpc_get_client_buffer(Device.getDeviceId()))
+    return plugin::Plugin::error(
+        "Attempt to get an RPC device while not initialized");
+
+  return Instances[Device.getDeviceId()].get();
+#else
+  return plugin::Plugin::error(
+      "Attempt to get an RPC device while not availible");
+#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();
+}
+
+Error RPCServerTy::deinitDevice(plugin::GenericDeviceTy &Device) {
+#ifdef LIBOMPTARGET_RPC_SUPPORT
+  auto Dealloc = [](void *Ptr, void *Data) {
+    plugin::GenericDeviceTy &Device =
+        *reinterpret_cast<plugin::GenericDeviceTy *>(Data);
+    Device.free(Ptr, TARGET_ALLOC_HOST);
+  };
+  if (rpc_status_t Err =
+          rpc_server_shutdown(Device.getDeviceId(), Dealloc, &Device))
+    return plugin::Plugin::error(
+        "Failed to shut down RPC server for device %d: %d",
+        Device.getDeviceId(), Err);
+#endif
+  return Error::success();
+}
+
+RPCServerTy::~RPCServerTy() {
+#ifdef LIBOMPTARGET_RPC_SUPPORT
+  rpc_shutdown();
+#endif
+}

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

struct CUDADeviceTy : public GenericDeviceTy {
   }
 
   /// Set the current context to this device's context.
   Error setContext() override {
     CUresult Res = cuCtxSetCurrent(Context);
     return Plugin::check(Res, "Error in cuCtxSetCurrent: %s");
   }
 
+  /// We want to set up the RPC server for host services to the GPU if it is
+  /// availible.
+  bool shouldSetupRPCServer() const override {
+    return libomptargetSupportsRPC();
+  }
+
   /// Get the stream of the asynchronous info sructure or get a new one.
   CUstream getStream(AsyncInfoWrapperTy &AsyncInfoWrapper) {
     CUstream &Stream = AsyncInfoWrapper.getQueueAs<CUstream>();
     if (!Stream)
       Stream = CUDAStreamManager.getResource();
     return Stream;
   }
 
if (auto Err = Plugin::check(Res, "Error in cuMemFree[Host]: %s")) {
       return OFFLOAD_FAIL;
     }
     return OFFLOAD_SUCCESS;
   }
 
   /// 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 = cuStreamSynchronize(Stream);
+    CUresult Res;
+    if (getRPCInstance()) {
+      do {
+        Res = cuStreamQuery(Stream);
+        if (auto Err = getRPCInstance()->runServer())
+          return Err;
+      } while (Res == CUDA_ERROR_NOT_READY);
+    } else {
+      Res = cuStreamSynchronize(Stream);
+    }
 
     // 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);
     AsyncInfo.Queue = nullptr;
 
     return Plugin::check(Res, "Error in cuStreamSynchronize: %s");