HIP workgroup scratch launch property

libclc/libspirv/lib/amdgpu/synchronization/barrier.cl

@@ -57,11 +57,22 @@ _CLC_OVERLOAD _CLC_DEF void __spirv_MemoryBarrier(int scope_memory,
   __mem_fence(scope_memory, semantics);
 }
 
+// Grid-wide synchronization provided by the ROCm device library (ockl). It is
+// only valid when the kernel is launched cooperatively (all work-groups are
+// co-resident), which the HIP UR adapter guarantees for root-group barriers.
+extern _CLC_CONVERGENT void __ockl_grid_sync(void);
+
 _CLC_OVERLOAD _CLC_DEF _CLC_CONVERGENT void
 __spirv_ControlBarrier(int scope_execution, int scope_memory,
                        int semantics) {
   if (semantics) {
     __mem_fence(scope_memory, semantics);
   }
-  __builtin_amdgcn_s_barrier();
+  // A Device/CrossDevice execution scope corresponds to a root-group barrier
+  // that must synchronize every work-item in the grid, not just the work-group.
+  if (scope_execution == Device || scope_execution == CrossDevice) {
+    __ockl_grid_sync();
+  } else {
+    __builtin_amdgcn_s_barrier();
+  }
 }

unified-runtime/source/adapters/hip/enqueue.cpp

@@ -260,18 +260,30 @@ static ur_result_t urEnqueueKernelLaunch(
 
   ur_kernel_launch_ext_properties_t *_launchPropList =
       const_cast<ur_kernel_launch_ext_properties_t *>(launchPropList);
-  // Adapters that don't support cooperative kernels are currently expected
-  // to ignore COOPERATIVE launch properties. Ideally we should avoid passing
-  // these at the SYCL RT level instead, see
-  // https://github.com/intel/llvm/issues/18421
-  if (_launchPropList &&
-      _launchPropList->flags & ~UR_KERNEL_LAUNCH_FLAG_COOPERATIVE) {
-    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
-  }
+  // HIP supports cooperative kernel launches through
+  // hipModuleLaunchCooperativeKernel and dynamic work-group local memory
+  // requested via a work-group launch property. Any other launch property is
+  // unsupported, see https://github.com/intel/llvm/issues/18421
+  bool UseCooperativeLaunch = false;
+  size_t WorkGroupMemory = 0;
 
   while (_launchPropList != nullptr) {
-    if (_launchPropList->stype !=
-        as_stype<ur_kernel_launch_ext_properties_t>()) {
+    switch (_launchPropList->stype) {
+    case UR_STRUCTURE_TYPE_KERNEL_LAUNCH_EXT_PROPERTIES: {
+      if (_launchPropList->flags & ~UR_KERNEL_LAUNCH_FLAG_COOPERATIVE)
+        return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+      if (_launchPropList->flags & UR_KERNEL_LAUNCH_FLAG_COOPERATIVE)
+        UseCooperativeLaunch = true;
+      break;
+    }
+    case UR_STRUCTURE_TYPE_KERNEL_LAUNCH_WORKGROUP_PROPERTY: {
+      WorkGroupMemory =
+          reinterpret_cast<ur_kernel_launch_workgroup_property_t *>(
+              _launchPropList)
+              ->workgroup_mem_size;
+      break;
+    }
+    default:
       return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
     }
     _launchPropList = static_cast<ur_kernel_launch_ext_properties_t *>(
@@ -331,10 +343,20 @@ static ur_result_t urEnqueueKernelLaunch(
       UR_CHECK_ERROR(RetImplEvent->start());
     }
 
-    UR_CHECK_ERROR(hipModuleLaunchKernel(
-        HIPFunc, BlocksPerGrid[0], BlocksPerGrid[1], BlocksPerGrid[2],
-        ThreadsPerBlock[0], ThreadsPerBlock[1], ThreadsPerBlock[2],
-        hKernel->getLocalSize(), HIPStream, ArgPointers.data(), nullptr));
+    // Dynamic work-group local memory requested through a launch property is
+    // added on top of the static local memory of the kernel arguments.
+    uint32_t SharedMemBytes = hKernel->getLocalSize() + WorkGroupMemory;
+    if (UseCooperativeLaunch) {
+      UR_CHECK_ERROR(hipModuleLaunchCooperativeKernel(
+          HIPFunc, BlocksPerGrid[0], BlocksPerGrid[1], BlocksPerGrid[2],
+          ThreadsPerBlock[0], ThreadsPerBlock[1], ThreadsPerBlock[2],
+          SharedMemBytes, HIPStream, ArgPointers.data()));
+    } else {
+      UR_CHECK_ERROR(hipModuleLaunchKernel(
+          HIPFunc, BlocksPerGrid[0], BlocksPerGrid[1], BlocksPerGrid[2],
+          ThreadsPerBlock[0], ThreadsPerBlock[1], ThreadsPerBlock[2],
+          SharedMemBytes, HIPStream, ArgPointers.data(), nullptr));
+    }
 
     if (phEvent) {
       UR_CHECK_ERROR(RetImplEvent->record());

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,