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
|