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           52 : UbLocalNotify::UbLocalNotify(RdmaHandle rdmaHandle, bool devUsed)
      22           52 :     : BaseLocalNotify(RmaType::UB, devUsed), rdmaHandle(rdmaHandle)
      23              : {
      24           52 :     auto devType = HrtGetDeviceType(); // 先查询，避免后续失败资源泄露
      25           52 :     HrtDevResInfo devResInfo;
      26           52 :     devResInfo.dieId            = 0;
      27           52 :     devResInfo.procType         = HrtDevResProcType::PROCESS_HCCP;
      28           52 :     devResInfo.resType          = HrtDevResType::RES_TYPE_STARS_NOTIFY_RECORD;
      29           52 :     devResInfo.resId            = GetNotify()->GetId();
      30           52 :     devResInfo.flag             = 0;
      31           52 :     auto resAddrInfo            = HrtGetDevResAddress(devResInfo);
      32           52 :     addr                        = resAddrInfo.address;
      33           52 :     DevCapability::GetInstance().Init(devType); // 单例初始化
      34           52 :     size                        = DevCapability::GetInstance().GetNotifySize();
      35           52 :     auto tokenIdInfoPair        = RdmaHandleManager::GetInstance().GetTokenIdInfo(rdmaHandle);
      36           52 :     TokenIdHandle tokenIdHandle = tokenIdInfoPair.first;
      37           52 :     tokenId                     = tokenIdInfoPair.second;
      38          156 :     HCCL_INFO("[UbLocalNotify] tokenIdHandle=0x[%llx]", tokenIdHandle);
      39          156 :     HCCL_INFO("mapped addr=[%llx]", addr);
      40          156 :     HCCL_INFO("UB notify size=[%u]", size);
      41              : 
      42              :     // halNotifyMap 返回的地址不保证4K对齐，
      43              :     // notify的地址还是使用hal接口返回的addr，但是注册mem的时候我们需要自己做向下对齐
      44           52 :     tokenValue = GetUbToken();
      45           52 :     std::pair<u64, u64> alignBuf = BufAlign(addr, size);
      46           52 :     HrtRaUbLocMemRegParam lmemReg{alignBuf.first, alignBuf.second, tokenValue, tokenIdHandle, 1};
      47           52 :     reqReg = HrtRaUbLocalMemReg(rdmaHandle, lmemReg);
      48           52 :     keySize         = reqReg.keySize;
      49           52 :     memHandle       = reqReg.handle;
      50           52 :     (void)memcpy_s(key, HRT_UB_MEM_KEY_MAX_LEN, reqReg.key, HRT_UB_MEM_KEY_MAX_LEN);
      51           52 : }
      52              : 
      53           88 : string UbLocalNotify::Describe() const
      54              : {
      55              :     return StringFormat("UbLocalNotify:notify=%s, addr=0x%llx, keySize=%u, memHandle=0x%llx",
      56           88 :                         GetNotify()->Describe().c_str(), addr, keySize, memHandle);
      57              : }
      58              : 
      59            2 : void UbLocalNotify::Wait(const Stream &stream, u32 timeout) const
      60              : {
      61            2 :     GetNotify()->Wait(stream, timeout);
      62            2 : }
      63              : 
      64            1 : void UbLocalNotify::Post(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           52 : void UbLocalNotify::ReleaseResource() const
      79              : {
      80           52 :     if (rdmaHandle && memHandle != 0) {
      81            5 :         HrtRaUbLocalMemUnreg(rdmaHandle, memHandle);
      82              :     }
      83              : 
      84           52 :     HrtDevResInfo devResInfo;
      85           52 :     devResInfo.dieId    = 0;
      86           52 :     devResInfo.procType = HrtDevResProcType::PROCESS_HCCP;
      87           52 :     devResInfo.resType  = HrtDevResType::RES_TYPE_STARS_NOTIFY_RECORD;
      88           52 :     devResInfo.resId    = GetNotify()->GetId();
      89           52 :     devResInfo.flag     = 0;
      90           52 :     HrtReleaseDevResAddress(devResInfo);
      91           52 : }
      92              : 
      93           56 : UbLocalNotify::~UbLocalNotify()
      94              : {
      95           52 :     DECTOR_TRY_CATCH("UbLocalNotify", ReleaseResource());
      96           56 : }
      97              : } // namespace Hccl