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           17 : bool TaskMap::BatchAddTask(const AICPUSharderTaskInfo& taskInfo, const std::queue<aicpu::Closure>& queue)
      20              : {
      21           17 :     const std::lock_guard<std::mutex> lk(mapMutex_);
      22           17 :     const auto& iter = taskMap_.find(taskInfo);
      23           17 :     if (iter == taskMap_.end()) {
      24           15 :         (void)taskMap_.emplace(taskInfo, queue);
      25           15 :         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           17 : }
      40              : 
      41            8 : bool TaskMap::PopTask(const AICPUSharderTaskInfo& taskInfo, aicpu::Closure& closure)
      42              : {
      43            8 :     const std::lock_guard<std::mutex> lk(mapMutex_);
      44            8 :     const auto& iter = taskMap_.find(taskInfo);
      45            8 :     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            7 :     auto& taskQueue = iter->second;
      51            7 :     if (taskQueue.empty()) {
      52            1 :         aicpusd_run_warn("Pop task queue from empty.");
      53            1 :         return false;
      54              :     }
      55              : 
      56            6 :     closure = taskQueue.front();
      57            6 :     taskQueue.pop();
      58              : 
      59            6 :     if (taskQueue.empty()) {
      60            5 :         (void)taskMap_.erase(taskInfo);
      61              :     }
      62              : 
      63            6 :     return true;
      64            8 : }
      65              : 
      66            6 : void TaskMap::Clear()
      67              : {
      68            6 :     const std::lock_guard<std::mutex> lk(mapMutex_);
      69            7 :     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            6 :     taskMap_.clear();
      76            6 : }
      77              : 
      78            5 : std::string TaskMap::DebugString()
      79              : {
      80            5 :     const std::lock_guard<std::mutex> lk(mapMutex_);
      81            5 :     std::ostringstream oss;
      82            5 :     oss << "Split kernel TaskMapSize=" << taskMap_.size() << ". ";
      83            5 :     uint32_t i = 0U;
      84            8 :     for (const auto& iter : taskMap_) {
      85            3 :         oss << "task=" << i++ << ", parallelId=" << iter.first.parallelId << ", size=" << iter.second.size()
      86            3 :             << ", shardNum=" << iter.first.shardNum;
      87              :     }
      88           10 :     return oss.str();
      89            5 : }
      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            3 : bool TaskQueue::Dequeue(aicpu::Closure& closure)
     103              : {
     104            3 :     const std::lock_guard<std::mutex> queLock(mtxQue_);
     105            3 :     if (taskQueue_.empty()) {
     106            1 :         aicpusd_err("Dequeue from empty");
     107            1 :         return false;
     108              :     }
     109            2 :     closure = taskQueue_.front();
     110            2 :     taskQueue_.pop();
     111            2 :     return true;
     112            3 : }
     113              : 
     114            6 : void TaskQueue::Clear()
     115              : {
     116            6 :     const std::lock_guard<std::mutex> queLock(mtxQue_);
     117         1030 :     while (!taskQueue_.empty()) {
     118         1024 :         taskQueue_.pop();
     119              :     }
     120            6 : }
     121              : 
     122            4 : std::string TaskQueue::DebugString()
     123              : {
     124            4 :     const std::lock_guard<std::mutex> queLock(mtxQue_);
     125            4 :     std::ostringstream oss;
     126            4 :     oss << "queueSize=" << taskQueue_.size();
     127            8 :     return oss.str();
     128            4 : }
     129              : } // namespace AicpuSchedule