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_CCU_CONTEXT_ALL_TO_ALL_MESH_2D_H
      12              : #define HCCLV2_CCU_CONTEXT_ALL_TO_ALL_MESH_2D_H
      13              : 
      14              : #include <vector>
      15              : #include <ios>
      16              : #include "log.h"
      17              : #include "ccu_context_alg_base.h"
      18              : #include "ccu_assist.h"
      19              : #include "ccu_datatype.h"
      20              : #include "ccu_instruction_all_to_all_mesh2d.h"
      21              : 
      22              : namespace Hccl {
      23              : 
      24              : class CcuContextAlltoAllMesh2D : public CcuContextAlgBase {
      25              : public:
      26              :     CcuContextAlltoAllMesh2D(
      27              :         const CcuCtxArg& arg, const std::vector<CcuTransport*>& transports, const CcuTransportGroup& group);
      28           11 :     ~CcuContextAlltoAllMesh2D() override {}
      29              : 
      30              :     void Algorithm() override;
      31              :     std::vector<uint64_t> GeneArgs(const CcuTaskArg& arg) override;
      32              : 
      33              : private:
      34              :     void CalculateArgs(const CcuTaskArgAlltoAllMesh2D* taskArg);
      35              :     void InitResources();
      36              :     void LoadArgs();
      37              :     void ExchangeInfoAndSync();
      38              :     void RankSync(uint32_t signalIndex);
      39              :     void AxisSync(uint32_t signalIndex);
      40              :     void FirstStep();
      41              :     void FirstStepOneSlice(uint16_t sliceId);
      42              :     void SecondStep();
      43              :     void CreateLocalCopyLoop();
      44              :     void LocalCopyByLoopGroup(CcuRep::Memory dst, CcuRep::Memory src, GroupOpSize& goPara);
      45              : 
      46              :     std::vector<uint32_t> dimSize;
      47              :     uint32_t axisId{0};
      48              : 
      49              :     std::vector<uint32_t> dimId; // 本rank所在行或列的编号
      50              :     uint32_t localId{0};         // 本chip所在行或列的编号
      51              :     uint32_t localSize{0};       // 本rank所在行或列的总rank数
      52              :     uint32_t anotherId{0};       // 本rank在另一个轴上的Id
      53              :     uint32_t anotherSize{0};
      54              : 
      55              :     GroupOpSize goSize_;
      56              : 
      57              :     // 从外部获取的参数
      58              :     CcuRep::Variable input;
      59              :     std::vector<CcuRep::Variable> bufferA; // 第一轮的目的地，需要交换
      60              :     CcuRep::Variable bufferB;              // 第二轮的起始，不需要交换
      61              :     std::vector<CcuRep::Variable> output;  // 第二轮的目的地，需要交换
      62              :     std::vector<CcuRep::Variable> token;
      63              : 
      64              :     // mem资源准备，记录实际访问的带偏移地址
      65              :     std::vector<CcuRep::Memory> inputAddrs;
      66              :     std::vector<CcuRep::Memory> bufferAddrs;
      67              :     std::vector<CcuRep::Memory> outputAddrs;
      68              : 
      69              :     // 在本地的搬运完成标记
      70              :     std::vector<CcuRep::MaskSignal> firstSignal;
      71              :     std::vector<CcuRep::MaskSignal> secondSignal;
      72              : 
      73              :     CcuRep::Variable sliceSize_;
      74              :     CcuRep::Variable baseOffset; // 多轮搬运时的每轮基础偏移
      75              :     // 地址计算：srcStride=sendRecvSize+sendStride，dstStride=sendRecvSize+recvStride
      76              :     CcuRep::Variable firstTransportSize; // a/b块大小
      77              :     CcuRep::Variable firstChunkOffset;   // 0/a块大小的偏移
      78              :     // 第一轮中，从inputMem，每次循环向每个对端发自己的一片，多次循环中的偏移+步进：
      79              :     // (die0--baseOffset+srcStride+D0*srcStride, die1--baseOffset+D0*srcStride+srcStride)
      80              :     CcuRep::Variable firstInputStrideLocal;
      81              :     CcuRep::Variable firstInputStrideAnother;
      82              :     // 第一轮中，写到对端buffer，die0写到对端bufferY，die1写到对端bufferX；共localSize-1个对端，每个对端写anotherSize-1片
      83              :     // anotherSize-1次循环的偏移+步进：(localId*sliceSize+localSize*sliceSize)
      84              :     CcuRep::Variable firstBufferOffset;
      85              :     CcuRep::Variable firstBufferStride;
      86              :     // 第一轮中，写到对端output，自身rankId对应到输出偏移：(baseOffset+rankId*dstStride)
      87              :     CcuRep::Variable firstOutputOffset;
      88              : 
      89              :     CcuRep::Variable secondTransportSize; // b/a块大小
      90              :     CcuRep::Variable secondChunkOffset;   // a/0块大小的偏移
      91              :     // 第二轮中，从inputMem，在某一次循环中，给每个对端发送自己的一片，共localSize片；给多个对端的偏移+步进：
      92              :     // (die0--baseOffset+yId*D0*srcStride+srcStride, die1--baseOffset+xId*srcStride+D0*srcStride)
      93              :     CcuRep::Variable secondInputOffset;
      94              :     CcuRep::Variable secondInputStride;
      95              :     // 第二轮中，从buffer读取，die0读bufferX，die1读bufferY，共localSize-1个对端，为每个对端读anotherSize-1片；偏移+步进：
      96              :     // 沿本方向每个rank（dst）步进，die0--D1*sliceSize，die1--D0*sliceSize；
      97              :     // 沿另一方向每个rank（src）步进，die0/die1--sliceSize；每次循环中向每个对端发一片
      98              :     CcuRep::Variable secondBufferStrideLocal;
      99              :     CcuRep::Variable secondBufferStrideAnother;
     100              :     // 第二轮中，写到对端output，地址与分片来源相对应，偏移+步进：
     101              :     // (die0--baseOffset+xId*dstStride+D0*dstStride, die1--baseOffset+yId*D0*dstStride+dstStride)
     102              :     CcuRep::Variable secondOutputOffset;
     103              :     CcuRep::Variable secondOutputStride;
     104              : 
     105              :     // 跨轴同步信号
     106              :     std::string localAxisSignalName;
     107              :     std::string anotherAxisSignalName;
     108              :     CcuRep::MaskSignal localAxisSignal;
     109              :     CcuRep::MaskSignal anotherAxisSignal;
     110              : 
     111              :     // 在geneArgs中使用
     112              :     uint64_t firstTransportSizeValue{0};
     113              :     uint64_t firstChunkOffsetValue{0};
     114              :     uint64_t firstInputStrideLocalValue{0};
     115              :     uint64_t firstInputStrideAnotherValue{0};
     116              :     uint64_t firstBufferOffsetValue{0};
     117              :     uint64_t firstBufferStrideValue{0};
     118              :     uint64_t firstOutputOffsetValue{0};
     119              :     uint64_t secondTransportSizeValue{0};
     120              :     uint64_t secondChunkOffsetValue{0};
     121              :     uint64_t secondInputOffsetValue{0};
     122              :     uint64_t secondInputStrideValue{0};
     123              :     uint64_t secondBufferStrideLocalValue{0};
     124              :     uint64_t secondBufferStrideAnotherValue{0};
     125              :     uint64_t secondOutputOffsetValue{0};
     126              :     uint64_t secondOutputStrideValue{0};
     127              : };
     128              : } // namespace Hccl
     129              : 
     130              : #endif // HCCLV2_CCU_CONTEXT_ALL_TO_ALL_MESH_2D_H