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