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