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              : #ifndef MEM_TRANSPORT_LITE_H
      11              : #define MEM_TRANSPORT_LITE_H
      12              : #include <vector>
      13              : #include <memory>
      14              : #include <unordered_map>
      15              : #include <functional>
      16              : #include "task_param.h"
      17              : #include "base_transport_lite_impl.h"
      18              : #include "stream_lite.h"
      19              : #include "rma_buffer_lite.h"
      20              : #include "buffer.h"
      21              : #include "kernel_param_lite.h"
      22              : 
      23              : namespace Hccl {
      24              : 
      25              : class MemTransportLite {
      26              : public:
      27              :     explicit MemTransportLite(
      28              :         std::vector<char>& uniqueId,
      29              :         std::function<void(u32 streamId, u32 taskId, const TaskParam& taskParam)> callback);
      30              : 
      31              :     std::string Describe() const;
      32              : 
      33              :     using TransferOp = struct TransferOp;
      34              : 
      35           11 :     Buffer GetRmtBuffer(u32 index) { return impl->GetRmtBuffer(index); }
      36              : 
      37            6 :     void Post(u32 index, const StreamLite& stream) { impl->Post(index, stream); }
      38              : 
      39            4 :     void Wait(u32 index, const StreamLite& stream) { impl->Wait(index, stream); }
      40              : 
      41            2 :     void Read(const RmaBufferLite& loc, const Buffer& rmt, const StreamLite& stream) { impl->Read(loc, rmt, stream); }
      42              : 
      43            3 :     void Write(const RmaBufferLite& loc, const Buffer& rmt, const StreamLite& stream) { impl->Write(loc, rmt, stream); }
      44              : 
      45            2 :     void ReadReduce(const RmaBufferLite& loc, const Buffer& rmt, const ReduceIn& reduceIn, const StreamLite& stream)
      46              :     {
      47            2 :         impl->ReadReduce(loc, rmt, reduceIn, stream);
      48            2 :     }
      49              : 
      50            2 :     void WriteReduce(const RmaBufferLite& loc, const Buffer& rmt, const ReduceIn& reduceIn, const StreamLite& stream)
      51              :     {
      52            2 :         impl->WriteReduce(loc, rmt, reduceIn, stream);
      53            2 :     }
      54              : 
      55            3 :     void WriteWithNotify(
      56              :         const RmaBufferLite& loc, const Buffer& rmt, const WithNotifyIn& withNotify, const StreamLite& stream)
      57              :     {
      58            3 :         impl->WriteWithNotify(loc, rmt, withNotify, stream);
      59            3 :     }
      60              : 
      61            1 :     void WriteReduceWithNotify(
      62              :         const RmaBufferLite& loc, const Buffer& rmt, const ReduceIn& reduceIn, const WithNotifyIn& withNotify,
      63              :         const StreamLite& stream)
      64              :     {
      65            1 :         impl->WriteReduceWithNotify(loc, rmt, reduceIn, withNotify, stream);
      66            1 :     }
      67              : 
      68            1 :     void BatchOneSidedRead(
      69              :         const std::vector<RmaBufSliceLite>& loc, const std::vector<RmtRmaBufSliceLite>& rmt, const StreamLite& stream)
      70              :     {
      71            1 :         impl->BatchOneSidedRead(loc, rmt, stream);
      72            1 :     }
      73              : 
      74            1 :     void BatchOneSidedWrite(
      75              :         const std::vector<RmaBufSliceLite>& loc, const std::vector<RmtRmaBufSliceLite>& rmt, const StreamLite& stream)
      76              :     {
      77            1 :         impl->BatchOneSidedWrite(loc, rmt, stream);
      78            1 :     }
      79              : 
      80            3 :     void BatchTransfer(
      81              :         const std::vector<RmaBufferLite>& loc, const std::vector<Buffer>& rmt,
      82              :         const std::vector<BaseTransportLiteImpl::TransferOp>& transferOp, const StreamLite& stream)
      83              :     {
      84            3 :         impl->BatchTransfer(loc, rmt, transferOp, stream);
      85            3 :     }
      86              : 
      87              : private:
      88              :     TransportType type;
      89              : 
      90              :     std::unique_ptr<BaseTransportLiteImpl> impl;
      91              : };
      92              : 
      93              : } // namespace Hccl
      94              : #endif