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_INSTRUCTION_REDUCE_SCATTER_MESH_1D_MEM2MEM_H_
      12              : #define HCCLV2_CCU_INSTRUCTION_REDUCE_SCATTER_MESH_1D_MEM2MEM_H_
      13              : 
      14              : #include "template_utils.h"
      15              : #include "instruction.h"
      16              : #include "ins_queue.h"
      17              : #include "ccu_context_utils.h"
      18              : #include "ccu_ctx_signature.h"
      19              : #include "ccu_ins.h"
      20              : #include "ccu_rank_group.h"
      21              : 
      22              : namespace Hccl {
      23              : 
      24              : class CcuCtxArgReduceScatterMeshMem2Mem1D : public CcuCtxArg {
      25              : public:
      26            0 :     explicit CcuCtxArgReduceScatterMeshMem2Mem1D(const std::vector<uint64_t> &dimSize, uint32_t rankId,
      27            0 :         const CollAlgOperator &op, const std::vector<std::vector<RankId>> &tempVTopo) :
      28            0 :             dimSize_(dimSize), rankId_(rankId),  op_(op), tempVTopo_(tempVTopo) {}
      29            0 :     CcuCtxSignature GetCtxSignature() const override
      30              :     {
      31            0 :         CcuCtxSignature signature;
      32            0 :         GenerateCcuCtxSignature(signature, CcuInstType::CCU_REDUCE_SCATTER_MESH_1D_MEM2MEM, op_, tempVTopo_);
      33            0 :         return signature;
      34            0 :     }
      35              :     std::vector<uint64_t>               dimSize_;
      36              :     uint32_t                            rankId_;
      37              :     CollAlgOperator                     op_;
      38              :     std::vector<std::vector<RankId>>    tempVTopo_;
      39              : };
      40              : 
      41              : class CcuTaskArgReduceScatterMeshMem2Mem1D : public CcuTaskArg {
      42              : public:
      43            0 :     explicit CcuTaskArgReduceScatterMeshMem2Mem1D(uint64_t inputAddr, uint64_t outputAddr, uint64_t token,
      44              :                                                         uint64_t scratchAddr,
      45              :                                                         uint64_t inputSliceStride,
      46              :                                                         uint64_t inputRepeatStride, uint64_t outputRepeatStride,
      47              :                                                         uint64_t normalSliceSize, uint64_t lastSliceSize, uint64_t repeatNum)
      48            0 :         : inputAddr_(inputAddr), outputAddr_(outputAddr), token_(token), scratchAddr_(scratchAddr),
      49            0 :         inputSliceStride_(inputSliceStride), inputRepeatStride_(inputRepeatStride), outputRepeatStride_(outputRepeatStride),
      50            0 :         normalSliceSize_(normalSliceSize), lastSliceSize_(lastSliceSize), repeatNum_(repeatNum) {}
      51              : 
      52              :     uint64_t inputAddr_;
      53              :     uint64_t outputAddr_;
      54              :     uint64_t token_;
      55              :     uint64_t scratchAddr_;
      56              :     uint64_t inputSliceStride_;
      57              :     uint64_t inputRepeatStride_;
      58              :     uint64_t outputRepeatStride_;
      59              :     uint64_t normalSliceSize_;
      60              :     uint64_t lastSliceSize_;
      61              :     uint64_t repeatNum_;
      62              : };
      63              : 
      64              : class CcuInstructionReduceScatterMeshMem2Mem1D : public CcuInstruction {
      65              : public:
      66            0 :     CcuInstructionReduceScatterMeshMem2Mem1D() : CcuInstruction()
      67              :     {
      68            0 :     }
      69              : 
      70            0 :     void Init(uint32_t rankId, uint64_t repeatNum, const CollAlgOperator &op,
      71              :               const std::vector<std::vector<RankId>> &tempVTopo, uint64_t inputAddr, uint64_t outputAddr,
      72              :               uint64_t token, uint64_t scratchAddr, uint64_t inputSliceStride, uint64_t inputRepeatStride,
      73              :               uint64_t outputRepeatStride, uint64_t normalSliceSize, uint64_t lastSliceSize)
      74              :     {
      75            0 :         dimSize_.push_back(tempVTopo[0].size());
      76            0 :         rankId_ = rankId;
      77            0 :         repeatNum_              = repeatNum;
      78            0 :         op_                     = op;
      79            0 :         tempVTopo_              = tempVTopo;
      80            0 :         inputAddr_              = inputAddr;
      81            0 :         outputAddr_             = outputAddr;
      82            0 :         token_                  = token;
      83            0 :         scratchAddr_            = scratchAddr;
      84            0 :         inputSliceStride_       = inputSliceStride;
      85            0 :         inputRepeatStride_      = inputRepeatStride;
      86            0 :         outputRepeatStride_     = outputRepeatStride;
      87            0 :         normalSliceSize_        = normalSliceSize;
      88            0 :         lastSliceSize_          = lastSliceSize;
      89              : 
      90            0 :         return;
      91              :     }
      92              : 
      93            0 :     CcuInstType GetInstType() const override
      94              :     {
      95            0 :         HCCL_INFO("CcuInstructionReduceScatterMeshMem2Mem1D instype is CCU_REDUCE_SCATTER_MESH_1D_MEM2MEM.");
      96            0 :         return instType_;
      97              :     }
      98              : 
      99            0 :     std::string Describe() const override
     100              :     {
     101            0 :         return StringFormat("CcuInstructionReduceScatterMeshMem2Mem1D rankId [%u], instType[%s]", rankId_, instType_.Describe().c_str());
     102              :     }
     103              : 
     104              :     void SetInstType(CcuInstType instType) 
     105              :     { 
     106              :         instType_ = instType; 
     107              :     }
     108              : 
     109            0 :     std::unique_ptr<CcuCtxArg> GetCtxArg() const override
     110              :     {
     111            0 :         return std::make_unique<CcuCtxArgReduceScatterMeshMem2Mem1D>(dimSize_, rankId_, op_, tempVTopo_);
     112              :     }
     113              : 
     114            0 :     std::unique_ptr<CcuTaskArg> GetTaskArg() const override
     115              :     {
     116            0 :         return std::make_unique<CcuTaskArgReduceScatterMeshMem2Mem1D>(inputAddr_, outputAddr_, token_, scratchAddr_,
     117            0 :         inputSliceStride_, inputRepeatStride_, outputRepeatStride_, normalSliceSize_, lastSliceSize_, repeatNum_);
     118              :     }
     119              : 
     120              : private:
     121              :         CcuInstType                      instType_ = CcuInstType::CCU_REDUCE_SCATTER_MESH_1D_MEM2MEM;
     122              :         std::vector<uint64_t>            dimSize_;
     123              :         uint32_t                         rankId_{0};
     124              :         uint64_t                         repeatNum_{0};
     125              :         CollAlgOperator                  op_;
     126              :         std::vector<std::vector<RankId>> tempVTopo_;
     127              :         uint64_t                         inputAddr_{0};
     128              :         uint64_t                         outputAddr_{0};
     129              :         uint64_t                         scratchAddr_{0};
     130              :         uint64_t                         token_{0};
     131              :         uint64_t                         inputSliceStride_{0};
     132              :         uint64_t                         inputRepeatStride_{0};
     133              :         uint64_t                         outputRepeatStride_{0};
     134              :         uint64_t                         normalSliceSize_{0};
     135              :         uint64_t                         lastSliceSize_{0};
     136              : };
     137              : }// namespace Hccl
     138              : #endif // HCCLV2_CCU_INSTRUCTION_REDUCE_SCATTER_MESH_1D_MEM2MEM