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              : #pragma once
      12              : 
      13              : #include <unordered_set>
      14              : #include <mutex>
      15              : #include <memory>
      16              : #include <string>
      17              : 
      18              : #include <hccl/hccl_comm.h>
      19              : #include <hccl/hccl_inner.h>
      20              : #include "hccl_common.h"
      21              : #include "hccl_ip_address.h"
      22              : #include "hccl_mem_defs.h"
      23              : 
      24              : namespace hccl {
      25              : 
      26              : // 进程粒度的内存记录
      27              : class GlobalMemRecord {
      28              : public:
      29              :     // 不使用默认构造函数
      30              :     GlobalMemRecord() = delete;
      31              : 
      32              :     explicit GlobalMemRecord(const HcclMem* mem);
      33              :     explicit GlobalMemRecord(const HcclMem& mem);
      34              : 
      35              :     // 因为成员中有mutex，单独实现移动构造函数，并禁用拷贝构造函数
      36              :     GlobalMemRecord(const GlobalMemRecord& other) = delete;
      37              :     GlobalMemRecord(GlobalMemRecord&& other) noexcept;
      38              : 
      39              :     GlobalMemRecord& operator=(const GlobalMemRecord& other) = delete;
      40              :     GlobalMemRecord& operator=(GlobalMemRecord&& other) = delete;
      41            5 :     inline HcclMemType GetMemType() const { return type_; }
      42              : 
      43            4 :     inline const void* GetAddr() const { return addr_; }
      44              : 
      45            4 :     inline u64 GetSize() const { return size_; }
      46              : 
      47            4 :     inline bool IsBeingBound() const { return !boundComm_.empty(); }
      48              : 
      49            1 :     inline std::vector<std::string> GetBoundComm() const
      50              :     {
      51            2 :         return std::vector<std::string>(boundComm_.cbegin(), boundComm_.cend());
      52              :     }
      53              : 
      54           20 :     bool operator<(const GlobalMemRecord& other) const
      55              :     {
      56              :         // 先比较type
      57           20 :         if (type_ != other.type_) {
      58            5 :             return type_ < other.type_;
      59              :         }
      60              :         // 再比较地址
      61           15 :         if (addr_ != other.addr_) {
      62            7 :             return addr_ < other.addr_;
      63              :         }
      64              :         // 最后比较size
      65            8 :         return size_ < other.size_;
      66              :     }
      67              : 
      68            3 :     bool operator==(const GlobalMemRecord& other) const
      69              :     {
      70            3 :         return (type_ == other.type_) && (addr_ == other.addr_) && (size_ == other.size_);
      71              :     }
      72              : 
      73              :     bool HasOverlap(const GlobalMemRecord& other) const;          // 判断传入的内存是否有重叠
      74              :     HcclResult BindToComm(const std::string& commIdentifier);     // 绑定一个通信域
      75              :     HcclResult UnbindFromComm(const std::string& commIdentifier); // 与一个通信域解绑
      76              :     std::string PrintInfo() const;                                // 获取内存信息字符串
      77            0 :     inline void SaveRegBufInfo(HcclNetDevCtx& ctx, HcclBuf& buf)
      78              :     {
      79            0 :         HCCL_INFO("[GlobalMemRecord][SaveRegBufInfo] ctx[%p], addr[%p], len[%llu].", ctx, buf.addr, buf.len);
      80            0 :         std::lock_guard<std::mutex> lock(regBufInfoMtx_);
      81            0 :         regBufInfo_.insert(std::make_pair(ctx, buf));
      82            0 :     }
      83              : 
      84            3 :     inline std::unordered_map<HcclNetDevCtx, HcclBuf> GetAllRegBufInfo() const { return regBufInfo_; }
      85              : 
      86              : private:
      87              :     const HcclMemType type_;                      // 内存块类型, host或device
      88              :     const void* addr_;                            // 内存块地址
      89              :     const u64 size_;                              // 内存块大小字节数
      90              :     std::unique_ptr<std::mutex> pLock_;           // 锁保证多线程访问安全
      91              :     std::unordered_set<std::string> boundComm_{}; // 绑定了的通信域
      92              :     std::mutex regBufInfoMtx_;
      93              :     std::unordered_map<HcclNetDevCtx, HcclBuf> regBufInfo_{};
      94              : };
      95              : 
      96              : } // namespace hccl