From cc7a197ff8830397c834898dafca79b364f3d2a7 Mon Sep 17 00:00:00 2001
From: Sean Keely <Sean.Keely@amd.com>
Date: Wed, 24 Jun 2020 17:04:43 -0500
Subject: [PATCH] Use correct queue id for gws assignment.

Change-Id: I9157951abce65b60dfe0e1db5819510a7703547f
---
 runtime/hsa-runtime/core/inc/amd_aql_queue.h       | 3 +++
 runtime/hsa-runtime/core/runtime/amd_aql_queue.cpp | 9 +++++++++
 runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp | 7 ++-----
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/runtime/hsa-runtime/core/inc/amd_aql_queue.h b/runtime/hsa-runtime/core/inc/amd_aql_queue.h
index 744bade58b..17f68b2b59 100644
--- a/runtime/hsa-runtime/core/inc/amd_aql_queue.h
+++ b/runtime/hsa-runtime/core/inc/amd_aql_queue.h
@@ -196,6 +196,9 @@ class AqlQueue : public core::Queue, private core::LocalSignal, public core::Doo
   /// @brief Update signal value using Release semantics
   void StoreRelease(hsa_signal_value_t value) override;
 
+  /// @brief Enable use of GWS from this queue.
+  hsa_status_t EnableGWS(int gws_slot_count);
+
  protected:
   bool _IsA(Queue::rtti_t id) const override { return id == &rtti_id_; }
 
diff --git a/runtime/hsa-runtime/core/runtime/amd_aql_queue.cpp b/runtime/hsa-runtime/core/runtime/amd_aql_queue.cpp
index 65ac0c07de..ca5e3deb64 100644
--- a/runtime/hsa-runtime/core/runtime/amd_aql_queue.cpp
+++ b/runtime/hsa-runtime/core/runtime/amd_aql_queue.cpp
@@ -1119,5 +1119,14 @@ void AqlQueue::InitScratchSRD() {
   amd_queue_.compute_tmpring_size = tmpring_size.u32All;
   return;
 }
+
+hsa_status_t AqlQueue::EnableGWS(int gws_slot_count) {
+  uint32_t discard;
+  auto status = hsaKmtAllocQueueGWS(queue_id_, gws_slot_count, &discard);
+  if (status != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  amd_queue_.hsa_queue.type = HSA_QUEUE_TYPE_COOPERATIVE;
+  return HSA_STATUS_SUCCESS;
+}
+
 }  // namespace amd
 }  // namespace rocr
diff --git a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
index 78cd7aa459..1cb61fc7fa 100644
--- a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
+++ b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
@@ -605,12 +605,9 @@ void GpuAgent::InitGWS() {
       throw AMD::hsa_exception(HSA_STATUS_ERROR_OUT_OF_RESOURCES,
                                "Internal queue creation failed.");
 
-    uint32_t discard;
-    auto status = hsaKmtAllocQueueGWS(queue->amd_queue_.hsa_queue.id, 1, &discard);
-    if (status != HSAKMT_STATUS_SUCCESS)
-      throw AMD::hsa_exception(HSA_STATUS_ERROR_OUT_OF_RESOURCES, "GWS allocation failed.");
+    auto err = static_cast<AqlQueue*>(queue.get())->EnableGWS(1);
+    if (err != HSA_STATUS_SUCCESS) throw AMD::hsa_exception(err, "GWS allocation failed.");
 
-    queue->amd_queue_.hsa_queue.type = HSA_QUEUE_TYPE_COOPERATIVE;
     gws_queue_.ref_ct_ = 0;
     return queue.release();
   });