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 "offload_stream_manager.h"
      12              : 
      13              : namespace hccl {
      14          530 : OffloadStreamManager::OffloadStreamManager() = default;
      15          530 : OffloadStreamManager::~OffloadStreamManager() = default;
      16              : 
      17            3 : HcclResult OffloadStreamManager::RegisterMaster(const std::string &tag, Stream &stream)
      18              : {
      19            3 :     std::unique_lock<std::mutex> lock(masterMapMutex_);
      20            3 :     masterMap_[tag] = stream;
      21            3 :     HCCL_DEBUG("[OffloadStreamManager][RegisterMaster]register master stream[%p] success, tag[%s].",
      22              :         stream.ptr(), tag.c_str());
      23            3 :     return HCCL_SUCCESS;
      24            3 : }
      25              : 
      26          131 : HcclResult OffloadStreamManager::RegisterSlaves(const std::string &tag, std::vector<Stream> &stream)
      27              : {
      28          131 :     HCCL_DEBUG("[OffloadStreamManager][RegisterSlaves]start register slaves stream, tag[%s], size[%u].",
      29              :         tag.c_str(), stream.size());
      30              : 
      31          131 :     std::unique_lock<std::mutex> lock(slavesMapMutex_);
      32          131 :     auto iter = slavesMap_.find(tag);
      33          131 :     if (iter != slavesMap_.end()) {
      34            1 :         HCCL_ERROR("[OffloadStreamManager][RegisterSlaves]in offload stream manager, register slaves fail,"
      35              :             "tag[%s] has existed", tag.c_str());
      36            1 :         return HCCL_E_PARA;
      37              :     }
      38          130 :     slavesMap_.insert(std::make_pair(tag, stream));
      39          130 :     HCCL_INFO("[OffloadStreamManager][RegisterSlaves]register slaves stream success, tag[%s], size[%u].",
      40              :         tag.c_str(), stream.size());
      41          130 :     return HCCL_SUCCESS;
      42          131 : }
      43              : 
      44            2 : Stream OffloadStreamManager::GetMaster(const std::string &tag)
      45              : {
      46            2 :     std::unique_lock<std::mutex> lock(masterMapMutex_);
      47            2 :     auto iter = masterMap_.find(tag);
      48            2 :     if (iter == masterMap_.end()) {
      49            1 :         HCCL_ERROR("[OffloadStreamManager][GetMaster]can't find tag[%s]", tag.c_str());
      50            1 :         return Stream();
      51              :     }
      52            1 :     return iter->second;
      53            2 : }
      54              : 
      55           23 : std::vector<Stream> OffloadStreamManager::GetSlaves(const std::string &tag, u32 num)
      56              : {
      57           23 :     HCCL_DEBUG("[OffloadStreamManager][GetSlaves]requesting for [%u] slaves, tag[%s].", num, tag.c_str());
      58           22 :     if (num == 0) {
      59            2 :         HCCL_WARNING("[OffloadStreamManager][GetSlaves]requesting for 0 slaves, return empty vector.");
      60            4 :         return std::vector<Stream>(0);
      61              :     }
      62              : 
      63           20 :     std::unique_lock<std::mutex> lock(slavesMapMutex_);
      64           21 :     auto iter = slavesMap_.find(tag);
      65           19 :     if (iter == slavesMap_.end()) {
      66            1 :         HCCL_ERROR("[OffloadStreamManager][GetSlaves]can't find tag[%s]", tag.c_str());
      67            1 :         return std::vector<Stream>();
      68              :     }
      69              : 
      70           17 :     if (iter->second.size() < num) {
      71            1 :         HCCL_ERROR("[OffloadStreamManager][GetSlaves]" \
      72              :             "trying to get [%u] slaves fail, only [%u] slaves available, tag[%s].",
      73              :             num, iter->second.size(), tag.c_str());
      74            1 :         return std::vector<Stream>();
      75              :     }
      76              : 
      77           16 :     std::vector<Stream> res(iter->second.begin(), iter->second.begin() + num);
      78           17 :     iter->second.erase(iter->second.begin(), iter->second.begin() + num);
      79           15 :     HCCL_INFO("[OffloadStreamManager][GetSlaves]get [%u] slaves success, returning.", res.size());
      80           19 :     return res;
      81           21 : }
      82              : 
      83           18 : HcclResult OffloadStreamManager::ClearSlaves(const std::string &tag)
      84              : {
      85           18 :     std::unique_lock<std::mutex> lock(slavesMapMutex_);
      86           16 :     auto iter = slavesMap_.find((tag));
      87           17 :     if (iter != slavesMap_.end()) {
      88           17 :         slavesMap_.erase(tag);
      89              :     }
      90           18 :     HCCL_DEBUG("[OffloadStreamManager][ClearSlaves]Destroy slaves stream success, tag[%s]", tag.c_str());
      91           18 :     return HCCL_SUCCESS;
      92           18 : }
      93              : 
      94            6 : HcclResult OffloadStreamManager::ClearSlaves()
      95              : {
      96            6 :     std::unique_lock<std::mutex> lock(slavesMapMutex_);
      97            6 :     slavesMap_.clear();
      98            6 :     HCCL_DEBUG("[OffloadStreamManager][ClearSlaves]clear all slave streams.");
      99            6 :     return HCCL_SUCCESS;
     100            6 : }
     101              : 
     102              : }  // namespace hccl