DeepNorm
Function Usage
During the training process of a deep neural network, DeepNorm can be used as a replacement for LayerNorm normalization in order to improve the stability of Transformers by expanding residual connections.
This API applies DeepNorm normalization to input data with a shape size of [B, S, H]. The calculation formula is as follows:
DeepNorm(x) = LayerNorm(α * X + SubLayer(X))
SubLayer(X) typically refers to a sub-layer in the DeepNorm model, and is used to implement the self-attention mechanism. This API will be passed as an input tensor.
Prototype
- Pass the temporary space through the sharedTmpBuffer input parameter.
1 2
template <typename T, bool isReuseSrc = false, bool isBasicBlock = false> __aicore__ inline void DeepNorm(const LocalTensor<T>& dstLocal, const LocalTensor<T>& meanLocal, const LocalTensor<T>& rstdLocal, const LocalTensor<T>& srcLocal, const LocalTensor<T>& gxLocal, const LocalTensor<T>& betaLocal, const LocalTensor<T>& gammaLocal, const LocalTensor<uint8_t>& sharedTmpBuffer, const T alpha, const T epsilon, DeepNormTiling& tiling)
- Allocate the temporary space through the API framework.
1 2
template <typename T, bool isReuseSrc = false, bool isBasicBlock = false> __aicore__ inline void DeepNorm(const LocalTensor<T>& dstLocal, const LocalTensor<T>& meanLocal, const LocalTensor<T>& rstdLocal, const LocalTensor<T>& srcLocal, const LocalTensor<T>& gxLocal, const LocalTensor<T>& betaLocal, const LocalTensor<T>& gammaLocal, const T alpha, const T epsilon, DeepNormTiling& tiling)
Parameters
Parameter |
Description |
|---|---|
T |
Data type of the operand. |
isReuseSrc |
Whether the source operand can be modified. The default value is false. If developers permit the modification of the source operand, enable this parameter, and this can save certain buffer space. If this parameter is set to true, the buffer space of srcLocal is reused during internal computation of this API to save the buffer space. If this parameter is set to false, the buffer space of srcLocal is not reused during internal computation of this API. This parameter can be enabled for float data inputs but cannot be enabled for half data inputs. For details about how to use isReuseSource, see Example 3. |
isBasicBlock |
If the shape information of srcTensor meets the basic block requirements, this parameter can be enabled to improve performance. By default, this parameter is disabled. For basic blocks, the shape of srcTensor must meet the following requirements:
|
Parameter |
Input/Output |
Description |
|---|---|---|
dstLocal |
Output |
Destination operand, with a type of LocalTensor and a shape of [B, S, H]. The length of H cannot exceed 2040. |
meanLocal |
Output |
Mean, destination operand, with a type of LocalTensor and a shape of [B, S]. The data type of meanLocal must be the same as that of dstLocal. |
rstdLocal |
Output |
Variance, destination operand, with a type of LocalTensor and a shape of [B, S]. The data type of rstdLocal must be the same as that of dstLocal. |
srcLocal |
Input |
Source operand, with a type of LocalTensor and a shape of [B, S, H]. The data type of srcLocal must be the same as that of the destination operand, and the last axis length must be 32-byte aligned. The length of H cannot exceed 2040. |
gxLocal |
Input |
Source operand, with a type of LocalTensor and a shape of [B, S, H]. The data type of gxLocal must be the same as that of the destination operand, and the last axis length must be 32-byte aligned. The length of H cannot exceed 2040. This parameter corresponds to the calculation result of SubLayer(X) in the formula. |
betaLocal |
Input |
Source operand, with a type of LocalTensor and a shape of [H]. The data type of betaLocal must be the same as that of the destination operand, and the length must be 32-byte aligned. The length of H cannot exceed 2040. |
gammaLocal |
Input |
Source operand, with a type of LocalTensor and a shape of [H]. The data type of gammaLocal must be the same as that of the destination operand, and the length must be 32-byte aligned. The length of H cannot exceed 2040. |
sharedTmpBuffer |
Input |
This parameter is used to store intermediate variables during complex internal API computation and is provided by developers. For details about how to obtain the temporary space size (BufferSize), see DeepNorm Tiling. |
alpha |
Input |
Weight coefficient. The data type must be the same as that of the destination operand. |
epsilon |
Input |
Weight coefficient, which is used to prevent division by zero errors. The data type must be the same as that of the destination operand. |
tiling |
Input |
Tiling information required for DeepNorm computation. For details about how to obtain the tiling information, see DeepNorm Tiling. |
Returns
None
Availability
Constraints
- For details about the alignment requirements of the operand address offset, see General Restrictions.
- The tensor space of srcLocal and dstLocal can be reused.
- The input shape must be in ND format.
- If the input data does not meet the alignment requirements, developers need to pad the data. The padded data should be set to 0 to prevent abnormal values from affecting network computation.
Example
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | #include "kernel_operator.h" template <typename dataType, bool isReuseSrc = false, bool isBasicBlock = false> class KernelDeepNorm { public: __aicore__ inline KernelDeepNorm() {} __aicore__ inline void Init(GM_ADDR inputGm, GM_ADDR inputGxGm, GM_ADDR betaGm, GM_ADDR gammaGm, GM_ADDR outputGm, GM_ADDR outputMeanGm, GM_ADDR outputVarianceGm, const DeepNormCustomTiling &customTiling) { this->tiling = customTiling.tiling; // DeepNormTiling alpha = customTiling.alpha; epsilon = customTiling.epsilon; const uint32_t bLength = tiling.bLength; const uint32_t sLength = tiling.sLength; hLength = tiling.hLength; bshLength = bLength * sLength * hLength; bsLength = bLength * sLength; inputXGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(inputGm), bshLength); inputGxGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(inputGxGm), bshLength); betaGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(betaGm), hLength); gammaGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(gammaGm), hLength); outputGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(outputGm), bshLength); outputMeanGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(outputMeanGm), bsLength); outputVarianceGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(outputVarianceGm), bsLength); constexpr uint32_t typeSize = sizeof(dataType); pipe.InitBuffer(inQueueX, 1, bshLength * typeSize); pipe.InitBuffer(inQueueGx, 1, bshLength * typeSize); pipe.InitBuffer(inQueueBeta, 1, hLength * typeSize); pipe.InitBuffer(inQueueGamma, 1, hLength * typeSize); pipe.InitBuffer(outQueue, 1, bshLength * typeSize); pipe.InitBuffer(outMeanQueue, 1, bsLength * typeSize); pipe.InitBuffer(outVarianceQueue, 1, bsLength * typeSize); } __aicore__ inline void Process() { CopyIn(); Compute(); CopyOut(); } private: __aicore__ inline void CopyIn() { AscendC::LocalTensor<dataType> inputXLocal = inQueueX.AllocTensor<dataType>(); AscendC::LocalTensor<dataType> inputGxLocal = inQueueGx.AllocTensor<dataType>(); AscendC::LocalTensor<dataType> betaLocal = inQueueBeta.AllocTensor<dataType>(); AscendC::LocalTensor<dataType> gammaLocal = inQueueGamma.AllocTensor<dataType>(); AscendC::DataCopy(inputXLocal, inputXGlobal, bshLength); AscendC::DataCopy(inputGxLocal, inputGxGlobal, bshLength); AscendC::DataCopy(betaLocal, betaGlobal, hLength); AscendC::DataCopy(gammaLocal, gammaGlobal, hLength); inQueueX.EnQue(inputXLocal); inQueueGx.EnQue(inputGxLocal); inQueueBeta.EnQue(betaLocal); inQueueGamma.EnQue(gammaLocal); } __aicore__ inline void Compute() { AscendC::LocalTensor<dataType> inputXLocal = inQueueX.DeQue<dataType>(); AscendC::LocalTensor<dataType> inputGxLocal = inQueueGx.DeQue<dataType>(); AscendC::LocalTensor<dataType> betaLocal = inQueueBeta.DeQue<dataType>(); AscendC::LocalTensor<dataType> gammaLocal = inQueueGamma.DeQue<dataType>(); AscendC::LocalTensor<dataType> outputLocal = outQueue.AllocTensor<dataType>(); AscendC::LocalTensor<dataType> outputMeanLocal = outMeanQueue.AllocTensor<dataType>(); AscendC::LocalTensor<dataType> outputVarianceLocal = outVarianceQueue.AllocTensor<dataType>(); AscendC::DeepNorm<dataType, isReuseSrc, isBasicBlock>(outputLocal, outputMeanLocal, outputVarianceLocal, inputXLocal, inputGxLocal, betaLocal, gammaLocal, alpha, epsilon, tiling); inQueueX.FreeTensor(inputXLocal); inQueueGx.FreeTensor(inputGxLocal); inQueueBeta.FreeTensor(betaLocal); inQueueGamma.FreeTensor(gammaLocal); outQueue.EnQue(outputLocal); outMeanQueue.EnQue(outputMeanLocal); outVarianceQueue.EnQue(outputVarianceLocal); } __aicore__ inline void CopyOut() { AscendC::LocalTensor<dataType> outputLocal = outQueue.DeQue<dataType>(); AscendC::LocalTensor<dataType> outputMeanLocal = outMeanQueue.DeQue<dataType>(); AscendC::LocalTensor<dataType> outputVarianceLocal = outVarianceQueue.DeQue<dataType>(); AscendC::DataCopy(outputGlobal, outputLocal, bshLength); AscendC::DataCopy(outputMeanGlobal, outputMeanLocal, bsLength); AscendC::DataCopy(outputVarianceGlobal, outputVarianceLocal, bsLength); outQueue.FreeTensor(outputLocal); outMeanQueue.FreeTensor(outputMeanLocal); outVarianceQueue.FreeTensor(outputVarianceLocal); } private: AscendC::GlobalTensor<dataType> inputXGlobal; AscendC::GlobalTensor<dataType> inputGxGlobal; AscendC::GlobalTensor<dataType> betaGlobal; AscendC::GlobalTensor<dataType> gammaGlobal; AscendC::GlobalTensor<dataType> outputGlobal; AscendC::GlobalTensor<dataType> outputMeanGlobal; AscendC::GlobalTensor<dataType> outputVarianceGlobal; AscendC::TPipe pipe; AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueX; AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueGx; AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueBeta; AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueGamma; AscendC::TQue<AscendC::QuePosition::VECOUT, 1> outQueue; AscendC::TQue<AscendC::QuePosition::VECOUT, 1> outMeanQueue; AscendC::TQue<AscendC::QuePosition::VECOUT, 1> outVarianceQueue; DeepNormTiling tiling; uint32_t bshLength; uint32_t bsLength; uint32_t hLength; dataType alpha; dataType epsilon; }; template <typename dataType, bool isReuseSrc = false, bool isBasicBlock = false> __aicore__ inline void kernel_deepnorm_operator(GM_ADDR inputGm, GM_ADDR inputGxGm, GM_ADDR betaGm, GM_ADDR gammaGm, GM_ADDR outputGm, GM_ADDR outputMeanGm, GM_ADDR outputVarianceGm, GM_ADDR customTiling) { GET_TILING_DATA(tilingData, customTiling) KernelDeepNorm<dataType, isReuseSrc, isBasicBlock> op; op.Init(inputGm, inputGxGm, betaGm, gammaGm, outputGm, outputMeanGm, outputVarianceGm, tilingData); op.Process(); } |