Merged
refactor(segsum/exp_segsum_grad): 重命名op_kernel头文件区分arch22/arch35 #5615
songkai111Add on Sep 12
refactor(segsum/exp_segsum_grad): 重命名op_kernel头文件区分arch22/arch35 #5615
Merged
songkai111Add on Sep 12
Total: 6 files changed+39-43
@@ -14,7 +14,7 @@
14 * Reuses the ported A2 compute logic; tiling data comes from a POD struct.14 * Reuses the ported A2 compute logic; tiling data comes from a POD struct.
15 */15 */
16 16 
17-#include "exp_segsum_grad.h"17+#include "exp_segsum_grad_arch35.h"
18 18 
19using namespace ExpSegsumGradArch35;19using namespace ExpSegsumGradArch35;
20 20 
Rmath/exp_segsum_grad/op_kernel/arch35/exp_segsum_grad.h→math/exp_segsum_grad/op_kernel/arch35/exp_segsum_grad_arch35.h+1-1
@@ -9,7 +9,7 @@
9 */9 */
10 10 
11/*!11/*!
12- * \file exp_segsum_grad.h (arch35 / Ascend950)12+ * \file exp_segsum_grad_arch35.h (arch35 / Ascend950)
13 * \brief arch35 port of ExpSegsumGrad.13 * \brief arch35 port of ExpSegsumGrad.
14 */14 */
15 15 
@@ -1,10 +1,10 @@
1/**1/**
2- * This program is free software, you can redistribute it and/or modify it.
3 * Copyright (c) 2025 Huawei Technologies Co., Ltd.2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4- * This file is a part of the CANN Open Software.3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5- * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").4+ * CANN Open Software License Agreement Version 2.0 (the "License").
6 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
7- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.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.8 * See LICENSE in the root of the software repository for the full text of the License.
9 */9 */
10 10 
@@ -13,19 +13,20 @@
13 * \brief13 * \brief
14 */14 */
15 15 
16-#include "exp_segsum_grad.h"16+#include "exp_segsum_grad_arch22.h"
17 17 
18using namespace ExpSegsumGrad;18using namespace ExpSegsumGrad;
19 19 
20extern "C" __global__ __aicore__ void exp_segsum_grad(GM_ADDR grad_output, GM_ADDR grad_self, GM_ADDR grad_input,20extern "C" __global__ __aicore__ void exp_segsum_grad(GM_ADDR grad_output, GM_ADDR grad_self, GM_ADDR grad_input,
21- GM_ADDR workspace, GM_ADDR tiling) {21+ GM_ADDR workspace, GM_ADDR tiling)
22+{
22 GET_TILING_DATA(tilingData, tiling);23 GET_TILING_DATA(tilingData, tiling);
23 24 
24 GM_ADDR userWS = GetUserWorkspace(workspace);25 GM_ADDR userWS = GetUserWorkspace(workspace);
25 if (userWS == nullptr) {26 if (userWS == nullptr) {
26 return;27 return;
27 }28 }
28- 29+ 
29 if (TILING_KEY_IS(0)) {30 if (TILING_KEY_IS(0)) {
30 ExpSegsumGradND<DTYPE_GRAD_OUTPUT, 0> op;31 ExpSegsumGradND<DTYPE_GRAD_OUTPUT, 0> op;
31 op.Init(grad_output, grad_self, grad_input, workspace, &tilingData);32 op.Init(grad_output, grad_self, grad_input, workspace, &tilingData);
@@ -37,4 +38,4 @@ extern "C" __global__ __aicore__ void exp_segsum_grad(GM_ADDR grad_output, GM_AD
37 } else {38 } else {
38 return;39 return;
39 }40 }
40-}41+}
Rmath/exp_segsum_grad/op_kernel/exp_segsum_grad.h→math/exp_segsum_grad/op_kernel/exp_segsum_grad_arch22.h+3-3
@@ -9,12 +9,12 @@
9 */9 */
10 10 
11/*!11/*!
12- * \file exp_segsum_grad.h12+ * \file exp_segsum_grad_arch22.h
13 * \brief13 * \brief
14 */14 */
15 15 
16-#ifndef EXP_SEGSUM_GRAD_H16+#ifndef EXP_SEGSUM_GRAD_ARCH22_H
17-#define EXP_SEGSUM_GRAD_H17+#define EXP_SEGSUM_GRAD_ARCH22_H
18 18 
19#include <type_traits>19#include <type_traits>
20#include "kernel_operator.h"20#include "kernel_operator.h"
@@ -13,7 +13,7 @@
13 * \brief13 * \brief
14 */14 */
15 15 
16-#include "segsum.h"16+#include "segsum_arch22.h"
17 17 
18using namespace Segsum;18using namespace Segsum;
19 19 
@@ -65,4 +65,4 @@ extern "C" __global__ __aicore__ void segsum(GM_ADDR x, GM_ADDR y, GM_ADDR works
65 }65 }
66#endif66#endif
67 }67 }
68-}68+}
Rmath/segsum/op_kernel/segsum.h→math/segsum/op_kernel/segsum_arch22.h+23-28
@@ -9,12 +9,12 @@
9 */9 */
10 10 
11/*!11/*!
12- * \file segsum.h12+ * \file segsum_arch22.h
13 * \brief13 * \brief
14 */14 */
15 15 
16-#ifndef SEGSUM16+#ifndef SEGSUM_ARCH22_H
17-#define SEGSUM17+#define SEGSUM_ARCH22_H
18 18 
19#include <type_traits>19#include <type_traits>
20#include "kernel_operator.h"20#include "kernel_operator.h"
@@ -29,8 +29,7 @@ constexpr int32_t NO_BUFFER_NUM = 1;
29 29 
30constexpr float INF_FLOAT = -INFINITY;30constexpr float INF_FLOAT = -INFINITY;
31template <typename T, int32_t MODE>31template <typename T, int32_t MODE>
32-class SegsumND32+class SegsumND {
33-{
34public:33public:
35 TPipe pipe;34 TPipe pipe;
36 35 
@@ -69,12 +68,10 @@ private:
69 __aicore__ inline void Compute(int64_t batchIdx, int64_t colIdx);68 __aicore__ inline void Compute(int64_t batchIdx, int64_t colIdx);
70 __aicore__ inline void CopyInBatch(int64_t offset, int64_t batchesEachCopy, uint32_t calSize);69 __aicore__ inline void CopyInBatch(int64_t offset, int64_t batchesEachCopy, uint32_t calSize);
71 __aicore__ inline void CopyOutBatch(int64_t offset, int64_t batchesEachCopy, int64_t calNum, uint32_t calSize);70 __aicore__ inline void CopyOutBatch(int64_t offset, int64_t batchesEachCopy, int64_t calNum, uint32_t calSize);
72- __aicore__ inline void processY(71+ __aicore__ inline void processY(int64_t outProcessedNum, int64_t calCount, LocalTensor<float> currentRow,
73- int64_t outProcessedNum, int64_t calCount, LocalTensor<float> currentRow, int64_t batchesEachCopy,72+ int64_t batchesEachCopy, int64_t blockNumData);
74- int64_t blockNumData);73+ __aicore__ inline void processCurrent(int64_t dataIdx, LocalTensor<float> currentRow, LocalTensor<T> xTensor,
75- __aicore__ inline void processCurrent(74+ int64_t batchesEachCopy, int64_t blockNumData);
76- int64_t dataIdx, LocalTensor<float> currentRow, LocalTensor<T> xTensor, int64_t batchesEachCopy,
77- int64_t blockNumData);
78 75 
79private:76private:
80 TBuf<QuePosition::VECCALC> lastQueue;77 TBuf<QuePosition::VECCALC> lastQueue;
@@ -98,8 +95,8 @@ private:
98};95};
99 96 
100template <typename T, int32_t MODE>97template <typename T, int32_t MODE>
101-__aicore__ inline void SegsumND<T, MODE>::Init(98+__aicore__ inline void SegsumND<T, MODE>::Init(GM_ADDR input, GM_ADDR output, GM_ADDR workspace,
102- GM_ADDR input, GM_ADDR output, GM_ADDR workspace, const SegsumTilingData* tilingData)99+ const SegsumTilingData* tilingData)
103{100{
104 blockIdx = GetBlockIdx();101 blockIdx = GetBlockIdx();
105 ParseTilingData(tilingData);102 ParseTilingData(tilingData);
@@ -191,9 +188,9 @@ __aicore__ inline void SegsumND<T, MODE>::Compute(int64_t batchIdx, int64_t colI
191}188}
192 189 
193template <typename T, int32_t MODE>190template <typename T, int32_t MODE>
194-__aicore__ inline void SegsumND<T, MODE>::processY(191+__aicore__ inline void SegsumND<T, MODE>::processY(int64_t outProcessedNum, int64_t calCount,
195- int64_t outProcessedNum, int64_t calCount, LocalTensor<float> currentRow, int64_t batchesEachCopy,192+ LocalTensor<float> currentRow, int64_t batchesEachCopy,
196- int64_t blockNumData)193+ int64_t blockNumData)
197{194{
198 if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) {195 if constexpr (std::is_same<T, bfloat16_t>::value || std::is_same<T, half>::value) {
199 LocalTensor<T> yTensor = outQueue.AllocTensor<T>();196 LocalTensor<T> yTensor = outQueue.AllocTensor<T>();
@@ -228,9 +225,9 @@ __aicore__ inline void SegsumND<T, MODE>::processY(
228}225}
229 226 
230template <typename T, int32_t MODE>227template <typename T, int32_t MODE>
231-__aicore__ inline void SegsumND<T, MODE>::processCurrent(228+__aicore__ inline void SegsumND<T, MODE>::processCurrent(int64_t dataIdx, LocalTensor<float> currentRow,
232- int64_t dataIdx, LocalTensor<float> currentRow, LocalTensor<T> xTensor, int64_t batchesEachCopy,229+ LocalTensor<T> xTensor, int64_t batchesEachCopy,
233- int64_t blockNumData)230+ int64_t blockNumData)
234{231{
235 if constexpr (std::is_same<T, bfloat16_t>::value) {232 if constexpr (std::is_same<T, bfloat16_t>::value) {
236 for (int64_t i = 0; i < batchesEachCopy; i++) {233 for (int64_t i = 0; i < batchesEachCopy; i++) {
@@ -239,9 +236,8 @@ __aicore__ inline void SegsumND<T, MODE>::processCurrent(
239 }236 }
240 } else {237 } else {
241 for (int64_t i = 0; i < batchesEachCopy; i++) {238 for (int64_t i = 0; i < batchesEachCopy; i++) {
242- Duplicate(239+ Duplicate(currentRow[i * blockNumData], static_cast<float>(xTensor.GetValue(dataIdx + i * blockNumData)),
243- currentRow[i * blockNumData], static_cast<float>(xTensor.GetValue(dataIdx + i * blockNumData)),240+ dataIdx);
244- dataIdx);
245 PipeBarrier<PIPE_V>();241 PipeBarrier<PIPE_V>();
246 }242 }
247 }243 }
@@ -295,13 +291,12 @@ __aicore__ inline void SegsumND<T, MODE>::CopyInBatch(int64_t offset, int64_t ba
295 inQueue.EnQue(xTensor);291 inQueue.EnQue(xTensor);
296}292}
297template <typename T, int32_t MODE>293template <typename T, int32_t MODE>
298-__aicore__ inline void SegsumND<T, MODE>::CopyOutBatch(294+__aicore__ inline void SegsumND<T, MODE>::CopyOutBatch(int64_t offset, int64_t batchesEachCopy, int64_t calNum,
299- int64_t offset, int64_t batchesEachCopy, int64_t calNum, uint32_t calSize)295+ uint32_t calSize)
300{296{
301 LocalTensor<T> yTensor = outQueue.DeQue<T>();297 LocalTensor<T> yTensor = outQueue.DeQue<T>();
302- DataCopyExtParams copyParams{298+ DataCopyExtParams copyParams{static_cast<uint16_t>(batchesEachCopy), calSize, 0,
303- static_cast<uint16_t>(batchesEachCopy), calSize, 0, static_cast<uint32_t>((calNum - 1) * calNum * sizeof(T)),299+ static_cast<uint32_t>((calNum - 1) * calNum * sizeof(T)), 0};
304- 0};
305 DataCopyPad(outTensorsGM[offset], yTensor, copyParams);300 DataCopyPad(outTensorsGM[offset], yTensor, copyParams);
306 outQueue.FreeTensor(yTensor);301 outQueue.FreeTensor(yTensor);
307}302}
@@ -318,4 +313,4 @@ __aicore__ inline void SegsumND<T, MODE>::ParseTilingData(const SegsumTilingData
318 blockSize = 32 / sizeof(T);313 blockSize = 32 / sizeof(T);
319}314}
320} // namespace Segsum315} // namespace Segsum
321-#endif316+#endif