[UR] Add dynamic shared mem support for cooperative group count suggestion (#22964)

Add dynamic shared memory support for cooperative group count suggestion

fixes: #22936

Author: bratpiorka

unified-runtime/source/adapters/level_zero/common/helpers/kernel_helpers.cpp

@@ -62,6 +62,32 @@ ur_result_t getSuggestedLocalWorkSize(ur_device_handle_t hDevice,
   return UR_RESULT_SUCCESS;
 }
 
+uint32_t getMaxCooperativeGroupCountWithDynamicSharedMemory(
+    uint32_t suggestedGroupCount, uint32_t computeUnitCount,
+    uint32_t maxSharedLocalMemory, uint32_t staticSharedLocalMemory,
+    size_t dynamicSharedLocalMemory) {
+
+  if (dynamicSharedLocalMemory == 0)
+    return suggestedGroupCount;
+
+  // Reject a per-group allocation that cannot fit on one compute unit
+  // and avoid overflowing the combined allocation size.
+  if (staticSharedLocalMemory > maxSharedLocalMemory ||
+      dynamicSharedLocalMemory > maxSharedLocalMemory - staticSharedLocalMemory)
+    return 0;
+
+  const size_t sharedLocalMemoryPerGroup =
+      staticSharedLocalMemory + dynamicSharedLocalMemory;
+
+  // Cap the driver's estimate by the number of groups that can reside given
+  // the requested local memory per group.
+  const uint64_t sharedLocalMemoryGroupCount =
+      (uint64_t)computeUnitCount *
+      (maxSharedLocalMemory / sharedLocalMemoryPerGroup);
+  return static_cast<uint32_t>(
+      std::min((uint64_t)suggestedGroupCount, sharedLocalMemoryGroupCount));
+}
+
 ur_result_t setKernelGlobalOffset(ur_context_handle_t Context,
                                   ze_kernel_handle_t Kernel, uint32_t WorkDim,
                                   const size_t *GlobalWorkOffset) {

unified-runtime/source/adapters/level_zero/common/helpers/kernel_helpers.hpp

@@ -59,6 +59,11 @@ ur_result_t getSuggestedLocalWorkSize(ur_device_handle_t hDevice,
                                       size_t GlobalWorkSize3D[3],
                                       uint32_t SuggestedLocalWorkSize3D[3]);
 
+uint32_t getMaxCooperativeGroupCountWithDynamicSharedMemory(
+    uint32_t suggestedGroupCount, uint32_t computeUnitCount,
+    uint32_t maxSharedLocalMemory, uint32_t staticSharedLocalMemory,
+    size_t dynamicSharedLocalMemory);
+
 /**
  * Handle uncommon conditions after kernel submission.
  * Resets the offset to {0, 0, 0} if one was supplied.

unified-runtime/source/adapters/level_zero/kernel.cpp

@@ -1094,7 +1094,6 @@ ur_result_t urKernelSuggestMaxCooperativeGroupCount(
     size_t dynamicSharedMemorySize, uint32_t *pGroupCountRet) {
   auto hKernel = v1_cast(hKernelOpque);
   auto hDevice = common_cast(hDeviceOpque);
-  (void)dynamicSharedMemorySize;
   std::shared_lock<ur_shared_mutex> Guard(hKernel->Mutex);
 
   ze_kernel_handle_t ZeKernel = nullptr;
@@ -1109,7 +1108,12 @@ ur_result_t urKernelSuggestMaxCooperativeGroupCount(
   uint32_t TotalGroupCount = 0;
   ZE2UR_CALL(zeKernelSuggestMaxCooperativeGroupCount,
              (ZeKernel, &TotalGroupCount));
-  *pGroupCountRet = TotalGroupCount;
+  *pGroupCountRet = getMaxCooperativeGroupCountWithDynamicSharedMemory(
+      TotalGroupCount,
+      hDevice->ZeDeviceProperties->numSubslicesPerSlice *
+          hDevice->ZeDeviceProperties->numSlices,
+      hDevice->ZeDeviceComputeProperties->maxSharedLocalMemory,
+      hKernel->ZeKernelProperties->localMemSize, dynamicSharedMemorySize);
   return UR_RESULT_SUCCESS;
 }
 

unified-runtime/source/adapters/level_zero/v2/kernel.cpp

@@ -666,7 +666,6 @@ ur_result_t urKernelSuggestMaxCooperativeGroupCount(
     size_t dynamicSharedMemorySize, uint32_t *pGroupCountRet) {
   auto hKernel = v2_cast(hKernelOpque);
   auto hDevice = common_cast(hDeviceOpque);
-  (void)dynamicSharedMemorySize;
 
   uint32_t wg[3];
   wg[0] = ur_cast<uint32_t>(pLocalWorkSize[0]);
@@ -678,7 +677,12 @@ ur_result_t urKernelSuggestMaxCooperativeGroupCount(
   uint32_t totalGroupCount = 0;
   ZE2UR_CALL(zeKernelSuggestMaxCooperativeGroupCount,
              (hKernel->getZeHandle(hDevice), &totalGroupCount));
-  *pGroupCountRet = totalGroupCount;
+  *pGroupCountRet = getMaxCooperativeGroupCountWithDynamicSharedMemory(
+      totalGroupCount,
+      hDevice->ZeDeviceProperties->numSubslicesPerSlice *
+          hDevice->ZeDeviceProperties->numSlices,
+      hDevice->ZeDeviceComputeProperties->maxSharedLocalMemory,
+      hKernel->getProperties(hDevice).localMemSize, dynamicSharedMemorySize);
   return UR_RESULT_SUCCESS;
 }
 

unified-runtime/test/conformance/kernel/urKernelSuggestMaxCooperativeGroupCount.cpp

@@ -5,6 +5,8 @@
 
 #include <uur/fixtures.h>
 
+#include <array>
+
 struct urKernelSuggestMaxCooperativeGroupCountTest
     : uur::urKernelExecutionTest {
   void SetUp() override {
@@ -36,6 +38,36 @@ TEST_P(urKernelSuggestMaxCooperativeGroupCountTest, Success) {
   ASSERT_GE(suggested_work_groups, 0);
 }
 
+TEST_P(urKernelSuggestMaxCooperativeGroupCountTest, DynamicSharedMemory) {
+  uint64_t local_memory_size = 0;
+  ASSERT_SUCCESS(urDeviceGetInfo(device, UR_DEVICE_INFO_LOCAL_MEM_SIZE,
+                                 sizeof(local_memory_size), &local_memory_size,
+                                 nullptr));
+  uint32_t compute_unit_count = 0;
+  ASSERT_SUCCESS(urDeviceGetInfo(device, UR_DEVICE_INFO_NUM_COMPUTE_UNITS,
+                                 sizeof(compute_unit_count),
+                                 &compute_unit_count, nullptr));
+
+  ASSERT_SUCCESS(urKernelSuggestMaxCooperativeGroupCount(
+      kernel, device, n_dimensions, &local_size, 0, &suggested_work_groups));
+
+  // Check that increasing the dynamic shared memory size reduces
+  // the number of suggested work groups.
+  constexpr std::array<uint32_t, 4> divs = {8, 4, 2, 1};
+  uint32_t previous_suggested_work_groups = suggested_work_groups;
+  for (const uint32_t i : divs) {
+    const size_t dynamic_shared_memory_size = local_memory_size / i;
+    ASSERT_SUCCESS(urKernelSuggestMaxCooperativeGroupCount(
+        kernel, device, n_dimensions, &local_size, dynamic_shared_memory_size,
+        &suggested_work_groups));
+
+    ASSERT_LE(suggested_work_groups, previous_suggested_work_groups);
+    ASSERT_LE((uint64_t)suggested_work_groups,
+              (uint64_t)(compute_unit_count)*i);
+    previous_suggested_work_groups = suggested_work_groups;
+  }
+}
+
 TEST_P(urKernelSuggestMaxCooperativeGroupCountTest, InvalidNullHandleKernel) {
   ASSERT_EQ_RESULT(urKernelSuggestMaxCooperativeGroupCount(
                        nullptr, device, n_dimensions, &local_size, 0,