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_local_notify.h"
      12              : 
      13              : #include "dev_capability.h"
      14              : #include "not_support_exception.h"
      15              : #include "exchange_ub_buffer_dto.h"
      16              : #include "rdma_handle_manager.h"
      17              : #include "local_ub_rma_buffer.h"
      18              : 
      19              : namespace Hccl {
      20              : 
      21           59 : UbLocalNotify::UbLocalNotify(RdmaHandle rdmaHandle, bool devUsed)
      22              :     : BaseLocalNotify(RmaType::UB, devUsed),
      23           59 :       rdmaHandle(rdmaHandle)
      24              : {
      25           59 :     auto devType = HrtGetDeviceType(); // 先查询，避免后续失败资源泄露
      26           59 :     HrtDevResInfo devResInfo;
      27           59 :     devResInfo.dieId = 0;
      28           59 :     devResInfo.procType = HrtDevResProcType::PROCESS_HCCP;
      29           59 :     devResInfo.resType = HrtDevResType::RES_TYPE_STARS_NOTIFY_RECORD;
      30           59 :     devResInfo.resId = GetNotify()->GetId();
      31           59 :     devResInfo.flag = 0;
      32           59 :     auto resAddrInfo = HrtGetDevResAddress(devResInfo);
      33           59 :     addr = resAddrInfo.address;
      34           59 :     DevCapability::GetInstance().Init(devType); // 单例初始化
      35           59 :     size = DevCapability::GetInstance().GetNotifySize();
      36           59 :     std::pair<u64, u64> alignBuf = BufAlign(addr, size);
      37           59 :     bufKey_ = BufferKey<uintptr_t, u64>{alignBuf.first, alignBuf.second};
      38           59 :     auto tokenIdInfoPair = RdmaHandleManager::GetInstance().GetTokenIdInfo(rdmaHandle, bufKey_);
      39           59 :     tokenIdHandle_ = tokenIdInfoPair.first;
      40           59 :     tokenId = tokenIdInfoPair.second;
      41          177 :     HCCL_INFO("[UbLocalNotify] tokenIdHandle=0x[%llx]", tokenIdHandle_);
      42          177 :     HCCL_INFO("mapped addr=[%llx]", addr);
      43          177 :     HCCL_INFO("UB notify size=[%u]", size);
      44              : 
      45              :     // halNotifyMap 返回的地址不保证4K对齐，
      46              :     // notify的地址还是使用hal接口返回的addr，但是注册mem的时候我们需要自己做向下对齐
      47           59 :     tokenValue = GetUbToken();
      48           59 :     HrtRaUbLocMemRegParam lmemReg{alignBuf.first, alignBuf.second, tokenValue, tokenIdHandle_, 1};
      49           59 :     reqReg = HrtRaUbLocalMemReg(rdmaHandle, lmemReg);
      50           59 :     keySize = reqReg.keySize;
      51           59 :     memHandle = reqReg.handle;
      52           59 :     (void)memcpy_s(key, HRT_UB_MEM_KEY_MAX_LEN, reqReg.key, HRT_UB_MEM_KEY_MAX_LEN);
      53           59 : }
      54              : 
      55           96 : string UbLocalNotify::Describe() const
      56              : {
      57              :     return StringFormat(
      58          192 :         "UbLocalNotify:notify=%s, addr=0x%llx, keySize=%u, memHandle=0x%llx", GetNotify()->Describe().c_str(), addr,
      59          288 :         keySize, memHandle);
      60              : }
      61              : 
      62            2 : void UbLocalNotify::Wait(const Stream& stream, u32 timeout) const { GetNotify()->Wait(stream, timeout); }
      63              : 
      64            1 : void UbLocalNotify::Post([[maybe_unused]] const Stream& stream) const
      65              : {
      66            1 :     std::string msg = "UbLocalNotify does not support submitting record task";
      67            4 :     MACRO_THROW(NotSupportException, msg);
      68            1 : }
      69              : 
      70            2 : std::unique_ptr<Serializable> UbLocalNotify::GetExchangeDto()
      71              : {
      72              :     std::unique_ptr<ExchangeUbBufferDto> dto
      73            2 :         = make_unique<ExchangeUbBufferDto>(addr, size, tokenValue, tokenId, keySize, GetNotify()->GetId());
      74            2 :     (void)memcpy_s(dto->key, HRT_UB_MEM_KEY_MAX_LEN, key, HRT_UB_MEM_KEY_MAX_LEN);
      75            4 :     return std::unique_ptr<Serializable>(dto.release());
      76            2 : }
      77              : 
      78           61 : void UbLocalNotify::ReleaseResource()
      79              : {
      80           61 :     const bool ctxValid = rdmaHandle != nullptr && RdmaHandleManager::GetInstance().IsHandleValid(rdmaHandle);
      81              : 
      82           61 :     if (rdmaHandle && memHandle != 0) {
      83            8 :         if (!ctxValid) {
      84           21 :             HCCL_WARNING(
      85              :                 "[UbLocalNotify][%s] skip HrtRaUbLocalMemUnreg, "
      86              :                 "rdmaHandle=%p invalid (DeInit/DestroyAll done), memHandle=0x%llx",
      87              :                 __func__, rdmaHandle, static_cast<unsigned long long>(memHandle));
      88            7 :             memHandle = 0;
      89              :         } else {
      90            1 :             HrtRaUbLocalMemUnreg(rdmaHandle, memHandle);
      91            1 :             memHandle = 0;
      92              :         }
      93              :     }
      94              : 
      95           61 :     if (ctxValid) {
      96            4 :         RdmaHandleManager::GetInstance().PutTokenIdInfo(rdmaHandle, bufKey_, tokenIdHandle_);
      97           57 :     } else if (rdmaHandle != nullptr) {
      98          171 :         HCCL_WARNING("[UbLocalNotify][%s] skip PutTokenIdInfo, rdmaHandle=%p invalid", __func__, rdmaHandle);
      99              :     }
     100              : 
     101           61 :     HrtDevResInfo devResInfo;
     102           61 :     devResInfo.dieId = 0;
     103           61 :     devResInfo.procType = HrtDevResProcType::PROCESS_HCCP;
     104           61 :     devResInfo.resType = HrtDevResType::RES_TYPE_STARS_NOTIFY_RECORD;
     105           61 :     devResInfo.resId = GetNotify()->GetId();
     106           61 :     devResInfo.flag = 0;
     107           61 :     HrtReleaseDevResAddress(devResInfo);
     108           61 : }
     109              : 
     110           63 : UbLocalNotify::~UbLocalNotify() { DECTOR_TRY_CATCH("UbLocalNotify", ReleaseResource()); }
     111              : } // namespace Hccl