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