From 3af54da059d9383087e048203bcaf234ac74e3a6 Mon Sep 17 00:00:00 2001
From: ROCm CI Service Account <66695075+rocm-ci@users.noreply.github.com>
Date: Fri, 23 Sep 2022 12:37:42 +0530
Subject: [PATCH] SWDEV-287351 - Fix hipStreamWithCUMask hang on Navi21 (#2939)

- For gfx >= 10, one work group processor encompasses 2 CUs & hence the CUs need to be enabled in pair

Change-Id: I359df1b221b4400b260b02201b7b0385054784f3
---
 catch/unit/stream/hipStreamGetCUMask.cc  | 14 +++++++++++++-
 catch/unit/stream/hipStreamWithCUMask.cc | 22 ++++++++++++++--------
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/catch/unit/stream/hipStreamGetCUMask.cc b/catch/unit/stream/hipStreamGetCUMask.cc
index 49f21bdbda..f63acc212a 100644
--- a/catch/unit/stream/hipStreamGetCUMask.cc
+++ b/catch/unit/stream/hipStreamGetCUMask.cc
@@ -117,9 +117,21 @@ TEST_CASE("Unit_hipExtStreamGetCUMask_verifyDefaultAndCustomMask") {
   }
 
   SECTION("Verify with custom mask set") {
+    hipDeviceProp_t props;
     std::vector<uint32_t> customMask(defaultCUMask);
     hipStream_t stream;
-    customMask[0] = 0xe;
+    int deviceId;
+
+    HIP_CHECK(hipGetDevice(&deviceId));
+    HIP_CHECK(hipGetDeviceProperties(&props, deviceId));
+
+    if (props.major >= 10) {
+      // For gfx >= 10, one work group processor encompasses 2 CUs &
+      // hence the CUs need to be enabled in pair
+      customMask[0] = 0xc;
+    } else {
+      customMask[0] = 0xe;
+    }
 
     HIP_CHECK(hipExtStreamCreateWithCUMask(&stream, customMask.size(),
                                                       customMask.data()));
diff --git a/catch/unit/stream/hipStreamWithCUMask.cc b/catch/unit/stream/hipStreamWithCUMask.cc
index ea6b033d9c..adf1db3e80 100644
--- a/catch/unit/stream/hipStreamWithCUMask.cc
+++ b/catch/unit/stream/hipStreamWithCUMask.cc
@@ -204,21 +204,27 @@ TEST_CASE("Unit_hipExtStreamCreateWithCUMask_Functionality") {
     hA[i] = CONSTANT + i;
   }
 
+  int cuCountPerGroup = 1;
+  if (props.major >= 10) {
+    cuCountPerGroup = 2; // For gfx >= 10, one work group processor encompasses 2 CUs
+  }
+
+  unsigned long mask = pow(2, cuCountPerGroup) - 1;
   for (int np = 0; np < KNumPartition; np++) {
     HIP_CHECK(hipMalloc(&dA[np], Nbytes));
     HIP_CHECK(hipMalloc(&dC[np], Nbytes));
     // make unique CU masks in the multiple of dwords for each stream
     uint32_t temp = 0;
-    uint32_t bit_index = np;
-    for (int i = np; i < props.multiProcessorCount; i = i + 4) {
-      temp |= 1UL << bit_index;
+    uint32_t bit_index = cuCountPerGroup * np;
+    for (int i = np; i < props.multiProcessorCount; i = i + cuCountPerGroup * 4) {
+      temp |= mask << bit_index;
       if (bit_index >= 32) {
         cuMasks[np].push_back(temp);
         temp = 0;
-        bit_index = np;
-        temp |= 1UL << bit_index;
+        bit_index = cuCountPerGroup * np;
+        temp |= mask << bit_index;
       }
-      bit_index += 4;
+      bit_index += cuCountPerGroup * 4;
     }
     if (bit_index != 0) {
       cuMasks[np].push_back(temp);
@@ -229,9 +235,9 @@ TEST_CASE("Unit_hipExtStreamCreateWithCUMask_Functionality") {
 
     HIP_CHECK(hipMemcpy(dA[np], hA, Nbytes, hipMemcpyHostToDevice));
 
-    ss[np] << std::hex;
+    ss[np] << std::hex << std::setfill('0');
     for (int i = cuMasks[np].size() - 1; i >= 0; i--) {
-      ss[np] << cuMasks[np][i];
+      ss[np] << std::setw(8) << cuMasks[np][i];
     }
   }