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