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 "ub_mem.h"
      12              : #include <algorithm>
      13              : #include "log.h"
      14              : 
      15              : namespace hcomm {
      16              : 
      17           13 : UbMemRegedMemMgr::UbMemRegedMemMgr() { localIpcRmaBufferMgr_ = std::make_unique<LocalIpcRmaBufferMgr>(); }
      18              : 
      19           27 : HcclResult UbMemRegedMemMgr::RegisterMemory(HcommMem mem, const char* memTag, void** memHandle)
      20              : {
      21           27 :     HCCL_INFO("[%s] Begin", __func__);
      22           27 :     CHK_PTR_NULL(localIpcRmaBufferMgr_);
      23           27 :     std::lock_guard<std::mutex> lock(memMtx_);
      24           54 :     return RegisterMemoryImpl(
      25           27 :         mem, memTag, memHandle, localIpcRmaBufferMgr_, allRegisteredBuffers_,
      26              :         static_cast<std::vector<std::shared_ptr<Hccl::LocalIpcRmaBuffer>>*>(nullptr), "UbMemRegedMemMgr",
      27            0 :         [&](auto& bufPtr, auto& parent) {
      28            5 :             return std::make_shared<Hccl::LocalIpcRmaBuffer>(bufPtr, *parent);
      29              :         },
      30            0 :         [&](auto& bufPtr) {
      31           16 :             return std::make_shared<Hccl::LocalIpcRmaBuffer>(bufPtr);
      32           27 :         });
      33           27 : }
      34              : 
      35           25 : HcclResult UbMemRegedMemMgr::UnregisterMemory(void* memHandle)
      36              : {
      37           25 :     HCCL_INFO("[%s] Begin", __func__);
      38           25 :     CHK_PTR_NULL(localIpcRmaBufferMgr_);
      39           25 :     std::lock_guard<std::mutex> lock(memMtx_);
      40           50 :     return UnregisterMemoryImpl(
      41           25 :         memHandle, localIpcRmaBufferMgr_, allRegisteredBuffers_,
      42              :         static_cast<std::vector<std::shared_ptr<Hccl::LocalIpcRmaBuffer>>*>(nullptr),
      43            0 :         [](auto* b) {
      44           12 :             return b->GetIpcPtr();
      45              :         },
      46            0 :         [](auto a, auto b) {
      47            6 :             return a == b;
      48           25 :         });
      49           25 : }
      50              : 
      51            0 : HcclResult UbMemRegedMemMgr::MemoryExport(
      52              :     [[maybe_unused]] const EndpointDesc endpointDesc, [[maybe_unused]] void* memHandle, [[maybe_unused]] void** memDesc,
      53              :     [[maybe_unused]] uint32_t* memDescLen)
      54              : {
      55            0 :     HCCL_INFO("UbMemRegedMemMgr MemoryExport is not supported.");
      56            0 :     return HCCL_SUCCESS;
      57              : }
      58              : 
      59            0 : HcclResult UbMemRegedMemMgr::MemoryImport(
      60              :     [[maybe_unused]] const void* memDesc, [[maybe_unused]] uint32_t descLen, [[maybe_unused]] HcommMem* outMem)
      61              : {
      62            0 :     HCCL_INFO("UbMemRegedMemMgr MemoryImport is not supported.");
      63            0 :     return HCCL_SUCCESS;
      64              : }
      65              : 
      66            0 : HcclResult UbMemRegedMemMgr::MemoryUnimport([[maybe_unused]] const void* memDesc, [[maybe_unused]] uint32_t descLen)
      67              : {
      68            0 :     HCCL_INFO("UbMemRegedMemMgr MemoryUnimport is not supported.");
      69            0 :     return HCCL_SUCCESS;
      70              : }
      71              : 
      72              : HcclResult
      73            0 : UbMemRegedMemMgr::GetAllMemHandles([[maybe_unused]] void** memHandles, [[maybe_unused]] uint32_t* memHandleNum)
      74              : {
      75            0 :     HCCL_INFO("UbMemRegedMemMgr GetAllMemHandles is not supported.");
      76            0 :     return HCCL_SUCCESS;
      77              : }
      78              : } // namespace hcomm