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