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