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 "mirror_task_manager_lite.h"
      11              : 
      12              : namespace Hccl {
      13              : 
      14          128 : MirrorTaskManagerLite::MirrorTaskManagerLite()
      15              : {
      16          128 : }
      17              : 
      18            0 : void MirrorTaskManagerLite::RegFullyCallBack(std::function<void()> callBack)
      19              : {
      20            0 :     fullyCallBack_ = callBack;
      21            0 :     return;
      22              : }
      23              : 
      24            0 : void MirrorTaskManagerLite::RegGetRemoteRankCallBack(std::function<u32(u64)> callBack)
      25              : {
      26            0 :     getRemoteRankCallback_ = callBack;
      27            0 :     return;
      28              : }
      29              : 
      30            0 : HcclResult MirrorTaskManagerLite::AddTaskInfo(u32 streamId, u32 taskId, const Hccl::TaskParam &taskParam, u64 handle)
      31              : {
      32            0 :     u32 remoteRankId = getRemoteRankCallback_ ? getRemoteRankCallback_(handle) : INVALID_U32;
      33            0 :     PrintTaskLog(streamId, taskId, taskParam, remoteRankId);
      34              : 
      35            0 :     auto it = streamQueues_.find(streamId);
      36            0 :     if (UNLIKELY(it == streamQueues_.end())) {
      37            0 :         auto cq = std::make_unique<CircularQueue<std::unique_ptr<TaskInfo>>>(MAX_AICPU_CIRCULAR_QUEUE_LENGTH);
      38            0 :         auto entry = StreamQueueEntry{std::move(cq), MAX_AICPU_CIRCULAR_QUEUE_LENGTH, 0};
      39            0 :         it = streamQueues_.emplace(streamId, std::move(entry)).first;
      40            0 :     }
      41              : 
      42            0 :     auto &entry = it->second;
      43            0 :     if (UNLIKELY(entry.taskNum == entry.capacity)) {
      44            0 :         fullyCallBack_();
      45            0 :         entry.taskNum = 0;
      46              :     }
      47              : 
      48            0 :     auto& taskInfo = entry.queue->GetAndUpdate();
      49            0 :     if (taskInfo == nullptr) {
      50            0 :         taskInfo = std::make_unique<Hccl::TaskInfo>(
      51            0 :             streamId, taskId, INVALID_U32, taskParam, currDfxOpInfo_, taskParam.isMaster);
      52              :     } else {
      53            0 :         taskInfo->streamId_ = streamId;
      54            0 :         taskInfo->taskId_ = taskId;
      55            0 :         taskInfo->taskParam_ = taskParam;
      56            0 :         taskInfo->dfxOpInfo_ = currDfxOpInfo_;
      57            0 :         taskInfo->remoteRank_ = INVALID_U32;
      58            0 :         taskInfo->isMaster_ = taskParam.isMaster;
      59              :     }
      60              : 
      61            0 :     taskInfo->channelHandle_ = handle;
      62            0 :     taskInfo->getRemoteRankByHandle_ = getRemoteRankCallback_;
      63            0 :     entry.taskNum++;
      64            0 :     return HCCL_SUCCESS;
      65              : }
      66              : 
      67           20 : void MirrorTaskManagerLite::AddTaskInfo(std::unique_ptr<TaskInfo> &&taskInfo)
      68              : {
      69           20 :     if (UNLIKELY(taskInfo == nullptr)) {
      70            0 :         THROW<InternalException>(
      71            0 :             StringFormat("MirrorTaskManagerLite::AddTaskInfo taskInfo is nullptr"));
      72              :     }
      73              : 
      74           20 :     auto it = streamQueues_.find(taskInfo->streamId_);
      75           20 :     if (UNLIKELY(it == streamQueues_.end())) {
      76           14 :         auto cq = std::make_unique<CircularQueue<std::unique_ptr<TaskInfo>>>(MAX_AICPU_CIRCULAR_QUEUE_LENGTH);
      77           14 :         auto entry = StreamQueueEntry{std::move(cq), MAX_AICPU_CIRCULAR_QUEUE_LENGTH, 0};
      78           14 :         it = streamQueues_.emplace(taskInfo->streamId_, std::move(entry)).first;
      79           14 :     }
      80              : 
      81           20 :     auto &entry = it->second;
      82           20 :     if (UNLIKELY(entry.taskNum == entry.capacity)) {
      83            0 :         fullyCallBack_();
      84            0 :         entry.taskNum = 0;
      85              :     }
      86              : 
      87           20 :     entry.queue->Append(std::move(taskInfo));
      88           20 :     entry.taskNum++;
      89           40 :     return;
      90              : }
      91              : 
      92            1 : HcclResult MirrorTaskManagerLite::SetCurrDfxOpInfo(std::shared_ptr<DfxOpInfo> dfxOpInfo)
      93              : {
      94            1 :     CHK_PTR_NULL(dfxOpInfo);
      95            1 :     currDfxOpInfo_ = std::move(dfxOpInfo);
      96            1 :     return HCCL_SUCCESS;
      97              : }
      98              : 
      99            6 : std::shared_ptr<DfxOpInfo> MirrorTaskManagerLite::GetCurrDfxOpInfo() const
     100              : {
     101            6 :     return currDfxOpInfo_;
     102              : }
     103              : 
     104            0 : TaskInfoQueue *MirrorTaskManagerLite::GetQueue(u32 streamId) const
     105              : {
     106            0 :     auto it = streamQueues_.find(streamId);
     107            0 :     if (it == streamQueues_.end()) {
     108            0 :         HCCL_ERROR("MirrorTaskManagerLite::GetQueue streamId(sqId)[%u] out of range", streamId);
     109            0 :         return nullptr;
     110              :     }
     111            0 :     return it->second.queue.get();
     112              : }
     113              : 
     114            0 : TaskInfo* MirrorTaskManagerLite::GetTaskInfo(u32 streamId, u32 taskId) const
     115              : {
     116            0 :     TaskInfoQueue *queue = nullptr;
     117              :     try {
     118            0 :         queue = GetQueue(streamId);
     119            0 :     } catch (HcclException &e) {
     120            0 :         HCCL_ERROR("Hccl exception %s was caught.", e.what());
     121            0 :         return nullptr;
     122            0 :     }
     123              : 
     124            0 :     auto FindTask = [taskId](const std::unique_ptr<TaskInfo> &taskInfo) {
     125            0 :         return taskInfo->taskId_ == taskId;
     126            0 :     };
     127              : 
     128            0 :     auto task = *queue->Find(FindTask);
     129            0 :     if (task == *queue->End()) {
     130            0 :         return nullptr;
     131              :     };
     132              : 
     133            0 :     HCCL_INFO("[MirrorTaskManagerLite][GetTaskInfo]find streamdId(sqId)[%u] taskId(sqeId)[%u]", streamId, taskId);
     134              : 
     135            0 :     return (*task).get();
     136            0 : }
     137              : 
     138            2 : std::unordered_map<u32, StreamQueueEntry>::iterator MirrorTaskManagerLite::Begin()
     139              : {
     140            2 :     return streamQueues_.begin();
     141              : }
     142              : 
     143            6 : std::unordered_map<u32, StreamQueueEntry>::iterator MirrorTaskManagerLite::End()
     144              : {
     145            6 :     return streamQueues_.end();
     146              : }
     147              : 
     148          128 : MirrorTaskManagerLite::~MirrorTaskManagerLite()
     149              : {
     150          128 : }
     151              : 
     152              : } // namespace Hccl