Fix a crash when printf used in a kernel launched on a stream with custom CU mask

SWDEV-249719 - root cause: queues with custom CU mask are not inserted
into queuePool_ (i.e., queue of reusable HSA queues) of ROC device class
causing a crash when creating hostcall buffers for printf

Change-Id: Ieee7005d9a5a30b3113394ce23ee65927126d0d6

Author: AryanSalmanpour

device/rocm/rocdevice.cpp

@@ -143,6 +143,7 @@ Device::Device(hsa_agent_t bkendDevice)
     , hsa_exclusive_gpu_access_(false)
     , queuePool_(QueuePriority::Total)
     , coopHostcallBuffer_(nullptr)
+    , queueWithCUMaskPool_(QueuePriority::Total)
     , numOfVgpus_(0) {
   group_segment_.handle = 0;
   system_segment_.handle = 0;
@@ -2300,7 +2301,13 @@ hsa_queue_t* Device::acquireQueue(uint32_t queue_size_hint, bool coop_queue,
       hsa_queue_destroy(queue);
       return nullptr;
     }
-    // Skip queue recycling for queues with custom CU mask
+    // add queues with custom CU mask into their special pool to keep track
+    // of mapping of these queues to their associated queueInfo (i.e., hostcall buffers)
+    auto result = queueWithCUMaskPool_[qIndex].emplace(std::make_pair(queue, QueueInfo()));
+    assert(result.second && "QueueInfo already exists");
+    auto &qInfo = result.first->second;
+    qInfo.refCount = 1;
+
     return queue;
   }
   if (coop_queue) {
@@ -2315,8 +2322,8 @@ hsa_queue_t* Device::acquireQueue(uint32_t queue_size_hint, bool coop_queue,
   return queue;
 }
 
-void Device::releaseQueue(hsa_queue_t* queue) {
-  for (auto& it : queuePool_) {
+void Device::releaseQueue(hsa_queue_t* queue, const std::vector<uint32_t>& cuMask) {
+  for (auto& it : cuMask.size() == 0 ? queuePool_ : queueWithCUMaskPool_) {
     auto qIter = it.find(queue);
     if (qIter != it.end()) {
       auto &qInfo = qIter->second;
@@ -2345,18 +2352,22 @@ void Device::releaseQueue(hsa_queue_t* queue) {
   hsa_queue_destroy(queue);
 }
 
-void* Device::getOrCreateHostcallBuffer(hsa_queue_t* queue, bool coop_queue) {
+void* Device::getOrCreateHostcallBuffer(hsa_queue_t* queue, bool coop_queue,
+                                        const std::vector<uint32_t>& cuMask) {
   decltype(queuePool_)::value_type::iterator qIter;
 
   if (!coop_queue) {
-    for (auto &it : queuePool_) {
+    for (auto &it : cuMask.size() == 0 ? queuePool_ : queueWithCUMaskPool_) {
       qIter = it.find(queue);
       if (qIter != it.end()) {
         break;
       }
     }
-
-    assert(qIter != queuePool_[QueuePriority::High].end());
+    if (cuMask.size() == 0) {
+      assert(qIter != queuePool_[QueuePriority::High].end());
+    } else {
+      assert(qIter != queueWithCUMaskPool_[QueuePriority::High].end());
+    }
 
     if (qIter->second.hostcallBuffer_) {
       return qIter->second.hostcallBuffer_;

device/rocm/rocdevice.hpp

@@ -459,11 +459,12 @@ class Device : public NullDevice {
                             amd::CommandQueue::Priority priority = amd::CommandQueue::Priority::Normal);
 
   //! Release HSA queue
-  void releaseQueue(hsa_queue_t*);
+  void releaseQueue(hsa_queue_t*, const std::vector<uint32_t>& cuMask = {});
 
   //! For the given HSA queue, return an existing hostcall buffer or create a
   //! new one. queuePool_ keeps a mapping from HSA queue to hostcall buffer.
-  void* getOrCreateHostcallBuffer(hsa_queue_t* queue, bool coop_queue = false);
+  void* getOrCreateHostcallBuffer(hsa_queue_t* queue, bool coop_queue = false,
+                                  const std::vector<uint32_t>& cuMask = {});
 
   //! Return multi GPU grid launch sync buffer
   address MGSync() const { return mg_sync_; }
@@ -524,7 +525,7 @@ class Device : public NullDevice {
     void* hostcallBuffer_;
   };
 
-  //!< a vector for keeping Pool of HSA queues with low, normal and high priorities for recycling
+  //! a vector for keeping Pool of HSA queues with low, normal and high priorities for recycling
   std::vector<std::map<hsa_queue_t*, QueueInfo>> queuePool_;
 
   //! returns a hsa queue from queuePool with least refCount and updates the refCount as well
@@ -535,6 +536,9 @@ class Device : public NullDevice {
   virtual bool findLinkInfo(const hsa_amd_memory_pool_t& pool,
                             std::vector<LinkAttrType>* link_attr);
 
+  //! Pool of HSA queues with custom CU masks
+  std::vector<std::map<hsa_queue_t*, QueueInfo>> queueWithCUMaskPool_;
+
  public:
   std::atomic<uint> numOfVgpus_;  //!< Virtual gpu unique index
 

device/rocm/rocvirtual.cpp

@@ -761,7 +761,7 @@ VirtualGPU::~VirtualGPU() {
   }
 
   if (gpu_queue_) {
-    roc_device_.releaseQueue(gpu_queue_);
+    roc_device_.releaseQueue(gpu_queue_, cuMask_);
   }
 }
 
@@ -2238,7 +2238,7 @@ bool VirtualGPU::submitKernelInternal(const amd::NDRangeContainer& sizes, const
         }
         case amd::KernelParameterDescriptor::HiddenHostcallBuffer: {
           if (amd::IS_HIP) {
-            auto buffer = roc_device_.getOrCreateHostcallBuffer(gpu_queue_, coopGroups);
+            auto buffer = roc_device_.getOrCreateHostcallBuffer(gpu_queue_, coopGroups, cuMask_);
             if (!buffer) {
               ClPrint(amd::LOG_ERROR, amd::LOG_KERN,
                       "Kernel expects a hostcall buffer, but none found");