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
      42              :     {
      43            5 :         return type_;
      44              :     }
      45              : 
      46            4 :     inline const void* GetAddr() const
      47              :     {
      48            4 :         return addr_;
      49              :     }
      50              : 
      51            4 :     inline u64 GetSize() const
      52              :     {
      53            4 :         return size_;
      54              :     }
      55              : 
      56            4 :     inline bool IsBeingBound() const
      57              :     {
      58            4 :         return !boundComm_.empty();
      59              :     }
      60              : 
      61            1 :     inline std::vector<std::string> GetBoundComm() const
      62              :     {
      63            2 :         return std::vector<std::string>(boundComm_.cbegin(), boundComm_.cend());
      64              :     }
      65              : 
      66           20 :     bool operator < (const GlobalMemRecord& other) const
      67              :     {
      68              :         // 先比较type
      69           20 :         if (type_ != other.type_) {
      70            5 :             return type_ < other.type_;
      71              :         }
      72              :         // 再比较地址
      73           15 :         if (addr_ != other.addr_) {
      74            7 :             return addr_ < other.addr_;
      75              :         }
      76              :         // 最后比较size
      77            8 :         return size_ < other.size_;
      78              :     }
      79              : 
      80            3 :     bool operator == (const GlobalMemRecord& other) const
      81              :     {
      82            3 :         return (type_ == other.type_) && (addr_ == other.addr_) && (size_ == other.size_);
      83              :     }
      84              : 
      85              :     bool HasOverlap(const GlobalMemRecord& other) const;  // 判断传入的内存是否有重叠
      86              :     HcclResult BindToComm(const std::string &commIdentifier);       // 绑定一个通信域
      87              :     HcclResult UnbindFromComm(const std::string &commIdentifier);   // 与一个通信域解绑
      88              :     std::string PrintInfo() const; // 获取内存信息字符串
      89            0 :     inline void SaveRegBufInfo(HcclNetDevCtx& ctx, HcclBuf& buf)
      90              :     {
      91            0 :         HCCL_INFO("[GlobalMemRecord][SaveRegBufInfo] ctx[%p], addr[%p], len[%llu].", ctx, buf.addr, buf.len);
      92            0 :         std::lock_guard<std::mutex> lock(regBufInfoMtx_);
      93            0 :         regBufInfo_.insert(std::make_pair(ctx, buf));
      94            0 :     }
      95              : 
      96            3 :     inline std::unordered_map<HcclNetDevCtx, HcclBuf> GetAllRegBufInfo() const
      97              :     {
      98            3 :         return regBufInfo_;
      99              :     }
     100              : 
     101              : private:
     102              :     const HcclMemType type_; // 内存块类型, host或device
     103              :     const void* addr_; // 内存块地址
     104              :     const u64 size_; // 内存块大小字节数
     105              :     std::unique_ptr<std::mutex> pLock_; //锁保证多线程访问安全
     106              :     std::unordered_set<std::string> boundComm_{}; // 绑定了的通信域
     107              :     std::mutex regBufInfoMtx_;
     108              :     std::unordered_map<HcclNetDevCtx, HcclBuf> regBufInfo_{};
     109              : };
     110              : 
     111              : } // namespace hccl