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 HCCLV2_COLL_OPERATOR_H
      12              : #define HCCLV2_COLL_OPERATOR_H
      13              : #include <memory>
      14              : #include <string>
      15              : #include <vector>
      16              : #include "types.h"
      17              : #include "op_type.h"
      18              : #include "op_mode.h"
      19              : #include "data_type.h"
      20              : #include "reduce_op.h"
      21              : #include "buffer_type.h"
      22              : #include "buffer.h"
      23              : #include "const_val.h"
      24              : namespace Hccl {
      25              : using BaseCollOperator = struct BaseCollOperatorDef {
      26              :     OpMode   opMode{OpMode::INVALID};
      27              :     OpType   opType{OpType::DEBUGCASE};  //A5的类型
      28              :     u32      oldOpType{0}; // A3的类型
      29              :     ReduceOp reduceOp{ReduceOp::INVALID};
      30              :     u32 oldReduceOp{0}; // A3的类型
      31              :     DataType dataType{DataType::INVALID};
      32              :     u32     oldDataType{0}; // 仅用于A3的场景
      33              :     DataType outputDataType{DataType::INVALID}; // 低精度场景，存在指定输出数据类型
      34              :     u64      dataCount{0};
      35              :     u32      root{0};
      36              :     u32      numBlocksLimit{0};
      37              :     RankId   sendRecvRemoteRank{0};
      38              :     u64 inputAddr{0};
      39              :     u64 inputSize{0};
      40              :     u64 outputAddr{0};
      41              :     u64 outputSize{0};
      42              :     std::shared_ptr<Buffer> inputMem{nullptr};
      43              :     std::shared_ptr<Buffer> outputMem{nullptr};
      44              :     std::shared_ptr<Buffer> scratchMem{nullptr};
      45              :     union {
      46              :         struct {
      47              :             u64 dataCount;
      48              :             DataType dataType;
      49              :             u64 strideCount;
      50              :         } dataDes;
      51              :         struct {
      52              :             void* counts;
      53              :             void* displs;
      54              :             DataType dataType;
      55              :         } vDataDes;
      56              :         struct {
      57              :             DataType sendType;
      58              :             DataType recvType;
      59              :             u64 sendCount;
      60              :             u64 recvCount;
      61              :         } all2AllDataDes;
      62              :         struct {
      63              :             DataType sendType;
      64              :             DataType recvType;
      65              :             void* sendCounts;
      66              :             void* recvCounts;
      67              :             void* sdispls;
      68              :             void* rdispls;
      69              :         } all2AllVDataDes;
      70              :         struct {
      71              :             DataType sendType;
      72              :             DataType recvType;
      73              :             void* sendCountMatrix;
      74              :         } all2AllVCDataDes;
      75              :         struct {
      76              :             void* sendRecvItemsPtr;
      77              :             u32 itemNum;
      78              :         } batchSendRecvDataDes;
      79              :     };
      80              :     // 使用初始化列表
      81         1150 :     BaseCollOperatorDef() : opMode(), opType(), reduceOp(), dataType(), dataCount(0), root(0), sendRecvRemoteRank() {
      82              :     // 显式初始化 union 的默认成员
      83         1150 :         dataDes = {0, DataType::INVALID, 0}; // 假设 dataDes 是默认使用的成员
      84         1150 :     }
      85            4 :     Buffer *GetBuffer(const BufferType type)
      86              :     {
      87            4 :         if (type == BufferType::INPUT) {
      88            1 :             return inputMem.get();
      89            3 :         } else if (type == BufferType::OUTPUT) {
      90            1 :             return outputMem.get();
      91            2 :         } else if (type == BufferType::SCRATCH) {
      92            2 :             return scratchMem.get();
      93              :         } else {
      94            0 :             return nullptr;
      95              :         }
      96              :     }
      97              : };
      98              : 
      99              : using CollAlgOperator = BaseCollOperator;
     100              : 
     101              : using CollOperator = struct CollOperatorDef : public BaseCollOperator {
     102              :     std::string             opTag;
     103              :     bool                    staticAddr{false};
     104              :     bool                    staticShape{false};
     105              :     RankId                  myRank{INVALID_RANKID};
     106              :     std::vector<char>       GetUniqueId() const;
     107              :     static CollOperatorDef         GetPackedData(std::vector<char> &byteVector);
     108              : };
     109              : 
     110              : std::string MemBufferDesc(const BaseCollOperator &collOp);
     111              : std::string OpDesc(const BaseCollOperator &collOp);
     112              : std::string DescReduceScatter(const BaseCollOperator &collOp);
     113              : std::string DescAllreduce(const BaseCollOperator &collOp);
     114              : std::string DescAllgather(const BaseCollOperator &collOp);
     115              : std::string DescScatter(const BaseCollOperator &collOp);
     116              : std::string DescAlltoall(const BaseCollOperator &collOp);
     117              : std::string DescAlltoallV(const BaseCollOperator &collOp);
     118              : std::string DescAlltoallVC(const BaseCollOperator &collOp);
     119              : std::string DescSend(const BaseCollOperator &collOp);
     120              : std::string DescRecv(const BaseCollOperator &collOp);
     121              : std::string DescReduce(const BaseCollOperator &collOp);
     122              : std::string DescBroadcast(const BaseCollOperator &collOp);
     123              : std::string DescBatchSendRecv(const BaseCollOperator &collOp);
     124              : std::string DescHalfAlltoAllV(const BaseCollOperator &collOp);
     125              : std::string DescReduceScatterV(const BaseCollOperator &collOp);
     126              : std::string DescAllGatherV(const BaseCollOperator &collOp);
     127              : 
     128              : std::string CollOpToString(const BaseCollOperator &collOp);
     129              : 
     130              : } // namespace Hccl
     131              : 
     132              : #endif // !HCCLV2_COLL_OPERATOR_H