include/ck/host_utility/stream_utility.hpp Source File

include/ck/host_utility/stream_utility.hpp Source File#

Composable Kernel: include/ck/host_utility/stream_utility.hpp Source File

 // SPDX-License-Identifier: MIT
 // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
  
 #pragma once
  
 #include <hip/hip_runtime.h>
  
 #include "ck/stream_config.hpp"
 #include "ck/host_utility/hip_check_error.hpp"
  
 static inline int getAvailableComputeUnitCount(const StreamConfig& stream_config)
 {
     constexpr int MAX_MASK_DWORDS = 64;
  
     // assume at most 64*32 = 2048 CUs
     uint32_t cuMask[MAX_MASK_DWORDS];
  
     for(int i = 0; i < MAX_MASK_DWORDS; i++)
         cuMask[i] = 0;
  
     auto countSetBits = [](uint32_t dword) {
         int count = 0;
  
         while(dword != 0)
         {
             if(dword & 0x1)
                 count++;
  
             dword = dword >> 1;
         };
  
         return (count);
     };
  
     hip_check_error(hipExtStreamGetCUMask(stream_config.stream_id_, MAX_MASK_DWORDS, &cuMask[0]));
  
     int ret = 0;
  
     for(int i = 0; i < MAX_MASK_DWORDS; i++)
         ret += countSetBits(cuMask[i]);
  
     return (ret);
 };