Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2025 Huawei Technologies Co., Ltd.
       3              :  * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
       4              :  * CANN Open Software License Agreement Version 2.0 (the "License").
       5              :  * Please refer to the License for details. You may not use this file except in compliance with the License.
       6              :  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
       7              :  * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
       8              :  * See LICENSE in the root of the software repository for the full text of the License.
       9              :  */
      10              : #include <sstream>
      11              : #include "aicpusd_status.h"
      12              : #include "aicpusd_task_queue.h"
      13              : namespace {
      14              : // list queue max size 1024
      15              : constexpr uint32_t MAX_TASK_QUEUE_SIZE = 1024U;
      16              : } // namespace
      17              : 
      18              : namespace AicpuSchedule {
      19           49 : bool TaskMap::BatchAddTask(const AICPUSharderTaskInfo& taskInfo, const std::queue<aicpu::Closure>& queue)
      20              : {
      21           49 :     const std::lock_guard<std::mutex> lk(mapMutex_);
      22           49 :     const auto& iter = taskMap_.find(taskInfo);
      23           49 :     if (iter == taskMap_.end()) {
      24           47 :         (void)taskMap_.emplace(taskInfo, queue);
      25           47 :         return true;
      26              :     }
      27              : 
      28            2 :     if (!iter->second.empty()) {
      29            1 :         aicpusd_err(
      30              :             "Try to add new task queue, but last queue is not been consumed. parallelId=%u, "
      31              :             "size=%lu, shardNum=%ld",
      32              :             taskInfo.parallelId, iter->second.size(), taskInfo.shardNum);
      33            1 :         return false;
      34              :     }
      35              : 
      36            1 :     iter->second = queue;
      37              : 
      38            1 :     return true;
      39           49 : }
      40              : 
      41          132 : bool TaskMap::PopTask(const AICPUSharderTaskInfo& taskInfo, aicpu::Closure& closure)
      42              : {
      43          132 :     const std::lock_guard<std::mutex> lk(mapMutex_);
      44          132 :     const auto& iter = taskMap_.find(taskInfo);
      45          132 :     if (iter == taskMap_.end()) {
      46            1 :         aicpusd_run_warn("Get task from map failed. parallelId=%u", taskInfo.parallelId);
      47            1 :         return false;
      48              :     }
      49              : 
      50          131 :     auto& taskQueue = iter->second;
      51          131 :     if (taskQueue.empty()) {
      52            1 :         aicpusd_run_warn("Pop task queue from empty. parallelId=%u", taskInfo.parallelId);
      53            1 :         return false;
      54              :     }
      55              : 
      56          130 :     closure = taskQueue.front();
      57          130 :     taskQueue.pop();
      58              : 
      59          130 :     if (taskQueue.empty()) {
      60           36 :         (void)taskMap_.erase(taskInfo);
      61              :     }
      62              : 
      63          130 :     return true;
      64          132 : }
      65              : 
      66            9 : void TaskMap::Clear()
      67              : {
      68            9 :     const std::lock_guard<std::mutex> lk(mapMutex_);
      69           10 :     for (auto iter = taskMap_.begin(); iter != taskMap_.end(); ++iter) {
      70            1 :         auto& taskQueue = iter->second;
      71            2 :         while (!taskQueue.empty()) {
      72            1 :             taskQueue.pop();
      73              :         }
      74              :     }
      75            9 :     taskMap_.clear();
      76            9 : }
      77              : 
      78            8 : std::string TaskMap::DebugString()
      79              : {
      80            8 :     const std::lock_guard<std::mutex> lk(mapMutex_);
      81            8 :     std::ostringstream oss;
      82            8 :     oss << "Split kernel TaskMapSize=" << taskMap_.size() << ". ";
      83            8 :     uint32_t i = 0U;
      84           12 :     for (const auto& iter : taskMap_) {
      85            4 :         oss << "task=" << i++ << ", parallelId=" << iter.first.parallelId << ", size=" << iter.second.size()
      86            4 :             << ", shardNum=" << iter.first.shardNum;
      87              :     }
      88           16 :     return oss.str();
      89            8 : }
      90              : 
      91         3077 : bool TaskQueue::Enqueue(const aicpu::Closure& closure)
      92              : {
      93         3077 :     const std::lock_guard<std::mutex> queLock(mtxQue_);
      94         3077 :     if (taskQueue_.size() >= MAX_TASK_QUEUE_SIZE) {
      95            1 :         aicpusd_err("Queue is too large");
      96            1 :         return false;
      97              :     }
      98         3076 :     taskQueue_.push(closure);
      99         3076 :     return true;
     100         3077 : }
     101              : 
     102            6 : bool TaskQueue::Dequeue(aicpu::Closure& closure)
     103              : {
     104            6 :     const std::lock_guard<std::mutex> queLock(mtxQue_);
     105            6 :     if (taskQueue_.empty()) {
     106            4 :         aicpusd_err("Dequeue from empty");
     107            4 :         return false;
     108              :     }
     109            2 :     closure = taskQueue_.front();
     110            2 :     taskQueue_.pop();
     111            2 :     return true;
     112            6 : }
     113              : 
     114            9 : void TaskQueue::Clear()
     115              : {
     116            9 :     const std::lock_guard<std::mutex> queLock(mtxQue_);
     117         1033 :     while (!taskQueue_.empty()) {
     118         1024 :         taskQueue_.pop();
     119              :     }
     120            9 : }
     121              : 
     122            7 : std::string TaskQueue::DebugString()
     123              : {
     124            7 :     const std::lock_guard<std::mutex> queLock(mtxQue_);
     125            7 :     std::ostringstream oss;
     126            7 :     oss << "queueSize=" << taskQueue_.size();
     127           14 :     return oss.str();
     128            7 : }
     129              : } // namespace AicpuSchedule