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              : #ifndef USER_REMOTE_MEM_GETTER_H
      12              : #define USER_REMOTE_MEM_GETTER_H
      13              : 
      14              : #include <cstdint>
      15              : #include <vector>
      16              : #include <array>
      17              : #include <string>
      18              : #include <functional>
      19              : #include "hccl_mem_defs.h"
      20              : #include "hccl/hccl_types.h"
      21              : #include "log.h"
      22              : 
      23              : constexpr uint32_t MAX_BUFFER_NUM = 30000;
      24              : 
      25              : namespace Hccl {
      26              : 
      27              : template <typename T>
      28              : struct RemoteMemCtx {
      29              :     bool& cacheValid;
      30              :     std::vector<T>& rmtBufferVec;
      31              :     std::vector<CommMem>& remoteUserMems;
      32              :     std::vector<std::string>& memInfoCopies;
      33              :     std::vector<char*>& memInfoPointers;
      34              :     CommMem** remoteMem;
      35              :     char*** memInfos;
      36              :     uint32_t* memNum;
      37              : 
      38           28 :     RemoteMemCtx(
      39              :         bool& cacheValid, std::vector<T>& rmtBufferVec, std::vector<CommMem>& remoteUserMems,
      40              :         std::vector<std::string>& memInfoCopies, std::vector<char*>& memInfoPointers, CommMem** remoteMem,
      41              :         char*** memInfos, uint32_t* memNum)
      42           28 :         : cacheValid(cacheValid),
      43           28 :           rmtBufferVec(rmtBufferVec),
      44           28 :           remoteUserMems(remoteUserMems),
      45           28 :           memInfoCopies(memInfoCopies),
      46           28 :           memInfoPointers(memInfoPointers),
      47           28 :           remoteMem(remoteMem),
      48           28 :           memInfos(memInfos),
      49           28 :           memNum(memNum) {};
      50              : };
      51              : 
      52            0 : inline HcclMemType CommMemTypeToHcclMemType(CommMemType type)
      53              : {
      54            0 :     switch (type) {
      55            0 :         case COMM_MEM_TYPE_DEVICE:
      56            0 :             return HCCL_MEM_TYPE_DEVICE;
      57            0 :         case COMM_MEM_TYPE_HOST:
      58            0 :             return HCCL_MEM_TYPE_HOST;
      59            0 :         default:
      60            0 :             return HCCL_MEM_TYPE_NUM;
      61              :     }
      62              : }
      63              : 
      64           19 : inline CommMemType HcclMemTypeToCommMemType(HcclMemType type)
      65              : {
      66           19 :     switch (type) {
      67           10 :         case HCCL_MEM_TYPE_DEVICE:
      68           10 :             return COMM_MEM_TYPE_DEVICE;
      69            4 :         case HCCL_MEM_TYPE_HOST:
      70            4 :             return COMM_MEM_TYPE_HOST;
      71            5 :         default:
      72            5 :             return COMM_MEM_TYPE_INVALID;
      73              :     }
      74              : }
      75              : 
      76              : template <typename T>
      77           28 : HcclResult GetRemoteUserMems(RemoteMemCtx<T>& remoteMemCtx)
      78              : {
      79           28 :     CHK_PRT_RET(!remoteMemCtx.remoteMem, HCCL_ERROR("[GetRemoteUserMems] remoteMem is nullptr"), HCCL_E_PARA);
      80           25 :     CHK_PRT_RET(!remoteMemCtx.memInfos, HCCL_ERROR("[GetRemoteUserMems] memInfos is nullptr"), HCCL_E_PARA);
      81           22 :     CHK_PRT_RET(!remoteMemCtx.memNum, HCCL_ERROR("[GetRemoteUserMems] memNum is nullptr"), HCCL_E_PARA);
      82           18 :     *(remoteMemCtx.remoteMem) = nullptr;
      83           18 :     *(remoteMemCtx.memInfos) = nullptr;
      84           18 :     *(remoteMemCtx.memNum) = 0;
      85           18 :     uint32_t userMemCount = remoteMemCtx.rmtBufferVec.size();
      86           18 :     if (userMemCount == 0) {
      87            8 :         HCCL_WARNING("[%s] bufferNum is 0.", __func__);
      88            6 :         return HCCL_SUCCESS;
      89              :     }
      90              :     // 检查是否有缓存
      91           12 :     if (!remoteMemCtx.cacheValid) {
      92           11 :         remoteMemCtx.remoteUserMems.clear();
      93           11 :         remoteMemCtx.memInfoCopies.clear();
      94           11 :         remoteMemCtx.memInfoCopies.reserve(userMemCount);
      95           11 :         remoteMemCtx.memInfoPointers.clear();
      96           11 :         remoteMemCtx.memInfoPointers.reserve(userMemCount);
      97           32 :         for (uint32_t i = 0; i < userMemCount; ++i) {
      98           19 :             auto& rmtBuffer = remoteMemCtx.rmtBufferVec[i];
      99           19 :             if (rmtBuffer == nullptr) {
     100            0 :                 return HCCL_E_PTR;
     101              :             }
     102           19 :             CommMem mem{};
     103           19 :             mem.type = HcclMemTypeToCommMemType(rmtBuffer->GetMemType());
     104           19 :             mem.addr = reinterpret_cast<void*>(rmtBuffer->GetAddr());
     105           19 :             mem.size = rmtBuffer->GetSize();
     106           19 :             remoteMemCtx.remoteUserMems.push_back(mem);
     107           19 :             std::string memInfoCopy = rmtBuffer->GetMemInfo();
     108           19 :             remoteMemCtx.memInfoCopies.push_back(std::move(memInfoCopy));
     109           19 :             remoteMemCtx.memInfoPointers.push_back(const_cast<char*>(remoteMemCtx.memInfoCopies.back().c_str()));
     110           23 :             HCCL_INFO(
     111              :                 "[%s] Found buffer[addr:%p, size:%llu, memInfo:%s]", __func__, mem.addr, mem.size,
     112              :                 remoteMemCtx.memInfoCopies.back().c_str());
     113              :         }
     114           11 :         remoteMemCtx.cacheValid = true;
     115              :     }
     116           12 :     *(remoteMemCtx.remoteMem) = remoteMemCtx.remoteUserMems.data();
     117           12 :     *(remoteMemCtx.memInfos) = remoteMemCtx.memInfoPointers.data();
     118           12 :     *(remoteMemCtx.memNum) = userMemCount;
     119           12 :     return HCCL_SUCCESS;
     120              : }
     121              : } // namespace Hccl
     122              : 
     123              : #endif // USER_REMOTE_MEM_GETTER_H