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
Merged fromsongkai111/ops-math_6421masterintocann/ops-mathmaster
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 | ||
| 19 | using namespace ExpSegsumGradArch35; | 19 | using 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 | * \brief | 13 | * \brief |
| 14 | */ | 14 | */ |
| 15 | 15 | ||
| 16 | -#include "exp_segsum_grad.h" | 16 | +#include "exp_segsum_grad_arch22.h" |
| 17 | 17 | ||
| 18 | using namespace ExpSegsumGrad; | 18 | using namespace ExpSegsumGrad; |
| 19 | 19 | ||
| 20 | extern "C" __global__ __aicore__ void exp_segsum_grad(GM_ADDR grad_output, GM_ADDR grad_self, GM_ADDR grad_input, | 20 | extern "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.h | 12 | + * \file exp_segsum_grad_arch22.h |
| 13 | * \brief | 13 | * \brief |
| 14 | */ | 14 | */ |
| 15 | 15 | ||
| 16 | -#ifndef EXP_SEGSUM_GRAD_H | 16 | +#ifndef EXP_SEGSUM_GRAD_ARCH22_H |
| 17 | -#define EXP_SEGSUM_GRAD_H | 17 | +#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 | * \brief | 13 | * \brief |
| 14 | */ | 14 | */ |
| 15 | 15 | ||
| 16 | -#include "segsum.h" | 16 | +#include "segsum_arch22.h" |
| 17 | 17 | ||
| 18 | using namespace Segsum; | 18 | using 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 | #endif | 66 | #endif |
| 67 | } | 67 | } |
| 68 | -} | 68 | +} |
| @@ -9,12 +9,12 @@ | |||
| 9 | */ | 9 | */ |
| 10 | 10 | ||
| 11 | /*! | 11 | /*! |
| 12 | - * \file segsum.h | 12 | + * \file segsum_arch22.h |
| 13 | * \brief | 13 | * \brief |
| 14 | */ | 14 | */ |
| 15 | 15 | ||
| 16 | -#ifndef SEGSUM | 16 | +#ifndef SEGSUM_ARCH22_H |
| 17 | -#define SEGSUM | 17 | +#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 | ||
| 30 | constexpr float INF_FLOAT = -INFINITY; | 30 | constexpr float INF_FLOAT = -INFINITY; |
| 31 | template <typename T, int32_t MODE> | 31 | template <typename T, int32_t MODE> |
| 32 | -class SegsumND | 32 | +class SegsumND { |
| 33 | -{ | ||
| 34 | public: | 33 | public: |
| 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 | ||
| 79 | private: | 76 | private: |
| 80 | TBuf<QuePosition::VECCALC> lastQueue; | 77 | TBuf<QuePosition::VECCALC> lastQueue; |
| @@ -98,8 +95,8 @@ private: | |||
| 98 | }; | 95 | }; |
| 99 | 96 | ||
| 100 | template <typename T, int32_t MODE> | 97 | template <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 | ||
| 193 | template <typename T, int32_t MODE> | 190 | template <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 | ||
| 230 | template <typename T, int32_t MODE> | 227 | template <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 | } |
| 297 | template <typename T, int32_t MODE> | 293 | template <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 Segsum | 315 | } // namespace Segsum |
| 321 | -#endif | 316 | +#endif |