From ec23f3605785a7b632d8fa31e722076b0004f011 Mon Sep 17 00:00:00 2001
From: Zheming Jin <zjin-lcf@users.noreply.github.com>
Date: Sat, 18 Jul 2026 07:51:14 -0700
Subject: [PATCH] [UR][HIP] Implement urKernelSuggestMaxCooperativeGroupCount

Implement the cooperative-group-count occupancy query for the HIP (AMD)
adapter, which previously returned UR_RESULT_ERROR_UNSUPPORTED_FEATURE and
forced the SYCL runtime onto a coarse 0/1 host-side fallback for
kernel::ext_oneapi_get_info<max_num_work_groups>.

As noted in intel/llvm#21803, a direct port of the CUDA implementation does
not work: hipOccupancyMaxActiveBlocksPerMultiprocessor expects a host
function pointer, whereas our kernel handle (hKernel->get()) is a
module-loaded hipFunction_t. Use the module-occupancy entry point
hipModuleOccupancyMaxActiveBlocksPerMultiprocessor instead, mirroring the
existing hipModuleOccupancyMaxPotentialBlockSize usage in the adapter.

The total suggested group count is the per-CU occupancy multiplied by the
device multiprocessor (compute-unit) count. When the per-CU occupancy is 0,
fall back to reporting 1 group if the requested work-group size, dynamic
local memory, and per-block register usage all fit within device limits,
matching the CUDA adapter's semantics.

Closes: https://github.com/intel/llvm/issues/21803

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .../source/adapters/hip/kernel.cpp            | 74 ++++++++++++++++++-
 1 file changed, 70 insertions(+), 4 deletions(-)

diff --git a/unified-runtime/source/adapters/hip/kernel.cpp b/unified-runtime/source/adapters/hip/kernel.cpp
index de20e31699a66..604aa6193abd6 100644
--- a/unified-runtime/source/adapters/hip/kernel.cpp
+++ b/unified-runtime/source/adapters/hip/kernel.cpp
@@ -156,10 +156,76 @@ urKernelGetNativeHandle(ur_kernel_handle_t, ur_native_handle_t *) {
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urKernelSuggestMaxCooperativeGroupCount(
-    ur_kernel_handle_t /*hKernel*/, ur_device_handle_t /*hDevice*/,
-    uint32_t /*workDim*/, const size_t * /*pLocalWorkSize*/,
-    size_t /*dynamicSharedMemorySize*/, uint32_t * /*pGroupCountRet*/) {
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+    ur_kernel_handle_t hKernel, ur_device_handle_t /*hDevice*/,
+    uint32_t workDim, const size_t *pLocalWorkSize,
+    size_t dynamicSharedMemorySize, uint32_t *pGroupCountRet) {
+  UR_ASSERT(hKernel, UR_RESULT_ERROR_INVALID_KERNEL);
+
+  size_t localWorkSize = pLocalWorkSize[0];
+  localWorkSize *= (workDim >= 2 ? pLocalWorkSize[1] : 1);
+  localWorkSize *= (workDim == 3 ? pLocalWorkSize[2] : 1);
+
+  // We need to set the active current device for this kernel explicitly here,
+  // because the occupancy querying API does not take a device parameter.
+  ur_device_handle_t Device = hKernel->getProgram()->getDevice();
+  ScopedDevice Active(Device);
+  try {
+    // We need to calculate max num of work-groups using per-device semantics.
+
+    // Unlike CUDA's cuOccupancyMaxActiveBlocksPerMultiprocessor, the
+    // hipOccupancy* variant expects a host function pointer; our kernel is a
+    // module-loaded hipFunction_t (hKernel->get()), so we must use the
+    // hipModule* occupancy entry point instead (see issue #21803).
+    int MaxNumActiveGroupsPerCU{0};
+    UR_CHECK_ERROR(hipModuleOccupancyMaxActiveBlocksPerMultiprocessor(
+        &MaxNumActiveGroupsPerCU, hKernel->get(),
+        static_cast<int>(localWorkSize), dynamicSharedMemorySize));
+    assert(MaxNumActiveGroupsPerCU >= 0);
+    // Handle the case where we can't have all CUs active with at least 1 group
+    // per CU. In that case, the device is still able to run 1 work-group, hence
+    // we will manually check if it is possible with the available HW resources.
+    if (MaxNumActiveGroupsPerCU == 0) {
+      size_t MaxWorkGroupSize{};
+      urKernelGetGroupInfo(
+          hKernel, Device, UR_KERNEL_GROUP_INFO_WORK_GROUP_SIZE,
+          sizeof(MaxWorkGroupSize), &MaxWorkGroupSize, nullptr);
+      size_t MaxLocalSizeBytes{};
+      urDeviceGetInfo(Device, UR_DEVICE_INFO_LOCAL_MEM_SIZE,
+                      sizeof(MaxLocalSizeBytes), &MaxLocalSizeBytes, nullptr);
+      // Check whether the kernel's per-thread register usage would exceed the
+      // device's per-block register budget for the requested work-group size.
+      int NumRegs{0};
+      UR_CHECK_ERROR(hipFuncGetAttribute(&NumRegs, HIP_FUNC_ATTRIBUTE_NUM_REGS,
+                                         hKernel->get()));
+      int MaxRegsPerBlock{0};
+      UR_CHECK_ERROR(hipDeviceGetAttribute(
+          &MaxRegsPerBlock, hipDeviceAttributeMaxRegistersPerBlock,
+          Device->get()));
+      const bool HasExceededMaxRegistersPerBlock =
+          localWorkSize * static_cast<size_t>(NumRegs) >
+          static_cast<size_t>(MaxRegsPerBlock);
+      if (localWorkSize > MaxWorkGroupSize ||
+          dynamicSharedMemorySize > MaxLocalSizeBytes ||
+          HasExceededMaxRegistersPerBlock)
+        *pGroupCountRet = 0;
+      else
+        *pGroupCountRet = 1;
+    } else {
+      // Multiply by the number of CUs (compute units = streaming
+      // multiprocessors) on the device in order to retrieve the total number
+      // of groups/blocks that can be launched.
+      int NumComputeUnits{0};
+      UR_CHECK_ERROR(hipDeviceGetAttribute(
+          &NumComputeUnits, hipDeviceAttributeMultiprocessorCount,
+          Device->get()));
+      assert(NumComputeUnits >= 0);
+      *pGroupCountRet =
+          static_cast<uint32_t>(NumComputeUnits) * MaxNumActiveGroupsPerCU;
+    }
+  } catch (ur_result_t Err) {
+    return Err;
+  }
+  return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urKernelGetInfo(ur_kernel_handle_t hKernel,