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              : 
      11              : #ifndef CORE_AICPUSD_TASK_QUEUE_H
      12              : #define CORE_AICPUSD_TASK_QUEUE_H
      13              : 
      14              : #include <mutex>
      15              : #include <unordered_map>
      16              : #include <queue>
      17              : #include "aicpusd_info.h"
      18              : #include "aicpu_sharder.h"
      19              : 
      20              : namespace AicpuSchedule {
      21              : struct HashKey {
      22           30 :     std::size_t operator() (const AICPUSharderTaskInfo &sharderTaskInfo) const noexcept
      23              :     {
      24           30 :         std::size_t h1 = std::hash<uint32_t>()(sharderTaskInfo.parallelId);
      25           30 :         return h1;
      26              :     }
      27              : };
      28              : 
      29              : class TaskMap {
      30              : public:
      31              :     void Clear();
      32              :     std::string DebugString();
      33              :     bool BatchAddTask(const AICPUSharderTaskInfo &taskInfo, const std::queue<aicpu::Closure> &queue);
      34              :     bool PopTask(const AICPUSharderTaskInfo &taskInfo, aicpu::Closure &closure);
      35              : 
      36           23 :     TaskMap() = default;
      37           23 :     ~TaskMap() = default;
      38              : 
      39              : private:
      40              :     TaskMap(const TaskMap &) = delete;
      41              :     TaskMap &operator=(const TaskMap &) = delete;
      42              :     TaskMap(TaskMap &&) = delete;
      43              :     TaskMap &operator=(TaskMap &&) = delete;
      44              : 
      45              :     std::mutex mapMutex_;
      46              :     std::unordered_map<AICPUSharderTaskInfo, std::queue<aicpu::Closure>, HashKey> taskMap_;
      47              : };
      48              : 
      49              : class TaskQueue {
      50              : public:
      51              :     void Clear();
      52              :     std::string DebugString();
      53              :     bool Enqueue(const aicpu::Closure &closure);
      54              :     bool Dequeue(aicpu::Closure &closure);
      55              : 
      56           21 :     TaskQueue() = default;
      57           21 :     ~TaskQueue() = default;
      58              : 
      59              : private:
      60              :     TaskQueue(const TaskQueue &) = delete;
      61              :     TaskQueue &operator=(const TaskQueue &) = delete;
      62              :     TaskQueue(TaskQueue&&) = delete;
      63              :     TaskQueue& operator=(TaskQueue&&) = delete;
      64              : 
      65              :     std::mutex mtxQue_;
      66              :     std::queue<aicpu::Closure> taskQueue_;
      67              : };
      68              : }
      69              : #endif