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              : #ifndef TRANSPORT_DEVICE_ROCE_MEM_H
      12              : #define TRANSPORT_DEVICE_ROCE_MEM_H
      13              : 
      14              : #include "transport_mem.h"
      15              : #include <atomic>
      16              : #include <chrono>
      17              : 
      18              : namespace hccl {
      19              : class TransportDeviceRoceMem : public TransportMem {
      20              : public:
      21              :     TransportDeviceRoceMem(
      22              :         const std::unique_ptr<NotifyPool>& notifyPool, const HcclNetDevCtx& netDevCtx, const HcclDispatcher& dispatcher,
      23              :         AttrInfo& attrInfo, bool aicpuUnfoldMode, const HcclQpInfoV2& qpInfo);
      24              :     ~TransportDeviceRoceMem() override;
      25              : 
      26              :     HcclResult
      27              :     ExchangeMemDesc(const RmaMemDescs& localMemDescs, RmaMemDescs& remoteMemDescs, u32& actualNumOfRemote) override;
      28              :     HcclResult EnableMemAccess(const RmaMemDesc& remoteMemDesc, RmaMem& remoteMem) override;
      29              :     HcclResult DisableMemAccess(const RmaMemDesc& remoteMemDesc) override;
      30              :     HcclResult SetSocket(const std::shared_ptr<HcclSocket>& socket) override;
      31              :     HcclResult Connect(s32 timeoutSec) override;
      32              :     HcclResult Write(const HcclBuf& remoteMem, const HcclBuf& localMem, const rtStream_t& stream) override;
      33              :     HcclResult Read(const HcclBuf& localMem, const HcclBuf& remoteMem, const rtStream_t& stream) override;
      34              :     HcclResult Write(const RmaOpMem& remoteMem, const RmaOpMem& localMem, const rtStream_t& stream) override;
      35              :     HcclResult Read(const RmaOpMem& localMem, const RmaOpMem& remoteMem, const rtStream_t& stream) override;
      36              :     HcclResult AddOpFence(const rtStream_t& stream) override;
      37              :     HcclResult
      38              :     GetTransInfo(HcclQpInfoV2& qpInfo, u32* lkey, u32* rkey, HcclBuf* localMem, HcclBuf* remoteMem, u32 num) override;
      39              :     HcclResult WaitOpFence(const rtStream_t& stream) override;
      40              : 
      41              :     HcclResult BatchWrite(
      42              :         const std::vector<MemDetails>& remoteMems, const std::vector<MemDetails>& localMems, Stream& stream) override;
      43              :     HcclResult BatchRead(
      44              :         const std::vector<MemDetails>& localMems, const std::vector<MemDetails>& remoteMems, Stream& stream) override;
      45              :     HcclResult AddOpFence(const MemDetails& localFenceMem, const MemDetails& remoteFenceMem, Stream& stream) override;
      46              : 
      47              : private:
      48              :     enum class RdmaOp { OP_WRITE = 0, OP_READ = 4 };
      49              : 
      50              :     template <typename T>
      51            0 :     inline T CeilDiv(T left, T right)
      52              :     {
      53            0 :         if (right == 0) {
      54            0 :             return left;
      55              :         }
      56            0 :         return (left + right - 1) / right;
      57              :     }
      58              : 
      59              :     HcclResult BatchOp(
      60              :         Stream& stream, const std::vector<MemDetails>& localMems, const std::vector<MemDetails>& remoteMems,
      61              :         bool isRead, bool fence);
      62              :     HcclResult FillMemDetails(
      63              :         std::vector<MemDetails>& localMemList, std::vector<MemDetails>& remoteMemList, MemDetails& localMem,
      64              :         MemDetails& remoteMem);
      65              :     HcclResult DoorBellSend(Stream& stream, u64 dbInfo, u32 wrDataLen, bool fence);
      66              :     HcclResult BatchPostSend(
      67              :         Stream& stream, u64& dbInfo, std::vector<MemDetails>& localMemList, std::vector<MemDetails>& remoteMemList,
      68              :         bool isRead, bool fence, u32& wqeCount, u64& wrDataLen);
      69              :     HcclResult PostSend(
      70              :         Stream& stream, u64& dbInfo, MemDetails* localMems, MemDetails* remoteMems, u32 memNum, bool isRead, bool fence,
      71              :         u32& wqeCount, u64& wrDataLen);
      72              :     HcclResult
      73              :     RdmaPostSend(u64& dbInfo, MemDetails* localMems, MemDetails* remoteMems, u32 memNum, RdmaOp opCode, bool fence);
      74              : 
      75              :     static std::atomic<u64> wrIdOffset_;
      76              : 
      77              :     const u64 MAX_RDMA_WQE_SIZE = 2ULL * 1024 * 1024 * 1024; // RDMA最大WQE限制, 2G限制是RDMA导致
      78              :     const u32 SEND_WR_LEN = 64;
      79              :     const std::chrono::microseconds timeout_;
      80              :     HcclQpInfoV2 qpInfo_{};
      81              : };
      82              : } // namespace hccl
      83              : #endif