Star1141
2752
代码
代码
Issues240
Pull Requests393
Actions
项目讨论
Wiki
分析
Star1141
2752
已合并
mhc_pre_sinkhorn算子性能优化增加切M分支 #5033
何宇航创建于 5月8日
mhc_pre_sinkhorn算子性能优化增加切M分支 #5033
已合并
何宇航创建于 5月8日
从weixin_44156099/ops-transformerfix_mhc_pre_sinkhorn合入到cann/ops-transformermaster
共 13 个文件变更+1081-321
Mmhc/mhc_pre_sinkhorn/examples/test_aclnn_mhc_pre_sinkhorn.cpp+1-14
@@ -275,7 +275,7 @@ int main()
275 std::vector<int64_t> alpha_shape = {3};275 std::vector<int64_t> alpha_shape = {3};
276 std::vector<int64_t> bias_shape = {hc_mix};276 std::vector<int64_t> bias_shape = {hc_mix};
277 277 
278- std::vector<int64_t> hin_shape = {bs, seq_len, n, c};278+ std::vector<int64_t> hin_shape = {bs, seq_len, c};
279 std::vector<int64_t> h_post_shape = {bs, seq_len, n};279 std::vector<int64_t> h_post_shape = {bs, seq_len, n};
280 std::vector<int64_t> h_res_shape = {bs, seq_len, n * n};280 std::vector<int64_t> h_res_shape = {bs, seq_len, n * n};
281 std::vector<int64_t> h_pre_shape = {bs, seq_len, n};281 std::vector<int64_t> h_pre_shape = {bs, seq_len, n};
@@ -319,19 +319,6 @@ int main()
319 CHECK_RET(aclrtSynchronizeStream(stream) == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed\n"); return -1);319 CHECK_RET(aclrtSynchronizeStream(stream) == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed\n"); return -1);
320 320 
321 LOG_PRINT("MhcPreSinkhorn compute success!\n");321 LOG_PRINT("MhcPreSinkhorn compute success!\n");
322- LOG_PRINT("Required outputs:\n");
323- PrintTensorDataBfloat16(hin_shape, tensors.hin_addr);
324- PrintTensorDataFloat(h_post_shape, tensors.h_post_addr);
325- PrintTensorDataFloat(h_res_shape, tensors.h_res_addr);
326- 
327- if (need_backward) {
328- LOG_PRINT("Optional outputs (needBackward=true):\n");
329- PrintTensorDataFloat(h_pre_shape, tensors.h_pre_addr);
330- PrintTensorDataFloat(hc_before_norm_shape, tensors.hc_before_norm_addr);
331- PrintTensorDataFloat(inv_rms_shape, tensors.inv_rms_addr);
332- PrintTensorDataFloat(sum_out_shape, tensors.sum_out_addr);
333- PrintTensorDataFloat(norm_out_shape, tensors.norm_out_addr);
334- }
335 322 
336 DestroyTensors(tensors);323 DestroyTensors(tensors);
337 FreeDeviceMemory(tensors);324 FreeDeviceMemory(tensors);
Mmhc/mhc_pre_sinkhorn/op_host/mhc_pre_sinkhorn_def.cpp+36-36
@@ -22,64 +22,64 @@ public:
22 {22 {
23 this->Input("x")23 this->Input("x")
24 .ParamType(REQUIRED)24 .ParamType(REQUIRED)
25- .DataType({ge::DT_BF16})25+ .DataType({ge::DT_BF16, ge::DT_FLOAT16})
26- .Format({ge::FORMAT_ND})26+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
27- .UnknownShapeFormat({ge::FORMAT_ND});27+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
28 this->Input("phi")28 this->Input("phi")
29 .ParamType(REQUIRED)29 .ParamType(REQUIRED)
30- .DataType({ge::DT_FLOAT})30+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
31- .Format({ge::FORMAT_ND})31+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
32- .UnknownShapeFormat({ge::FORMAT_ND});32+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
33 this->Input("alpha")33 this->Input("alpha")
34 .ParamType(REQUIRED)34 .ParamType(REQUIRED)
35- .DataType({ge::DT_FLOAT})35+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
36- .Format({ge::FORMAT_ND})36+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
37- .UnknownShapeFormat({ge::FORMAT_ND});37+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
38 this->Input("bias")38 this->Input("bias")
39 .ParamType(REQUIRED)39 .ParamType(REQUIRED)
40- .DataType({ge::DT_FLOAT})40+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
41- .Format({ge::FORMAT_ND})41+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
42- .UnknownShapeFormat({ge::FORMAT_ND});42+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
43 this->Output("hin")43 this->Output("hin")
44 .ParamType(REQUIRED)44 .ParamType(REQUIRED)
45- .DataType({ge::DT_BF16})45+ .DataType({ge::DT_BF16, ge::DT_FLOAT})
46- .Format({ge::FORMAT_ND})46+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
47- .UnknownShapeFormat({ge::FORMAT_ND});47+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
48 this->Output("hPost")48 this->Output("hPost")
49 .ParamType(REQUIRED)49 .ParamType(REQUIRED)
50- .DataType({ge::DT_FLOAT})50+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
51- .Format({ge::FORMAT_ND})51+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
52- .UnknownShapeFormat({ge::FORMAT_ND});52+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
53 this->Output("hRes")53 this->Output("hRes")
54 .ParamType(REQUIRED)54 .ParamType(REQUIRED)
55- .DataType({ge::DT_FLOAT})55+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
56- .Format({ge::FORMAT_ND})56+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
57- .UnknownShapeFormat({ge::FORMAT_ND});57+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
58 this->Output("hPre")58 this->Output("hPre")
59 .ParamType(REQUIRED)59 .ParamType(REQUIRED)
60- .DataType({ge::DT_FLOAT})60+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
61- .Format({ge::FORMAT_ND})61+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
62- .UnknownShapeFormat({ge::FORMAT_ND});62+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
63 this->Output("hcBeforeNorm")63 this->Output("hcBeforeNorm")
64 .ParamType(REQUIRED)64 .ParamType(REQUIRED)
65- .DataType({ge::DT_FLOAT})65+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
66- .Format({ge::FORMAT_ND})66+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
67- .UnknownShapeFormat({ge::FORMAT_ND});67+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
68 this->Output("invRms")68 this->Output("invRms")
69 .ParamType(REQUIRED)69 .ParamType(REQUIRED)
70- .DataType({ge::DT_FLOAT})70+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
71- .Format({ge::FORMAT_ND})71+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
72- .UnknownShapeFormat({ge::FORMAT_ND});72+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
73 this->Output("sumOut")73 this->Output("sumOut")
74 .ParamType(REQUIRED)74 .ParamType(REQUIRED)
75- .DataType({ge::DT_FLOAT})75+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
76- .Format({ge::FORMAT_ND})76+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
77- .UnknownShapeFormat({ge::FORMAT_ND});77+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
78 this->Output("normOut")78 this->Output("normOut")
79 .ParamType(REQUIRED)79 .ParamType(REQUIRED)
80- .DataType({ge::DT_FLOAT})80+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
81- .Format({ge::FORMAT_ND})81+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
82- .UnknownShapeFormat({ge::FORMAT_ND});82+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
83 83 
84 this->Attr("hc_mult").AttrType(OPTIONAL).Int(4);84 this->Attr("hc_mult").AttrType(OPTIONAL).Int(4);
85 this->Attr("num_iters").AttrType(OPTIONAL).Int(20);85 this->Attr("num_iters").AttrType(OPTIONAL).Int(20);
Mmhc/mhc_pre_sinkhorn/op_host/op_api/aclnn_mhc_pre_sinkhorn.cpp+1-1
@@ -245,7 +245,7 @@ static inline aclnnStatus CheckParams(const aclTensor *x, const aclTensor *phi,
245 245 
246aclnnStatus aclnnMhcPreSinkhornGetWorkspaceSize(246aclnnStatus aclnnMhcPreSinkhornGetWorkspaceSize(
247 const aclTensor *x, const aclTensor *phi, const aclTensor *alpha, const aclTensor *bias,247 const aclTensor *x, const aclTensor *phi, const aclTensor *alpha, const aclTensor *bias,
248- int64_t hcMult, int64_t numIters, float hcEps, float normEps, bool needBackward,248+ int64_t hcMult, int64_t numIters, double hcEps, double normEps, bool needBackward,
249 aclTensor *hin, aclTensor *hPost, aclTensor *hRes,249 aclTensor *hin, aclTensor *hPost, aclTensor *hRes,
250 aclTensor *hPre, aclTensor *hcBeforeNorm, aclTensor *invRms,250 aclTensor *hPre, aclTensor *hcBeforeNorm, aclTensor *invRms,
251 aclTensor *sumOut, aclTensor *normOut,251 aclTensor *sumOut, aclTensor *normOut,
Mmhc/mhc_pre_sinkhorn/op_host/op_api/aclnn_mhc_pre_sinkhorn.h+1-1
@@ -35,7 +35,7 @@ extern "C" {
35 */35 */
36aclnnStatus aclnnMhcPreSinkhornGetWorkspaceSize(36aclnnStatus aclnnMhcPreSinkhornGetWorkspaceSize(
37 const aclTensor *x, const aclTensor *phi, const aclTensor *alpha, const aclTensor *bias,37 const aclTensor *x, const aclTensor *phi, const aclTensor *alpha, const aclTensor *bias,
38- int64_t hcMult, int64_t numIters, float hcEps, float normEps, bool needBackward,38+ int64_t hcMult, int64_t numIters, double hcEps, double normEps, bool needBackward,
39 aclTensor *hin, aclTensor *hPost, aclTensor *hRes,39 aclTensor *hin, aclTensor *hPost, aclTensor *hRes,
40 aclTensor *hPre, aclTensor *hcBeforeNorm, aclTensor *invRms,40 aclTensor *hPre, aclTensor *hcBeforeNorm, aclTensor *invRms,
41 aclTensor *sumOut, aclTensor *normOut,41 aclTensor *sumOut, aclTensor *normOut,
Mmhc/mhc_pre_sinkhorn/op_host/op_api/mhc_pre_sinkhorn.cpp+5-5
@@ -29,8 +29,8 @@ namespace l0op {
29OP_TYPE_REGISTER(MhcPreSinkhorn);29OP_TYPE_REGISTER(MhcPreSinkhorn);
30 30 
31static const aclTensor *MhcPreSinkhornAiCore(const aclTensor *x, const aclTensor *phi, const aclTensor *alpha,31static const aclTensor *MhcPreSinkhornAiCore(const aclTensor *x, const aclTensor *phi, const aclTensor *alpha,
32- const aclTensor *bias, int64_t hcMult, int64_t numIters, float hcEps,32+ const aclTensor *bias, int64_t hcMult, int64_t numIters, double hcEps,
33- float normEps, bool needBackward,33+ double normEps, bool needBackward,
34 aclTensor *hin, aclTensor *hPost, aclTensor *hRes,34 aclTensor *hin, aclTensor *hPost, aclTensor *hRes,
35 aclTensor *hPre, aclTensor *hcBeforeNorm, aclTensor *invRms,35 aclTensor *hPre, aclTensor *hcBeforeNorm, aclTensor *invRms,
36 aclTensor *sumOut, aclTensor *normOut,36 aclTensor *sumOut, aclTensor *normOut,
@@ -43,7 +43,7 @@ static const aclTensor *MhcPreSinkhornAiCore(const aclTensor *x, const aclTensor
43 MhcPreSinkhorn,43 MhcPreSinkhorn,
44 OP_INPUT(x, phi, alpha, bias),44 OP_INPUT(x, phi, alpha, bias),
45 OP_OUTPUT(hin, hPost, hRes, hPre, hcBeforeNorm, invRms, sumOut, normOut),45 OP_OUTPUT(hin, hPost, hRes, hPre, hcBeforeNorm, invRms, sumOut, normOut),
46- OP_ATTR(hcMult, numIters, hcEps, normEps, needBackward));46+ OP_ATTR(hcMult, numIters, static_cast<float>(hcEps), static_cast<float>(normEps), needBackward));
47 47
48 OP_CHECK(ret == ACLNN_SUCCESS,48 OP_CHECK(ret == ACLNN_SUCCESS,
49 OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "MhcPreSinkhorn ADD_TO_LAUNCHER_LIST_AICORE failed."),49 OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "MhcPreSinkhorn ADD_TO_LAUNCHER_LIST_AICORE failed."),
@@ -52,8 +52,8 @@ static const aclTensor *MhcPreSinkhornAiCore(const aclTensor *x, const aclTensor
52}52}
53 53 
54const aclTensor *MhcPreSinkhorn(const aclTensor *x, const aclTensor *phi, const aclTensor *alpha,54const aclTensor *MhcPreSinkhorn(const aclTensor *x, const aclTensor *phi, const aclTensor *alpha,
55- const aclTensor *bias, int64_t hcMult, int64_t numIters, float hcEps,55+ const aclTensor *bias, int64_t hcMult, int64_t numIters, double hcEps,
56- float normEps, bool needBackward,56+ double normEps, bool needBackward,
57 aclTensor *hin, aclTensor *hPost, aclTensor *hRes,57 aclTensor *hin, aclTensor *hPost, aclTensor *hRes,
58 aclTensor *hPre, aclTensor *hcBeforeNorm, aclTensor *invRms,58 aclTensor *hPre, aclTensor *hcBeforeNorm, aclTensor *invRms,
59 aclTensor *sumOut, aclTensor *normOut,59 aclTensor *sumOut, aclTensor *normOut,
Mmhc/mhc_pre_sinkhorn/op_host/op_api/mhc_pre_sinkhorn.h+2-2
@@ -46,8 +46,8 @@ namespace l0op {
46 * @return aclTensor*: Output tensor hin46 * @return aclTensor*: Output tensor hin
47 */47 */
48const aclTensor *MhcPreSinkhorn(const aclTensor *x, const aclTensor *phi, const aclTensor *alpha,48const aclTensor *MhcPreSinkhorn(const aclTensor *x, const aclTensor *phi, const aclTensor *alpha,
49- const aclTensor *bias, int64_t hcMult, int64_t numIters, float hcEps,49+ const aclTensor *bias, int64_t hcMult, int64_t numIters, double hcEps,
50- float normEps, bool needBackward,50+ double normEps, bool needBackward,
51 aclTensor *hin, aclTensor *hPost, aclTensor *hRes,51 aclTensor *hin, aclTensor *hPost, aclTensor *hRes,
52 aclTensor *hPre, aclTensor *hcBeforeNorm, aclTensor *invRms,52 aclTensor *hPre, aclTensor *hcBeforeNorm, aclTensor *invRms,
53 aclTensor *sumOut, aclTensor *normOut,53 aclTensor *sumOut, aclTensor *normOut,
Mmhc/mhc_pre_sinkhorn/op_host/op_tiling/mhc_pre_sinkhorn_tiling.cpp+126-139
@@ -54,6 +54,8 @@ constexpr int64_t NORM_EPS_ATTR_IDX = 3;
54constexpr int64_t NEED_BACKWARD_ATTR_IDX = 4;54constexpr int64_t NEED_BACKWARD_ATTR_IDX = 4;
55constexpr int64_t DEFAULT_ITER_TIMES = 20;55constexpr int64_t DEFAULT_ITER_TIMES = 20;
56constexpr int64_t BS_SPLIT_THRESHOLD = 128;56constexpr int64_t BS_SPLIT_THRESHOLD = 128;
57+constexpr int64_t MAX_BS_PER_LOOP = 32;
58+constexpr int64_t MAX_REDUCE_SIZE = 256;
57}59}
58 60 
59ge::graphStatus MhcPreSinkhornTiling::GetPlatformInfo()61ge::graphStatus MhcPreSinkhornTiling::GetPlatformInfo()
@@ -69,6 +71,9 @@ ge::graphStatus MhcPreSinkhornTiling::GetPlatformInfo()
69 auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);71 auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
70 aivCoreNum_ = ascendcPlatform.GetCoreNumAiv();72 aivCoreNum_ = ascendcPlatform.GetCoreNumAiv();
71 aicCoreNum_ = ascendcPlatform.GetCoreNumAic();73 aicCoreNum_ = ascendcPlatform.GetCoreNumAic();
74+ aivCoreNum_ = 40;
75+ aicCoreNum_ = 20;
76+ 
72 uint64_t ubSizePlatForm;77 uint64_t ubSizePlatForm;
73 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm);78 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm);
74 ubSize_ = ubSizePlatForm;79 ubSize_ = ubSizePlatForm;
@@ -144,21 +149,17 @@ ge::graphStatus MhcPreSinkhornTiling::GetShapeAttrsInfoInner()
144 return ge::GRAPH_SUCCESS;149 return ge::GRAPH_SUCCESS;
145}150}
146 151 
147- 
148ge::graphStatus MhcPreSinkhornTiling::CalcMKSplitCoreMembasePart2Tiling()152ge::graphStatus MhcPreSinkhornTiling::CalcMKSplitCoreMembasePart2Tiling()
149{153{
150- rowOfFormerBlock_ = CeilDiv(curBsSplit_, static_cast<int64_t>(aivCoreNum_));154+ rowOfFormerBlock_ = CeilDiv(bs_, static_cast<int64_t>(aivCoreNum_));
151- usedAivCoreNums_ = std::min(CeilDiv(curBsSplit_, rowOfFormerBlock_), static_cast<int64_t>(aivCoreNum_));155+ usedAivCoreNums_ = std::min(CeilDiv(bs_, rowOfFormerBlock_), static_cast<int64_t>(aivCoreNum_));
152- rowOfTailBlock_ = curBsSplit_ - (usedAivCoreNums_ - 1) * rowOfFormerBlock_;156+ rowOfTailBlock_ = bs_ - (usedAivCoreNums_ - 1) * rowOfFormerBlock_;
153 157 
154 int64_t minRowPerCore = 1;158 int64_t minRowPerCore = 1;
155 int64_t rowOnceLoop = std::min(rowOfFormerBlock_, minRowPerCore);159 int64_t rowOnceLoop = std::min(rowOfFormerBlock_, minRowPerCore);
156 int64_t kBlockNum = tilingData_.get_cubeBlockDimK();160 int64_t kBlockNum = tilingData_.get_cubeBlockDimK();
157 161 
158 hcMultAlign_ = RoundUp(hcMult_, BLOCK_SIZE / sizeof(float));162 hcMultAlign_ = RoundUp(hcMult_, BLOCK_SIZE / sizeof(float));
159- int64_t mix0OriginSize = kBlockNum * rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
160- int64_t mix1OriginSize = kBlockNum * rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
161- int64_t mix2OriginSize = kBlockNum * rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
162 int64_t mix0Size = rowOnceLoop * hcMultAlign_ * sizeof(float);163 int64_t mix0Size = rowOnceLoop * hcMultAlign_ * sizeof(float);
163 int64_t mix1Size = rowOnceLoop * hcMultAlign_ * sizeof(float);164 int64_t mix1Size = rowOnceLoop * hcMultAlign_ * sizeof(float);
164 int64_t mix2Size = rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float);165 int64_t mix2Size = rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float);
@@ -168,6 +169,7 @@ ge::graphStatus MhcPreSinkhornTiling::CalcMKSplitCoreMembasePart2Tiling()
168 int64_t ySize = rowOnceLoop * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER;169 int64_t ySize = rowOnceLoop * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER;
169 int64_t postSize = rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;170 int64_t postSize = rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
170 int64_t combFragSize = rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;171 int64_t combFragSize = rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
172+ int64_t combFragBufSize = rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER * 2;
171 int64_t base0Size = hcMultAlign_ * sizeof(float);173 int64_t base0Size = hcMultAlign_ * sizeof(float);
172 int64_t base1Size = hcMultAlign_ * sizeof(float);174 int64_t base1Size = hcMultAlign_ * sizeof(float);
173 int64_t base2Size = hcMult_ * hcMultAlign_ * sizeof(float);175 int64_t base2Size = hcMult_ * hcMultAlign_ * sizeof(float);
@@ -178,8 +180,8 @@ ge::graphStatus MhcPreSinkhornTiling::CalcMKSplitCoreMembasePart2Tiling()
178 int64_t reduceBufSize = rowOnceLoop * hcMultAlign_ * sizeof(float);180 int64_t reduceBufSize = rowOnceLoop * hcMultAlign_ * sizeof(float);
179 int64_t maskPatternSize = BLOCK_SIZE * 16;181 int64_t maskPatternSize = BLOCK_SIZE * 16;
180 182 
181- int64_t totalSize = mix0OriginSize + mix1OriginSize + mix2OriginSize + mix0Size + mix1Size + mix2Size + squareSumSize + rsqrtSize +183+ int64_t totalSize = mix0Size + mix1Size + mix2Size + squareSumSize + rsqrtSize + xSize + ySize + postSize +
182- xSize + ySize + postSize + combFragSize + base0Size + base1Size + base2Size + xCastSize + yCastSize +184+ combFragSize + combFragBufSize + base0Size + base1Size + base2Size + xCastSize + yCastSize +
183 rowBrcb0Size + hcBrcb1Size + reduceBufSize + maskPatternSize;185 rowBrcb0Size + hcBrcb1Size + reduceBufSize + maskPatternSize;
184 rowFactor_ = rowOnceLoop;186 rowFactor_ = rowOnceLoop;
185 if (totalSize <= ubSize_) {187 if (totalSize <= ubSize_) {
@@ -187,8 +189,9 @@ ge::graphStatus MhcPreSinkhornTiling::CalcMKSplitCoreMembasePart2Tiling()
187 dFactor_ = d_;189 dFactor_ = d_;
188 tailDFactor_ = dFactor_;190 tailDFactor_ = dFactor_;
189 } else {191 } else {
190- int64_t usedUbSize = mix0OriginSize + mix1OriginSize + mix2OriginSize + mix0Size + mix1Size + mix2Size + squareSumSize + rsqrtSize192+ int64_t usedUbSize = mix0Size + mix1Size + mix2Size + squareSumSize + rsqrtSize + postSize + combFragSize +
191- + postSize + combFragSize + base0Size + base1Size + base2Size + rowBrcb0Size + hcBrcb1Size + reduceBufSize + maskPatternSize;193+ base0Size + base1Size + base2Size + rowBrcb0Size + hcBrcb1Size + reduceBufSize +
194+ maskPatternSize;
192 int64_t ubRemain = ubSize_ - usedUbSize;195 int64_t ubRemain = ubSize_ - usedUbSize;
193 dFactor_ = d_;196 dFactor_ = d_;
194 int64_t base = 2;197 int64_t base = 2;
@@ -213,9 +216,7 @@ ge::graphStatus MhcPreSinkhornTiling::CalcMKSplitCoreMembasePart2Tiling()
213 216 
214 if (dFactor_ == d_) {217 if (dFactor_ == d_) {
215 while (rowFactor_ <= rowOfFormerBlock_) {218 while (rowFactor_ <= rowOfFormerBlock_) {
216- mix0OriginSize = kBlockNum * rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;219+ break;
217- mix1OriginSize = kBlockNum * rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
218- mix2OriginSize = kBlockNum * rowFactor_ * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
219 mix0Size = rowFactor_ * hcMultAlign_ * sizeof(float);220 mix0Size = rowFactor_ * hcMultAlign_ * sizeof(float);
220 mix1Size = rowFactor_ * hcMultAlign_ * sizeof(float);221 mix1Size = rowFactor_ * hcMultAlign_ * sizeof(float);
221 mix2Size = rowFactor_ * hcMult_ * hcMultAlign_ * sizeof(float);222 mix2Size = rowFactor_ * hcMult_ * hcMultAlign_ * sizeof(float);
@@ -231,7 +232,7 @@ ge::graphStatus MhcPreSinkhornTiling::CalcMKSplitCoreMembasePart2Tiling()
231 hcBrcb1Size = RoundUp(rowFactor_ * hcMultAlign_, 8) * BLOCK_SIZE;232 hcBrcb1Size = RoundUp(rowFactor_ * hcMultAlign_, 8) * BLOCK_SIZE;
232 reduceBufSize = rowFactor_ * hcMultAlign_ * sizeof(float);233 reduceBufSize = rowFactor_ * hcMultAlign_ * sizeof(float);
233 maskPatternSize = BLOCK_SIZE;234 maskPatternSize = BLOCK_SIZE;
234- totalSize = mix0OriginSize + mix1OriginSize + mix2OriginSize + mix0Size + mix1Size + mix2Size + squareSumSize + rsqrtSize +235+ totalSize = mix0Size + mix1Size + mix2Size + squareSumSize + rsqrtSize +
235 xSize + ySize + postSize + combFragSize + base0Size + base1Size + base2Size + xCastSize + yCastSize +236 xSize + ySize + postSize + combFragSize + base0Size + base1Size + base2Size + xCastSize + yCastSize +
236 rowBrcb0Size + hcBrcb1Size + reduceBufSize + maskPatternSize;237 rowBrcb0Size + hcBrcb1Size + reduceBufSize + maskPatternSize;
237 if (totalSize > ubSize_) {238 if (totalSize > ubSize_) {
@@ -270,119 +271,7 @@ ge::graphStatus MhcPreSinkhornTiling::CalcMKSplitCoreMembasePart2Tiling()
270 return ge::GRAPH_SUCCESS;271 return ge::GRAPH_SUCCESS;
271}272}
272 273 
273-ge::graphStatus MhcPreSinkhornTiling::CalcMembaseOpTiling()274+// 预留功能，切K情况尾核单独计算
274-{
275- rowOfFormerBlock_ = CeilDiv(bs_, static_cast<int64_t>(aivCoreNum_));
276- usedAivCoreNums_ = std::min(CeilDiv(bs_, rowOfFormerBlock_), static_cast<int64_t>(aivCoreNum_));
277- rowOfTailBlock_ = bs_ - (usedAivCoreNums_ - 1) * rowOfFormerBlock_;
278- 
279- int64_t minRowPerCore = 1;
280- int64_t rowOnceLoop = std::min(rowOfFormerBlock_, minRowPerCore);
281- 
282- hcMultAlign_ = RoundUp(hcMult_, BLOCK_SIZE / sizeof(float));
283- int64_t mix0Size = rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
284- int64_t mix1Size = rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
285- int64_t mix2Size = rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
286- int64_t rsqrtSize = RoundUp(rowOnceLoop, BLOCK_SIZE / sizeof(float)) * sizeof(float) * DOUBLE_BUFFER;
287- int64_t xSize = rowOnceLoop * hcMult_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型
288- int64_t ySize = rowOnceLoop * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER;
289- int64_t postSize = rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
290- int64_t combFragSize = rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
291- int64_t base0Size = hcMultAlign_ * sizeof(float);
292- int64_t base1Size = hcMultAlign_ * sizeof(float);
293- int64_t base2Size = hcMult_ * hcMultAlign_ * sizeof(float);
294- int64_t xCastSize = rowOnceLoop * hcMult_ * RoundUp(d_, 8) * sizeof(float);
295- int64_t yCastSize = rowOnceLoop * RoundUp(d_, 8) * sizeof(float);
296- int64_t rowBrcb0Size = RoundUp(rowOnceLoop, 8) * BLOCK_SIZE;
297- int64_t hcBrcb1Size = RoundUp(rowOnceLoop * hcMultAlign_, 8) *2 * BLOCK_SIZE;
298- int64_t reduceBufSize = rowOnceLoop * hcMultAlign_ * sizeof(float);
299- 
300- int64_t totalSize = mix0Size + mix1Size + mix2Size + rsqrtSize + xSize + ySize + postSize + combFragSize +
301- base0Size + base1Size + base2Size + xCastSize + yCastSize + rowBrcb0Size + hcBrcb1Size + reduceBufSize;
302- rowFactor_ = rowOnceLoop;
303- if (totalSize <= ubSize_) {
304- dLoop_ = 1;
305- dFactor_ = d_;
306- tailDFactor_ = dFactor_;
307- } else {
308- int64_t usedUbSize = mix0Size + mix1Size + mix2Size + rsqrtSize + postSize + combFragSize +
309- base0Size + base1Size + base2Size + rowBrcb0Size + hcBrcb1Size + reduceBufSize;
310- int64_t ubRemain = ubSize_ - usedUbSize;
311- dFactor_ = d_;
312- int64_t base = 2;
313- while (1) {
314- dFactor_ = CeilDiv(d_, base);
315- xSize = rowOnceLoop * hcMult_ * RoundUp(dFactor_, 16) * 2 * DOUBLE_BUFFER;
316- ySize = rowOnceLoop * RoundUp(dFactor_, 16) * 2 * DOUBLE_BUFFER;
317- xCastSize = rowOnceLoop * hcMult_ * RoundUp(dFactor_, 8) * sizeof(float);
318- yCastSize = rowOnceLoop * RoundUp(dFactor_, 8) * sizeof(float);
319- int64_t targetSize = xSize + ySize + xCastSize + yCastSize;
320- if (targetSize <= ubRemain) {
321- break;
322- }
323- base++;
324- }
325- if (dFactor_ > 32) {
326- dFactor_ = DownAlign(dFactor_, 32);
327- }
328- dLoop_ = CeilDiv(d_, dFactor_);
329- tailDFactor_ = d_ % dFactor_ == 0 ? dFactor_ : d_ % dFactor_;
330- }
331- 
332- if (dFactor_ == d_) {
333- while (rowFactor_ <= rowOfFormerBlock_) {
334- mix0Size = rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
335- mix1Size = rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
336- mix2Size = rowFactor_ * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
337- rsqrtSize = RoundUp(rowFactor_, BLOCK_SIZE / sizeof(float)) * sizeof(float) * DOUBLE_BUFFER;
338- xSize = rowFactor_ * hcMult_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER;
339- ySize = rowFactor_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER;
340- postSize = rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
341- combFragSize = rowFactor_ * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER;
342- xCastSize = rowFactor_ * hcMult_ * RoundUp(d_, 8) * sizeof(float);
343- yCastSize = rowFactor_ * RoundUp(d_, 8) * sizeof(float);
344- rowBrcb0Size = RoundUp(rowFactor_, 8) * BLOCK_SIZE;
345- hcBrcb1Size = RoundUp(rowFactor_ * hcMultAlign_, 8) * 2 * BLOCK_SIZE;
346- reduceBufSize = rowFactor_ * hcMultAlign_ * sizeof(float);
347- totalSize = mix0Size + mix1Size + mix2Size + rsqrtSize + xSize + ySize + postSize + combFragSize +
348- base0Size + base1Size + base2Size + xCastSize + yCastSize + rowBrcb0Size + hcBrcb1Size + reduceBufSize;;
349- if (totalSize > ubSize_) {
350- rowFactor_ = rowFactor_ - 1;
351- break;
352- }
353- rowFactor_ = rowFactor_ + 1;
354- }
355- rowFactor_ = rowFactor_ > rowOfFormerBlock_ ? rowFactor_ - 1 : rowFactor_;
356- }
357- rowLoopOfFormerBlock_ = CeilDiv(rowOfFormerBlock_, rowFactor_);
358- rowLoopOfTailBlock_ = CeilDiv(rowOfTailBlock_, rowFactor_);
359- tailRowFactorOfFormerBlock_ = rowOfFormerBlock_ % rowFactor_ == 0 ? rowFactor_ : rowOfFormerBlock_ % rowFactor_;
360- tailRowFactorOfTailBlock_ = rowOfTailBlock_ % rowFactor_ == 0 ? rowFactor_ : rowOfTailBlock_ % rowFactor_;
361- 
362- tilingData_.set_bs(bs_);
363- tilingData_.set_hcMix(hcMix_);
364- tilingData_.set_hcMult(hcMult_);
365- tilingData_.set_d(d_);
366- tilingData_.set_hcMultAlign(hcMultAlign_);
367- tilingData_.set_rowOfFormerBlock(rowOfFormerBlock_);
368- tilingData_.set_rowOfTailBlock(rowOfTailBlock_);
369- tilingData_.set_rowLoopOfFormerBlock(rowLoopOfFormerBlock_);
370- tilingData_.set_rowLoopOfTailBlock(rowLoopOfTailBlock_);
371- tilingData_.set_rowFactor(rowFactor_);
372- tilingData_.set_tailRowFactorOfFormerBlock(tailRowFactorOfFormerBlock_);
373- tilingData_.set_tailRowFactorOfTailBlock(tailRowFactorOfTailBlock_);
374- tilingData_.set_dLoop(dLoop_);
375- tilingData_.set_dFactor(dFactor_);
376- tilingData_.set_tailDFactor(tailDFactor_);
377- tilingData_.set_iterTimes(iterTimes_);
378- tilingData_.set_hcEps(hcEps_);
379- 
380- tilingData_.set_kBlockFactor(32);
381- tilingData_.set_kFactor(512);
382- return ge::GRAPH_SUCCESS;
383-}
384- 
385-// 尾核单独计算：
386ge::graphStatus MhcPreSinkhornTiling::CalcBsSplit()275ge::graphStatus MhcPreSinkhornTiling::CalcBsSplit()
387{276{
388 tailBs_ = bs_;277 tailBs_ = bs_;
@@ -420,8 +309,6 @@ ge::graphStatus MhcPreSinkhornTiling::CalcTailBsTiling()
420 tilingData_.set_tailBsCubeBlockDimM(tilingData_.get_cubeBlockDimM());309 tilingData_.set_tailBsCubeBlockDimM(tilingData_.get_cubeBlockDimM());
421 return ge::GRAPH_SUCCESS;310 return ge::GRAPH_SUCCESS;
422 }311 }
423- 
424- int64_t kSize = hcMult_ * d_;
425 312
426 uint64_t tailMDimNum = std::min(aicCoreNum_, static_cast<uint64_t>(CeilDiv(tailBs_, M_L1_MAX_SIZE)));313 uint64_t tailMDimNum = std::min(aicCoreNum_, static_cast<uint64_t>(CeilDiv(tailBs_, M_L1_MAX_SIZE)));
427 uint64_t tailSingleCoreM = RoundUp(CeilDiv(tailBs_, tailMDimNum), AscendC::BLOCK_CUBE);314 uint64_t tailSingleCoreM = RoundUp(CeilDiv(tailBs_, tailMDimNum), AscendC::BLOCK_CUBE);
@@ -468,8 +355,10 @@ ge::graphStatus MhcPreSinkhornTiling::CalcOpTiling()
468 }355 }
469 uint64_t kSize = hcMult_ * d_;356 uint64_t kSize = hcMult_ * d_;
470 tilingData_.set_k(kSize);357 tilingData_.set_k(kSize);
471- uint64_t mDimNum = std::min(aicCoreNum_, static_cast<uint64_t>(CeilDiv(curBsSplit_, M_L1_MAX_SIZE)));358+ // 切K计算stage1
472- uint64_t singleCoreM = RoundUp(CeilDiv(curBsSplit_, mDimNum), AscendC::BLOCK_CUBE);359+ uint64_t mDimNum = std::min(aicCoreNum_, static_cast<uint64_t>(CeilDiv(bs_, M_L1_MAX_SIZE)));
360+ uint64_t singleCoreM = RoundUp(CeilDiv(bs_, mDimNum), AscendC::BLOCK_CUBE);
361+ 
473 uint64_t kDimNum = aicCoreNum_ / mDimNum;362 uint64_t kDimNum = aicCoreNum_ / mDimNum;
474 uint64_t splitKSize = RoundUp(CeilDiv(kSize, kDimNum), K_MULIT_CORE_SPLIT_BASE_SIZE);363 uint64_t splitKSize = RoundUp(CeilDiv(kSize, kDimNum), K_MULIT_CORE_SPLIT_BASE_SIZE);
475 364
@@ -486,13 +375,9 @@ ge::graphStatus MhcPreSinkhornTiling::CalcOpTiling()
486 int64_t lineByteSize = (sizeof(int16_t) + sizeof(int32_t)) * DOUBLE_BUFFER * tilingData_.get_cvLoopKSize();375 int64_t lineByteSize = (sizeof(int16_t) + sizeof(int32_t)) * DOUBLE_BUFFER * tilingData_.get_cvLoopKSize();
487 int64_t stage1MFactorValue = ubSize_ / lineByteSize;376 int64_t stage1MFactorValue = ubSize_ / lineByteSize;
488 tilingData_.set_stage1MFactor(stage1MFactorValue);377 tilingData_.set_stage1MFactor(stage1MFactorValue);
489- if (kDimNum != 1) {
490- return CalcMKSplitCoreMembasePart2Tiling();
491- }
492 return CalcMKSplitCoreMembasePart2Tiling();378 return CalcMKSplitCoreMembasePart2Tiling();
493}379}
494 380 
495- 
496ge::graphStatus MhcPreSinkhornTiling::DoOpTiling()381ge::graphStatus MhcPreSinkhornTiling::DoOpTiling()
497{382{
498 if (GetPlatformInfo() == ge::GRAPH_FAILED) {383 if (GetPlatformInfo() == ge::GRAPH_FAILED) {
@@ -501,12 +386,18 @@ ge::graphStatus MhcPreSinkhornTiling::DoOpTiling()
501 if (GetShapeAttrsInfoInner() == ge::GRAPH_FAILED) {386 if (GetShapeAttrsInfoInner() == ge::GRAPH_FAILED) {
502 return ge::GRAPH_FAILED;387 return ge::GRAPH_FAILED;
503 }388 }
389+ // 预留参数计算（切K）与stage2计算
504 if (CalcOpTiling() == ge::GRAPH_FAILED) {390 if (CalcOpTiling() == ge::GRAPH_FAILED) {
505 return ge::GRAPH_FAILED;391 return ge::GRAPH_FAILED;
506 }392 }
393+ // 预留参数计算（切K）
507 if (CalcTailBsTiling() == ge::GRAPH_FAILED) {394 if (CalcTailBsTiling() == ge::GRAPH_FAILED) {
508 return ge::GRAPH_FAILED;395 return ge::GRAPH_FAILED;
509 }396 }
397+ // 目前stage1计算在用分支
398+ if (CalcStage1Tiling() == ge::GRAPH_FAILED) {
399+ return ge::GRAPH_FAILED;
400+ }
510 if (GetWorkspaceSize() == ge::GRAPH_FAILED) {401 if (GetWorkspaceSize() == ge::GRAPH_FAILED) {
511 return ge::GRAPH_FAILED;402 return ge::GRAPH_FAILED;
512 }403 }
@@ -587,20 +478,116 @@ static void PrintTilingData(const gert::TilingContext* context, MhcPreSinkhornTi
587 OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: tailBsKL1Size is %ld.", tiling.get_tailBsKL1Size());478 OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: tailBsKL1Size is %ld.", tiling.get_tailBsKL1Size());
588 OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: tailBsMultCoreSplitMSize is %ld.", tiling.get_tailBsMultCoreSplitMSize());479 OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: tailBsMultCoreSplitMSize is %ld.", tiling.get_tailBsMultCoreSplitMSize());
589 OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: tailBsCubeBlockDimM is %ld.", tiling.get_tailBsCubeBlockDimM());480 OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: tailBsCubeBlockDimM is %ld.", tiling.get_tailBsCubeBlockDimM());
481+ OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: stage1VecCoreNum is %ld.", tiling.get_stage1VecCoreNum());
482+ OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: stage1CubeCoreNum is %ld.", tiling.get_stage1CubeCoreNum());
483+ OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: stage1BsPerVecCore is %ld.", tiling.get_stage1BsPerVecCore());
484+ OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: stage1TailBsPerVecCore is %ld.", tiling.get_stage1TailBsPerVecCore());
485+ OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: stage1BsLoop is %ld.", tiling.get_stage1BsLoop());
486+ OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: stage1BsFactor is %ld.", tiling.get_stage1BsFactor());
487+ OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: stage1NcLoop is %ld.", tiling.get_stage1NcLoop());
488+ OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: stage1NcFactor is %ld.", tiling.get_stage1NcFactor());
489+ OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: stage1TailNcFactor is %ld.", tiling.get_stage1TailNcFactor());
490+ OP_LOGD(context->GetNodeName(), "MhcPreSinkhorn_tiling: stage1XCastWsSize is %ld.", tiling.get_stage1XCastWsSize());
491+}
492+ 
493+ge::graphStatus MhcPreSinkhornTiling::CalcStage1Tiling()
494+{
495+ // 1. 计算核心数量
496+ int64_t vecCoreNum = static_cast<int64_t>(aivCoreNum_);
497+ int64_t cubeCoreNum = static_cast<int64_t>(aicCoreNum_);
498+
499+ tilingData_.set_stage1VecCoreNum(vecCoreNum);
500+ tilingData_.set_stage1CubeCoreNum(cubeCoreNum);
501+
502+ // 2. 计算 BS 分核
503+ int64_t bsPerVecCore = CeilDiv(bs_, vecCoreNum);
504+ int64_t tailBsPerVecCore = bs_ - (vecCoreNum - 1) * bsPerVecCore;
505+ if (tailBsPerVecCore <= 0) {
506+ tailBsPerVecCore = bsPerVecCore;
507+ }
508+
509+ tilingData_.set_stage1BsPerVecCore(bsPerVecCore);
510+ tilingData_.set_stage1TailBsPerVecCore(tailBsPerVecCore);
511+
512+ // 3. 计算 BS 循环（每轮最多 32 个 BS）
513+ int64_t bsLoop = CeilDiv(bsPerVecCore, MAX_BS_PER_LOOP);
514+ int64_t bsFactor = MAX_BS_PER_LOOP;
515+
516+ tilingData_.set_stage1BsLoop(bsLoop);
517+ tilingData_.set_stage1BsFactor(bsFactor);
518+
519+ // 4. 计算 ReduceSum 切分（每次最多 256 个元素）
520+ int64_t ncSize = hcMult_ * d_;
521+ int64_t ncLoop = CeilDiv(ncSize, MAX_REDUCE_SIZE);
522+ int64_t ncFactor = MAX_REDUCE_SIZE;
523+ int64_t tailNcFactor = ncSize % MAX_REDUCE_SIZE;
524+ if (tailNcFactor == 0) {
525+ tailNcFactor = MAX_REDUCE_SIZE;
526+ }
527+
528+ tilingData_.set_stage1NcLoop(ncLoop);
529+ tilingData_.set_stage1NcFactor(ncFactor);
530+ tilingData_.set_stage1TailNcFactor(tailNcFactor);
531+ 
532+ int64_t mm1M = bsFactor * 2;
533+ int64_t mm1N = hcMult_ * hcMult_ + 2 * hcMult_;
534+ int64_t mm1K = hcMult_ * d_;
535+ 
536+ auto platformInfo = context_->GetPlatformInfo();
537+ OP_CHECK_IF(platformInfo == nullptr, OP_LOGE("TilingForMhcPreSinkhorn", "Tiling platformInfo is null"),
538+ return ge::GRAPH_FAILED);
539+ auto ascendPlatformInfo = platform_ascendc::PlatformAscendC(platformInfo);
540+ 
541+ auto featureDataType = matmul_tiling::DataType::DT_FLOAT;
542+ matmul_tiling::MatmulApiTiling mm1Tiling(ascendPlatformInfo);
543+ 
544+ mm1Tiling.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, featureDataType);
545+ mm1Tiling.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, featureDataType, true);
546+ mm1Tiling.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, featureDataType);
547+ mm1Tiling.SetOrgShape(mm1M, mm1N, mm1K);
548+ mm1Tiling.SetShape(mm1M, mm1N, mm1K);
549+ mm1Tiling.SetBias(false);
550+ mm1Tiling.SetBufferSpace(-1, -1, -1);
551+ if (mm1Tiling.GetTiling(tilingData_.mm1TilingData) == -1) {
552+ OP_LOGE(context_, "mm1Tiling.GetTiling failed, M=%ld, N=%ld, K=%ld", mm1M, mm1N, mm1K);
553+ return ge::GRAPH_FAILED;
554+ }
555+ 
556+ // 5. 计算 Workspace 大小
557+ // X_cast: cubeCoreNum * maxBsPerCore * n * c * sizeof(float)
558+ int64_t xCastWsSize = bsFactor * vecCoreNum * 2 * ncSize * sizeof(float);
559+ workspaceSize_ += xCastWsSize;
560+ tilingData_.set_stage1XCastWsSize(xCastWsSize);
561+
562+ // 6. 计算额外的 workspace 大小（needGrad=false 时使用）
563+ if (!needGrad_) {
564+ // invRms: bs * sizeof(float)
565+ int64_t invRmsWsSize = bs_ * sizeof(float);
566+ // hcBeforeNorm: bs * hcMix * sizeof(float)
567+ int64_t hcMix = hcMult_ * hcMult_ + 2 * hcMult_;
568+ int64_t hcBeforeNormWsSize = bs_ * hcMix * sizeof(float);
569+
570+ workspaceSize_ += invRmsWsSize + hcBeforeNormWsSize;
571+ }
572+
573+ int64_t kBlockNum = 1;
574+ tilingData_.set_cubeBlockDimK(kBlockNum);
575+ 
576+ return ge::GRAPH_SUCCESS;
590}577}
591 578 
592ge::graphStatus MhcPreSinkhornTiling::GetWorkspaceSize()579ge::graphStatus MhcPreSinkhornTiling::GetWorkspaceSize()
593{580{
594- workspaceSize_ = 16 * 1024 * 1024 + 128 * 1024 * 1024;581+ workspaceSize_ += 16 * 1024 * 1024 + 128 * 1024 * 1024;
595 return ge::GRAPH_SUCCESS;582 return ge::GRAPH_SUCCESS;
596}583}
597 584 
598ge::graphStatus MhcPreSinkhornTiling::PostTiling()585ge::graphStatus MhcPreSinkhornTiling::PostTiling()
599{586{
600- context_->SetTilingKey(0);587+ context_->SetTilingKey(0); // 0: 切M分支；预留1：切M K分支
601 context_->SetBlockDim(aicCoreNum_);588 context_->SetBlockDim(aicCoreNum_);
602 size_t* workspaces = context_->GetWorkspaceSizes(1);589 size_t* workspaces = context_->GetWorkspaceSizes(1);
603- workspaces[0] = workspaceSize_ + tilingData_.get_curBsSplit() * d_ * 4 * 16;590+ workspaces[0] = workspaceSize_;
604 591
605 PrintTilingData(context_, tilingData_);592 PrintTilingData(context_, tilingData_);
606 593
Mmhc/mhc_pre_sinkhorn/op_host/op_tiling/mhc_pre_sinkhorn_tiling.h+83-70
@@ -41,84 +41,96 @@ struct TilingOptionalParaInfo {
41};41};
42 42 
43BEGIN_TILING_DATA_DEF(MhcPreSinkhornTilingData)43BEGIN_TILING_DATA_DEF(MhcPreSinkhornTilingData)
44-TILING_DATA_FIELD_DEF(int64_t, bs);44+TILING_DATA_FIELD_DEF(int64_t, bs); // batch size
45-TILING_DATA_FIELD_DEF(int64_t, hcMix);45+TILING_DATA_FIELD_DEF(int64_t, hcMix); // hcMult * hcMult + 2 * hcMult
46-TILING_DATA_FIELD_DEF(int64_t, hcMult);46+TILING_DATA_FIELD_DEF(int64_t, hcMult); // n维度大小
47-TILING_DATA_FIELD_DEF(int64_t, d);47+TILING_DATA_FIELD_DEF(int64_t, d); // 特征维度大小
48-TILING_DATA_FIELD_DEF(int64_t, hcMultAlign);48+TILING_DATA_FIELD_DEF(int64_t, hcMultAlign); // 对齐后的hcMult
49-TILING_DATA_FIELD_DEF(int64_t, rowOfFormerBlock);49+TILING_DATA_FIELD_DEF(int64_t, rowOfFormerBlock); // 前核处理BS总数
50-TILING_DATA_FIELD_DEF(int64_t, rowOfTailBlock);50+TILING_DATA_FIELD_DEF(int64_t, rowOfTailBlock); // 尾核处理BS总数
51-TILING_DATA_FIELD_DEF(int64_t, rowLoopOfFormerBlock);51+TILING_DATA_FIELD_DEF(int64_t, rowLoopOfFormerBlock); // 前核处理行循环次数
52-TILING_DATA_FIELD_DEF(int64_t, rowLoopOfTailBlock);52+TILING_DATA_FIELD_DEF(int64_t, rowLoopOfTailBlock); // 尾核处理行循环次数
53-TILING_DATA_FIELD_DEF(int64_t, rowFactor);53+TILING_DATA_FIELD_DEF(int64_t, rowFactor); // 每次处理的行数
54-TILING_DATA_FIELD_DEF(int64_t, tailRowFactorOfFormerBlock);54+TILING_DATA_FIELD_DEF(int64_t, tailRowFactorOfFormerBlock); // 前核尾次处理的行数
55-TILING_DATA_FIELD_DEF(int64_t, tailRowFactorOfTailBlock);55+TILING_DATA_FIELD_DEF(int64_t, tailRowFactorOfTailBlock); // 尾核尾次处理的行数
56-TILING_DATA_FIELD_DEF(int64_t, dLoop);56+TILING_DATA_FIELD_DEF(int64_t, dLoop); // d维度循环次数
57-TILING_DATA_FIELD_DEF(int64_t, dFactor);57+TILING_DATA_FIELD_DEF(int64_t, dFactor); // 每次处理的d维度大小
58-TILING_DATA_FIELD_DEF(int64_t, tailDFactor);58+TILING_DATA_FIELD_DEF(int64_t, tailDFactor); // 尾次处理的d维度大小
59-TILING_DATA_FIELD_DEF(int64_t, iterTimes);59+TILING_DATA_FIELD_DEF(int64_t, iterTimes); // Sinkhorn迭代轮数
60-TILING_DATA_FIELD_DEF(float, hcEps);60+TILING_DATA_FIELD_DEF(float, hcEps); // Sinkhorn算法的epsilon参数
61-TILING_DATA_FIELD_DEF(float, normEps);61+TILING_DATA_FIELD_DEF(float, normEps); // 归一化的epsilon参数
62-TILING_DATA_FIELD_DEF(int64_t, kBlockFactor);62+TILING_DATA_FIELD_DEF(int64_t, kBlockFactor); // K维度块因子
63-TILING_DATA_FIELD_DEF(int64_t, kFactor);63+TILING_DATA_FIELD_DEF(int64_t, kFactor); // K维度每次处理大小
64-TILING_DATA_FIELD_DEF(int64_t, tailKFactor);64+TILING_DATA_FIELD_DEF(int64_t, tailKFactor); // K维度尾次处理大小
65-TILING_DATA_FIELD_DEF(int64_t, kLoop);65+TILING_DATA_FIELD_DEF(int64_t, kLoop); // K维度循环次数
66-TILING_DATA_FIELD_DEF(int64_t, stage2RowFactor);66+TILING_DATA_FIELD_DEF(int64_t, stage2RowFactor); // 阶段2每次处理的行数
67+// 预留参数，切K优化使用
68+TILING_DATA_FIELD_DEF(int64_t, k); // K维度总大小
69+TILING_DATA_FIELD_DEF(int64_t, kLoopOfFormerBlock); // 前核K维度循环次数
70+TILING_DATA_FIELD_DEF(int64_t, kLoopOfTailBlock); // 尾核K维度循环次数
71+TILING_DATA_FIELD_DEF(int64_t, stage1KFactor); // 阶段1 K维度每次处理大小
72+TILING_DATA_FIELD_DEF(int64_t, kFactorOfFormerBlock); // 前核K维度每次处理大小
73+TILING_DATA_FIELD_DEF(int64_t, kL1Size); // K维度L1缓冲大小
74+TILING_DATA_FIELD_DEF(int64_t, mL1Size); // M维度L1缓冲大小
75+TILING_DATA_FIELD_DEF(int64_t, cubeBlockDimM); // Cube核M维度切分数量
76+TILING_DATA_FIELD_DEF(int64_t, cubeBlockDimK); // Cube核K维度切分数量
77+TILING_DATA_FIELD_DEF(int64_t, kUbSize); // K维度UB缓冲大小
78+TILING_DATA_FIELD_DEF(int64_t, cvLoopKSize); // CV循环K维度大小
79+TILING_DATA_FIELD_DEF(int64_t, multCoreSplitKSize); // 多核切分K维度大小
80+TILING_DATA_FIELD_DEF(int64_t, multCoreSplitMSize); // 多核切分M维度大小
81+TILING_DATA_FIELD_DEF(int64_t, tailKSizeOfFormerBlock); // 前核K维度尾次大小
82+TILING_DATA_FIELD_DEF(int64_t, tailKSizeOfTailBlock); // 尾核K维度尾次大小
67 83 
68-TILING_DATA_FIELD_DEF(int64_t, k);84+TILING_DATA_FIELD_DEF(int64_t, mLoopOfFormerBlock); // 前核M维度循环次数
69-TILING_DATA_FIELD_DEF(int64_t, kLoopOfFormerBlock);85+TILING_DATA_FIELD_DEF(int64_t, mLoopOfTailBlock); // 尾核M维度循环次数
70-TILING_DATA_FIELD_DEF(int64_t, kLoopOfTailBlock);86+TILING_DATA_FIELD_DEF(int64_t, formerMSize); // 前核M维度大小
71-TILING_DATA_FIELD_DEF(int64_t, stage1KFactor);87+TILING_DATA_FIELD_DEF(int64_t, tailMSizeOfFormerBlock); // 前核M维度尾次大小
72-TILING_DATA_FIELD_DEF(int64_t, kFactorOfFormerBlock);88+TILING_DATA_FIELD_DEF(int64_t, tailMSizeOfTailBlock); // 尾核M维度尾次大小
73-TILING_DATA_FIELD_DEF(int64_t, kL1Size);
74-TILING_DATA_FIELD_DEF(int64_t, mL1Size);
75-TILING_DATA_FIELD_DEF(int64_t, cubeBlockDimM);
76-TILING_DATA_FIELD_DEF(int64_t, cubeBlockDimK);
77-TILING_DATA_FIELD_DEF(int64_t, kUbSize);
78-TILING_DATA_FIELD_DEF(int64_t, cvLoopKSize);
79-TILING_DATA_FIELD_DEF(int64_t, multCoreSplitKSize);
80-TILING_DATA_FIELD_DEF(int64_t, multCoreSplitMSize);
81-TILING_DATA_FIELD_DEF(int64_t, tailKSizeOfFormerBlock);
82-TILING_DATA_FIELD_DEF(int64_t, tailKSizeOfTailBlock);
83 89 
84-TILING_DATA_FIELD_DEF(int64_t, mLoopOfFormerBlock);90+TILING_DATA_FIELD_DEF(int64_t, firstUsedCoreNum); // 阶段1使用的核心数
85-TILING_DATA_FIELD_DEF(int64_t, mLoopOfTailBlock);91+TILING_DATA_FIELD_DEF(int64_t, secondUsedCoreNum); // 阶段2使用的核心数
86-TILING_DATA_FIELD_DEF(int64_t, formerMSize);
87-TILING_DATA_FIELD_DEF(int64_t, tailMSizeOfFormerBlock);
88-TILING_DATA_FIELD_DEF(int64_t, tailMSizeOfTailBlock);
89 92 
90-TILING_DATA_FIELD_DEF(int64_t, firstUsedCoreNum);93+TILING_DATA_FIELD_DEF(int64_t, rowInnerFactor); // 行内部因子
91-TILING_DATA_FIELD_DEF(int64_t, secondUsedCoreNum);
92 94 
93-TILING_DATA_FIELD_DEF(int64_t, rowInnerFactor);95+TILING_DATA_FIELD_DEF(int64_t, cubeCoreNum); // Cube核心数
96+TILING_DATA_FIELD_DEF(int64_t, stage1MFactor); // 阶段1 M维度每次处理大小
94 97 
95-TILING_DATA_FIELD_DEF(int64_t, cubeCoreNum);98+TILING_DATA_FIELD_DEF(int64_t, bufferPool0Size); // 缓冲池0大小
96-TILING_DATA_FIELD_DEF(int64_t, stage1MFactor);99+TILING_DATA_FIELD_DEF(int64_t, bufferPool1Size); // 缓冲池1大小
100+TILING_DATA_FIELD_DEF(int64_t, mUbSize); // M维度UB缓冲大小
97 101 
98-TILING_DATA_FIELD_DEF(int64_t, bufferPool0Size);102+TILING_DATA_FIELD_DEF(int64_t, needGrad); // 是否需要梯度
99-TILING_DATA_FIELD_DEF(int64_t, bufferPool1Size);103+TILING_DATA_FIELD_DEF(int64_t, bsSplitThreshold); // BS切分阈值
100-TILING_DATA_FIELD_DEF(int64_t, mUbSize);104+TILING_DATA_FIELD_DEF(int64_t, bsLoop); // BS循环次数
105+TILING_DATA_FIELD_DEF(int64_t, tailBs); // 尾批次BS数
106+TILING_DATA_FIELD_DEF(int64_t, curBsSplit); // 当前BS切分大小
101 107 
102-TILING_DATA_FIELD_DEF(int64_t, needGrad);108+// 预留参数，切K优化情况 尾批次专用参数 (tailBs != BS_SPLIT_THRESHOLD 时使用)
103-TILING_DATA_FIELD_DEF(int64_t, bsSplitThreshold);109+TILING_DATA_FIELD_DEF(int64_t, tailBsRowOfFormerBlock); // 尾批次前核处理BS总数
104-TILING_DATA_FIELD_DEF(int64_t, bsLoop);110+TILING_DATA_FIELD_DEF(int64_t, tailBsRowOfTailBlock); // 尾批次尾核处理BS总数
105-TILING_DATA_FIELD_DEF(int64_t, tailBs);111+TILING_DATA_FIELD_DEF(int64_t, tailBsRowLoopOfFormerBlock); // 尾批次前核处理行循环次数
106-TILING_DATA_FIELD_DEF(int64_t, curBsSplit);112+TILING_DATA_FIELD_DEF(int64_t, tailBsRowLoopOfTailBlock); // 尾批次尾核处理行循环次数
107- 113+TILING_DATA_FIELD_DEF(int64_t, tailBsUsedCoreNum); // 尾批次使用的核心数
108-// 尾批次专用参数 (tailBs != BS_SPLIT_THRESHOLD 时使用)114+TILING_DATA_FIELD_DEF(int64_t, tailBsRowFactor); // 尾批次每次处理的行数
109-TILING_DATA_FIELD_DEF(int64_t, tailBsRowOfFormerBlock);115+TILING_DATA_FIELD_DEF(int64_t, tailBsTailRowFactorOfFormerBlock); // 尾批次前核尾次处理的行数
110-TILING_DATA_FIELD_DEF(int64_t, tailBsRowOfTailBlock);116+TILING_DATA_FIELD_DEF(int64_t, tailBsTailRowFactorOfTailBlock); // 尾批次尾核尾次处理的行数
111-TILING_DATA_FIELD_DEF(int64_t, tailBsRowLoopOfFormerBlock);117+TILING_DATA_FIELD_DEF(int64_t, tailBsML1Size); // 尾批次M维度L1缓冲大小
112-TILING_DATA_FIELD_DEF(int64_t, tailBsRowLoopOfTailBlock);118+TILING_DATA_FIELD_DEF(int64_t, tailBsKL1Size); // 尾批次K维度L1缓冲大小
113-TILING_DATA_FIELD_DEF(int64_t, tailBsUsedCoreNum);119+TILING_DATA_FIELD_DEF(int64_t, tailBsMultCoreSplitMSize); // 尾批次多核切分M维度大小
114-TILING_DATA_FIELD_DEF(int64_t, tailBsRowFactor);120+TILING_DATA_FIELD_DEF(int64_t, tailBsCubeBlockDimM); // 尾批次Cube核M维度切分数量
115-TILING_DATA_FIELD_DEF(int64_t, tailBsTailRowFactorOfFormerBlock);
116-TILING_DATA_FIELD_DEF(int64_t, tailBsTailRowFactorOfTailBlock);
117-TILING_DATA_FIELD_DEF(int64_t, tailBsML1Size);
118-TILING_DATA_FIELD_DEF(int64_t, tailBsKL1Size);
119-TILING_DATA_FIELD_DEF(int64_t, tailBsMultCoreSplitMSize);
120-TILING_DATA_FIELD_DEF(int64_t, tailBsCubeBlockDimM);
121 121 
122+// 阶段1专用 Tiling 字段
123+TILING_DATA_FIELD_DEF(int64_t, stage1VecCoreNum); // 阶段1使用的 Vector Core 数量
124+TILING_DATA_FIELD_DEF(int64_t, stage1CubeCoreNum); // 阶段1使用的 Cube Core 数量
125+TILING_DATA_FIELD_DEF(int64_t, stage1BsPerVecCore); // 每个 Vector Core 处理的 BS 数
126+TILING_DATA_FIELD_DEF(int64_t, stage1TailBsPerVecCore); // 尾核处理的 BS 数
127+TILING_DATA_FIELD_DEF(int64_t, stage1BsLoop); // BS 循环轮数
128+TILING_DATA_FIELD_DEF(int64_t, stage1BsFactor); // 每轮 BS 数
129+TILING_DATA_FIELD_DEF(int64_t, stage1NcLoop); // n*c 切分段数
130+TILING_DATA_FIELD_DEF(int64_t, stage1NcFactor); // 每段大小
131+TILING_DATA_FIELD_DEF(int64_t, stage1TailNcFactor); // 尾段大小
132+TILING_DATA_FIELD_DEF(int64_t, stage1XCastWsSize); // X_cast workspace 大小
133+TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, mm1TilingData)
122END_TILING_DATA_DEF;134END_TILING_DATA_DEF;
123 135 
124REGISTER_TILING_DATA_CLASS(MhcPreSinkhorn, MhcPreSinkhornTilingData)136REGISTER_TILING_DATA_CLASS(MhcPreSinkhorn, MhcPreSinkhornTilingData)
@@ -144,6 +156,7 @@ public:
144 ge::graphStatus CalcMKSplitCoreMembasePart2Tiling();156 ge::graphStatus CalcMKSplitCoreMembasePart2Tiling();
145 ge::graphStatus CalcBsSplit();157 ge::graphStatus CalcBsSplit();
146 ge::graphStatus CalcTailBsTiling();158 ge::graphStatus CalcTailBsTiling();
159+ ge::graphStatus CalcStage1Tiling();
147 160 
148private:161private:
149 gert::TilingContext *context_ = nullptr;162 gert::TilingContext *context_ = nullptr;
Mmhc/mhc_pre_sinkhorn/op_kernel/mhc_pre_sinkhorn.cpp+23-8
@@ -15,9 +15,9 @@
15 15 
16#include "kernel_operator.h"16#include "kernel_operator.h"
17#include "kernel_operator_intf.h"17#include "kernel_operator_intf.h"
18- 
19#include "lib/matmul_intf.h"18#include "lib/matmul_intf.h"
20#include "mhc_pre_sinkhorn_m_k_split_core.h"19#include "mhc_pre_sinkhorn_m_k_split_core.h"
20+#include "mhc_pre_sinkhorn_m_split_core.h"
21#include "mhc_pre_sinkhorn_base.h"21#include "mhc_pre_sinkhorn_base.h"
22using namespace MhcPreSinkhorn;22using namespace MhcPreSinkhorn;
23 23 
@@ -38,21 +38,37 @@ extern "C" __global__ __aicore__ void mhc_pre_sinkhorn(GM_ADDR x, GM_ADDR phi, G
38 if (userWs == nullptr) {38 if (userWs == nullptr) {
39 return;39 return;
40 }40 }
41- // A341+ 
42+ GET_TILING_DATA_WITH_STRUCT(MhcPreSinkhornTilingData, tiling_data_in, tiling);
43+ const MhcPreSinkhornTilingData *__restrict tilingData = &tiling_data_in;
44+ 
42 if (TILING_KEY_IS(0)) {45 if (TILING_KEY_IS(0)) {
43- GET_TILING_DATA_WITH_STRUCT(MhcPreSinkhornTilingData, tiling_data_in, tiling);46+ // -----------split m--------------
44- const MhcPreSinkhornTilingData *__restrict tilingData = &tiling_data_in;47+ TPipe pipe;
45- 48+ MhcPreSinkhorn::MhcPreSinkhornStage1<DTYPE_X> op1;
49+ op1.cubeCompute_.mm1_.Init(&tilingData->mm1TilingData, &pipe);
50+ op1.cubeCompute_.mm1_.SetSubBlockIdx(0);
51+ op1.Init(x, phi, invRms, hcBeforeNorm, userWs, tilingData, &pipe);
52+ op1.Process();
53+ pipe.Destroy();
54+
55+ TPipe pipeStage2;
56+ MhcPreSinkhorn::MhcPreSinkhornStage2<DTYPE_X> op2;
57+ op2.Init(x, alpha, bias, hin, hPost, hRes,
58+ hPre, hcBeforeNorm, invRms, sumOut, normOut,
59+ userWs, tilingData, &pipeStage2);
60+ op2.Process(false);
61+ pipeStage2.Destroy();
62+ } else if (TILING_KEY_IS(1)) {
63+ // -----------split m k------------
46 for (int64_t bsLoopIdx = 0; bsLoopIdx < tilingData->bsLoop; bsLoopIdx++) {64 for (int64_t bsLoopIdx = 0; bsLoopIdx < tilingData->bsLoop; bsLoopIdx++) {
47 int64_t curBsOffset = bsLoopIdx * tilingData->curBsSplit;65 int64_t curBsOffset = bsLoopIdx * tilingData->curBsSplit;
48 int64_t curBs = (bsLoopIdx == tilingData->bsLoop - 1) ? tilingData->tailBs : tilingData->curBsSplit;66 int64_t curBs = (bsLoopIdx == tilingData->bsLoop - 1) ? tilingData->tailBs : tilingData->curBsSplit;
49 bool isTailBsLoop = (bsLoopIdx == tilingData->bsLoop - 1) && (tilingData->tailBs != tilingData->curBsSplit);67 bool isTailBsLoop = (bsLoopIdx == tilingData->bsLoop - 1) && (tilingData->tailBs != tilingData->curBsSplit);
50-
51 TPipe pipe;68 TPipe pipe;
52 MhcPreSinkhorn::MhcPreSinkhornMembaseKSplitCorePart1<DTYPE_X> op;69 MhcPreSinkhorn::MhcPreSinkhornMembaseKSplitCorePart1<DTYPE_X> op;
53 op.Init(x, phi, userWs, tilingData, &pipe, curBsOffset, curBs, isTailBsLoop);70 op.Init(x, phi, userWs, tilingData, &pipe, curBsOffset, curBs, isTailBsLoop);
54 op.Process(curBsOffset, curBs, isTailBsLoop);71 op.Process(curBsOffset, curBs, isTailBsLoop);
55- 
56 pipe.Destroy();72 pipe.Destroy();
57 73 
58 TPipe pipeStage2;74 TPipe pipeStage2;
@@ -61,7 +77,6 @@ extern "C" __global__ __aicore__ void mhc_pre_sinkhorn(GM_ADDR x, GM_ADDR phi, G
61 hPre, hcBeforeNorm, invRms, sumOut, normOut,77 hPre, hcBeforeNorm, invRms, sumOut, normOut,
62 userWs, tilingData, &pipeStage2, curBsOffset, isTailBsLoop);78 userWs, tilingData, &pipeStage2, curBsOffset, isTailBsLoop);
63 op2.Process(curBsOffset, curBs, isTailBsLoop);79 op2.Process(curBsOffset, curBs, isTailBsLoop);
64- 
65 pipeStage2.Destroy();80 pipeStage2.Destroy();
66 SyncAll<false>(); // cv全部同步81 SyncAll<false>(); // cv全部同步
67 }82 }
Mmhc/mhc_pre_sinkhorn/op_kernel/mhc_pre_sinkhorn_base.h+5-5
@@ -38,6 +38,8 @@ constexpr uint64_t NUM_EIGHT = 8;
38constexpr uint64_t NUM_SIXTEEN = 16;38constexpr uint64_t NUM_SIXTEEN = 16;
39constexpr uint64_t SQUARE_SUM_SIZE = 16;39constexpr uint64_t SQUARE_SUM_SIZE = 16;
40constexpr uint64_t MM_CACHE_LINE_BYTES = 512;40constexpr uint64_t MM_CACHE_LINE_BYTES = 512;
41+constexpr uint64_t MAX_BS_PER_LOOP = 32;
42+constexpr uint64_t ELEMENTS_SIZE_PER_BLOCK = BLOCK_SIZE / sizeof(float);
41 43 
42__aicore__ inline int32_t CeilDiv(int32_t a, int32_t b)44__aicore__ inline int32_t CeilDiv(int32_t a, int32_t b)
43{45{
@@ -171,13 +173,11 @@ __aicore__ inline void MulABLastDimBrcInline2(const LocalTensor<T> &output, cons
171 const int32_t curRowNum, const int32_t curColNum, const int32_t numN)173 const int32_t curRowNum, const int32_t curColNum, const int32_t numN)
172{174{
173 uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T);175 uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T);
174- 
175 if constexpr (needBrc) {176 if constexpr (needBrc) {
176 uint32_t repeatTimes = CeilDiv(curRowNum * CeilDiv(elemInOneBlock, numN), ONE_REPEAT_BLOCK_NUMS);177 uint32_t repeatTimes = CeilDiv(curRowNum * CeilDiv(elemInOneBlock, numN), ONE_REPEAT_BLOCK_NUMS);
177- Brcb(tmpBuffer, input1, repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE});178+ Brcb(tmpBuffer, input1, repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE}); // n-> n c
178 }179 }
179 180 
180- 
181 PipeBarrier<PIPE_V>();181 PipeBarrier<PIPE_V>();
182 uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T);182 uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T);
183 uint32_t curColNumAlign = RoundUp<T>(curColNum);183 uint32_t curColNumAlign = RoundUp<T>(curColNum);
@@ -216,7 +216,7 @@ __aicore__ inline void MulABLastDimBrcInline2(const LocalTensor<T> &output, cons
216 }216 }
217 }217 }
218 }218 }
219- 219+ // 非128对齐场景;
220 if (numRemainPerLine > 0) {220 if (numRemainPerLine > 0) {
221 if (dstRepStridePerLine > MAX_REPEAT_STRIDE) {221 if (dstRepStridePerLine > MAX_REPEAT_STRIDE) {
222 instrParams.dstBlkStride = 1;222 instrParams.dstBlkStride = 1;
@@ -579,7 +579,7 @@ __aicore__ inline void CastTwoDim(const LocalTensor<T0> &output, const LocalTens
579 PipeBarrier<PIPE_V>();579 PipeBarrier<PIPE_V>();
580}580}
581 581 
582- 582+// dim0:bs dim1:n dim2:c
583template <typename T>583template <typename T>
584__aicore__ void inline ProcessY(const LocalTensor<T> &yLocal, const LocalTensor<T> &xLocal,584__aicore__ void inline ProcessY(const LocalTensor<T> &yLocal, const LocalTensor<T> &xLocal,
585 const LocalTensor<float> &mix01Local, const LocalTensor<float> &hcBrcbLocal1,585 const LocalTensor<float> &mix01Local, const LocalTensor<float> &hcBrcbLocal1,
Mmhc/mhc_pre_sinkhorn/op_kernel/mhc_pre_sinkhorn_cube_compute.h+106-38
@@ -18,6 +18,8 @@
18#include "kernel_operator.h"18#include "kernel_operator.h"
19#include "kernel_operator_intf.h"19#include "kernel_operator_intf.h"
20#include "mhc_pre_sinkhorn_base.h"20#include "mhc_pre_sinkhorn_base.h"
21+#include "lib/matmul_intf.h"
22+#include <cstdint>
21 23 
22using AscendC::BLOCK_CUBE;24using AscendC::BLOCK_CUBE;
23using AscendC::GlobalTensor;25using AscendC::GlobalTensor;
@@ -32,6 +34,9 @@ using AscendC::WaitFlag;
32using namespace AscendC;34using namespace AscendC;
33 35 
34namespace MhcPreSinkhorn {36namespace MhcPreSinkhorn {
37+ 
38+constexpr MatmulConfig MHC_PRE_GRAD_MM1_CFG = GetMDLConfig(false, false, 0, false, false, false, true);
39+ 
35struct MmParams {40struct MmParams {
36 uint64_t curML1;41 uint64_t curML1;
37 uint64_t curKL1;42 uint64_t curKL1;
@@ -45,16 +50,18 @@ struct MmParams {
45 bool isLastK;50 bool isLastK;
46 bool isFirstK;51 bool isFirstK;
47};52};
48-#define MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM template <bool enableSquareSum>53+#define MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM template <bool enableSquareSum>
54+#define MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS HcCubeComputeSplitK<enableSquareSum>
55+#define MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM template<typename T>
56+#define MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS HcCubeCompute<T>
49 57 
50-#define MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS HcCubeCompute<enableSquareSum>58+// 切K使用的矩阵乘实现
51- 59+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
52-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM60+class HcCubeComputeSplitK {
53-class HcCubeCompute {
54public:61public:
55- __aicore__ inline HcCubeCompute(){};62+ __aicore__ inline HcCubeComputeSplitK(){};
56 63 
57- __aicore__ inline void Init(const GlobalTensor<float>& xGm, const GlobalTensor<float>& fnGm, TPipe *tpipe);64+ __aicore__ inline void Init(const GlobalTensor<float>& xGm, const GlobalTensor<float>& phiGm, TPipe *tpipe);
58 __aicore__ inline void ComputeDecode(const AscendC::GlobalTensor<float> &xGm, const AscendC::GlobalTensor<float> &workspaceGlobalA2,65 __aicore__ inline void ComputeDecode(const AscendC::GlobalTensor<float> &xGm, const AscendC::GlobalTensor<float> &workspaceGlobalA2,
59 const AscendC::GlobalTensor<float> &workspaceGlobalAB, const MmParams &mmParams);66 const AscendC::GlobalTensor<float> &workspaceGlobalAB, const MmParams &mmParams);
60 __aicore__ inline void CopyInB1(67 __aicore__ inline void CopyInB1(
@@ -93,7 +100,7 @@ private:
93 static constexpr int32_t ONE_BLOCK_SIZE = 32;100 static constexpr int32_t ONE_BLOCK_SIZE = 32;
94 int32_t perBlock32 = ONE_BLOCK_SIZE / sizeof(float);101 int32_t perBlock32 = ONE_BLOCK_SIZE / sizeof(float);
95 102 
96- GlobalTensor<float> fnGm_;103+ GlobalTensor<float> phiGm_;
97 GlobalTensor<float> yGm_;104 GlobalTensor<float> yGm_;
98 105 
99 static constexpr uint64_t MM1_MTE2_MTE1_EVENT = 2;106 static constexpr uint64_t MM1_MTE2_MTE1_EVENT = 2;
@@ -139,10 +146,10 @@ private:
139 uint64_t n_ = 0;146 uint64_t n_ = 0;
140};147};
141 148 
142-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM149+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
143-__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::Init(const GlobalTensor<float>& xGm, const GlobalTensor<float>& fnGm, TPipe *tpipe)150+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS::Init(const GlobalTensor<float>& xGm, const GlobalTensor<float>& phiGm, TPipe *tpipe)
144{151{
145- fnGm_ = fnGm;152+ phiGm_ = phiGm;
146 153 
147 TBuf<TPosition::A1> l1aBuffer;154 TBuf<TPosition::A1> l1aBuffer;
148 tpipe->InitBuffer(l1aBuffer, 256 * 1024);155 tpipe->InitBuffer(l1aBuffer, 256 * 1024);
@@ -175,8 +182,8 @@ __aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::Init(const
175 }182 }
176}183}
177 184 
178-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM185+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
179-__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::CopyInA1(186+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS::CopyInA1(
180 uint64_t kL1Size,187 uint64_t kL1Size,
181 const GlobalTensor<float> &aGlobal, const LocalTensor<float> &al1Local, const MmParams &mmParams)188 const GlobalTensor<float> &aGlobal, const LocalTensor<float> &al1Local, const MmParams &mmParams)
182{189{
@@ -192,8 +199,8 @@ __aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::CopyInA1(
192 DataCopy(al1Local, aGlobal, nd2nzParams);199 DataCopy(al1Local, aGlobal, nd2nzParams);
193}200}
194 201 
195-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM202+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
196-__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::CopyInB1(203+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS::CopyInB1(
197 uint64_t mGmOffset, uint64_t kGmOffset, uint64_t kL1Size, const MmParams &mmParams)204 uint64_t mGmOffset, uint64_t kGmOffset, uint64_t kL1Size, const MmParams &mmParams)
198{205{
199 AscendC::Nd2NzParams nd2nzParams;206 AscendC::Nd2NzParams nd2nzParams;
@@ -205,11 +212,11 @@ __aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::CopyInB1(
205 nd2nzParams.dstNzC0Stride = (mmParams.curNL1 + BLOCK_CUBE - 1) / BLOCK_CUBE * BLOCK_CUBE;212 nd2nzParams.dstNzC0Stride = (mmParams.curNL1 + BLOCK_CUBE - 1) / BLOCK_CUBE * BLOCK_CUBE;
206 nd2nzParams.dstNzNStride = 1;213 nd2nzParams.dstNzNStride = 1;
207 nd2nzParams.dstNzMatrixStride = 1;214 nd2nzParams.dstNzMatrixStride = 1;
208- DataCopy(l1b_[(l1bLoopIdx_ % L1_BUF_NUM) * L1_BUF_OFFSET], fnGm_[mGmOffset * mmParams.kGmSize + kGmOffset], nd2nzParams);215+ DataCopy(l1b_[(l1bLoopIdx_ % L1_BUF_NUM) * L1_BUF_OFFSET], phiGm_[mGmOffset * mmParams.kGmSize + kGmOffset], nd2nzParams);
209}216}
210 217 
211-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM218+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
212-__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::CopyOut(const AscendC::GlobalTensor<float> &workspaceGlobal,219+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS::CopyOut(const AscendC::GlobalTensor<float> &workspaceGlobal,
213 const AscendC::LocalTensor<float> &c1Local, uint64_t baseM, uint64_t baseN, bool enableNz2Nd, uint64_t N)220 const AscendC::LocalTensor<float> &c1Local, uint64_t baseM, uint64_t baseN, bool enableNz2Nd, uint64_t N)
214{221{
215 AscendC::DataCopyCO12DstParams intriParams;222 AscendC::DataCopyCO12DstParams intriParams;
@@ -231,8 +238,8 @@ __aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::CopyOut(con
231 AscendC::DataCopy(workspaceGlobal, c1Local, intriParams);238 AscendC::DataCopy(workspaceGlobal, c1Local, intriParams);
232}239}
233 240 
234-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM241+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
235-__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::Fixp(const AscendC::GlobalTensor<float> &workspaceGlobalA2,242+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS::Fixp(const AscendC::GlobalTensor<float> &workspaceGlobalA2,
236 const AscendC::GlobalTensor<float> &workspaceGlobalAB, const MmParams &mmParams)243 const AscendC::GlobalTensor<float> &workspaceGlobalAB, const MmParams &mmParams)
237{244{
238 // Copy MmadA2245 // Copy MmadA2
@@ -251,21 +258,21 @@ __aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::Fixp(const
251 mmParams.nOutSize); // nd m,n 512B对齐258 mmParams.nOutSize); // nd m,n 512B对齐
252}259}
253 260 
254-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM261+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
255-__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::SetBL1Mte1ToMte2Flag()262+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS::SetBL1Mte1ToMte2Flag()
256{263{
257 SetFlag<HardEvent::MTE1_MTE2>(B_MTE1_MTE2_EVENT + l1bLoopIdx_ % L1_BUF_NUM);264 SetFlag<HardEvent::MTE1_MTE2>(B_MTE1_MTE2_EVENT + l1bLoopIdx_ % L1_BUF_NUM);
258 l1bLoopIdx_++;265 l1bLoopIdx_++;
259}266}
260 267 
261-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM268+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
262-__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::WaitBL1Mte1ToMte2Flag()269+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS::WaitBL1Mte1ToMte2Flag()
263{270{
264 WaitFlag<HardEvent::MTE1_MTE2>(B_MTE1_MTE2_EVENT + l1bLoopIdx_ % L1_BUF_NUM);271 WaitFlag<HardEvent::MTE1_MTE2>(B_MTE1_MTE2_EVENT + l1bLoopIdx_ % L1_BUF_NUM);
265}272}
266 273 
267-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM274+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
268-__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::ComputeDecode(275+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS::ComputeDecode(
269 const AscendC::GlobalTensor<float> &xGm,276 const AscendC::GlobalTensor<float> &xGm,
270 const AscendC::GlobalTensor<float> &workspaceGlobalA2, const AscendC::GlobalTensor<float> &workspaceGlobalAB,277 const AscendC::GlobalTensor<float> &workspaceGlobalA2, const AscendC::GlobalTensor<float> &workspaceGlobalAB,
271 const MmParams &mmParams)278 const MmParams &mmParams)
@@ -342,8 +349,8 @@ __aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::ComputeDeco
342 Fixp(workspaceGlobalA2, workspaceGlobalAB, mmParams); // l0cLoopIdx_++;349 Fixp(workspaceGlobalA2, workspaceGlobalAB, mmParams); // l0cLoopIdx_++;
343}350}
344 351 
345-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM352+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
346-__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::End()353+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS::End()
347{354{
348 for (int i = 0; i < L0A_BUF_NUM; i++) {355 for (int i = 0; i < L0A_BUF_NUM; i++) {
349 WaitFlag<HardEvent::M_MTE1>(M_MTE1_EVENT_L0A + i);356 WaitFlag<HardEvent::M_MTE1>(M_MTE1_EVENT_L0A + i);
@@ -355,8 +362,8 @@ __aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::End()
355 }362 }
356}363}
357 364 
358-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM365+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
359-__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::LoadAToL0A(366+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS::LoadAToL0A(
360 uint64_t kL1Offset, uint64_t kL0Size, uint64_t l1LoopIdx, const MmParams &mmParams)367 uint64_t kL1Offset, uint64_t kL0Size, uint64_t l1LoopIdx, const MmParams &mmParams)
361{368{
362 static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true};369 static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true};
@@ -393,8 +400,8 @@ __aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::LoadAToL0A(
393 loadData3DParams);400 loadData3DParams);
394}401}
395 402 
396-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM403+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
397-__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::LoadAToL0B(404+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS::LoadAToL0B(
398 uint64_t kL1Offset, uint64_t kL0Size, uint64_t l1LoopIdx, const MmParams &mmParams)405 uint64_t kL1Offset, uint64_t kL0Size, uint64_t l1LoopIdx, const MmParams &mmParams)
399{406{
400 // mk nz -> m,k zz407 // mk nz -> m,k zz
@@ -413,8 +420,8 @@ __aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::LoadAToL0B(
413 }420 }
414}421}
415 422 
416-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM423+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
417-__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::MmadA2(424+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS::MmadA2(
418 uint64_t kGmOffset, uint64_t kL0Size, bool isLastK, const MmParams &mmParams)425 uint64_t kGmOffset, uint64_t kL0Size, bool isLastK, const MmParams &mmParams)
419{426{
420 SetFlag<HardEvent::MTE1_M>(MTE1_M_EVENT);427 SetFlag<HardEvent::MTE1_M>(MTE1_M_EVENT);
@@ -435,8 +442,8 @@ __aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::MmadA2(
435 }442 }
436}443}
437 444 
438-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM445+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
439-__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::LoadBToL0B(446+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS::LoadBToL0B(
440 uint64_t kL1Offset, uint64_t kL0Size, uint64_t l1LoopIdx, const MmParams &mmParams)447 uint64_t kL1Offset, uint64_t kL0Size, uint64_t l1LoopIdx, const MmParams &mmParams)
441{448{
442 LoadData2DParams l1ToL0bParams;449 LoadData2DParams l1ToL0bParams;
@@ -452,8 +459,8 @@ __aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::LoadBToL0B(
452 l1ToL0bParams);459 l1ToL0bParams);
453}460}
454 461 
455-MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM462+MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_PARAM
456-__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::MmadAB(463+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_SPLIT_K_TEMPLATE_CLASS::MmadAB(
457 uint64_t kGmOffset, uint64_t kL0Size, bool isLastK, const MmParams &mmParams)464 uint64_t kGmOffset, uint64_t kL0Size, bool isLastK, const MmParams &mmParams)
458{465{
459 SetFlag<HardEvent::MTE1_M>(MTE1_M_EVENT);466 SetFlag<HardEvent::MTE1_M>(MTE1_M_EVENT);
@@ -474,6 +481,67 @@ __aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::MmadAB(
474 mmadParams);481 mmadParams);
475 AscendC::SetHF32Mode(0);482 AscendC::SetHF32Mode(0);
476}483}
484+ 
485+// 不切K使用的矩阵乘实现
486+MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM
487+class HcCubeCompute {
488+public:
489+ __aicore__ inline HcCubeCompute(){};
490+ __aicore__ inline void Init(const GlobalTensor<float> &xGm, const GlobalTensor<float> &phiGm,
491+ const GlobalTensor<float> &workspaceGlobalAB, int64_t bs, int64_t n, int64_t c, int64_t vecCoreNum);
492+ __aicore__ inline void ProcessMatmulXPhi(const int32_t taskOffset, const int32_t mm1M);
493+ 
494+public:
495+ GlobalTensor<float> xGm_;
496+ GlobalTensor<float> phiGm_;
497+ GlobalTensor<float> workspaceGlobalAB_;
498+ using AType = matmul::MatmulType<TPosition::GM, CubeFormat::ND, T>;
499+ using BType = matmul::MatmulType<TPosition::GM, CubeFormat::ND, T, true>;
500+ using CType = matmul::MatmulType<TPosition::GM, CubeFormat::ND, T>;
501+
502+ matmul::MatmulImpl<AType, BType, CType, CType, MHC_PRE_GRAD_MM1_CFG> mm1_;
503+ 
504+ int64_t n_, c_, curBs_;
505+ int64_t mm1K_, mm1M_, mm1N_;
506+ int64_t ping4Cub = 1;
507+ uint64_t vecCoreNum_ = 0;
508+ uint64_t blockIdx_ = 0;
509+};
510+ 
511+MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM
512+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::Init(const GlobalTensor<float> &xGm,
513+ const GlobalTensor<float> &phiGm,
514+ const GlobalTensor<float> &workspaceGlobalAB,
515+ int64_t bs, int64_t n, int64_t c, int64_t vecCoreNum)
516+{
517+ blockIdx_ = GetBlockIdx();
518+ xGm_ = xGm;
519+ phiGm_ = phiGm;
520+ workspaceGlobalAB_ = workspaceGlobalAB;
521+ vecCoreNum_ = vecCoreNum;
522+ n_ = n;
523+ c_ = c;
524+ curBs_ = bs;
525+ 
526+ mm1N_= n_ * n_ + 2 * n_;
527+ mm1K_ = n_ * c_;
528+}
529+ 
530+MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_PARAM
531+__aicore__ inline void MHC_PRE_SINKHORN_CUBE_COMPUTE_TEMPLATE_CLASS::ProcessMatmulXPhi(const int32_t taskOffset, const int32_t mm1M)
532+{
533+ if (mm1M <= 0)
534+ return;
535+ uint64_t xxoffset = blockIdx_ * curBs_* 2 * mm1K_ + ping4Cub * vecCoreNum_ * curBs_* mm1K_;
536+ mm1_.SetTensorA(xGm_[blockIdx_ * curBs_* 2 * mm1K_ + ping4Cub * vecCoreNum_ * curBs_* mm1K_]);
537+ mm1_.SetTensorB(phiGm_, true);
538+ mm1_.SetHF32(true, 1);
539+ mm1_.SetOrgShape(mm1M, mm1N_, mm1K_);
540+ mm1_.SetSingleShape(mm1M, mm1N_, mm1K_);
541+ mm1_.template IterateAll<false> (workspaceGlobalAB_[taskOffset * mm1N_]);
542+ mm1_.End();
543+ ping4Cub = 1 - ping4Cub;
544+}
477} // namespace MhcPreSinkhorn545} // namespace MhcPreSinkhorn
478 546 
479#endif // MHC_PRE_SINKHORN_CUBE_COMPUTE_H547#endif // MHC_PRE_SINKHORN_CUBE_COMPUTE_H
Mmhc/mhc_pre_sinkhorn/op_kernel/mhc_pre_sinkhorn_m_k_split_core.h+3-2
@@ -181,7 +181,7 @@ private:
181 LocalTensor<float> yCastLocal;181 LocalTensor<float> yCastLocal;
182 LocalTensor<float> mmOutLocal;182 LocalTensor<float> mmOutLocal;
183 183
184- HcCubeCompute<false> cubeCompute_;184+ HcCubeComputeSplitK<false> cubeCompute_;
185 static constexpr uint64_t SYNC_AIV_TO_AIC_FLAG = 8;185 static constexpr uint64_t SYNC_AIV_TO_AIC_FLAG = 8;
186 static constexpr uint64_t SYNC_AIC_TO_AIV_FLAG = 9;186 static constexpr uint64_t SYNC_AIC_TO_AIV_FLAG = 9;
187 static constexpr uint64_t SYNC_MODE2 = NUM_TWO;187 static constexpr uint64_t SYNC_MODE2 = NUM_TWO;
@@ -230,7 +230,7 @@ public:
230 }230 }
231 231 
232 // InQue232 // InQue
233- int64_t stage1UsedCoreNum = tilingData->cubeBlockDimK; // todo check 此处不应该写死32233+ int64_t stage1UsedCoreNum = tilingData->cubeBlockDimK;
234 int64_t mixesQue01Size = stage1UsedCoreNum * tilingData->stage2RowFactor * tilingData->hcMultAlign * NUM_TWO * sizeof(float);234 int64_t mixesQue01Size = stage1UsedCoreNum * tilingData->stage2RowFactor * tilingData->hcMultAlign * NUM_TWO * sizeof(float);
235 pipe->InitBuffer(mixesQue01, NUM_TWO, mixesQue01Size);235 pipe->InitBuffer(mixesQue01, NUM_TWO, mixesQue01Size);
236 pipe->InitBuffer(mixesQue2, NUM_TWO,236 pipe->InitBuffer(mixesQue2, NUM_TWO,
@@ -353,6 +353,7 @@ public:
353 mxies01Local[stage1UsedCoreNum * curRowFactor * tilingData->hcMultAlign],353 mxies01Local[stage1UsedCoreNum * curRowFactor * tilingData->hcMultAlign],
354 stage1UsedCoreNum, curRowFactor, tilingData->hcMult, curBs,354 stage1UsedCoreNum, curRowFactor, tilingData->hcMult, curBs,
355 CeilAlign(tilingData->hcMix, WORKSPACE_ALIGN_SIZE / sizeof(float)));355 CeilAlign(tilingData->hcMix, WORKSPACE_ALIGN_SIZE / sizeof(float)));
356+ 
356 // wk:[2, kcorenum, bs, n^2 +2n]357 // wk:[2, kcorenum, bs, n^2 +2n]
357 // mx0[2,kcorenum, curRowfator, n]358 // mx0[2,kcorenum, curRowfator, n]
358 mixesQue01.EnQue(mxies01Local);359 mixesQue01.EnQue(mxies01Local);
Amhc/mhc_pre_sinkhorn/op_kernel/mhc_pre_sinkhorn_m_split_core.h+689-0
@@ -0,0 +1,689 @@
1+/**
2+ * Copyright (c) 2026 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+/*!
12+ * \file mhc_pre_sinkhorn_m_split_core.h
13+ * \brief
14+ */
15+ 
16+#ifndef MHC_PRE_SINKHORN_M_SPLIT_A3_CORE_H
17+#define MHC_PRE_SINKHORN_M_SPLIT_A3_CORE_H
18+ 
19+#include "kernel_operator.h"
20+#include "mhc_pre_sinkhorn_base.h"
21+#include "mhc_pre_sinkhorn_cube_compute.h"
22+ 
23+namespace MhcPreSinkhorn {
24+using namespace AscendC;
25+template <typename T>
26+class MhcPreSinkhornStage1 {
27+public:
28+ __aicore__ inline MhcPreSinkhornStage1() {}
29+ 
30+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR phi, GM_ADDR invRms, GM_ADDR hcBeforeNorm,
31+ GM_ADDR workspace, const MhcPreSinkhornTilingData *tilingDataPtr, TPipe *pipePtr)
32+ {
33+ pipe = pipePtr;
34+ tilingData = tilingDataPtr;
35+ needGrad_ = tilingData->needGrad;
36+ 
37+ // 获取当前核心 ID
38+ uint64_t blockIdx = GetBlockIdx();
39+ uint64_t cubeCoreId = blockIdx;
40+ uint64_t vecCoreId = blockIdx;
41+
42+ // 计算 workspace 偏移
43+ int64_t ncSize = tilingData->hcMult * tilingData->d;
44+ 
45+ // 计算 Cube Core ID（Vector Core ID / 2）
46+ if ASCEND_IS_AIV {
47+ cubeCoreId = blockIdx / CV_RATIO;
48+ int64_t maxBsPerLoop = 32;
49+ int64_t xQueSize = maxBsPerLoop * tilingData->stage1NcFactor;
50+ pipe->InitBuffer(xQue, NUM_TWO, xQueSize * sizeof(T));
51+ pipe->InitBuffer(xCastQue, NUM_TWO, xQueSize * sizeof(float));
52+ pipe->InitBuffer(tmpQue, maxBsPerLoop * tilingData->stage1NcFactor * sizeof(float));
53+ pipe->InitBuffer(sumQue, maxBsPerLoop * sizeof(float));
54+ LocalTensor<float> sumLocal = sumQue.Get<float>();
55+ pipe->InitBuffer(invRmsQue, NUM_TWO, maxBsPerLoop * sizeof(float));
56+ }
57+ 
58+ // 设置 GM 地址
59+ xGm.SetGlobalBuffer((__gm__ T *)x);
60+ phiGm.SetGlobalBuffer((__gm__ float *)phi);
61+ invRmsGm.SetGlobalBuffer((__gm__ float *)invRms);
62+ hcBeforeNormGm.SetGlobalBuffer((__gm__ float *)hcBeforeNorm);
63+ xCastWsGm.SetGlobalBuffer((__gm__ float *)workspace);
64+ 
65+ // 设置额外的 workspace 地址（needGrad=false 时使用）
66+ if (!needGrad_) {
67+ int64_t xCastWsTotalSize = tilingData->stage1XCastWsSize / sizeof(float);
68+ invRmsWsGm.SetGlobalBuffer((__gm__ float *)workspace + xCastWsTotalSize);
69+ int64_t invRmsWsTotalSize = tilingData->bs * sizeof(float) / sizeof(float);
70+ hcBeforeNormWsGm.SetGlobalBuffer((__gm__ float *)workspace + xCastWsTotalSize + invRmsWsTotalSize);
71+ }
72+ 
73+ // 初始化 Cube Compute
74+ if ASCEND_IS_AIC {
75+ // 计算当前 Cube Core 的输出偏移
76+ int64_t mmOutSize = tilingData->hcMult * tilingData->hcMult + 2 * tilingData->hcMult;
77+ // 根据 needGrad 决定输出地址
78+ if (needGrad_) {
79+ cubeCompute_.Init(xCastWsGm, phiGm, hcBeforeNormGm,
80+ tilingData->stage1BsFactor, tilingData->hcMult, tilingData->d, tilingData->stage1VecCoreNum);
81+ } else {
82+ cubeCompute_.Init(xCastWsGm, phiGm, hcBeforeNormWsGm,
83+ tilingData->stage1BsFactor, tilingData->hcMult, tilingData->d, tilingData->stage1VecCoreNum);
84+ }
85+ return;
86+ }
87+ }
88+ 
89+ __aicore__ inline void Process()
90+ {
91+ uint64_t blockIdx = GetBlockIdx();
92+ uint64_t vecCoreId = blockIdx;
93+ uint64_t cubeCoreId = blockIdx;
94+ 
95+ if ASCEND_IS_AIC {
96+ CrossCoreSetFlag<SYNC_MODE2, PIPE_FIX>(SYNC_AIC_TO_AIV_FLAG);
97+ CrossCoreSetFlag<SYNC_MODE2, PIPE_FIX>(SYNC_AIC_TO_AIV_FLAG);
98+ }
99+ 
100+ // 循环处理 BS（每轮最多 32 个）
101+ // ASCEND_IS_AIC 直接根据总BS数，计算当前处理的BS数
102+ int32_t totalTasksAligned_ = AlignUp(tilingData->bs, tilingData->stage1VecCoreNum * tilingData->stage1BsFactor);
103+ 
104+ // Cube Core 流程
105+ if ASCEND_IS_AIC {
106+ for (int32_t taskOffset = blockIdx * 2 * tilingData->stage1BsFactor; taskOffset < totalTasksAligned_;
107+ taskOffset += tilingData->stage1CubeCoreNum * 2 * tilingData->stage1BsFactor) {
108+ int32_t tileTaskCount = min(static_cast<int32_t>(2 * tilingData->stage1BsFactor),
109+ static_cast<int32_t>(tilingData->bs - taskOffset));
110+ // 等待 Vector Core 完成 Cast
111+ CrossCoreWaitFlag(SYNC_AIV_TO_AIC_FLAG);
112+ // 执行矩阵乘
113+ if (tileTaskCount > 0) {
114+ cubeCompute_.ProcessMatmulXPhi(taskOffset, tileTaskCount);
115+ }
116+ // 通知 Vector Core 完成
117+ CrossCoreSetFlag<SYNC_MODE2, PIPE_FIX>(SYNC_AIC_TO_AIV_FLAG);
118+ }
119+ }
120+ 
121+ if ASCEND_IS_AIV {
122+ int64_t ncSize = tilingData->hcMult * tilingData->d;
123+ for (int32_t taskOffset = blockIdx * tilingData->stage1BsFactor; taskOffset < totalTasksAligned_;
124+ taskOffset += tilingData->stage1VecCoreNum * tilingData->stage1BsFactor) {
125+ CrossCoreWaitFlag(SYNC_AIC_TO_AIV_FLAG);
126+ int32_t curBs = min(static_cast<int32_t>(tilingData->stage1BsFactor),
127+ static_cast<int32_t>(tilingData->bs - taskOffset));
128+ if (curBs > 0) {
129+ int64_t curBsOffset = taskOffset;
130+ // 初始化 sum 为 0
131+ LocalTensor<float> invRmsLocal = invRmsQue.AllocTensor<float>();
132+ LocalTensor<float> squareLocal = tmpQue.Get<float>();
133+ LocalTensor<float> sumLocal = sumQue.Get<float>();
134+ Duplicate(sumLocal, 0.0f, curBs);
135+ 
136+ // 分段 ReduceSum
137+ for (int64_t ncIdx = 0; ncIdx < tilingData->stage1NcLoop; ncIdx++) {
138+ int64_t curNcSize = (ncIdx == tilingData->stage1NcLoop - 1)
139+ ? tilingData->stage1TailNcFactor
140+ : tilingData->stage1NcFactor;
141+ int64_t ncOffset = ncIdx * tilingData->stage1NcFactor;
142+ int64_t curNcSizeAlign = AlignUp(static_cast<uint64_t>(curNcSize), ELEMENTS_SIZE_PER_BLOCK);
143+ 
144+ // 1. CopyIn X (BF16)
145+ LocalTensor<T> xLocal = xQue.AllocTensor<T>();
146+ CopyIn(xGm[curBsOffset * ncSize + ncOffset], xLocal, curBs, curNcSize, ncSize - curNcSize);
147+ xQue.EnQue(xLocal);
148+ 
149+ // 2. Cast X (BF16) → X_cast (FP32)
150+ xLocal = xQue.DeQue<T>();
151+ LocalTensor<float> xCastLocal = xCastQue.AllocTensor<float>();
152+ Cast(xCastLocal, xLocal, RoundMode::CAST_NONE, curBs * curNcSizeAlign);
153+ xQue.FreeTensor(xLocal);
154+ xCastQue.EnQue(xCastLocal);
155+ 
156+ // 3. CopyOut X_cast to workspace
157+ xCastLocal = xCastQue.DeQue<float>();
158+ int64_t xCastoffset = blockIdx * tilingData->stage1BsFactor * ncSize + ping4vec * tilingData->stage1VecCoreNum * tilingData->stage1BsFactor * ncSize;
159+ CopyOut(xCastLocal, xCastWsGm[xCastoffset + ncOffset], curBs, curNcSize, ncSize - curNcSize);
160+ xCastQue.FreeTensor(xCastLocal);
161+ 
162+ // 4. 计算 inv_rms
163+ // 4.1 Square
164+ PipeBarrier<PIPE_V>();
165+ Mul(squareLocal, xCastLocal, xCastLocal, curBs * curNcSizeAlign);
166+ PipeBarrier<PIPE_V>();
167+ 
168+ // 对当前段进行 ReduceSum
169+ uint32_t srcShape[2] = {static_cast<uint32_t>(curBs), static_cast<uint32_t>(curNcSize)};
170+ AscendC::ReduceSum<float, Pattern::Reduce::AR, true>(squareLocal, squareLocal, srcShape, true);
171+ PipeBarrier<PIPE_V>();
172+
173+ // 累加到总和
174+ Add(sumLocal, sumLocal, squareLocal, curBs);
175+ PipeBarrier<PIPE_V>();
176+ }
177+ // 4.3 Compute inv_rms
178+ // inv_rms = 1 / sqrt(sum / (n*c) + eps)
179+ float invNc = 1.0f / static_cast<float>(ncSize);
180+ Muls(invRmsLocal, sumLocal, invNc, curBs);
181+ PipeBarrier<PIPE_V>();
182+ Adds(invRmsLocal, invRmsLocal, tilingData->normEps, curBs);
183+ PipeBarrier<PIPE_V>();
184+ Sqrt(invRmsLocal, invRmsLocal, curBs);
185+ PipeBarrier<PIPE_V>();
186+ Duplicate(sumLocal, 1.0f, curBs);
187+ PipeBarrier<PIPE_V>();
188+ Div(invRmsLocal, sumLocal, invRmsLocal, curBs);
189+ invRmsQue.EnQue(invRmsLocal);
190+ 
191+ // 5. CopyOut inv_rms to output
192+ invRmsLocal = invRmsQue.DeQue<float>();
193+ if (needGrad_) {
194+ CopyOut(invRmsLocal, invRmsGm[curBsOffset], 1, curBs);
195+ } else {
196+ CopyOut(invRmsLocal, invRmsWsGm[curBsOffset], 1, curBs);
197+ }
198+ invRmsQue.FreeTensor(invRmsLocal);
199+ ping4vec = 1 - ping4vec;
200+ }
201+ CrossCoreSetFlag<SYNC_MODE2, PIPE_MTE3>(SYNC_AIV_TO_AIC_FLAG);
202+ }
203+ CrossCoreWaitFlag(SYNC_AIC_TO_AIV_FLAG);
204+ CrossCoreWaitFlag(SYNC_AIC_TO_AIV_FLAG);
205+ }
206+ }
207+ 
208+public:
209+ TPipe *pipe;
210+ const MhcPreSinkhornTilingData *tilingData;
211+ 
212+ // GM 地址
213+ GlobalTensor<T> xGm; // 输入 X (BF16)
214+ GlobalTensor<float> phiGm; // phi 矩阵 (FP32)
215+ GlobalTensor<float> invRmsGm; // inv_rms 输出 (FP32)
216+ GlobalTensor<float> hcBeforeNormGm; // hcBeforeNorm 输出 (FP32)
217+ GlobalTensor<float> xCastWsGm; // X_cast workspace
218+ GlobalTensor<float> invRmsWsGm; // invRms workspace (needGrad=false)
219+ GlobalTensor<float> hcBeforeNormWsGm; // hcBeforeNorm workspace (needGrad=false)
220+ 
221+ // Queue
222+ TQue<QuePosition::VECIN, NUM_TWO> xQue; // X 输入队列
223+ TQue<QuePosition::VECOUT, NUM_TWO> xCastQue; // X_cast 输出队列
224+ TBuf<QuePosition::VECCALC> tmpQue;
225+ TBuf<QuePosition::VECCALC> sumAddQue; // sum 队列
226+ TBuf<QuePosition::VECCALC> sumQue; // sum 队列
227+ TQue<QuePosition::VECOUT, NUM_TWO> invRmsQue; // inv_rms 队列
228+ 
229+ // Cube Compute
230+ HcCubeCompute<float> cubeCompute_;
231+
232+ // 标志
233+ bool needGrad_ = true;
234+ uint64_t ping4vec = 1;
235+ 
236+ // 同步标志
237+ static constexpr uint64_t SYNC_AIV_TO_AIC_FLAG = 8;
238+ static constexpr uint64_t SYNC_AIC_TO_AIV_FLAG = 9;
239+ static constexpr uint64_t SYNC_MODE2 = NUM_TWO;
240+ static constexpr uint64_t CV_RATIO = 2;
241+};
242+ 
243+template <typename T>
244+class MhcPreSinkhornStage2 {
245+public:
246+ __aicore__ inline MhcPreSinkhornStage2()
247+ {
248+ }
249+ 
250+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR hcScale, GM_ADDR hcBase, GM_ADDR y,
251+ GM_ADDR post, GM_ADDR combFrag,
252+ GM_ADDR hPre, GM_ADDR hcBeforeNorm, GM_ADDR invRms,
253+ GM_ADDR sumOut, GM_ADDR normOut,
254+ GM_ADDR workspace,
255+ const MhcPreSinkhornTilingData *tilingDataPtr, TPipe *pipePtr)
256+ {
257+ pipe = pipePtr;
258+ tilingData = tilingDataPtr;
259+ needGrad_ = tilingData->needGrad;
260+ 
261+ xGm.SetGlobalBuffer((__gm__ T *)x);
262+ hcScaleGm.SetGlobalBuffer((__gm__ float *)hcScale);
263+ hcBaseGm.SetGlobalBuffer((__gm__ float *)hcBase);
264+ yGm.SetGlobalBuffer((__gm__ T *)y);
265+ postGm.SetGlobalBuffer((__gm__ float *)post);
266+ combFragGm.SetGlobalBuffer((__gm__ float *)combFrag);
267+ workspaceGm.SetGlobalBuffer((__gm__ float *)workspace);
268+ 
269+ if (needGrad_) {
270+ hPreGm.SetGlobalBuffer((__gm__ float *)hPre);
271+ hcBeforeNormGm.SetGlobalBuffer((__gm__ float *)hcBeforeNorm);
272+ invRmsGm.SetGlobalBuffer((__gm__ float *)invRms);
273+ sumOutGm.SetGlobalBuffer((__gm__ float *)sumOut);
274+ normOutGm.SetGlobalBuffer((__gm__ float *)normOut);
275+ }
276+ // InQue
277+ int64_t mixesQue01Size = tilingData->stage2RowFactor * tilingData->hcMultAlign * NUM_TWO * sizeof(float);
278+ pipe->InitBuffer(mixesQue01, NUM_TWO, mixesQue01Size);
279+ pipe->InitBuffer(mixesQue2, NUM_TWO, tilingData->stage2RowFactor * tilingData->hcMult * tilingData->hcMultAlign * sizeof(float));
280+ pipe->InitBuffer(squareSumQue, NUM_TWO, tilingData->stage2RowFactor * SQUARE_SUM_SIZE * sizeof(float));
281+ int64_t xQueNum2 = tilingData->stage2RowFactor * tilingData->hcMult * RoundUp<T>(tilingData->dFactor);
282+ pipe->InitBuffer(xQue, NUM_TWO, xQueNum2 * sizeof(T));
283+ // OutQue
284+ pipe->InitBuffer(yQue, NUM_TWO, tilingData->stage2RowFactor * RoundUp<T>(tilingData->dFactor) * sizeof(T));
285+ pipe->InitBuffer(postQue, NUM_TWO, tilingData->stage2RowFactor * tilingData->hcMultAlign * sizeof(float));
286+ pipe->InitBuffer(combFragQue, NUM_TWO,
287+ tilingData->stage2RowFactor * tilingData->hcMult * tilingData->hcMultAlign * sizeof(float));
288+ pipe->InitBuffer(reduceQue, NUM_TWO, tilingData->stage2RowFactor * tilingData->hcMultAlign * sizeof(float));
289+ // TBuf
290+ pipe->InitBuffer(combFragBuf,
291+ tilingData->stage2RowFactor * tilingData->hcMult * tilingData->hcMultAlign * NUM_TWO * sizeof(float));
292+ pipe->InitBuffer(hcBaseBuf0, tilingData->hcMultAlign * sizeof(float));
293+ pipe->InitBuffer(hcBaseBuf1, tilingData->hcMultAlign * sizeof(float));
294+ pipe->InitBuffer(hcBaseBuf2, tilingData->hcMult * tilingData->hcMultAlign * sizeof(float));
295+ pipe->InitBuffer(rowBrcbBuf0, RoundUp<float>(tilingData->stage2RowFactor) * BLOCK_SIZE);
296+ pipe->InitBuffer(hcBrcbBuf1, RoundUp<float>(tilingData->stage2RowFactor * tilingData->hcMultAlign * NUM_TWO) * BLOCK_SIZE);
297+ pipe->InitBuffer(reduceBuf, tilingData->stage2RowFactor * tilingData->hcMultAlign * sizeof(float));
298+ pipe->InitBuffer(mxies01ReduceBuf, tilingData->stage2RowFactor * tilingData->hcMultAlign * NUM_TWO * sizeof(float));
299+ pipe->InitBuffer(mxies02ReduceBuf, tilingData->stage2RowFactor * tilingData->hcMultAlign * tilingData->hcMult * sizeof(float));
300+ pipe->InitBuffer(xCastBuf, xQueNum2 * sizeof(float));
301+ pipe->InitBuffer(yCastBuf, tilingData->stage2RowFactor * RoundUp<T>(tilingData->dFactor) * sizeof(float));
302+ 
303+ hcBase0Local = hcBaseBuf0.Get<float>();
304+ hcBase1Local = hcBaseBuf1.Get<float>();
305+ hcBase2Local = hcBaseBuf2.Get<float>();
306+ rowBrcbLocal0 = rowBrcbBuf0.Get<float>();
307+ hcBrcbLocal1 = hcBrcbBuf1.Get<float>();
308+ reduceLocal = reduceBuf.Get<float>();
309+ mxies01ReduceLocal = mxies01ReduceBuf.Get<float>();
310+ mxies02ReduceLocal = mxies02ReduceBuf.Get<float>();
311+ xCastLocal = xCastBuf.Get<float>();
312+ yCastLocal = yCastBuf.Get<float>();
313+ combFragBufLocal = combFragBuf.Get<float>();
314+ }
315+ 
316+ __aicore__ inline void Process(bool isTailBsLoop = false)
317+ {
318+ SyncAll();
319+ isTailBsLoop_ = isTailBsLoop;
320+
321+ int64_t curRowOfFormerBlock = isTailBsLoop ? tilingData->tailBsRowOfFormerBlock : tilingData->rowOfFormerBlock;
322+ int64_t curRowLoopOfFormerBlock = isTailBsLoop ? tilingData->tailBsRowLoopOfFormerBlock : tilingData->rowLoopOfFormerBlock;
323+ int64_t curRowLoopOfTailBlock = isTailBsLoop ? tilingData->tailBsRowLoopOfTailBlock : tilingData->rowLoopOfTailBlock;
324+ int64_t curSecondUsedCoreNum = isTailBsLoop ? tilingData->tailBsUsedCoreNum : tilingData->secondUsedCoreNum;
325+ int64_t curStage2RowFactor = isTailBsLoop ? tilingData->tailBsRowFactor : tilingData->stage2RowFactor;
326+ int64_t curTailRowFactorOfFormerBlock = isTailBsLoop ? tilingData->tailBsTailRowFactorOfFormerBlock : tilingData->tailRowFactorOfFormerBlock;
327+ int64_t curTailRowFactorOfTailBlock = isTailBsLoop ? tilingData->tailBsTailRowFactorOfTailBlock : tilingData->tailRowFactorOfTailBlock;
328+ int64_t curML1Size = isTailBsLoop ? tilingData->tailBsML1Size : tilingData->mL1Size;
329+
330+ if ASCEND_IS_AIV {
331+ int64_t stage1UsedCoreNum = tilingData->cubeBlockDimK;
332+ int64_t stage2BlockIdx = GetBlockIdx();
333+ int64_t stage2UsedCoreNum = curSecondUsedCoreNum;
334+ if (stage2BlockIdx >= stage2UsedCoreNum) {
335+ return;
336+ }
337+ int64_t mmLastAxisSize = CeilAlign(tilingData->hcMix, MM_CACHE_LINE_BYTES / sizeof(float));
338+ int64_t xCastFp32BufSize = curML1Size * CeilAlign(tilingData->cvLoopKSize, MM_CACHE_LINE_BYTES / sizeof(float)) * sizeof(float);
339+ int64_t workspaceSize1 = (tilingData->cubeCoreNum * DOUBLE_BUFFER * xCastFp32BufSize) / sizeof(float);
340+ int64_t workspaceSize2 = CeilAlign(stage1UsedCoreNum * tilingData->bs * mmLastAxisSize * sizeof(float), WORKSPACE_ALIGN_SIZE) / sizeof(float);
341+ CopyIn(hcBaseGm, hcBase0Local, 1, tilingData->hcMult);
342+ CopyIn(hcBaseGm[tilingData->hcMult], hcBase1Local, 1, tilingData->hcMult);
343+ CopyIn(hcBaseGm[tilingData->hcMult * NUM_TWO], hcBase2Local, tilingData->hcMult, tilingData->hcMult);
344+ event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
345+ SetFlag<HardEvent::MTE2_V>(eventId);
346+ WaitFlag<HardEvent::MTE2_V>(eventId);
347+ 
348+ int64_t rowOuterLoop =
349+ (stage2BlockIdx == stage2UsedCoreNum - 1) ? curRowLoopOfTailBlock : curRowLoopOfFormerBlock;
350+ int64_t tailRowFactor = (stage2BlockIdx == stage2UsedCoreNum - 1) ? curTailRowFactorOfTailBlock :
351+ curTailRowFactorOfFormerBlock;
352+ int64_t xGmBlockBaseOffsetPart2 = stage2BlockIdx * curRowOfFormerBlock * tilingData->hcMult * tilingData->d;
353+ uint64_t mixBaseOffset = 0;
354+ for (int64_t rowOuterIdx = 0; rowOuterIdx < rowOuterLoop; rowOuterIdx++) {
355+ int64_t xGmBsBaseOffsetPart2 = rowOuterIdx * curStage2RowFactor * tilingData->hcMult * tilingData->d;
356+ int64_t curRowFactor = (rowOuterIdx == rowOuterLoop - 1) ? tailRowFactor : curStage2RowFactor;
357+
358+ squareSumOutLocal = squareSumQue.AllocTensor<float>();
359+ Duplicate(rowBrcbLocal0, static_cast<float>(1.0f), curRowFactor);
360+ if (needGrad_) {
361+ CopyIn(invRmsGm[stage2BlockIdx * curRowOfFormerBlock + rowOuterIdx * curStage2RowFactor],
362+ squareSumOutLocal, 1, curRowFactor);
363+ } else {
364+ int64_t invRmsOffset = tilingData->bs * tilingData->hcMult * tilingData->d;
365+ CopyIn(workspaceGm[invRmsOffset + stage2BlockIdx * curRowOfFormerBlock + rowOuterIdx * curStage2RowFactor],
366+ squareSumOutLocal, 1, curRowFactor);
367+ }
368+
369+ squareSumQue.EnQue(squareSumOutLocal);
370+ squareSumOutLocal = squareSumQue.DeQue<float>();
371+ // 搬运矩阵乘的前两段结果--> 内存格式的变更： 连续地址改为 bs *n ; bs *n ,n为4 非对齐，则需要pad出对齐的8
372+ mxies01Local = mixesQue01.AllocTensor<float>();
373+ 
374+ if (needGrad_) {
375+ mixBaseOffset = stage2BlockIdx * curRowOfFormerBlock * tilingData->hcMix +
376+ rowOuterIdx * curStage2RowFactor * tilingData->hcMix;
377+ CopyInWithOuterFor(hcBeforeNormGm[mixBaseOffset], mxies01Local, 1, curRowFactor, tilingData->hcMult,
378+ tilingData->bs, tilingData->hcMix);
379+ CopyInWithOuterFor(hcBeforeNormGm[mixBaseOffset + tilingData->hcMult],
380+ mxies01Local[curRowFactor * tilingData->hcMultAlign],
381+ 1, curRowFactor, tilingData->hcMult, tilingData->bs,
382+ tilingData->hcMix);
383+ } else {
384+ int64_t invRmsOffset = tilingData->bs * tilingData->hcMult * tilingData->d;
385+ mixBaseOffset = workspaceSize1 + stage2BlockIdx * curRowOfFormerBlock * tilingData->hcMix +
386+ rowOuterIdx * curStage2RowFactor * tilingData->hcMix;
387+ CopyInWithOuterFor(workspaceGm[mixBaseOffset], mxies01Local, 1, curRowFactor, tilingData->hcMult,
388+ tilingData->bs, tilingData->hcMix);
389+ CopyInWithOuterFor(workspaceGm[mixBaseOffset + tilingData->hcMult],
390+ mxies01Local[curRowFactor * tilingData->hcMultAlign],
391+ 1, curRowFactor, tilingData->hcMult, tilingData->bs,
392+ tilingData->hcMix);
393+ }
394+ // // wk:[2, kcorenum, bs, n^2 +2n]
395+ // // mx0[2,kcorenum, curRowfator, n]
396+ mixesQue01.EnQue(mxies01Local);
397+ mxies01Local = mixesQue01.DeQue<float>();
398+ 
399+ ProcessPre(mxies01ReduceLocal, mxies01Local, hcBase0Local, squareSumOutLocal, rowBrcbLocal0, hcBrcbLocal1,
400+ hcScaleGm.GetValue(0), tilingData->hcEps, curRowFactor, tilingData->hcMult);
401+ 
402+ if (needGrad_) {
403+ int64_t hPreBaseOffset = stage2BlockIdx * curRowOfFormerBlock * tilingData->hcMult +
404+ rowOuterIdx * curStage2RowFactor * tilingData->hcMult;
405+ VToMTE3Sync();
406+ CopyOut(mxies01ReduceLocal, hPreGm[hPreBaseOffset], curRowFactor, tilingData->hcMult);
407+ MTE3ToVSync();
408+ }
409+ // --- pre --
410+ for (int64_t dLoopIdx = 0; dLoopIdx < tilingData->dLoop; dLoopIdx++) {
411+ int64_t curDFactor =
412+ (dLoopIdx == tilingData->dLoop - 1) ? tilingData->tailDFactor : tilingData->dFactor;
413+ xLocal = xQue.template AllocTensor<T>();
414+ CopyIn(xGm[xGmBlockBaseOffsetPart2 + xGmBsBaseOffsetPart2 +
415+ dLoopIdx * tilingData->dFactor],
416+ xLocal, curRowFactor * tilingData->hcMult, curDFactor, tilingData->d - curDFactor);
417+ xQue.template EnQue(xLocal);
418+ xLocal = xQue.template DeQue<T>();
419+ yLocal = yQue.template AllocTensor<T>();
420+ 
421+ ProcessY(yLocal, xLocal, mxies01ReduceLocal, hcBrcbLocal1, xCastLocal, yCastLocal, curRowFactor,
422+ tilingData->hcMult, curDFactor);
423+ xQue.template FreeTensor(xLocal);
424+ yQue.template EnQue(yLocal);
425+ yLocal = yQue.template DeQue<T>();
426+
427+ CopyOut(yLocal,
428+ yGm[stage2BlockIdx * curRowOfFormerBlock * tilingData->d +
429+ rowOuterIdx * curStage2RowFactor * tilingData->d + dLoopIdx * tilingData->dFactor],
430+ curRowFactor, curDFactor, tilingData->d - curDFactor);
431+ yQue.template FreeTensor(yLocal);
432+ }
433+ // post
434+ postLocal = postQue.AllocTensor<float>();
435+ ProcessPost(postLocal, mxies01Local[curRowFactor * tilingData->hcMultAlign], hcBase1Local,
436+ squareSumOutLocal, rowBrcbLocal0, hcBrcbLocal1, hcScaleGm.GetValue(1), curRowFactor,
437+ tilingData->hcMult);
438+ mixesQue01.FreeTensor(mxies01Local); // 这里对应的申请在上面
439+ postQue.EnQue(postLocal);
440+ postLocal = postQue.DeQue<float>();
441+ 
442+ CopyOut(postLocal,
443+ postGm[stage2BlockIdx * curRowOfFormerBlock * tilingData->hcMult +
444+ rowOuterIdx * curStage2RowFactor * tilingData->hcMult],
445+ curRowFactor, tilingData->hcMult);
446+ postQue.FreeTensor(postLocal);
447+ 
448+ // combFrag
449+ mixes2Local = mixesQue2.AllocTensor<float>();
450+ mixBaseOffset = stage2BlockIdx * curRowOfFormerBlock * tilingData->hcMix +
451+ rowOuterIdx * curStage2RowFactor * tilingData->hcMix;
452+ for (int64_t j = 0; j < curRowFactor; ++j) {
453+ CopyIn(hcBeforeNormGm[mixBaseOffset + j * tilingData->hcMix + tilingData->hcMult * NUM_TWO],
454+ mixes2Local[j * tilingData->hcMult * tilingData->hcMultAlign], tilingData->hcMult,
455+ tilingData->hcMult);
456+ }
457+ mixesQue2.EnQue(mixes2Local);
458+ mixes2Local = mixesQue2.DeQue<float>();
459+ 
460+ combFragLocal = combFragQue.AllocTensor<float>();
461+ MulABLastDimBrcInline<float, false>(mxies02ReduceLocal, mixes2Local, rsqrtLocal, rowBrcbLocal0, curRowFactor,
462+ tilingData->hcMult * tilingData->hcMultAlign);
463+ mixesQue2.FreeTensor(mixes2Local);
464+ Muls(mxies02ReduceLocal, mxies02ReduceLocal, hcScaleGm.GetValue(NUM_TWO),
465+ curRowFactor * tilingData->hcMult * tilingData->hcMultAlign);
466+ PipeBarrier<PIPE_V>();
467+ AddBAFirstDimBrcInline<float>(mxies02ReduceLocal, mxies02ReduceLocal, hcBase2Local, curRowFactor,
468+ tilingData->hcMult * tilingData->hcMultAlign);
469+ SoftmaxFP32Perf(mxies02ReduceLocal, mxies02ReduceLocal, reduceLocal, hcBrcbLocal1, curRowFactor * tilingData->hcMult,
470+ tilingData->hcMult, tilingData->hcEps);
471+
472+ if (needGrad_) {
473+ int64_t normOutBaseOffset = (stage2BlockIdx * curRowOfFormerBlock +
474+ rowOuterIdx * curStage2RowFactor) * tilingData->hcMult * tilingData->hcMult;
475+ VToMTE3Sync();
476+ CopyOut(mxies02ReduceLocal, normOutGm[normOutBaseOffset], curRowFactor * tilingData->hcMult, tilingData->hcMult);
477+ MTE3ToVSync();
478+ }
479+ 
480+ if (needGrad_) {
481+ VToMTE3Sync();
482+ int64_t sumOutBaseOffset = (stage2BlockIdx * curRowOfFormerBlock +
483+ rowOuterIdx * curStage2RowFactor) * tilingData->hcMult;
484+ CopyOut(reduceLocal, sumOutGm[sumOutBaseOffset], curRowFactor, tilingData->hcMult);
485+ MTE3ToVSync();
486+ }
487+ 
488+ ReduceSumARAPerf(reduceLocal, mxies02ReduceLocal, curRowFactor, tilingData->hcMult, tilingData->hcMult);
489+ if (needGrad_) {
490+ VToMTE3Sync();
491+ int64_t sumOutBaseOffset = (tilingData->bs +
492+ stage2BlockIdx * curRowOfFormerBlock +
493+ rowOuterIdx * curStage2RowFactor) * tilingData->hcMult;
494+ CopyOut(reduceLocal, sumOutGm[sumOutBaseOffset], curRowFactor, tilingData->hcMult);
495+ MTE3ToVSync();
496+ }
497+ Adds(reduceLocal, reduceLocal, tilingData->hcEps, curRowFactor * tilingData->hcMult);
498+ PipeBarrier<PIPE_V>();
499+ DivABABrcInline(combFragLocal, mxies02ReduceLocal, reduceLocal, curRowFactor, tilingData->hcMult,
500+ tilingData->hcMult);
501+ 
502+ if (needGrad_) {
503+ VToMTE3Sync();
504+ int64_t normOutBaseOffset = (tilingData->bs +
505+ stage2BlockIdx * curRowOfFormerBlock +
506+ rowOuterIdx * curStage2RowFactor) * tilingData->hcMult * tilingData->hcMult;
507+ CopyOut(combFragLocal, normOutGm[normOutBaseOffset], curRowFactor * tilingData->hcMult,
508+ tilingData->hcMult);
509+ MTE3ToVSync();
510+ }
511+ 
512+ // 循环内norm输出通过拼接,减少搬运
513+ for (int64_t iter = 1; iter < tilingData->iterTimes; iter++) {
514+ auto reduceLocal2 = reduceQue.AllocTensor<float>();
515+ LastDimReduceSumPerf(reduceLocal2, combFragLocal, curRowFactor * tilingData->hcMult, tilingData->hcMult);
516+ Adds(reduceLocal2, reduceLocal2, tilingData->hcEps, curRowFactor * tilingData->hcMult);
517+ PipeBarrier<PIPE_V>();
518+ DivABLastDimBrcInline<float, true>(combFragLocal, combFragLocal, reduceLocal2, hcBrcbLocal1,
519+ curRowFactor * tilingData->hcMult, tilingData->hcMult);
520+ if (needGrad_) {
521+ reduceQue.EnQue(reduceLocal2);
522+ reduceLocal2 = reduceQue.DeQue<float>();
523+ int64_t sumOutBaseOffset = ((iter * 2) * tilingData->bs +
524+ stage2BlockIdx * curRowOfFormerBlock +
525+ rowOuterIdx * curStage2RowFactor) * tilingData->hcMult;
526+ CopyOut(reduceLocal2, sumOutGm[sumOutBaseOffset], curRowFactor, tilingData->hcMult);
527+ }
528+ reduceQue.template FreeTensor(reduceLocal2);
529+ 
530+ if (needGrad_) {
531+ int64_t tileRepeatTimes_ = CeilDiv(tilingData->hcMult * curStage2RowFactor * 2 * tilingData->hcMult, 64);
532+ PipeBarrier<PIPE_V>();
533+ Copy(combFragBufLocal, combFragLocal, curStage2RowFactor * tilingData->hcMult * tilingData->hcMultAlign, 1, {1, 1, 0, 0});
534+ PipeBarrier<PIPE_V>();
535+ Cast(yCastLocal.template ReinterpretCast<int64_t>(),
536+ combFragBufLocal.template ReinterpretCast<int32_t>(), RoundMode::CAST_NONE,
537+ curStage2RowFactor * tilingData->hcMult * 8);
538+ PipeBarrier<PIPE_V>();
539+ 
540+ Copy(combFragBufLocal, yCastLocal, 64, tileRepeatTimes_, {1, 2, 8, 16});
541+ PipeBarrier<PIPE_V>();
542+ 
543+ Cast(combFragBufLocal.template ReinterpretCast<int32_t>(),
544+ combFragBufLocal.template ReinterpretCast<int64_t>(), RoundMode::CAST_NONE,
545+ curStage2RowFactor * tilingData->hcMult * tilingData->hcMult);
546+ PipeBarrier<PIPE_V>();
547+ 
548+ VToMTE3Sync();
549+ int64_t normOutBaseOffset = ((iter * 2) * tilingData->bs +
550+ stage2BlockIdx * curRowOfFormerBlock +
551+ rowOuterIdx * curStage2RowFactor) * tilingData->hcMult * tilingData->hcMult;
552+ CopyOut(combFragBufLocal, normOutGm[normOutBaseOffset], 1, curRowFactor * tilingData->hcMult *tilingData->hcMult);
553+ MTE3ToVSync();
554+ }
555+ auto reduceLocal3 = reduceQue.AllocTensor<float>();
556+ ReduceSumARAPerf(reduceLocal3, combFragLocal, curRowFactor, tilingData->hcMult, tilingData->hcMult);
557+ Adds(reduceLocal3, reduceLocal3, tilingData->hcEps, curRowFactor * tilingData->hcMult);
558+ PipeBarrier<PIPE_V>();
559+ 
560+ if (needGrad_) {
561+ reduceQue.EnQue(reduceLocal3);
562+ reduceLocal3 = reduceQue.DeQue<float>();
563+ int64_t sumOutBaseOffset = ((iter * 2 + 1) * tilingData->bs +
564+ stage2BlockIdx * curRowOfFormerBlock +
565+ rowOuterIdx * curStage2RowFactor) * tilingData->hcMult;
566+ CopyOut(reduceLocal3, sumOutGm[sumOutBaseOffset], curRowFactor, tilingData->hcMult);
567+ }
568+ reduceQue.template FreeTensor(reduceLocal3);
569+ DivABABrcInline(combFragLocal, combFragLocal, reduceLocal3, curRowFactor, tilingData->hcMult,
570+ tilingData->hcMult);
571+ 
572+ if (needGrad_) {
573+ int64_t tileRepeatTimes_ = CeilDiv(tilingData->hcMult * curStage2RowFactor * 2 * tilingData->hcMult, 64);
574+ PipeBarrier<PIPE_V>();
575+ Copy(combFragBufLocal, combFragLocal, curStage2RowFactor * tilingData->hcMult * tilingData->hcMultAlign, 1, {1, 1, 0, 0});
576+ PipeBarrier<PIPE_V>();
577+ Cast(yCastLocal.template ReinterpretCast<int64_t>(),
578+ combFragBufLocal.template ReinterpretCast<int32_t>(), RoundMode::CAST_NONE,
579+ curStage2RowFactor * tilingData->hcMult * 8);
580+ PipeBarrier<PIPE_V>();
581+ 
582+ Copy(combFragBufLocal, yCastLocal, 64, tileRepeatTimes_, {1, 2, 8, 16});
583+ PipeBarrier<PIPE_V>();
584+ 
585+ Cast(combFragBufLocal.template ReinterpretCast<int32_t>(),
586+ combFragBufLocal.template ReinterpretCast<int64_t>(), RoundMode::CAST_NONE,
587+ curStage2RowFactor * tilingData->hcMult * tilingData->hcMult);
588+ PipeBarrier<PIPE_V>();
589+ 
590+ VToMTE3Sync();
591+ int64_t normOutBaseOffset = ((iter * 2 + 1) * tilingData->bs +
592+ stage2BlockIdx * curRowOfFormerBlock +
593+ rowOuterIdx * curStage2RowFactor) * tilingData->hcMult * tilingData->hcMult;
594+ CopyOut(combFragBufLocal, normOutGm[normOutBaseOffset], 1, curRowFactor * tilingData->hcMult *tilingData->hcMult);
595+ MTE3ToVSync();
596+ }
597+ }
598+
599+ squareSumQue.template FreeTensor(squareSumOutLocal);
600+ 
601+ combFragQue.EnQue(combFragLocal);
602+ combFragLocal = combFragQue.DeQue<float>();
603+ CopyOut(combFragLocal,
604+ combFragGm[stage2BlockIdx * curRowOfFormerBlock * tilingData->hcMult * tilingData->hcMult +
605+ rowOuterIdx * curStage2RowFactor * tilingData->hcMult * tilingData->hcMult],
606+ curRowFactor * tilingData->hcMult, tilingData->hcMult);
607+ combFragQue.FreeTensor(combFragLocal);
608+ }
609+ }
610+ }
611+ 
612+private:
613+ TPipe *pipe;
614+ const MhcPreSinkhornTilingData *tilingData;
615+ GlobalTensor<float> mixesGm;
616+ GlobalTensor<float> rsqrtGm;
617+ GlobalTensor<float> hcScaleGm;
618+ GlobalTensor<float> hcBaseGm;
619+ GlobalTensor<float> workspaceGm;
620+ GlobalTensor<T> xGm;
621+ GlobalTensor<T> yGm;
622+ GlobalTensor<float> postGm;
623+ GlobalTensor<float> combFragGm;
624+ 
625+ GlobalTensor<float> hPreGm;
626+ GlobalTensor<float> hcBeforeNormGm;
627+ GlobalTensor<float> invRmsGm;
628+ GlobalTensor<float> sumOutGm;
629+ GlobalTensor<float> normOutGm;
630+ 
631+ bool needGrad_ = false;
632+ bool isTailBsLoop_ = false;
633+ 
634+ TQue<QuePosition::VECIN, 1> mixesQue01;
635+ TQue<QuePosition::VECIN, 1> mixesQue2;
636+ TQue<QuePosition::VECIN, 1> xQue;
637+ TQue<QuePosition::VECOUT, 1> yQue;
638+ TQue<QuePosition::VECOUT, 1> postQue;
639+ TQue<QuePosition::VECOUT, 1> combFragQue;
640+ 
641+ TQue<QuePosition::VECOUT, 1> reduceQue;
642+ TQue<QuePosition::VECIN, 1> squareSumQue;
643+ 
644+ TBuf<QuePosition::VECCALC> hcBaseBuf0;
645+ TBuf<QuePosition::VECCALC> hcBaseBuf1;
646+ TBuf<QuePosition::VECCALC> hcBaseBuf2;
647+ 
648+ TBuf<QuePosition::VECCALC> rowBrcbBuf0;
649+ TBuf<QuePosition::VECCALC> hcBrcbBuf1;
650+ TBuf<QuePosition::VECCALC> reduceBuf;
651+ 
652+ TBuf<QuePosition::VECCALC> rsqrtBuf;
653+ TBuf<QuePosition::VECCALC> squareReduceBuf;
654+ TBuf<QuePosition::VECCALC> mxies01ReduceBuf;
655+ TBuf<QuePosition::VECCALC> mxies02ReduceBuf;
656+ 
657+ TBuf<QuePosition::VECCALC> xCastBuf;
658+ TBuf<QuePosition::VECCALC> yCastBuf;
659+ TBuf<QuePosition::VECCALC> maskPatternBuf;
660+ TBuf<QuePosition::VECCALC> combFragBuf;
661+ 
662+ LocalTensor<float> mxies01Local;
663+ LocalTensor<float> mixes2Local;
664+ LocalTensor<float> rsqrtLocal;
665+ LocalTensor<T> xLocal;
666+ LocalTensor<T> yLocal;
667+ LocalTensor<float> postLocal;
668+ LocalTensor<float> combFragLocal;
669+ LocalTensor<float> hcBase0Local;
670+ LocalTensor<float> hcBase1Local;
671+ LocalTensor<float> hcBase2Local;
672+ LocalTensor<float> rowBrcbLocal0;
673+ LocalTensor<float> hcBrcbLocal1;
674+ LocalTensor<float> reduceLocal;
675+ LocalTensor<float> squareReduceLocal;
676+ LocalTensor<float> mxies01ReduceLocal;
677+ LocalTensor<float> mxies02ReduceLocal;
678+ LocalTensor<float> xCastLocal;
679+ LocalTensor<float> yCastLocal;
680+ LocalTensor<float> squareSumOutLocal;
681+ LocalTensor<uint32_t> maskPatternLocal;
682+ 
683+ LocalTensor<float> reduceLocal2;
684+ LocalTensor<float> combFragBufLocal;
685+};
686+ 
687+} // namespace MhcPreSinkhornSinkhorn
688+ 
689+#endif