[OpenMP] Increase opportunity for parallel kernel launch in AMDGPUs: add multiple hsa queue's per device in plugin

[carlo.bertolli]

Summary

This patch extends the AMDGPU plugin for OpenMP target offloading from using a single HSA queue to multiple queues (four in this patch) per device. This enables concurrent threads to concurrently submit kernel launches to the same GPU.

[JonChesterfield]

LG, thanks!

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

 }
 
 struct HSALifetime {
   // Wrapper around HSA used to ensure it is constructed before other types
   // and destructed after, which means said other types can use raii for
   // cleanup without risking running outside of the lifetime of HSA
   const hsa_status_t S;
 
-  bool success() { return S == HSA_STATUS_SUCCESS; }
+  bool HSAInitSuccess() { return S == HSA_STATUS_SUCCESS; }
   HSALifetime() : S(hsa_init()) {}
 
   ~HSALifetime() {
     if (S == HSA_STATUS_SUCCESS) {
       hsa_status_t Err = hsa_shut_down();
       if (Err != HSA_STATUS_SUCCESS) {
         // Can't call into HSA to get a string from the integer
         DP("Shutting down HSA failed: %d\n", Err);
       }
     }
   }
 };
 
+// Handle scheduling of multiple hsa_queue's per device to
+// multiple threads (one scheduler per device)
+class HSAQueueScheduler {
+public:
+  HSAQueueScheduler() : current(0) {}
+
+  HSAQueueScheduler(const HSAQueueScheduler &) = delete;
+
+  HSAQueueScheduler(HSAQueueScheduler &&q) {
+    current = q.current.load();
+    for (uint8_t i = 0; i < NUM_QUEUES_PER_DEVICE; i++) {
+      HSAQueues[i] = q.HSAQueues[i];
+      q.HSAQueues[i] = nullptr;
+    }
+  }
+
+  // \return false if any HSA queue creation fails
+  bool CreateQueues(hsa_agent_t HSAAgent, uint32_t queue_size) {
+    for (uint8_t i = 0; i < NUM_QUEUES_PER_DEVICE; i++) {
+      hsa_queue_t *Q = nullptr;
+      hsa_status_t rc =
+          hsa_queue_create(HSAAgent, queue_size, HSA_QUEUE_TYPE_MULTI,
+                           callbackQueue, NULL, UINT32_MAX, UINT32_MAX, &Q);
+      if (rc != HSA_STATUS_SUCCESS) {
+        DP("Failed to create HSA queue %d\n", i);
+        return false;
+      }
+      HSAQueues[i] = Q;
+    }
+    return true;
+  }
+
+  ~HSAQueueScheduler() {
+    for (uint8_t i = 0; i < NUM_QUEUES_PER_DEVICE; i++) {
+      if (HSAQueues[i]) {
+        hsa_status_t err = hsa_queue_destroy(HSAQueues[i]);
+        if (err != HSA_STATUS_SUCCESS)
+          DP("Error destroying HSA queue");
+      }
+    }
+  }
+
+  // \return next queue to use for device
+  hsa_queue_t *Next() {
+    return HSAQueues[(current.fetch_add(1, std::memory_order_relaxed)) %
+                     NUM_QUEUES_PER_DEVICE];
+  }
+
+private:
+  // Number of queues per device
+  enum : uint8_t { NUM_QUEUES_PER_DEVICE = 4 };
+  hsa_queue_t *HSAQueues[NUM_QUEUES_PER_DEVICE] = {};
+  std::atomic<uint8_t> current;
+};
+
 /// Class containing all the device information
-class RTLDeviceInfoTy {
-  HSALifetime HSA; // First field => constructed first and destructed last
+class RTLDeviceInfoTy : HSALifetime {
   std::vector<std::list<FuncOrGblEntryTy>> FuncGblEntries;
 
   struct QueueDeleter {
     void operator()(hsa_queue_t *Q) {
       if (Q) {
         hsa_status_t Err = hsa_queue_destroy(Q);
         if (Err != HSA_STATUS_SUCCESS) {
           DP("Error destroying hsa queue: %s\n", get_error_string(Err));
         }
       }
     }
   };
 
 public:
   bool ConstructionSucceeded = false;
 
   // load binary populates symbol tables and mutates various global state
   // run uses those symbol tables
   std::shared_timed_mutex load_run_lock;
 
   int NumberOfDevices = 0;
 
   // GPU devices
   std::vector<hsa_agent_t> HSAAgents;
-  std::vector<std::unique_ptr<hsa_queue_t, QueueDeleter>>
-      HSAQueues; // one per gpu
+  std::vector<HSAQueueScheduler> HSAQueueSchedulers; // one per gpu
 
   // CPUs
   std::vector<hsa_agent_t> CPUAgents;
 
   // Device properties
   std::vector<int> ComputeUnits;
   std::vector<int> GroupsPerDevice;
   std::vector<int> ThreadsPerGroup;
RTLDeviceInfoTy() {
     DP("Start initializing " GETNAME(TARGET_NAME) "\n");
 
     // LIBOMPTARGET_KERNEL_TRACE provides a kernel launch trace to stderr
     // anytime. You do not need a debug library build.
     //  0 => no tracing
     //  1 => tracing dispatch only
     // >1 => verbosity increase
 
-    if (!HSA.success()) {
+    if (!HSAInitSuccess()) {
       DP("Error when initializing HSA in " GETNAME(TARGET_NAME) "\n");
       return;
     }
 
     if (char *envStr = getenv("LIBOMPTARGET_KERNEL_TRACE"))
       print_kernel_trace = atoi(envStr);
     else
       print_kernel_trace = 0;
RTLDeviceInfoTy() {
     if (NumberOfDevices == 0) {
       DP("There are no devices supporting HSA.\n");
       return;
     } else {
       DP("There are %d devices supporting HSA.\n", NumberOfDevices);
     }
 
     // Init the device info
-    HSAQueues.resize(NumberOfDevices);
+    HSAQueueSchedulers.reserve(NumberOfDevices);
     FuncGblEntries.resize(NumberOfDevices);
     ThreadsPerGroup.resize(NumberOfDevices);
     ComputeUnits.resize(NumberOfDevices);
     GPUName.resize(NumberOfDevices);
     GroupsPerDevice.resize(NumberOfDevices);
     WarpSize.resize(NumberOfDevices);
     NumTeams.resize(NumberOfDevices);
     NumThreads.resize(NumberOfDevices);
for (int i = 0; i < NumberOfDevices; i++) {
         }
         enum { MaxQueueSize = 4096 };
         if (queue_size > MaxQueueSize) {
           queue_size = MaxQueueSize;
         }
       }
 
       {
-        hsa_queue_t *Q = nullptr;
-        hsa_status_t rc =
-            hsa_queue_create(HSAAgents[i], queue_size, HSA_QUEUE_TYPE_MULTI,
-                             callbackQueue, NULL, UINT32_MAX, UINT32_MAX, &Q);
-        if (rc != HSA_STATUS_SUCCESS) {
-          DP("Failed to create HSA queue %d\n", i);
+        HSAQueueScheduler QSched;
+        if (!QSched.CreateQueues(HSAAgents[i], queue_size))
           return;
-        }
-        HSAQueues[i].reset(Q);
+        HSAQueueSchedulers.emplace_back(std::move(QSched));
       }
 
       deviceStateStore[i] = {nullptr, 0};
     }
 
     for (int i = 0; i < NumberOfDevices; i++) {
       ThreadsPerGroup[i] = RTLDeviceInfoTy::Default_WG_Size;
       GroupsPerDevice[i] = RTLDeviceInfoTy::DefaultNumTeams;
RTLDeviceInfoTy() {
     // Default state.
     RequiresFlags = OMP_REQ_UNDEFINED;
 
     ConstructionSucceeded = true;
   }
 
   ~RTLDeviceInfoTy() {
     DP("Finalizing the " GETNAME(TARGET_NAME) " DeviceInfo.\n");
-    if (!HSA.success()) {
+    if (!HSAInitSuccess()) {
       // Then none of these can have been set up and they can't be torn down
       return;
     }
     // Run destructors on types that use HSA before
     // impl_finalize removes access to it
     deviceStateStore.clear();
     KernelArgPoolMap.clear();
     // Terminate hostrpc before finalizing hsa
fprintf(traceToStdout ? stdout : stderr,
             device_id, KernelInfo->ExecutionMode, KernelInfo->ConstWGSize,
             arg_num, num_groups, WorkgroupSize, num_teams, thread_limit,
             group_segment_size, sgpr_count, vgpr_count, sgpr_spill_count,
             vgpr_spill_count, loop_tripcount, KernelInfo->Name);
   }
 
   // Run on the device.
   {
-    hsa_queue_t *queue = DeviceInfo.HSAQueues[device_id].get();
+    hsa_queue_t *queue = DeviceInfo.HSAQueueSchedulers[device_id].Next();
     if (!queue) {
       return OFFLOAD_FAIL;
     }
     uint64_t packet_id = acquire_available_packet_id(queue);
 
     const uint32_t mask = queue->size - 1; // size is a power of 2
     hsa_kernel_dispatch_packet_t *packet =
         (hsa_kernel_dispatch_packet_t *)queue->base_address +