[Libomptarget] Configure the RPC port count from the plugin

[jhuber6]

Summary

This patch allows us to configure the port count to what the specific
card would desire for parallelism. For AMDGPU we need to use the maximum
number of hardware parallelism to avoid deadlocks. For NVPTX we don't
have this problem due to the friendlier scheduler, so we use the number
of warps active on an SM times the number of SMs as a good guess.

Note that the max ports currently is going to be smaller than these
numbers. That will be improved in the future.

[JonChesterfield]

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

1660

This is valid if a wave opens at most one port at a time. I think an argument could be made that a wave could try to open one port per thread. Likely to be something we should add to the libc tests.

Async also comprises the sizing, though I currently think we should limit openmp to synchronous calls.

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

770

Typos. Also isn't num_teams the wrong constant here? Should be whatever openmp calls warps

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

64

This should probably be a hard error if the plugin wants more ports than it can have, especially on platforms where that implies deadlock risk

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

391

It'll affect performance.

Also I'm not totally confident Nvidia has a fair scheduler on SMs, that would be a good thing to check. Amdgpu does not have a fair scheduler on CUs

899

Typo. Also doesn't seem to match the comment

[jdoerfert]

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

2471

In a follow up, can you please expose this via ompx_get_hardware_num_processing_elements(int device) (or whatever the non device version is spelled)?

[jhuber6]

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

2471

Yeah I'm not sure how to calculate this on CUDA unfortunately.

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

64

Was putting that off considering that the limit is currently like 64, which is far below the like 2000 that most platforms will want.

[jhuber6]

Updating

[jhuber6]

ping

[jhuber6]

ping

[jdoerfert]

LG, fix comments.

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

Error initImpl(GenericPluginTy &Plugin) override {
 
     // Compute the default number of teams.
     uint32_t ComputeUnits = 0;
     if (auto Err =
             getDeviceAttr(HSA_AMD_AGENT_INFO_COMPUTE_UNIT_COUNT, ComputeUnits))
       return Err;
     GridValues.GV_Default_Num_Teams = ComputeUnits * OMPX_DefaultTeamsPerCU;
 
+    uint32_t WavesPerCU = 0;
+    if (auto Err =
+            getDeviceAttr(HSA_AMD_AGENT_INFO_MAX_WAVES_PER_CU, WavesPerCU))
+      return Err;
+    HardwareParallelism = ComputeUnits * WavesPerCU;
+
     // Get maximum size of any device queues and maximum number of queues.
     uint32_t MaxQueueSize;
     if (auto Err = getDeviceAttr(HSA_AGENT_INFO_QUEUE_MAX_SIZE, MaxQueueSize))
       return Err;
 
     uint32_t MaxQueues;
     if (auto Err = getDeviceAttr(HSA_AGENT_INFO_QUEUES_MAX, MaxQueues))
       return Err;
struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   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();
   }
 
+  /// AMDGPU returns the product of the number of compute units and the waves
+  /// per compute unit.
+  uint64_t requestedRPCPortCount() const override  {
+    return HardwareParallelism;
+  }
+
   /// 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;
   }
 
private:
   hsa_agent_t Agent;
 
   /// The GPU architecture.
   std::string ComputeUnitKind;
 
   /// The frequency of the steady clock inside the device.
   uint64_t ClockFrequency;
 
+  /// The total number of concurrent work items that can be running on the GPU.
+  uint64_t HardwareParallelism;
+
   /// Reference to the host device.
   AMDHostDeviceTy &HostDevice;
 
   /// List of device packet queues.
   std::vector<AMDGPUQueueTy> Queues;
 };
 
 Error AMDGPUDeviceImageTy::loadExecutable(const AMDGPUDeviceTy &Device) {

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

struct GenericDeviceTy : public DeviceAllocatorTy {
   /// @see OMPX_MinThreadsForLowTripCount
   virtual uint32_t getMinThreadsForLowTripCountLoop() {
     return OMPX_MinThreadsForLowTripCount;
   }
 
   /// Get the RPC server running on this device.
   RPCServerTy *getRPCServer() const { return RPCServer; }
 
+  /// The number of parallel RPC ports to use on the device. In general, this
+  /// should be roughly equivalent to the amount of hardware parallelism the
+  /// device can support. This is because GPUs in general do not have forward
+  /// progress guarantees, so we minimize thread level dependencies by
+  /// allocating enough space such that each device thread can have a port. This
+  /// is likely overly pessimistic in the average case, but guarantees no
+  /// deadlocks at the cost of memory. This must be overloaded by targets
+  /// expecting to use the RPC server.
+  virtual uint64_t requestedRPCPortCount() const {
+    assert(!shouldSetupRPCServer() && "Default implementation cannot be used");
+    return 0;
+  }
+
 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,
#define defineOmptCallback(Name, Type, Code) Name##_t Name##_fn = nullptr;
   FOREACH_OMPT_DEVICE_EVENT(defineOmptCallback)
 #undef defineOmptCallback
 
   /// Internal representation for OMPT device (initialize & finalize)
   std::atomic<bool> OmptInitialized;
 #endif
 
 private:
-
   /// Return the kernel environment object for kernel \p Name.
   Expected<KernelEnvironmentTy>
   getKernelEnvironmentForKernel(StringRef Name, DeviceImageTy &Image);
 };
 
 /// 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.

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

Error RPCServerTy::initDevice(plugin::GenericDeviceTy &Device,
                               plugin::DeviceImageTy &Image) {
 #ifdef LIBOMPTARGET_RPC_SUPPORT
   uint32_t DeviceId = Device.getDeviceId();
   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(DeviceId, RPC_MAXIMUM_PORT_COUNT,
+  uint64_t NumPorts =
+      std::min(Device.requestedRPCPortCount(), RPC_MAXIMUM_PORT_COUNT);
+  if (rpc_status_t Err = rpc_server_init(DeviceId, NumPorts,
                                          Device.getWarpSize(), Alloc, &Device))
     return plugin::Plugin::error(
         "Failed to initialize RPC server for device %d: %d", DeviceId, 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(

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

Error initImpl(GenericPluginTy &Plugin) override {
     if (auto Err = getDeviceAttr(CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
                                  ComputeCapability.Major))
       return Err;
 
     if (auto Err = getDeviceAttr(CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
                                  ComputeCapability.Minor))
       return Err;
 
+    uint32_t NumMuliprocessors = 0;
+    uint32_t MaxThreadsPerSM = 0;
+    uint32_t WarpSize = 0;
+    if (auto Err = getDeviceAttr(CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT,
+                                 NumMuliprocessors))
+      return Err;
+    if (auto Err = getDeviceAttr(CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_MULTIPROCESSOR,
+                                 MaxThreadsPerSM))
+      return Err;
+    if (auto Err = getDeviceAttr(CU_DEVICE_ATTRIBUTE_WARP_SIZE, WarpSize))
+      return Err;
+    HardwareParallelism = NumMuliprocessors * (MaxThreadsPerSM / WarpSize);
+
     return Plugin::success();
   }
 
   /// Deinitialize the device and release its resources.
   Error deinitImpl() override {
     if (Context) {
       if (auto Err = setContext())
         return Err;
struct CUDADeviceTy : public GenericDeviceTy {
   }
 
   /// We want to set up the RPC server for host services to the GPU if it is
   /// availible.
   bool shouldSetupRPCServer() const override {
     return libomptargetSupportsRPC();
   }
 
+  /// NVIDIA returns the produce number of SMs and the number of warps that fit
+  /// if the maximum number of threads were scheduled on each SM.
+  uint64_t requestedRPCPortCount() const override {
+    return HardwareParallelism;
+  }
+
   /// 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;
   }
 
private:
   /// The compute capability of the corresponding CUDA device.
   struct ComputeCapabilityTy {
     uint32_t Major;
     uint32_t Minor;
     std::string str() const {
       return "sm_" + std::to_string(Major * 10 + Minor);
     }
   } ComputeCapability;
+
+  /// The maximum number of warps that can be resident on all the SMs
+  /// simultaneously.
+  uint32_t HardwareParallelism = 0;
 };
 
 Error CUDAKernelTy::launchImpl(GenericDeviceTy &GenericDevice,
                                uint32_t NumThreads, uint64_t NumBlocks,
                                KernelArgsTy &KernelArgs, void *Args,
                                AsyncInfoWrapperTy &AsyncInfoWrapper) const {
   CUDADeviceTy &CUDADevice = static_cast<CUDADeviceTy &>(GenericDevice);