From 37942c982a45845ea639b81e29f996dceff17da2 Mon Sep 17 00:00:00 2001
From: Sean Keely <Sean.Keely@amd.com>
Date: Sat, 15 Jan 2022 10:39:38 -0600
Subject: [PATCH] Add HSA_AMD_AGENT_INFO_COOPERATIVE_COMPUTE_UNIT_COUNT.

On gfx90a only a reduced number of CUs must be used for cooperative
dispatches due to CWSR and launcher interactions with asymetric
harvest.  We must use one fewer CUs per SE than the lowest count of
CUs on any SE.

Also adds env var HSA_COOP_CU_COUNT which enables the cooperative
CU count computation.  Set to 1 to enable the new computation.
This is an opt-in feature that will become enabled by default (opt-out)
in a future release.

Change-Id: Ifbb75ced3bbc15876eef44922c6a4f6fde8c4c28
---
 runtime/hsa-runtime/CMakeLists.txt                 |  2 +-
 runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp | 11 +++++++++++
 runtime/hsa-runtime/core/util/flag.h               |  8 ++++++++
 runtime/hsa-runtime/inc/hsa_ext_amd.h              |  8 +++++++-
 4 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/runtime/hsa-runtime/CMakeLists.txt b/runtime/hsa-runtime/CMakeLists.txt
index 9f4ce63076..8fb02b14e0 100644
--- a/runtime/hsa-runtime/CMakeLists.txt
+++ b/runtime/hsa-runtime/CMakeLists.txt
@@ -85,7 +85,7 @@ if (ROCM_CCACHE_BUILD)
 endif() # if (ROCM_CCACHE_BUILD)
 
 ## Get version strings
-get_version ( "1.4.0" )
+get_version ( "1.5.0" )
 if ( ${ROCM_PATCH_VERSION} )
   set ( VERSION_PATCH ${ROCM_PATCH_VERSION})
 endif()
diff --git a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
index e5582787a0..1f8b70ab9e 100644
--- a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
+++ b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
@@ -1010,6 +1010,17 @@ hsa_status_t GpuAgent::GetInfo(hsa_agent_info_t attribute, void* value) const {
     case HSA_AMD_AGENT_INFO_SVM_DIRECT_HOST_ACCESS:
       assert(regions_.size() != 0 && "No device local memory found!");
       *((bool*)value) = properties_.Capability.ui32.CoherentHostAccess == 1;
+    case HSA_AMD_AGENT_INFO_COOPERATIVE_COMPUTE_UNIT_COUNT:
+      if (core::Runtime::runtime_singleton_->flag().coop_cu_count() &&
+          (isa_->GetMajorVersion() == 9) && (isa_->GetMinorVersion() == 0) &&
+          (isa_->GetStepping() == 10)) {
+        uint32_t count = 0;
+        hsa_status_t err = GetInfo((hsa_agent_info_t)HSA_AMD_AGENT_INFO_COMPUTE_UNIT_COUNT, &count);
+        assert(err == HSA_STATUS_SUCCESS && "CU count query failed.");
+        *((uint32_t*)value) = (count & 0xFFFFFFF8) - 8;  // value = floor(count/8)*8-8
+        break;
+      }
+      return GetInfo((hsa_agent_info_t)HSA_AMD_AGENT_INFO_COMPUTE_UNIT_COUNT, value);
     default:
       return HSA_STATUS_ERROR_INVALID_ARGUMENT;
       break;
diff --git a/runtime/hsa-runtime/core/util/flag.h b/runtime/hsa-runtime/core/util/flag.h
index cd57903ead..045a6d0c1a 100644
--- a/runtime/hsa-runtime/core/util/flag.h
+++ b/runtime/hsa-runtime/core/util/flag.h
@@ -148,6 +148,11 @@ class Flag {
 
     var = os::GetEnvVar("HSA_CU_MASK_SKIP_INIT");
     cu_mask_skip_init_ = (var == "1") ? true : false;
+
+    // Temporary opt-in for corrected HSA_AMD_AGENT_INFO_COOPERATIVE_COMPUTE_UNIT_COUNT behavior.
+    // Will become opt-out and possibly removed in future releases.
+    var = os::GetEnvVar("HSA_COOP_CU_COUNT");
+    coop_cu_count_ = (var == "1") ? true : false;
   }
 
   void parse_masks(uint32_t maxGpu, uint32_t maxCU) {
@@ -214,6 +219,8 @@ class Flag {
 
   bool cu_mask_skip_init() const { return cu_mask_skip_init_; }
 
+  bool coop_cu_count() const { return coop_cu_count_; }
+
  private:
   bool check_flat_scratch_;
   bool enable_vm_fault_message_;
@@ -233,6 +240,7 @@ class Flag {
   bool check_sramecc_validity_;
   bool debug_;
   bool cu_mask_skip_init_;
+  bool coop_cu_count_;
 
   SDMA_OVERRIDE enable_sdma_;
 
diff --git a/runtime/hsa-runtime/inc/hsa_ext_amd.h b/runtime/hsa-runtime/inc/hsa_ext_amd.h
index bb568bdc9d..6c1a936e5f 100644
--- a/runtime/hsa-runtime/inc/hsa_ext_amd.h
+++ b/runtime/hsa-runtime/inc/hsa_ext_amd.h
@@ -311,7 +311,13 @@ typedef enum hsa_amd_agent_info_s {
    * physically resident in the agent's local memory.
    * The type of this attribute is bool.
    */
-  HSA_AMD_AGENT_INFO_SVM_DIRECT_HOST_ACCESS = 0xA013
+  HSA_AMD_AGENT_INFO_SVM_DIRECT_HOST_ACCESS = 0xA013,
+  /**
+   * Some processors support more CUs than can reliably be used in a cooperative
+   * dispatch.  This queries the count of CUs which are fully enabled for
+   * cooperative dispatch.
+   */
+  HSA_AMD_AGENT_INFO_COOPERATIVE_COMPUTE_UNIT_COUNT = 0xA014
 } hsa_amd_agent_info_t;
 
 typedef struct hsa_amd_hdp_flush_s {