diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp
index 8dd76d2aeb..5ab779c93b 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp
@@ -539,12 +539,11 @@ hsa_status_t hsa_amd_queue_cu_set_mask(const hsa_queue_t* queue, uint32_t num_cu
                                        const uint32_t* cu_mask) {
   TRY;
   IS_OPEN();
-  IS_BAD_PTR(cu_mask);
 
   core::Queue* cmd_queue = core::Queue::Convert(queue);
   IS_VALID(cmd_queue);
-  if ((num_cu_mask_count == 0) || (num_cu_mask_count % 32 != 0))
-    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  if (num_cu_mask_count != 0) IS_BAD_PTR(cu_mask);
+  if (num_cu_mask_count % 32 != 0) return HSA_STATUS_ERROR_INVALID_ARGUMENT;
   return cmd_queue->SetCUMasking(num_cu_mask_count, cu_mask);
   CATCH;
 }
diff --git a/projects/rocr-runtime/runtime/hsa-runtime/inc/hsa_ext_amd.h b/projects/rocr-runtime/runtime/hsa-runtime/inc/hsa_ext_amd.h
index 96140dc22f..bb568bdc9d 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/inc/hsa_ext_amd.h
+++ b/projects/rocr-runtime/runtime/hsa-runtime/inc/hsa_ext_amd.h
@@ -792,6 +792,8 @@ hsa_status_t HSA_API hsa_amd_image_get_info_max_dim(hsa_agent_t agent,
  * @details Enables the queue to run on only selected CUs.  The given mask is
  * combined by bitwise AND with any device wide mask in HSA_CU_MASK before
  * being applied.
+ * If num_cu_mask_count is 0 then the request is interpreted as a request to
+ * enable all CUs and no cu_mask array need be given.
  *
  * @param[in] queue A pointer to HSA queue.
  *
@@ -810,8 +812,8 @@ hsa_status_t HSA_API hsa_amd_image_get_info_max_dim(hsa_agent_t agent,
  *
  * @retval ::HSA_STATUS_ERROR_INVALID_QUEUE @p queue is NULL or invalid.
  *
- * @retval ::HSA_STATUS_ERROR_INVALID_ARGUMENT @p num_cu_mask_count is 0, not
- * a multiple of 32 or @p cu_mask is NULL.
+ * @retval ::HSA_STATUS_ERROR_INVALID_ARGUMENT @p num_cu_mask_count is not
+ * a multiple of 32 or @p num_cu_mask_count is not 0 and cu_mask is NULL.
  *
  */
 hsa_status_t HSA_API hsa_amd_queue_cu_set_mask(const hsa_queue_t* queue,