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 :     std::pair<u64, u64> alignBuf = BufAlign(addr, size);
      36           52 :     bufKey_ = BufferKey<uintptr_t, u64>{alignBuf.first, alignBuf.second};
      37           52 :     auto tokenIdInfoPair        = RdmaHandleManager::GetInstance().GetTokenIdInfo(rdmaHandle, bufKey_);
      38           52 :     tokenIdHandle_              = tokenIdInfoPair.first;
      39           52 :     tokenId                     = tokenIdInfoPair.second;
      40          156 :     HCCL_INFO("[UbLocalNotify] tokenIdHandle=0x[%llx]", tokenIdHandle_);
      41          156 :     HCCL_INFO("mapped addr=[%llx]", addr);
      42          156 :     HCCL_INFO("UB notify size=[%u]", size);
      43              : 
      44              :     // halNotifyMap 返回的地址不保证4K对齐，
      45              :     // notify的地址还是使用hal接口返回的addr，但是注册mem的时候我们需要自己做向下对齐
      46           52 :     tokenValue = GetUbToken();
      47           52 :     HrtRaUbLocMemRegParam lmemReg{alignBuf.first, alignBuf.second, tokenValue, tokenIdHandle_, 1};
      48           52 :     reqReg = HrtRaUbLocalMemReg(rdmaHandle, lmemReg);
      49           52 :     keySize         = reqReg.keySize;
      50           52 :     memHandle       = reqReg.handle;
      51           52 :     (void)memcpy_s(key, HRT_UB_MEM_KEY_MAX_LEN, reqReg.key, HRT_UB_MEM_KEY_MAX_LEN);
      52           52 : }
      53              : 
      54           88 : string UbLocalNotify::Describe() const
      55              : {
      56              :     return StringFormat("UbLocalNotify:notify=%s, addr=0x%llx, keySize=%u, memHandle=0x%llx",
      57           88 :                         GetNotify()->Describe().c_str(), addr, keySize, memHandle);
      58              : }
      59              : 
      60            2 : void UbLocalNotify::Wait(const Stream &stream, u32 timeout) const
      61              : {
      62            2 :     GetNotify()->Wait(stream, timeout);
      63            2 : }
      64              : 
      65            1 : void UbLocalNotify::Post(const Stream &stream) const
      66              : {
      67            1 :     std::string msg = "UbLocalNotify does not support submitting record task";
      68            4 :     MACRO_THROW(NotSupportException, msg);
      69            1 : }
      70              : 
      71            2 : std::unique_ptr<Serializable> UbLocalNotify::GetExchangeDto()
      72              : {
      73              :     std::unique_ptr<ExchangeUbBufferDto> dto
      74            2 :         = make_unique<ExchangeUbBufferDto>(addr, size, tokenValue, tokenId, keySize, GetNotify()->GetId());
      75            2 :     (void)memcpy_s(dto->key, HRT_UB_MEM_KEY_MAX_LEN, key, HRT_UB_MEM_KEY_MAX_LEN);
      76            4 :     return std::unique_ptr<Serializable>(dto.release());
      77            2 : }
      78              : 
      79           52 : void UbLocalNotify::ReleaseResource() const
      80              : {
      81           52 :     if (rdmaHandle && memHandle != 0) {
      82            5 :         HrtRaUbLocalMemUnreg(rdmaHandle, memHandle);
      83              :     }
      84              : 
      85           52 :     if (rdmaHandle) {
      86           52 :         RdmaHandleManager::GetInstance().PutTokenIdInfo(rdmaHandle, bufKey_, tokenIdHandle_);
      87              :     }
      88              : 
      89           52 :     HrtDevResInfo devResInfo;
      90           52 :     devResInfo.dieId    = 0;
      91           52 :     devResInfo.procType = HrtDevResProcType::PROCESS_HCCP;
      92           52 :     devResInfo.resType  = HrtDevResType::RES_TYPE_STARS_NOTIFY_RECORD;
      93           52 :     devResInfo.resId    = GetNotify()->GetId();
      94           52 :     devResInfo.flag     = 0;
      95           52 :     HrtReleaseDevResAddress(devResInfo);
      96           52 : }
      97              : 
      98           56 : UbLocalNotify::~UbLocalNotify()
      99              : {
     100           52 :     DECTOR_TRY_CATCH("UbLocalNotify", ReleaseResource());
     101           56 : }
     102              : } // namespace Hccl