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