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_SCATTER_MESH_2D_H_
12 : #define HCCLV2_CCU_INSTRUCTION_SCATTER_MESH_2D_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 : // 为ScatterMesh2D实现的CCUIns、CCUCtxArg与CCUTaskArg
25 : class CcuCtxArgScatterMesh2D : public CcuCtxArg {
26 : public:
27 0 : explicit CcuCtxArgScatterMesh2D(const std::vector<uint64_t> &dSize, uint32_t rankSize, uint32_t rId, uint32_t axisId, uint32_t root, const CollAlgOperator &op,
28 0 : const std::vector<std::vector<RankId>> &tempVTopo) :
29 0 : dimSize_(dSize), rankSize_(rankSize), rankId_(rId), axisId_(axisId), root_(root), op_(op), tempVTopo_(tempVTopo) {}
30 :
31 0 : CcuCtxSignature GetCtxSignature() const override
32 : {
33 0 : CcuCtxSignature signature;
34 0 : GenerateCcuCtxSignature(signature, CcuInstType::CCU_SCATTER_MESH_2D_DIRECT, op_, tempVTopo_);
35 0 : return signature;
36 0 : }
37 :
38 : std::vector<uint64_t> dimSize_;
39 : uint32_t rankSize_;
40 : uint32_t rankId_;
41 : uint32_t axisId_;
42 : uint32_t root_;
43 : CollAlgOperator op_;
44 : std::vector<std::vector<RankId>> tempVTopo_;
45 : };
46 :
47 : class CcuTaskArgScatterMesh2D : public CcuTaskArg {
48 : public:
49 0 : explicit CcuTaskArgScatterMesh2D(uint64_t inputAddr, uint64_t outputAddr, uint64_t scratchAddr, uint64_t token,
50 : uint64_t sliceSize, uint64_t stride, uint64_t xSliceSize, uint64_t ySliceSize)
51 0 : : inputAddr_(inputAddr), outputAddr_(outputAddr), scratchAddr_(scratchAddr), token_(token),
52 0 : sliceSize_(sliceSize), stride_(stride), xSliceSize_(xSliceSize), ySliceSize_(ySliceSize)
53 : {
54 0 : }
55 :
56 : uint64_t inputAddr_;
57 : uint64_t outputAddr_;
58 : uint64_t scratchAddr_;
59 : uint64_t token_;
60 : uint64_t sliceSize_;
61 : uint64_t stride_;
62 : uint64_t xSliceSize_;
63 : uint64_t ySliceSize_;
64 : };
65 :
66 : class CcuInstructionScatterMesh2D : public CcuInstruction {
67 : public:
68 0 : CcuInstructionScatterMesh2D() : CcuInstruction()
69 : {
70 0 : }
71 :
72 0 : void Init(uint32_t rankId, uint32_t rankSize, uint32_t axisId, uint64_t inputAddr, uint64_t outputAddr, uint64_t scratchAddr, uint64_t token, uint64_t sliceSize,
73 : uint64_t stride, uint64_t xSliceSize, uint64_t ySliceSize, CollAlgOperator &op, std::vector<std::vector<RankId>> &tempVTopo)
74 : {
75 0 : dimSize_.push_back(tempVTopo[0].size());
76 0 : dimSize_.push_back(tempVTopo[1].size());
77 0 : rankId_ = rankId;
78 0 : rankSize_ = rankSize;
79 0 : axisId_ = axisId;
80 0 : inputAddr_ = inputAddr;
81 0 : outputAddr_ = outputAddr;
82 0 : scratchAddr_ = scratchAddr;
83 0 : sliceSize_ = sliceSize;
84 0 : token_ = token;
85 0 : stride_ = stride;
86 0 : xSliceSize_ = xSliceSize;
87 0 : ySliceSize_ = ySliceSize;
88 0 : op_ = op;
89 0 : tempVTopo_ = tempVTopo;
90 :
91 0 : if ( dimSize_[0]*dimSize_[1] != rankSize_ ) {
92 0 : THROW<InvalidParamsException>(
93 0 : StringFormat("[CcuInstructionScatterMesh2D] each DimSize and rankSize is NOT match. dimSize[0][%llu], dimSize[1][%llu], rankSize[%llu]", dimSize_[0], dimSize_[1], rankSize_) );
94 : }
95 :
96 0 : return;
97 : }
98 :
99 0 : CcuInstType GetInstType() const override
100 : {
101 0 : HCCL_INFO("CcuInstructionScatterMesh2D instype is CCU_SCATTER_MESH_2D_DIRECT.");
102 0 : return instType_;
103 : }
104 :
105 : void SetInstType(CcuInstType instType)
106 : {
107 : instType_ = instType;
108 : }
109 :
110 0 : std::string Describe() const override
111 : {
112 0 : return StringFormat("CcuInstructionScatterMesh2D rankId [%u], instType[%s]", rankId_, instType_.Describe().c_str());
113 : }
114 :
115 0 : std::unique_ptr<CcuCtxArg> GetCtxArg() const override
116 : {
117 0 : return std::make_unique<CcuCtxArgScatterMesh2D>(dimSize_, rankSize_, rankId_, axisId_, op_.root, op_, tempVTopo_);
118 : }
119 :
120 0 : std::unique_ptr<CcuTaskArg> GetTaskArg() const override
121 : {
122 0 : return std::make_unique<CcuTaskArgScatterMesh2D>(inputAddr_, outputAddr_, scratchAddr_,token_, sliceSize_, stride_, xSliceSize_, ySliceSize_);
123 : }
124 :
125 : private:
126 : CcuInstType instType_ = CcuInstType::CCU_SCATTER_MESH_2D_DIRECT;
127 : std::vector<uint64_t> dimSize_;
128 : uint32_t rankId_{0};
129 : uint32_t rankSize_{0};
130 : uint32_t axisId_{0};
131 : uint64_t inputAddr_{0};
132 : uint64_t outputAddr_{0};
133 : uint64_t scratchAddr_{0};
134 : uint64_t token_{0};
135 : uint64_t stride_{0};
136 : uint64_t sliceSize_{0};
137 : uint64_t xSliceSize_{0};
138 : uint64_t ySliceSize_{0};
139 : CollAlgOperator op_;
140 : std::vector<std::vector<RankId>> tempVTopo_;
141 : };
142 :
143 : }
144 : #endif // HCCLV2_CCU_INSTRUCTION_SCATTER_MESH_2D_H_
|