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              : }
      17              : 
      18              : namespace AicpuSchedule {
      19           48 : bool TaskMap::BatchAddTask(const AICPUSharderTaskInfo &taskInfo, const std::queue<aicpu::Closure> &queue)
      20              : {
      21           48 :     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("Try to add new task queue, but last queue is not been consumed. parallelId=%u, "
      30              :             "size=%lu, shardNum=%ld", taskInfo.parallelId, iter->second.size(), taskInfo.shardNum);
      31            1 :         return false;
      32              :     }
      33              : 
      34            1 :     iter->second = queue;
      35              : 
      36            1 :     return true;
      37           49 : }
      38              : 
      39          132 : bool TaskMap::PopTask(const AICPUSharderTaskInfo &taskInfo, aicpu::Closure &closure)
      40              : {
      41          132 :     const std::lock_guard<std::mutex> lk(mapMutex_);
      42          132 :     const auto &iter = taskMap_.find(taskInfo);
      43          132 :     if (iter == taskMap_.end()) {
      44            1 :         aicpusd_run_warn("Get task from map failed. parallelId=%u", taskInfo.parallelId);
      45            1 :         return false;
      46              :     }
      47              : 
      48          131 :     auto &taskQueue = iter->second;
      49          131 :     if (taskQueue.empty()) {
      50            1 :         aicpusd_run_warn("Pop task queue from empty. parallelId=%u", taskInfo.parallelId);
      51            1 :         return false;
      52              :     }
      53              : 
      54          130 :     closure = taskQueue.front();
      55          130 :     taskQueue.pop();
      56              : 
      57          130 :     if (taskQueue.empty()) {
      58           36 :         (void)taskMap_.erase(taskInfo);
      59              :     }
      60              : 
      61          130 :     return true;
      62          132 : }
      63              : 
      64            9 : void TaskMap::Clear()
      65              : {
      66            9 :     const std::lock_guard<std::mutex> lk(mapMutex_);
      67           10 :     for (auto iter = taskMap_.begin(); iter != taskMap_.end(); ++iter) {
      68            1 :         auto &taskQueue = iter->second;
      69            2 :         while (!taskQueue.empty()) {
      70            1 :             taskQueue.pop();
      71              :         }
      72              :     }
      73            9 :     taskMap_.clear();
      74            9 : }
      75              : 
      76            8 : std::string TaskMap::DebugString()
      77              : {
      78            8 :     const std::lock_guard<std::mutex> lk(mapMutex_);
      79            8 :     std::ostringstream oss;
      80            8 :     oss << "Split kernel TaskMapSize=" << taskMap_.size() << ". ";
      81            8 :     uint32_t i = 0U;
      82           12 :     for (const auto &iter : taskMap_) {
      83            4 :         oss << "task=" << i++ << ", parallelId=" << iter.first.parallelId
      84            4 :             << ", size=" << iter.second.size()
      85            4 :             << ", shardNum=" << iter.first.shardNum;
      86              :     }
      87           16 :     return oss.str();
      88            8 : }
      89              : 
      90         3077 : bool TaskQueue::Enqueue(const aicpu::Closure &closure)
      91              : {
      92         3077 :     const std::lock_guard<std::mutex> queLock(mtxQue_);
      93         3077 :     if (taskQueue_.size() >= MAX_TASK_QUEUE_SIZE) {
      94            1 :         aicpusd_err("Queue is too large");
      95            1 :         return false;
      96              :     }
      97         3076 :     taskQueue_.push(closure);
      98         3076 :     return true;
      99         3077 : }
     100              : 
     101            6 : bool TaskQueue::Dequeue(aicpu::Closure &closure)
     102              : {
     103            6 :     const std::lock_guard<std::mutex> queLock(mtxQue_);
     104            6 :     if (taskQueue_.empty()) {
     105            4 :         aicpusd_err("Dequeue from empty");
     106            4 :         return false;
     107              :     }
     108            2 :     closure = taskQueue_.front();
     109            2 :     taskQueue_.pop();
     110            2 :     return true;
     111            6 : }
     112              : 
     113            9 : void TaskQueue::Clear()
     114              : {
     115            9 :     const std::lock_guard<std::mutex> queLock(mtxQue_);
     116         1033 :     while (!taskQueue_.empty()) {
     117         1024 :         taskQueue_.pop();
     118              :     }
     119            9 : }
     120              : 
     121            7 : std::string TaskQueue::DebugString()
     122              : {
     123            7 :     const std::lock_guard<std::mutex> queLock(mtxQue_);
     124            7 :     std::ostringstream oss;
     125            7 :     oss << "queueSize=" << taskQueue_.size();
     126           14 :     return oss.str();
     127            7 : }
     128              : }