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 "local_rma_buf_manager.h"
      12              : #include "internal_exception.h"
      13              : #include "rdma_handle_manager.h"
      14              : #include "local_ipc_rma_buffer_v2.h"
      15              : #include "local_rdma_rma_buffer_v2.h"
      16              : #include "local_ub_rma_buffer.h"
      17              : 
      18              : #include "log.h"
      19              : #include "stl_util.h"
      20              : #include "exception_util.h"
      21              : #include "communicator_impl.h"
      22              : 
      23              : namespace Hccl {
      24              : 
      25          329 : LocalRmaBufManager::LocalRmaBufManager(const CommunicatorImpl& communicator)
      26          329 :     : comm(const_cast<CommunicatorImpl*>(&communicator))
      27          329 : {}
      28              : 
      29          329 : LocalRmaBufManager::~LocalRmaBufManager() { DECTOR_TRY_CATCH("LocalRmaBufManager", Destroy()); }
      30              : 
      31           33 : bool LocalRmaBufManager::IsExist(const string& opTag, const PortData& portData, BufferType bufferType)
      32              : {
      33           49 :     return bufs.find(opTag) != bufs.end() && bufs[opTag].find(portData) != bufs[opTag].end()
      34           49 :            && bufs[opTag][portData].find(bufferType) != bufs[opTag][portData].end();
      35              : }
      36              : 
      37           12 : LocalRmaBuffer* LocalRmaBufManager::Reg(
      38              :     const string& opTag, BufferType bufferType, std::shared_ptr<Buffer> buffer, const PortData& portData,
      39              :     LinkProtocol linkProtocol)
      40              : {
      41           12 :     if (buffer == nullptr) {
      42            0 :         HCCL_ERROR("input buffer is null");
      43            0 :         return nullptr;
      44              :     }
      45           36 :     HCCL_INFO(
      46              :         "LocalRmaBufManager::Reg, buffer[%s], opTag[%s], bufferType[%u], portData[%s]", buffer->Describe().c_str(),
      47              :         opTag.c_str(), bufferType, portData.Describe().c_str());
      48           12 :     if (IsExist(opTag, portData, bufferType)) {
      49              :         string msg = StringFormat(
      50              :             "opTag=%s bufferType=%s, buffer=%s already reg to portData=%s", opTag.c_str(),
      51            1 :             bufferType.Describe().c_str(), buffer->Describe().c_str(), portData.Describe().c_str());
      52            3 :         HCCL_DEBUG(msg.c_str());
      53            1 :         return bufs[opTag][portData][bufferType].get();
      54            1 :     }
      55           11 :     if (portData.GetType() == PortDeploymentType::P2P) {
      56            0 :         bufs[opTag][portData][bufferType] = make_unique<LocalIpcRmaBuffer>(buffer);
      57            0 :         return bufs[opTag][portData][bufferType].get();
      58              :     } else {
      59           11 :         if (portData.GetProto() == LinkProtoType::RDMA) {
      60              :             RdmaHandle rdmaHandle
      61            0 :                 = RdmaHandleManager::GetInstance().Get(comm->GetDevicePhyId(), portData, linkProtocol);
      62            0 :             bufs[opTag][portData][bufferType] = make_unique<LocalRdmaRmaBuffer>(buffer, rdmaHandle);
      63            0 :             return bufs[opTag][portData][bufferType].get();
      64           11 :         } else if (portData.GetProto() == LinkProtoType::UB) {
      65           30 :             HCCL_INFO("LocalRmaBufManager::Reg, comm->GetOpAiCpuTSFeatureFlag[%d]", comm->GetOpAiCpuTSFeatureFlag());
      66           10 :             if (comm->GetOpAiCpuTSFeatureFlag()) { // 算子粒度
      67            3 :                 bufs[opTag][portData][bufferType] = make_unique<LocalUbRmaBuffer>(buffer);
      68              :             } else {
      69              :                 RdmaHandle rdmaHandle
      70            7 :                     = RdmaHandleManager::GetInstance().Get(comm->GetDevicePhyId(), portData, linkProtocol);
      71            7 :                 bufs[opTag][portData][bufferType] = make_unique<LocalUbRmaBuffer>(buffer, rdmaHandle);
      72              :             }
      73           10 :             return bufs[opTag][portData][bufferType].get();
      74              :         }
      75              :         // 待修改: 仅支持 P2P 和 RDMA
      76            1 :         string msg = StringFormat("PortData=%s is error", portData.Describe().c_str());
      77            4 :         MACRO_THROW(InternalException, msg);
      78            1 :     }
      79              : }
      80              : 
      81           21 : LocalRmaBuffer* LocalRmaBufManager::Get(const string& opTag, const PortData& portData, BufferType bufferType)
      82              : {
      83           21 :     if (IsExist(opTag, portData, bufferType)) { // if localRmaBuffer exists
      84           21 :         HCCL_INFO(
      85              :             "[LocalRmaBufManager][%s] LocalUbRmaBuffer[%s]", __func__,
      86              :             bufs[opTag][portData][bufferType]->Describe().c_str());
      87            7 :         return bufs[opTag][portData][bufferType].get();
      88              :     }
      89           42 :     HCCL_WARNING(
      90              :         "LocalRmaBuffer doesn't exist:opTag[%s], bufferType[%u], portData[%s]", opTag.c_str(), bufferType,
      91              :         portData.Describe().c_str());
      92           14 :     return nullptr;
      93              : }
      94              : 
      95          330 : void LocalRmaBufManager::Destroy() { bufs.clear(); }
      96              : 
      97            0 : LocalRmaBuffer* LocalRmaBufManager::Get(const PortData& portData)
      98              : {
      99            0 :     if (Contain(ccuBufs, portData)) {
     100            0 :         return ccuBufs[portData].get();
     101              :     }
     102            0 :     HCCL_WARNING("LocalRmaBuffer doesn't exist at port[%s].", portData.Describe().c_str());
     103            0 :     return nullptr;
     104              : }
     105              : 
     106            1 : HcclResult LocalRmaBufManager::Dereg(const string& opTag)
     107              : {
     108            1 :     if (bufs.find(opTag) == bufs.end()) {
     109            3 :         HCCL_WARNING("[LocalRmaBufManager::%s] opTag[%s] Cannot find Buffer in bufs.", __func__, opTag.c_str());
     110            1 :         return HCCL_SUCCESS;
     111              :     }
     112            0 :     bufs.erase(opTag);
     113            0 :     return HCCL_SUCCESS;
     114              : }
     115              : 
     116              : } // namespace Hccl