wehub-resource-sync
2062 lines
78 KiB
C++
2062 lines
78 KiB
C++
/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||
|
||
Licensed under the Apache License, Version 2.0 (the "License");
|
||
you may not use this file except in compliance with the License.
|
||
You may obtain a copy of the License at
|
||
|
||
http://www.apache.org/licenses/LICENSE-2.0
|
||
|
||
Unless required by applicable law or agreed to in writing, software
|
||
distributed under the License is distributed on an "AS IS" BASIS,
|
||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
See the License for the specific language governing permissions and
|
||
limitations under the License. */
|
||
|
||
#pragma once
|
||
|
||
#include <iostream>
|
||
|
||
#include "glog/logging.h"
|
||
#include "paddle/common/ddim.h"
|
||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||
#include "paddle/phi/backends/gpu/gpu_device_function.h"
|
||
#include "paddle/phi/backends/gpu/gpu_dnn.h"
|
||
#include "paddle/phi/common/memory_utils.h"
|
||
#include "paddle/phi/kernels/funcs/aligned_vector.h"
|
||
#include "paddle/phi/kernels/funcs/cub.h"
|
||
#include "paddle/phi/kernels/funcs/fake_quantize_functor.h"
|
||
#include "paddle/phi/kernels/funcs/fast_ln_v1.h"
|
||
|
||
namespace phi {
|
||
namespace funcs {
|
||
|
||
template <typename T>
|
||
using CudnnDataType = phi::backends::gpu::CudnnDataType<T>;
|
||
template <typename T>
|
||
using LayerNormParamType = typename CudnnDataType<T>::BatchNormParamType;
|
||
|
||
inline static int GetDesiredBlockDim(int64_t block_dim) {
|
||
const int kMaxBlockDim = 512;
|
||
#ifdef __HIPCC__
|
||
const int lwarpSize = 64;
|
||
#else
|
||
const int lwarpSize = 32;
|
||
#endif
|
||
return block_dim >= kMaxBlockDim ? kMaxBlockDim : lwarpSize;
|
||
}
|
||
|
||
static dim3 GetDesiredGridDim(int64_t grid_size) {
|
||
dim3 grid_dim(1, 1);
|
||
int64_t grid_x = grid_size > 1 ? grid_size : 1;
|
||
int64_t grid_y = 1;
|
||
if (grid_x > 2147483648LL) {
|
||
grid_y = 1024;
|
||
grid_x = (grid_x + grid_y - 1) / grid_y;
|
||
PADDLE_ENFORCE_LE_INT_MAX(grid_x, "grid_x");
|
||
}
|
||
grid_dim.x = static_cast<uint32_t>(grid_x);
|
||
grid_dim.y = static_cast<uint32_t>(grid_y);
|
||
return grid_dim;
|
||
}
|
||
|
||
template <typename U>
|
||
static __forceinline__ __device__ U WarpReduceSum(U val) {
|
||
unsigned mask = 0u;
|
||
CREATE_SHFL_MASK(mask, true);
|
||
for (int offset = warpSize / 2; offset > 0; offset /= 2) {
|
||
val += phi::backends::gpu::CudaShuffleDownSync(mask, val, offset);
|
||
}
|
||
return val;
|
||
}
|
||
|
||
template <typename U>
|
||
__forceinline__ __device__ U BlockReduceSum(U val, U *shared) {
|
||
int lane = threadIdx.x % warpSize;
|
||
int wid = threadIdx.x / warpSize;
|
||
|
||
val = WarpReduceSum(val); // Each warp performs partial reduction
|
||
if (lane == 0) shared[wid] = val; // Write reduced value to shared memory
|
||
__syncthreads(); // Wait for all partial reductions
|
||
// read from shared memory only if that warp existed
|
||
val =
|
||
(threadIdx.x < blockDim.x / warpSize) ? shared[lane] : static_cast<U>(0);
|
||
|
||
if (wid == 0) val = WarpReduceSum(val); // Final reduce within first warp
|
||
|
||
return val;
|
||
}
|
||
|
||
#define FIXED_BLOCK_DIM_CASE_BASE(log2_block_dim, ...) \
|
||
case (1 << (log2_block_dim)): { \
|
||
constexpr auto kBlockDim = (1 << (log2_block_dim)); \
|
||
__VA_ARGS__; \
|
||
} break
|
||
|
||
#define FIXED_BLOCK_DIM_CASE(...) \
|
||
FIXED_BLOCK_DIM_CASE_BASE(9, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_CASE_BASE(8, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_CASE_BASE(7, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_CASE_BASE(6, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_CASE_BASE(5, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_CASE_BASE(4, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_CASE_BASE(3, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_CASE_BASE(2, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_CASE_BASE(1, ##__VA_ARGS__)
|
||
|
||
#define FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE_BASE( \
|
||
log2_block_dim, feature_size, kMaxBlockNum, ...) \
|
||
case (1 << (log2_block_dim)): { \
|
||
for (int64_t i = 0; i < std::ceil(feature_size / (1.0 * kMaxBlockNum)); \
|
||
i++) { \
|
||
int64_t col_offset = i * static_cast<int64_t>(kMaxBlockNum); \
|
||
int block_num = static_cast<int>(std::min( \
|
||
feature_size - col_offset, static_cast<int64_t>(kMaxBlockNum))); \
|
||
constexpr auto kBlockDim = (1 << (log2_block_dim)); \
|
||
__VA_ARGS__; \
|
||
} \
|
||
} break
|
||
|
||
#define FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE(feature_size, kMaxBlockNum, ...) \
|
||
FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE_BASE( \
|
||
9, feature_size, kMaxBlockNum, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE_BASE( \
|
||
8, feature_size, kMaxBlockNum, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE_BASE( \
|
||
7, feature_size, kMaxBlockNum, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE_BASE( \
|
||
6, feature_size, kMaxBlockNum, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE_BASE( \
|
||
5, feature_size, kMaxBlockNum, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE_BASE( \
|
||
4, feature_size, kMaxBlockNum, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE_BASE( \
|
||
3, feature_size, kMaxBlockNum, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE_BASE( \
|
||
2, feature_size, kMaxBlockNum, ##__VA_ARGS__); \
|
||
FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE_BASE( \
|
||
1, feature_size, kMaxBlockNum, ##__VA_ARGS__)
|
||
|
||
static __device__ __forceinline__ float real_sqrt(float x) { return sqrtf(x); }
|
||
static __device__ __forceinline__ double real_sqrt(double x) {
|
||
return ::sqrt(x);
|
||
}
|
||
|
||
template <typename T>
|
||
struct PairForLayerNorm {
|
||
__device__ __forceinline__ PairForLayerNorm() {}
|
||
__device__ __forceinline__ PairForLayerNorm(const T &first, const T &second)
|
||
: first_(first), second_(second) {}
|
||
|
||
T first_;
|
||
T second_;
|
||
};
|
||
|
||
template <typename T>
|
||
struct PairForLayerNormAddFunctor {
|
||
__device__ __forceinline__ PairForLayerNorm<T> operator()(
|
||
const PairForLayerNorm<T> &p1, const PairForLayerNorm<T> &p2) {
|
||
return PairForLayerNorm<T>(p1.first_ + p2.first_, p1.second_ + p2.second_);
|
||
}
|
||
};
|
||
|
||
template <typename T>
|
||
__inline__ __device__ T rsqrt_(const T val) {
|
||
return static_cast<T>(1) / sqrt(val);
|
||
}
|
||
|
||
template <>
|
||
__inline__ __device__ float rsqrt_(const float val) {
|
||
return rsqrtf(val);
|
||
}
|
||
|
||
template <>
|
||
__inline__ __device__ double rsqrt_(const double val) {
|
||
return ::rsqrt(val);
|
||
}
|
||
|
||
#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__) || defined(PADDLE_WITH_HIP)
|
||
template <>
|
||
__inline__ __device__ half rsqrt_(const half val) {
|
||
return hrsqrt(val);
|
||
}
|
||
#endif
|
||
|
||
// fast_ln_v1_fwd_kernel is moved to paddle/phi/kernels/funcs/fast_ln_v1.h
|
||
|
||
template <typename T>
|
||
__forceinline__ __device__ int8_t quant_helper(const T input,
|
||
const float scale,
|
||
const int round_type,
|
||
const float max_bound,
|
||
const float min_bound) {
|
||
float quant_value = max_bound * scale * static_cast<float>(input);
|
||
|
||
if (round_type == 0) {
|
||
quant_value = static_cast<float>(roundWithTiesToEven(quant_value));
|
||
} else {
|
||
quant_value = static_cast<float>(round(quant_value));
|
||
}
|
||
quant_value = quant_value > max_bound ? max_bound : quant_value;
|
||
quant_value = quant_value < min_bound ? min_bound : quant_value;
|
||
return static_cast<int8_t>(quant_value);
|
||
}
|
||
|
||
template <typename T, typename U, bool ScaleBiasWithSameTypeX>
|
||
using LayerNormScaleBiasT =
|
||
typename std::conditional<ScaleBiasWithSameTypeX, T, U>::type;
|
||
|
||
template <typename T,
|
||
typename U,
|
||
int BlockDim,
|
||
bool ScaleBiasWithSameTypeX = false,
|
||
typename InType = T,
|
||
typename OutType = T>
|
||
__global__ void LayerNormForward(
|
||
const InType *x,
|
||
const LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *scale,
|
||
const LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *bias,
|
||
OutType *y,
|
||
U *mean,
|
||
U *var,
|
||
float epsilon,
|
||
int64_t feature_size,
|
||
const float *dequant_out_scale_data = nullptr,
|
||
const int quant_out_scale_offset = 0,
|
||
const float quant_in_scale = 1.0,
|
||
const int quant_round_type = 1,
|
||
const float quant_max_bound = 127.0,
|
||
const float quant_min_bound = -127.0) {
|
||
__shared__ U mean_share;
|
||
__shared__ U var_share;
|
||
__shared__ U shared_mean[32]; // threadIdx.x / warpSize <= kMaxBlockDim /
|
||
// warpSize <= 1024/32 = 32;
|
||
__shared__ U shared_var[32];
|
||
|
||
int64_t beg_idx =
|
||
(static_cast<int64_t>(blockIdx.x) * gridDim.y + blockIdx.y) *
|
||
feature_size +
|
||
threadIdx.x;
|
||
int64_t end_idx =
|
||
(static_cast<int64_t>(blockIdx.x) * gridDim.y + blockIdx.y + 1) *
|
||
feature_size;
|
||
|
||
// Step 1: Reduce to calculate mean and var
|
||
U mean_val = 0;
|
||
U var_val = 0;
|
||
for (int64_t i = beg_idx; i < end_idx; i += BlockDim) {
|
||
U tmp = static_cast<U>(x[i]);
|
||
mean_val += tmp;
|
||
var_val += (tmp * tmp);
|
||
}
|
||
|
||
mean_val = BlockReduceSum<U>(mean_val, shared_mean);
|
||
var_val = BlockReduceSum<U>(var_val, shared_var);
|
||
|
||
if (threadIdx.x == 0) {
|
||
auto scale = static_cast<U>(static_cast<float>(1.) /
|
||
static_cast<float>(feature_size));
|
||
auto tmp = mean_val * scale;
|
||
mean[blockIdx.x] = mean_share = static_cast<U>(tmp);
|
||
var_share = static_cast<U>(var_val * scale - mean_share * mean_share);
|
||
var_share = var_share > U(0) ? var_share : U(0);
|
||
var[blockIdx.x] = var_share;
|
||
}
|
||
__syncthreads();
|
||
|
||
mean_val = mean_share;
|
||
U invvar = rsqrt_<U>(var_share + static_cast<U>(epsilon));
|
||
|
||
// Step 2: Calculate y
|
||
if (scale != nullptr) {
|
||
if (bias != nullptr) {
|
||
for (int64_t i = beg_idx, j = threadIdx.x; i < end_idx;
|
||
i += BlockDim, j += BlockDim) {
|
||
if (std::is_same<OutType, int8_t>::value) {
|
||
y[i] = quant_helper(
|
||
static_cast<T>(static_cast<U>(scale[j]) *
|
||
(static_cast<U>(x[i]) - mean_val) * invvar +
|
||
static_cast<U>(bias[j])),
|
||
quant_in_scale,
|
||
quant_round_type,
|
||
quant_max_bound,
|
||
quant_min_bound);
|
||
} else {
|
||
y[i] = static_cast<OutType>(static_cast<U>(scale[j]) *
|
||
(static_cast<U>(x[i]) - mean_val) *
|
||
invvar +
|
||
static_cast<U>(bias[j]));
|
||
}
|
||
}
|
||
} else {
|
||
for (int64_t i = beg_idx, j = threadIdx.x; i < end_idx;
|
||
i += BlockDim, j += BlockDim) {
|
||
if (std::is_same<OutType, int8_t>::value) {
|
||
y[i] = quant_helper(
|
||
static_cast<T>(static_cast<U>(scale[j]) *
|
||
(static_cast<U>(x[i]) - mean_val) * invvar),
|
||
quant_in_scale,
|
||
quant_round_type,
|
||
quant_max_bound,
|
||
quant_min_bound);
|
||
} else {
|
||
y[i] =
|
||
static_cast<OutType>(static_cast<U>(scale[j]) *
|
||
(static_cast<U>(x[i]) - mean_val) * invvar);
|
||
}
|
||
}
|
||
}
|
||
} else { // scale == nullptr
|
||
if (bias != nullptr) {
|
||
for (int64_t i = beg_idx, j = threadIdx.x; i < end_idx;
|
||
i += BlockDim, j += BlockDim) {
|
||
if (std::is_same<OutType, int8_t>::value) {
|
||
y[i] = quant_helper(
|
||
static_cast<T>((static_cast<U>(x[i]) - mean_val) * invvar +
|
||
static_cast<U>(bias[j])),
|
||
quant_in_scale,
|
||
quant_round_type,
|
||
quant_max_bound,
|
||
quant_min_bound);
|
||
} else {
|
||
y[i] =
|
||
static_cast<OutType>((static_cast<U>(x[i]) - mean_val) * invvar +
|
||
static_cast<U>(bias[j]));
|
||
}
|
||
}
|
||
} else {
|
||
for (int64_t i = beg_idx, j = threadIdx.x; i < end_idx;
|
||
i += BlockDim, j += BlockDim) {
|
||
if (std::is_same<OutType, int8_t>::value) {
|
||
y[i] = quant_helper(
|
||
static_cast<T>((static_cast<U>(x[i]) - mean_val) * invvar),
|
||
quant_in_scale,
|
||
quant_round_type,
|
||
quant_max_bound,
|
||
quant_min_bound);
|
||
} else {
|
||
y[i] =
|
||
static_cast<OutType>((static_cast<U>(x[i]) - mean_val) * invvar);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
template <typename T, typename U, int VPT>
|
||
__inline__ __device__ void cuLoadAddStridedInputs(const int64_t i1_block,
|
||
const int thr_load_row_off,
|
||
const int thr_load_col_off,
|
||
const int i2_off,
|
||
const int row_stride,
|
||
U *warp_buf1,
|
||
U *warp_buf2,
|
||
const T *__restrict__ input,
|
||
const T *__restrict__ dout,
|
||
const int64_t i1_end,
|
||
const int64_t n2,
|
||
const U *__restrict__ mean,
|
||
const U *__restrict__ var,
|
||
const float epsilon) {
|
||
const int64_t i1 = i1_block + thr_load_row_off;
|
||
if (i1 >= i1_end) return;
|
||
U curr_mean = mean[i1];
|
||
U curr_invvar = rsqrt_<U>(var[i1] + epsilon);
|
||
#pragma unroll
|
||
for (int k = 0; k < VPT; ++k) {
|
||
const int i2 = i2_off + k;
|
||
const int64_t load_idx = i1 * n2 + i2;
|
||
const int64_t write_idx =
|
||
static_cast<int64_t>(thr_load_row_off) * row_stride + thr_load_col_off +
|
||
k;
|
||
if (i2 < n2) {
|
||
U curr_input = static_cast<U>(input[load_idx]);
|
||
U curr_dout = static_cast<U>(dout[load_idx]);
|
||
warp_buf1[write_idx] += curr_dout;
|
||
warp_buf2[write_idx] +=
|
||
curr_dout * (curr_input - curr_mean) * curr_invvar;
|
||
}
|
||
}
|
||
}
|
||
|
||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||
template <bool IsFusedDropoutResidualLn,
|
||
bool NeedDDropoutSrcPtr,
|
||
typename T,
|
||
typename U,
|
||
typename ScaleT = U,
|
||
typename MaskType = uint8_t,
|
||
int VecSize = 8,
|
||
int WARPS_M = 4,
|
||
int WARPS_N = 1,
|
||
int BYTES_PER_LDG = 16,
|
||
int ELTS_PER_ROW = 1024,
|
||
int THREADS_PER_WARP = 32,
|
||
int THREADS_PER_ROW = WARPS_N *THREADS_PER_WARP,
|
||
int THREADS_PER_CTA = WARPS_M *THREADS_PER_ROW,
|
||
int ROWS_PER_CTA = WARPS_M,
|
||
int ELTS_PER_ROW_PER_CTA = THREADS_PER_ROW *VecSize,
|
||
int LDGS = ELTS_PER_ROW / ELTS_PER_ROW_PER_CTA>
|
||
__global__ __launch_bounds__(THREADS_PER_CTA) void fused_ln_bwd_fast_kernel(
|
||
const int rows,
|
||
float epsilon,
|
||
const T *__restrict__ x_ptr,
|
||
const ScaleT *__restrict__ gamma_ptr,
|
||
const U *__restrict__ mean_ptr,
|
||
const U *__restrict__ var_ptr,
|
||
const T *__restrict__ dout_ptr,
|
||
U *__restrict__ dgamma_temp_ptr,
|
||
U *__restrict__ dbeta_temp_ptr,
|
||
T *__restrict__ dx_ptr,
|
||
const MaskType *mask_ptr = nullptr,
|
||
T factor = static_cast<T>(0),
|
||
T *d_dropout_src_ptr = nullptr) {
|
||
static_assert(
|
||
!IsFusedDropoutResidualLn || NeedDDropoutSrcPtr,
|
||
"When IsFusedDropoutResidualLn = true, NeedDDropoutSrcPtr must be true.");
|
||
|
||
using Vec = AlignedVector<T, VecSize>;
|
||
using Vec_scale = AlignedVector<ScaleT, VecSize>;
|
||
using MaskLoadT = AlignedVector<MaskType, VecSize>;
|
||
|
||
const int64_t tidx = threadIdx.x;
|
||
const int64_t bidx = blockIdx.x;
|
||
const int64_t lane = tidx % THREADS_PER_WARP; // 0, 1, ..., 31
|
||
const int64_t warp = tidx / THREADS_PER_WARP; // 0, 1, 2, 3
|
||
const int64_t warp_m = warp / WARPS_N; // 0, 1, 2, 3
|
||
const int64_t warp_n = warp % WARPS_N; // 0
|
||
const int64_t tid_r = warp_n * THREADS_PER_WARP + lane; // 0, 1, ..., 31
|
||
|
||
const int64_t r = bidx * ROWS_PER_CTA + warp_m;
|
||
const int64_t c = warp_n * THREADS_PER_WARP + lane;
|
||
|
||
static_assert(ELTS_PER_ROW == THREADS_PER_ROW * LDGS * VecSize, "");
|
||
|
||
// smem for column reduction
|
||
__shared__ U smem_[ROWS_PER_CTA * ELTS_PER_ROW];
|
||
|
||
U dgamma_sum[LDGS * VecSize];
|
||
U dbeta_sum[LDGS * VecSize];
|
||
|
||
memset(dgamma_sum, 0, sizeof(U) * LDGS * VecSize);
|
||
memset(dbeta_sum, 0, sizeof(U) * LDGS * VecSize);
|
||
|
||
// Note: it is no use for WARP_N = 1
|
||
__shared__ U smem_sum_loss1[ROWS_PER_CTA * WARPS_N]; // 4
|
||
__shared__ U smem_sum_loss2[ROWS_PER_CTA * WARPS_N]; // 4
|
||
U *sum_loss1_shared = &smem_sum_loss1[warp_m * WARPS_N];
|
||
U *sum_loss2_shared = &smem_sum_loss2[warp_m * WARPS_N];
|
||
|
||
// step-1: compute dx and local results of dscale and dbias
|
||
constexpr float rn = 1.f / static_cast<float>(ELTS_PER_ROW);
|
||
Vec_scale gamma[LDGS];
|
||
int64_t col = c;
|
||
#pragma unroll
|
||
for (int64_t it = 0; it < LDGS; it++) {
|
||
phi::Load<ScaleT, VecSize>(gamma_ptr + col * VecSize, &gamma[it]);
|
||
col += THREADS_PER_ROW;
|
||
}
|
||
|
||
#pragma unroll 1
|
||
for (int64_t row = r; row < rows; row += gridDim.x * ROWS_PER_CTA) {
|
||
const U mean_cur_row = mean_ptr[row];
|
||
const U var_cur_row = rsqrt_<U>(var_ptr[row] + epsilon);
|
||
Vec dout[LDGS], x[LDGS];
|
||
MaskLoadT mask_vec[LDGS];
|
||
int64_t col = c;
|
||
#pragma unroll
|
||
for (int64_t it = 0; it < LDGS; it++) {
|
||
phi::Load<T, VecSize>(dout_ptr + row * ELTS_PER_ROW + col * VecSize,
|
||
&dout[it]);
|
||
phi::Load<T, VecSize>(x_ptr + row * ELTS_PER_ROW + col * VecSize, &x[it]);
|
||
if (IsFusedDropoutResidualLn) {
|
||
phi::Load<MaskType, VecSize>(
|
||
mask_ptr + row * ELTS_PER_ROW + col * VecSize, &mask_vec[it]);
|
||
}
|
||
|
||
col += THREADS_PER_ROW;
|
||
}
|
||
|
||
// local reductions
|
||
U dy[LDGS * VecSize];
|
||
U y[LDGS * VecSize];
|
||
|
||
U sum_loss1 = 0.f;
|
||
U sum_loss2 = 0.f;
|
||
#pragma unroll
|
||
for (int it = 0; it < LDGS; it++) {
|
||
#pragma unroll
|
||
for (int jt = 0; jt < VecSize; jt++) {
|
||
U x_tmp = static_cast<U>(x[it][jt]);
|
||
U y_tmp = var_cur_row * (x_tmp - mean_cur_row);
|
||
U dy_tmp = static_cast<U>(gamma[it][jt]) *
|
||
static_cast<U>(dout[it][jt]); // scale * dy
|
||
U dout_tmp = static_cast<U>(dout[it][jt]); // dy
|
||
|
||
// used for get dx (row reduction)
|
||
sum_loss1 += dy_tmp; // scale * dy, sum_1
|
||
sum_loss2 += dy_tmp * y_tmp; // scale * dy * y, sum_2
|
||
|
||
dy[it * VecSize + jt] = dy_tmp; // scale * dy
|
||
y[it * VecSize + jt] = y_tmp; // y
|
||
|
||
// used for get dscale and dbias (column reduction)
|
||
dgamma_sum[it * VecSize + jt] += dout_tmp * y_tmp; // dy * y
|
||
dbeta_sum[it * VecSize + jt] += dout_tmp; // dy
|
||
}
|
||
}
|
||
|
||
// reduction across row for sum_loss1, sum_loss2
|
||
if (WARPS_N == 1) {
|
||
#pragma unroll
|
||
// row reduction among 32 threads.
|
||
for (int it = 1; it < THREADS_PER_WARP; it *= 2) {
|
||
#ifdef PADDLE_WITH_HIP
|
||
sum_loss1 += __shfl_xor(sum_loss1, it);
|
||
sum_loss2 += __shfl_xor(sum_loss2, it);
|
||
#else
|
||
sum_loss1 += __shfl_xor_sync(uint32_t(-1), sum_loss1, it);
|
||
sum_loss2 += __shfl_xor_sync(uint32_t(-1), sum_loss2, it);
|
||
#endif
|
||
}
|
||
sum_loss1 *= rn;
|
||
sum_loss2 *= rn;
|
||
} else {
|
||
#pragma unroll
|
||
for (int it = 16; it > 0; it /= 2) {
|
||
#ifdef PADDLE_WITH_HIP
|
||
sum_loss1 += __shfl_down(sum_loss1, it);
|
||
sum_loss2 += __shfl_down(sum_loss2, it);
|
||
#else
|
||
sum_loss1 += __shfl_down_sync(uint32_t(-1), sum_loss1, it);
|
||
sum_loss2 += __shfl_down_sync(uint32_t(-1), sum_loss2, it);
|
||
#endif
|
||
}
|
||
|
||
if (lane == 0) {
|
||
sum_loss1_shared[warp_n] = sum_loss1;
|
||
sum_loss2_shared[warp_n] = sum_loss2;
|
||
}
|
||
|
||
__syncthreads();
|
||
if (warp_n == 0 && lane == 0) {
|
||
sum_loss1 = 0.f;
|
||
sum_loss2 = 0.f;
|
||
for (int it = 0; it < WARPS_N; it++) {
|
||
sum_loss1 += sum_loss1_shared[it];
|
||
sum_loss2 += sum_loss2_shared[it];
|
||
}
|
||
sum_loss1_shared[0] = sum_loss1;
|
||
sum_loss2_shared[0] = sum_loss2;
|
||
}
|
||
__syncthreads();
|
||
|
||
sum_loss1 = sum_loss1_shared[0] * rn;
|
||
sum_loss2 = sum_loss2_shared[0] * rn;
|
||
}
|
||
|
||
#pragma unroll
|
||
for (int it = 0; it < LDGS; it++) {
|
||
#pragma unroll
|
||
for (int jt = 0; jt < VecSize; jt++) {
|
||
U dy_tmp = dy[it * VecSize + jt]; // scale * dy
|
||
U y_tmp = y[it * VecSize + jt]; // y
|
||
// dx = var * (scale * dy - sum_loss2 * y - sum_loss1)
|
||
U dx_tmp = var_cur_row * (dy_tmp - sum_loss2 * y_tmp - sum_loss1);
|
||
// Note: reuse x and dout vec register to store dx and d_dropout_src.
|
||
x[it][jt] = static_cast<T>(dx_tmp);
|
||
if (IsFusedDropoutResidualLn) {
|
||
if (factor == static_cast<T>(1.0f)) { // no dropout
|
||
dout[it][jt] = x[it][jt] * factor;
|
||
} else {
|
||
dout[it][jt] =
|
||
x[it][jt] * static_cast<T>(mask_vec[it][jt]) * factor;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// store dx to global memory
|
||
col = c;
|
||
#pragma unroll
|
||
for (int it = 0; it < LDGS; it++) {
|
||
phi::Store<T, VecSize>(x[it],
|
||
dx_ptr + row * ELTS_PER_ROW + col * VecSize);
|
||
if (IsFusedDropoutResidualLn) {
|
||
phi::Store<T, VecSize>(
|
||
dout[it], d_dropout_src_ptr + row * ELTS_PER_ROW + col * VecSize);
|
||
} else if (NeedDDropoutSrcPtr) {
|
||
phi::Store<T, VecSize>(
|
||
x[it], d_dropout_src_ptr + row * ELTS_PER_ROW + col * VecSize);
|
||
}
|
||
col += THREADS_PER_ROW;
|
||
}
|
||
}
|
||
|
||
// step-2: column reduction of dscale and dbias for each thread block.
|
||
// each block's sum: [4 * 1024] -> [1 * 1024]
|
||
enum { NUM_RES = ELTS_PER_ROW / THREADS_PER_CTA }; // 1024/128 = 8
|
||
static_assert(NUM_RES * THREADS_PER_CTA == ELTS_PER_ROW, "");
|
||
|
||
U *smem_write;
|
||
|
||
smem_write = &smem_[warp_m * ELTS_PER_ROW + tid_r * VecSize]; // [4 * 1024]
|
||
#pragma unroll
|
||
for (int it = 0; it < LDGS; it++) {
|
||
#pragma unroll
|
||
for (int jt = 0; jt < VecSize; jt++) {
|
||
smem_write[jt] = dbeta_sum[it * VecSize + jt];
|
||
}
|
||
smem_write += THREADS_PER_ROW * VecSize; // 32*8
|
||
}
|
||
__syncthreads();
|
||
U cta_dbeta_sum[NUM_RES];
|
||
memset(cta_dbeta_sum, 0, sizeof(U) * NUM_RES);
|
||
// column reduction for elems in smem: 4*1024 -> 1*1024.
|
||
for (int it = 0; it < ROWS_PER_CTA; it++) {
|
||
for (int jt = 0; jt < NUM_RES; jt++) {
|
||
cta_dbeta_sum[jt] +=
|
||
smem_[it * ELTS_PER_ROW + tidx + jt * THREADS_PER_CTA];
|
||
}
|
||
}
|
||
__syncthreads();
|
||
|
||
smem_write = &smem_[warp_m * ELTS_PER_ROW + tid_r * VecSize];
|
||
#pragma unroll
|
||
for (int it = 0; it < LDGS; it++) {
|
||
#pragma unroll
|
||
for (int jt = 0; jt < VecSize; jt++) {
|
||
smem_write[jt] = dgamma_sum[it * VecSize + jt];
|
||
}
|
||
smem_write += THREADS_PER_ROW * VecSize;
|
||
}
|
||
__syncthreads();
|
||
U cta_dgamma_sum[NUM_RES];
|
||
memset(cta_dgamma_sum, 0, sizeof(U) * NUM_RES);
|
||
for (int it = 0; it < ROWS_PER_CTA; it++) {
|
||
for (int jt = 0; jt < NUM_RES; jt++) {
|
||
cta_dgamma_sum[jt] +=
|
||
smem_[it * ELTS_PER_ROW + tidx + jt * THREADS_PER_CTA];
|
||
}
|
||
}
|
||
|
||
// the shape of results：(#blocks, 1024)
|
||
U *dgamma_part =
|
||
static_cast<U *>(dgamma_temp_ptr) + bidx * ELTS_PER_ROW + tidx;
|
||
for (int jt = 0; jt < NUM_RES; jt++) {
|
||
*dgamma_part = cta_dgamma_sum[jt];
|
||
dgamma_part += THREADS_PER_CTA;
|
||
}
|
||
|
||
U *dbeta_part = static_cast<U *>(dbeta_temp_ptr) + bidx * ELTS_PER_ROW + tidx;
|
||
for (int jt = 0; jt < NUM_RES; jt++) {
|
||
*dbeta_part = cta_dbeta_sum[jt];
|
||
dbeta_part += THREADS_PER_CTA;
|
||
}
|
||
}
|
||
|
||
/* This function carry out column reduction whose input is [rows, 1024] and
|
||
* output is [1, 1024].
|
||
* #blocks: 32
|
||
* #threads: 512
|
||
*/
|
||
// todo(@limin29): to think if there are better impl strategies
|
||
template <typename U,
|
||
typename ScaleT = U,
|
||
int VecSize = 1,
|
||
int WARPS_M = 16,
|
||
int WARPS_N = 1,
|
||
int BYTES_PER_LDG = 4,
|
||
int ELTS_PER_ROW = 1024,
|
||
int THREADS_PER_WARP = 32,
|
||
int THREADS_PER_ROW = WARPS_N *THREADS_PER_WARP,
|
||
int THREADS_PER_CTA = WARPS_M *THREADS_PER_ROW,
|
||
int ROWS_PER_CTA = WARPS_M,
|
||
int ELTS_PER_ROW_PER_CTA = THREADS_PER_ROW *VecSize,
|
||
int LDGS = ELTS_PER_ROW / ELTS_PER_ROW_PER_CTA,
|
||
int VEC_COLS = ELTS_PER_ROW / VecSize>
|
||
__global__ __launch_bounds__(THREADS_PER_CTA) void ln_bwd_fast_final_kernel(
|
||
const int rows,
|
||
U *__restrict__ dg_part_,
|
||
U *__restrict__ db_part_,
|
||
ScaleT *__restrict__ dg_,
|
||
ScaleT *__restrict__ db_) {
|
||
using Vec = AlignedVector<U, VecSize>;
|
||
static_assert(VEC_COLS == ELTS_PER_ROW / VecSize, "");
|
||
|
||
const int tidx = threadIdx.x;
|
||
const int bidx = blockIdx.x;
|
||
const int lane = tidx % THREADS_PER_WARP;
|
||
const int warp = tidx / THREADS_PER_WARP;
|
||
const int warp_m = warp / WARPS_N;
|
||
const int warp_n = warp % WARPS_N;
|
||
const int tid_c = warp_n * THREADS_PER_WARP + lane;
|
||
|
||
const int c = bidx * THREADS_PER_ROW + tid_c;
|
||
const int r = warp_m;
|
||
|
||
__shared__ U smem_space[(WARPS_M - 1) * THREADS_PER_ROW * VecSize];
|
||
|
||
for (int64_t col = c; col < VEC_COLS; col += gridDim.x * THREADS_PER_ROW) {
|
||
const U *dg_part_ptr = (dg_part_) + r * ELTS_PER_ROW + col * VecSize;
|
||
const U *db_part_ptr = (db_part_) + r * ELTS_PER_ROW + col * VecSize;
|
||
|
||
U dg_sum[VecSize];
|
||
U db_sum[VecSize];
|
||
memset(dg_sum, 0, sizeof(U) * VecSize);
|
||
memset(db_sum, 0, sizeof(U) * VecSize);
|
||
#pragma unroll
|
||
for (int64_t row = r; row < rows; row += ROWS_PER_CTA) {
|
||
Vec dg;
|
||
Vec db;
|
||
phi::Load<U, VecSize>(dg_part_ptr, &dg);
|
||
phi::Load<U, VecSize>(db_part_ptr, &db);
|
||
dg_part_ptr += ROWS_PER_CTA * ELTS_PER_ROW;
|
||
db_part_ptr += ROWS_PER_CTA * ELTS_PER_ROW;
|
||
|
||
#pragma unroll
|
||
for (int jt = 0; jt < VecSize; jt++) {
|
||
dg_sum[jt] += dg[jt];
|
||
db_sum[jt] += db[jt];
|
||
}
|
||
}
|
||
|
||
// reduction across rows of the thread block
|
||
U *smem_write;
|
||
smem_write = smem_space + (warp_m - 1) * THREADS_PER_ROW * VecSize + tid_c;
|
||
|
||
if (warp_m > 0) {
|
||
#pragma unroll
|
||
for (int jt = 0; jt < VecSize; jt++) {
|
||
*smem_write = dg_sum[jt];
|
||
smem_write += THREADS_PER_ROW;
|
||
}
|
||
}
|
||
__syncthreads();
|
||
|
||
U *smem_read;
|
||
smem_read = smem_space + tid_c;
|
||
if (warp_m == 0) {
|
||
#pragma unroll
|
||
for (int it = 0; it < WARPS_M - 1; it++) {
|
||
#pragma unroll
|
||
for (int jt = 0; jt < VecSize; jt++) {
|
||
dg_sum[jt] += *smem_read;
|
||
smem_read += THREADS_PER_ROW;
|
||
}
|
||
}
|
||
}
|
||
|
||
__syncthreads();
|
||
|
||
smem_write = smem_space + (warp_m - 1) * THREADS_PER_ROW * VecSize + tid_c;
|
||
|
||
if (warp_m > 0) {
|
||
#pragma unroll
|
||
for (int jt = 0; jt < VecSize; jt++) {
|
||
*smem_write = db_sum[jt];
|
||
smem_write += THREADS_PER_ROW;
|
||
}
|
||
}
|
||
__syncthreads();
|
||
|
||
smem_read = smem_space + tid_c;
|
||
if (warp_m == 0) {
|
||
#pragma unroll
|
||
for (int it = 0; it < WARPS_M - 1; it++) {
|
||
#pragma unroll
|
||
for (int jt = 0; jt < VecSize; jt++) {
|
||
db_sum[jt] += *smem_read;
|
||
smem_read += THREADS_PER_ROW;
|
||
}
|
||
}
|
||
|
||
union {
|
||
ScaleT raw;
|
||
ScaleT elt[VecSize];
|
||
} dg_out, db_out;
|
||
|
||
#pragma unroll
|
||
for (int jt = 0; jt < VecSize; jt++) {
|
||
dg_out.elt[jt] = dg_sum[jt];
|
||
db_out.elt[jt] = db_sum[jt];
|
||
}
|
||
ScaleT *dg_ptr = reinterpret_cast<ScaleT *>(dg_) + col;
|
||
ScaleT *db_ptr = reinterpret_cast<ScaleT *>(db_) + col;
|
||
*dg_ptr = dg_out.raw;
|
||
*db_ptr = db_out.raw;
|
||
}
|
||
}
|
||
}
|
||
|
||
/* This function support two kinds of computations (only for float and fp16
|
||
* type):
|
||
*
|
||
* Case-1: compute layer_norm_grad for layernorm op by setting mask_ptr and
|
||
* d_dropout_src_ptr to nullptr. Here, d_x_ptr returns the grad of layernorm
|
||
* input.
|
||
*
|
||
* Case-2: compute layer_norm_grad + residual_grad + dropout_grad for
|
||
* fused_dropout_residual_layernorm op. Here, dx_ptr returns residual_grad.
|
||
*
|
||
*/
|
||
template <typename T,
|
||
typename U,
|
||
typename ScaleT = U,
|
||
typename MaskType = uint8_t>
|
||
void ln_bwd_fast_kernel_driver(const GPUContext &dev_ctx,
|
||
const int64_t rows,
|
||
const int64_t cols,
|
||
float epsilon,
|
||
const T *x_ptr,
|
||
const ScaleT *scale_ptr,
|
||
const U *mean_ptr,
|
||
const U *var_ptr,
|
||
const T *dout_ptr,
|
||
T *dx_ptr,
|
||
ScaleT *dscale_ptr,
|
||
ScaleT *dbias_ptr,
|
||
const MaskType *mask_ptr = nullptr,
|
||
T factor = static_cast<T>(0),
|
||
T *d_dropout_src_ptr = nullptr) {
|
||
auto stream = dev_ctx.stream();
|
||
if (cols == 1024 || cols == 384 || cols == 256) {
|
||
// step-1: compute dx and reduced part results of dscale and dbias.
|
||
const int WARPS_M = 4; // how many rows deal in a cta.
|
||
const int WARPS_N = 1; // how many warps to deal with a row.
|
||
const int BYTES_PER_LDG = 16;
|
||
const int VecSize = BYTES_PER_LDG / sizeof(T);
|
||
|
||
const int THREADS_PER_WARP = 32;
|
||
const int THREADS_PER_ROW = WARPS_N * THREADS_PER_WARP;
|
||
const int THREADS_PER_CTA = WARPS_M * THREADS_PER_ROW;
|
||
const int ROWS_PER_CTA = WARPS_M;
|
||
|
||
// 4 * 1024 * 4
|
||
const int SMEM_BYTES = ROWS_PER_CTA * cols * sizeof(U);
|
||
|
||
// #blocks = 2 * #SM
|
||
const int gridx = 2 * dev_ctx.GetSMCount();
|
||
|
||
// get temp space for dscale and dbias.
|
||
DenseTensor dscale_temp;
|
||
dscale_temp.Resize({gridx, cols});
|
||
dev_ctx.template Alloc<U>(&dscale_temp);
|
||
U *dscale_temp_ptr = dscale_temp.data<U>();
|
||
|
||
DenseTensor dbias_temp;
|
||
dbias_temp.Resize({gridx, cols});
|
||
dev_ctx.template Alloc<U>(&dbias_temp);
|
||
U *dbias_temp_ptr = dbias_temp.data<U>();
|
||
|
||
if (mask_ptr != nullptr) {
|
||
if (d_dropout_src_ptr == nullptr) {
|
||
PADDLE_THROW(common::errors::InvalidArgument(
|
||
"To compute fused_dropout_residual_ln grad, d_dropout_src_ptr "
|
||
"can't be null"));
|
||
}
|
||
#define LAUNCH_MASK_FUSED_LN_BWD_FAST_KERNEL(vec_size, ele_per_row) \
|
||
fused_ln_bwd_fast_kernel<true, \
|
||
true, \
|
||
T, \
|
||
U, \
|
||
ScaleT, \
|
||
MaskType, \
|
||
vec_size, \
|
||
WARPS_M, \
|
||
WARPS_N, \
|
||
BYTES_PER_LDG, \
|
||
ele_per_row> \
|
||
<<<gridx, THREADS_PER_CTA, 0, stream>>>(rows, \
|
||
epsilon, \
|
||
x_ptr, \
|
||
scale_ptr, \
|
||
mean_ptr, \
|
||
var_ptr, \
|
||
dout_ptr, \
|
||
dscale_temp_ptr, \
|
||
dbias_temp_ptr, \
|
||
dx_ptr, \
|
||
mask_ptr, \
|
||
factor, \
|
||
d_dropout_src_ptr);
|
||
|
||
if (cols == 1024) {
|
||
LAUNCH_MASK_FUSED_LN_BWD_FAST_KERNEL(VecSize, 1024);
|
||
} else {
|
||
switch (cols) {
|
||
case 384:
|
||
LAUNCH_MASK_FUSED_LN_BWD_FAST_KERNEL(1, 384);
|
||
break;
|
||
case 256:
|
||
LAUNCH_MASK_FUSED_LN_BWD_FAST_KERNEL(VecSize, 256);
|
||
break;
|
||
}
|
||
}
|
||
#undef LAUNCH_MASK_FUSED_LN_BWD_FAST_KERNEL
|
||
|
||
} else {
|
||
#define LAUNCH_FUSED_LN_BWD_FAST_KERNEL_BASE( \
|
||
vec_size, ele_per_row, need_d_dropout_src_ptr) \
|
||
fused_ln_bwd_fast_kernel<false, \
|
||
need_d_dropout_src_ptr, \
|
||
T, \
|
||
U, \
|
||
ScaleT, \
|
||
MaskType, \
|
||
vec_size, \
|
||
WARPS_M, \
|
||
WARPS_N, \
|
||
BYTES_PER_LDG, \
|
||
ele_per_row> \
|
||
<<<gridx, THREADS_PER_CTA, 0, stream>>>(rows, \
|
||
epsilon, \
|
||
x_ptr, \
|
||
scale_ptr, \
|
||
mean_ptr, \
|
||
var_ptr, \
|
||
dout_ptr, \
|
||
dscale_temp_ptr, \
|
||
dbias_temp_ptr, \
|
||
dx_ptr, \
|
||
nullptr, \
|
||
factor, \
|
||
d_dropout_src_ptr);
|
||
|
||
#define LAUNCH_FUSED_LN_BWD_FAST_KERNEL(vec_size, ele_per_row) \
|
||
do { \
|
||
if (d_dropout_src_ptr != nullptr) { \
|
||
LAUNCH_FUSED_LN_BWD_FAST_KERNEL_BASE(vec_size, ele_per_row, true); \
|
||
} else { \
|
||
LAUNCH_FUSED_LN_BWD_FAST_KERNEL_BASE(vec_size, ele_per_row, false); \
|
||
} \
|
||
} while (0)
|
||
|
||
if (cols == 1024) {
|
||
LAUNCH_FUSED_LN_BWD_FAST_KERNEL(VecSize, 1024);
|
||
} else {
|
||
switch (cols) {
|
||
case 384:
|
||
LAUNCH_FUSED_LN_BWD_FAST_KERNEL(1, 384);
|
||
break;
|
||
case 256:
|
||
LAUNCH_FUSED_LN_BWD_FAST_KERNEL(VecSize, 256);
|
||
break;
|
||
}
|
||
}
|
||
|
||
#undef LAUNCH_FUSED_LN_BWD_FAST_KERNEL
|
||
}
|
||
|
||
const int WARPS_M_2 = 16;
|
||
const int WARPS_N_2 = 1;
|
||
const int BYTES_PER_LDG_2 = 4;
|
||
const int VecSize_2 =
|
||
std::max(1, static_cast<int>(BYTES_PER_LDG_2 / sizeof(U))); // 1
|
||
|
||
const int THREADS_PER_WARP_2 = 32;
|
||
const int THREADS_PER_ROW_2 = WARPS_N_2 * THREADS_PER_WARP_2; // 32
|
||
const int THREADS_PER_CTA_2 =
|
||
WARPS_M_2 * THREADS_PER_ROW_2; // 16 * 32 = 512
|
||
const int ROWS_PER_CTA_2 = WARPS_M_2; // 16
|
||
|
||
// #blocks: 32, #threads_per_block: 512
|
||
// Note: it is not supported for double type.
|
||
if (sizeof(U) > 4) {
|
||
PADDLE_THROW(
|
||
common::errors::InvalidArgument("Only support float and fp16 type"));
|
||
} else {
|
||
int gridx_2 = 0;
|
||
|
||
#define LAUNCH_LN_BWD_BETA_GAMMMA_KERNEL(vec_size, ele_per_row) \
|
||
gridx_2 = static_cast<int>(std::ceil( \
|
||
ele_per_row / static_cast<float>(THREADS_PER_ROW_2 * vec_size))); \
|
||
ln_bwd_fast_final_kernel<U, \
|
||
ScaleT, \
|
||
vec_size, \
|
||
WARPS_M_2, \
|
||
WARPS_N_2, \
|
||
BYTES_PER_LDG_2, \
|
||
ele_per_row> \
|
||
<<<gridx_2, THREADS_PER_CTA_2, 0, stream>>>( \
|
||
gridx, dscale_temp_ptr, dbias_temp_ptr, dscale_ptr, dbias_ptr);
|
||
|
||
if (cols == 1024) {
|
||
LAUNCH_LN_BWD_BETA_GAMMMA_KERNEL(VecSize_2, 1024);
|
||
} else {
|
||
switch (cols) {
|
||
case 384:
|
||
LAUNCH_LN_BWD_BETA_GAMMMA_KERNEL(1, 384);
|
||
break;
|
||
case 256:
|
||
LAUNCH_LN_BWD_BETA_GAMMMA_KERNEL(VecSize_2, 256);
|
||
break;
|
||
}
|
||
}
|
||
|
||
#undef LAUNCH_LN_BWD_BETA_GAMMMA_KERNEL
|
||
}
|
||
} else {
|
||
PADDLE_THROW(common::errors::InvalidArgument(
|
||
"Fast layer_norm kernel is only used when feature_size is 1024"));
|
||
}
|
||
}
|
||
#endif
|
||
|
||
template <typename T, typename U, int BDIMX, int BDIMY, int VPTX>
|
||
__global__ void LayerNormBackwardPartGradGammaBeta(const T *__restrict__ dout,
|
||
const T *__restrict__ input,
|
||
const int64_t n1,
|
||
const int64_t n2,
|
||
const U *__restrict__ mean,
|
||
const U *__restrict__ var,
|
||
float epsilon,
|
||
U *part_grad_gamma,
|
||
U *part_grad_beta) {
|
||
// VPTX -> value per thread.x, BDIMX -> blockDim.x,
|
||
// BDIMY -> blockDim.y, template for compile time optimizations.
|
||
constexpr int RowStride = BDIMX + 1;
|
||
constexpr int BLOCK_SIZE = BDIMX * BDIMY;
|
||
constexpr int VPTX_MUL_BDIMY = VPTX * BDIMY;
|
||
constexpr int SharedSize = (BLOCK_SIZE > 2 * VPTX_MUL_BDIMY * RowStride)
|
||
? BLOCK_SIZE
|
||
: 2 * VPTX_MUL_BDIMY * RowStride;
|
||
|
||
const int thr_load_col_off = (threadIdx.x * VPTX) & (BDIMX - 1);
|
||
const int thr_load_row_off =
|
||
(threadIdx.x * VPTX) / BDIMX + threadIdx.y * BDIMY;
|
||
const int i2_off = blockIdx.x * BDIMX + thr_load_col_off;
|
||
|
||
__shared__ U buf[SharedSize];
|
||
U *warp_buf1 = reinterpret_cast<U *>(buf);
|
||
U *warp_buf2 = warp_buf1 + VPTX_MUL_BDIMY * RowStride;
|
||
|
||
for (int idx = threadIdx.y * BDIMX + threadIdx.x;
|
||
idx < 2 * VPTX_MUL_BDIMY * RowStride;
|
||
idx += BLOCK_SIZE) {
|
||
buf[idx] = U(0);
|
||
}
|
||
__syncthreads();
|
||
|
||
for (int64_t i1_block = static_cast<int64_t>(blockIdx.y) * BDIMY * VPTX;
|
||
i1_block < n1;
|
||
i1_block += VPTX_MUL_BDIMY * gridDim.y) {
|
||
cuLoadAddStridedInputs<T, U, VPTX>(i1_block,
|
||
thr_load_row_off,
|
||
thr_load_col_off,
|
||
i2_off,
|
||
RowStride,
|
||
warp_buf1,
|
||
warp_buf2,
|
||
input,
|
||
dout,
|
||
n1,
|
||
n2,
|
||
mean,
|
||
var,
|
||
epsilon);
|
||
}
|
||
__syncthreads();
|
||
|
||
// inter-warp reductions, sum within each warp
|
||
U acc1 = U(0);
|
||
U acc2 = U(0);
|
||
#pragma unroll
|
||
for (int k = 0; k < VPTX; ++k) {
|
||
int row1 = threadIdx.y + k * VPTX;
|
||
int idx1 = row1 * RowStride + threadIdx.x;
|
||
acc1 += warp_buf1[idx1];
|
||
acc2 += warp_buf2[idx1];
|
||
}
|
||
warp_buf1[threadIdx.y * RowStride + threadIdx.x] = acc1;
|
||
warp_buf2[threadIdx.y * RowStride + threadIdx.x] = acc2;
|
||
__syncthreads();
|
||
|
||
// sum all warps
|
||
#pragma unroll
|
||
for (int offset = VPTX >> 1; offset > 1; offset >>= 1) {
|
||
if (threadIdx.y < offset) {
|
||
int row1 = threadIdx.y;
|
||
int row2 = threadIdx.y + offset;
|
||
int idx1 = row1 * RowStride + threadIdx.x;
|
||
int idx2 = row2 * RowStride + threadIdx.x;
|
||
warp_buf1[idx1] += warp_buf1[idx2];
|
||
warp_buf2[idx1] += warp_buf2[idx2];
|
||
}
|
||
__syncthreads();
|
||
}
|
||
int64_t i2 = static_cast<int64_t>(blockIdx.x) * BDIMX + threadIdx.x;
|
||
if (threadIdx.y == 0 && i2 < n2) {
|
||
int row1 = threadIdx.y;
|
||
int row2 = threadIdx.y + 1;
|
||
int idx1 = row1 * RowStride + threadIdx.x;
|
||
int idx2 = row2 * RowStride + threadIdx.x;
|
||
part_grad_beta[blockIdx.y * n2 + i2] = warp_buf1[idx1] + warp_buf1[idx2];
|
||
part_grad_gamma[blockIdx.y * n2 + i2] = warp_buf2[idx1] + warp_buf2[idx2];
|
||
}
|
||
}
|
||
|
||
template <typename T, typename U, int BDIMX, int BDIMY, typename ScaleT>
|
||
__global__ void LayerNormBackwardSumGradGammaBeta(const U *part_grad_gamma,
|
||
const U *part_grad_beta,
|
||
const int part_size,
|
||
const int64_t n1,
|
||
const int64_t n2,
|
||
ScaleT *grad_gamma,
|
||
ScaleT *grad_beta) {
|
||
// sum partial gradients for gamma and beta
|
||
__shared__ U buf[BDIMX * BDIMY];
|
||
int64_t i2 = static_cast<int64_t>(blockIdx.x) * BDIMX + threadIdx.x;
|
||
if (i2 < n2) {
|
||
// each warp does sequential reductions until reduced part_size is num_warps
|
||
int num_warp_reductions = part_size / BDIMY;
|
||
U sum_gamma = U(0);
|
||
U sum_beta = U(0);
|
||
const U *part_grad_gamma_ptr =
|
||
part_grad_gamma + threadIdx.y * num_warp_reductions * n2 + i2;
|
||
const U *part_grad_beta_ptr =
|
||
part_grad_beta + threadIdx.y * num_warp_reductions * n2 + i2;
|
||
for (int warp_offset = 0; warp_offset < num_warp_reductions;
|
||
++warp_offset) {
|
||
sum_gamma += part_grad_gamma_ptr[warp_offset * n2];
|
||
sum_beta += part_grad_beta_ptr[warp_offset * n2];
|
||
}
|
||
// inter-warp reductions
|
||
constexpr int nbsize3 = BDIMX * BDIMY / 2;
|
||
for (int offset = BDIMY / 2; offset >= 1; offset /= 2) {
|
||
// top half write to shared memory
|
||
if (threadIdx.y >= offset && threadIdx.y < 2 * offset) {
|
||
const int write_idx = (threadIdx.y - offset) * blockDim.x + threadIdx.x;
|
||
buf[write_idx] = sum_gamma;
|
||
buf[write_idx + nbsize3] = sum_beta;
|
||
}
|
||
__syncthreads();
|
||
// bottom half sums
|
||
if (threadIdx.y < offset) {
|
||
const int read_idx = threadIdx.y * BDIMX + threadIdx.x;
|
||
sum_gamma += buf[read_idx];
|
||
sum_beta += buf[read_idx + nbsize3];
|
||
}
|
||
__syncthreads();
|
||
}
|
||
// write out fully summed gradients
|
||
if (threadIdx.y == 0) {
|
||
grad_gamma[i2] = static_cast<ScaleT>(sum_gamma);
|
||
grad_beta[i2] = static_cast<ScaleT>(sum_beta);
|
||
}
|
||
}
|
||
}
|
||
|
||
template <typename T, typename U, int BDIMX, int BDIMY, typename ScaleT>
|
||
__global__ void LayerNormBackwardComputeGradInput(const T *__restrict__ dout,
|
||
const T *__restrict__ input,
|
||
const int64_t n1,
|
||
const int64_t n2,
|
||
const U *__restrict__ mean,
|
||
const U *__restrict__ var,
|
||
const float epsilon,
|
||
const ScaleT *gamma,
|
||
T *grad_input) {
|
||
#ifdef __HIPCC__
|
||
for (int64_t i1 = hipBlockIdx_x; i1 < n1; i1 += hipGridDim_x) {
|
||
#else
|
||
for (int64_t i1 = blockIdx.x; i1 < n1; i1 += gridDim.x) {
|
||
#endif
|
||
U sum_loss1 = U(0);
|
||
U sum_loss2 = U(0);
|
||
const U c_mean = mean[i1];
|
||
const U c_invvar = rsqrt_<U>(var[i1] + epsilon);
|
||
const T *k_input = input + i1 * n2;
|
||
const T *k_dout = dout + i1 * n2;
|
||
constexpr int numx = BDIMX * BDIMY;
|
||
const int thrx = threadIdx.x + threadIdx.y * BDIMX;
|
||
if (gamma != NULL) {
|
||
int64_t l = 4 * thrx;
|
||
for (; l + 3 < n2; l += 4 * numx) {
|
||
for (int k = 0; k < 4; ++k) {
|
||
const U c_h = static_cast<U>(k_input[l + k]);
|
||
const U c_loss = static_cast<U>(k_dout[l + k]);
|
||
sum_loss1 += c_loss * static_cast<U>(gamma[l + k]);
|
||
sum_loss2 +=
|
||
c_loss * static_cast<U>(gamma[l + k]) * (c_h - c_mean) * c_invvar;
|
||
}
|
||
}
|
||
for (; l < n2; ++l) {
|
||
const U c_h = static_cast<U>(k_input[l]);
|
||
const U c_loss = static_cast<U>(k_dout[l]);
|
||
sum_loss1 += c_loss * static_cast<U>(gamma[l]);
|
||
sum_loss2 +=
|
||
c_loss * static_cast<U>(gamma[l]) * (c_h - c_mean) * c_invvar;
|
||
}
|
||
} else {
|
||
int64_t l = 4 * thrx;
|
||
for (; l + 3 < n2; l += 4 * numx) {
|
||
for (int k = 0; k < 4; ++k) {
|
||
const U c_h = static_cast<U>(k_input[l + k]);
|
||
const U c_loss = static_cast<U>(k_dout[l + k]);
|
||
sum_loss1 += c_loss;
|
||
sum_loss2 += c_loss * (c_h - c_mean) * c_invvar;
|
||
}
|
||
}
|
||
for (; l < n2; ++l) {
|
||
const U c_h = static_cast<U>(k_input[l]);
|
||
const U c_loss = static_cast<U>(k_dout[l]);
|
||
sum_loss1 += c_loss;
|
||
sum_loss2 += c_loss * (c_h - c_mean) * c_invvar;
|
||
}
|
||
}
|
||
// intra-warp reductions
|
||
#pragma unroll
|
||
for (int mask = BDIMX / 2; mask > 0; mask /= 2) {
|
||
#ifdef PADDLE_WITH_HIP
|
||
// WARP_SHFL_XOR(sum_loss, mask);
|
||
sum_loss1 += __shfl_xor(sum_loss1, mask, warpSize);
|
||
sum_loss2 += __shfl_xor(sum_loss2, mask, warpSize);
|
||
#else
|
||
// WARP_SHFL_XOR(sum_loss, mask);
|
||
sum_loss1 += __shfl_xor_sync(0xffffffff, sum_loss1, mask, warpSize);
|
||
sum_loss2 += __shfl_xor_sync(0xffffffff, sum_loss2, mask, warpSize);
|
||
#endif
|
||
}
|
||
// inter-warp reductions
|
||
if (BDIMY > 1) {
|
||
__shared__ U buf[BDIMX * BDIMY];
|
||
for (int offset = BDIMY / 2; offset > 0; offset /= 2) {
|
||
// upper half of warps write to shared
|
||
if (threadIdx.y >= offset && threadIdx.y < 2 * offset) {
|
||
const int wrt_i = (threadIdx.y - offset) * BDIMX + threadIdx.x;
|
||
buf[2 * wrt_i] = sum_loss1;
|
||
buf[2 * wrt_i + 1] = sum_loss2;
|
||
}
|
||
__syncthreads();
|
||
// lower half merges
|
||
if (threadIdx.y < offset) {
|
||
const int read_i = threadIdx.y * blockDim.x + threadIdx.x;
|
||
sum_loss1 += buf[2 * read_i];
|
||
sum_loss2 += buf[2 * read_i + 1];
|
||
}
|
||
__syncthreads();
|
||
}
|
||
if (threadIdx.y == 0) {
|
||
buf[2 * threadIdx.x] = sum_loss1;
|
||
buf[2 * threadIdx.x + 1] = sum_loss2;
|
||
}
|
||
__syncthreads();
|
||
if (threadIdx.y != 0) {
|
||
sum_loss1 = buf[2 * threadIdx.x];
|
||
sum_loss2 = buf[2 * threadIdx.x + 1];
|
||
}
|
||
}
|
||
// all threads now have the two sums over l
|
||
U fH = (U)n2;
|
||
U term1 = (U(1) / fH) * c_invvar;
|
||
T *k_grad_input = grad_input + i1 * n2;
|
||
if (gamma != NULL) {
|
||
for (int64_t l = thrx; l < n2; l += numx) {
|
||
const U c_h = static_cast<U>(k_input[l]);
|
||
const U c_loss = static_cast<U>(k_dout[l]);
|
||
U f_grad_input = fH * c_loss * static_cast<U>(gamma[l]);
|
||
f_grad_input -= sum_loss1;
|
||
f_grad_input -= (c_h - c_mean) * c_invvar * sum_loss2;
|
||
f_grad_input *= term1;
|
||
k_grad_input[l] = static_cast<T>(f_grad_input);
|
||
}
|
||
} else {
|
||
for (int64_t l = thrx; l < n2; l += numx) {
|
||
const U c_h = static_cast<U>(k_input[l]);
|
||
const U c_loss = static_cast<U>(k_dout[l]);
|
||
U f_grad_input = fH * c_loss;
|
||
f_grad_input -= sum_loss1;
|
||
f_grad_input -= (c_h - c_mean) * c_invvar * sum_loss2;
|
||
f_grad_input *= term1;
|
||
k_grad_input[l] = static_cast<T>(f_grad_input);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
template <typename T, typename U, typename ScaleT, int DataPerTid>
|
||
__global__ void LayerNormBackwardComputeGradInputWithSmallFeatureSize(
|
||
const T *__restrict__ dout,
|
||
const T *__restrict__ input,
|
||
const int n1,
|
||
const int n2,
|
||
const U *__restrict__ mean,
|
||
const U *__restrict__ var,
|
||
const float epsilon,
|
||
const ScaleT *__restrict__ gamma,
|
||
T *grad_input) {
|
||
constexpr int WarpSize = 32;
|
||
#ifdef __HIPCC__
|
||
for (int64_t bid = hipBlockIdx_x; bid < n1; bid += hipGridDim_x) {
|
||
#else
|
||
for (int64_t bid = blockIdx.x; bid < n1; bid += gridDim.x) {
|
||
#endif
|
||
U sum_loss1 = U(0);
|
||
U sum_loss2 = U(0);
|
||
const U c_mean = mean[bid];
|
||
const U c_invvar = rsqrt_<U>(var[bid] + epsilon);
|
||
|
||
const int64_t main_vec_n2 = n2 / DataPerTid;
|
||
const int64_t tid_num = WarpSize * blockDim.y;
|
||
const int64_t thrx = threadIdx.x + threadIdx.y * WarpSize;
|
||
|
||
// One feature-size per block.
|
||
const T *__restrict__ k_dout = dout + bid * n2;
|
||
const T *__restrict__ k_input = input + bid * n2;
|
||
T *k_grad_input = grad_input + bid * n2;
|
||
|
||
// Data storage location in local register.
|
||
using VecT = AlignedVector<T, DataPerTid>;
|
||
using VecScaleT = AlignedVector<ScaleT, DataPerTid>;
|
||
|
||
const VecT *__restrict__ v_k_dout =
|
||
reinterpret_cast<const VecT *__restrict__>(k_dout);
|
||
const VecT *__restrict__ v_k_input =
|
||
reinterpret_cast<const VecT *__restrict__>(k_input);
|
||
const VecScaleT *__restrict__ v_gamma =
|
||
reinterpret_cast<const VecScaleT *__restrict__>(gamma);
|
||
VecT *v_grad = reinterpret_cast<VecT *>(k_grad_input);
|
||
|
||
// Each thread shall deal with no more than 8 data.
|
||
U dout_data[8];
|
||
U input_data[8];
|
||
U gamma_data[8];
|
||
|
||
if (gamma != NULL) {
|
||
int64_t tid = thrx;
|
||
for (int64_t i = 0; tid < main_vec_n2; tid += tid_num, ++i) {
|
||
VecT v_tmp_dout = v_k_dout[tid];
|
||
VecT v_tmp_input = v_k_input[tid];
|
||
VecScaleT v_tmp_gamma = v_gamma[tid];
|
||
#pragma unroll
|
||
for (int64_t k = 0; k < DataPerTid; ++k) {
|
||
const int64_t idx = k + i * DataPerTid;
|
||
dout_data[idx] = static_cast<U>(v_tmp_dout[k]);
|
||
input_data[idx] = static_cast<U>(v_tmp_input[k]);
|
||
gamma_data[idx] = static_cast<U>(v_tmp_gamma[k]);
|
||
sum_loss1 += dout_data[idx] * gamma_data[idx];
|
||
sum_loss2 += dout_data[idx] * gamma_data[idx] *
|
||
(input_data[idx] - c_mean) * c_invvar;
|
||
}
|
||
}
|
||
} else {
|
||
int64_t tid = thrx;
|
||
for (int64_t i = 0; tid < main_vec_n2; tid += tid_num, ++i) {
|
||
VecT v_tmp_dout = v_k_dout[tid];
|
||
VecT v_tmp_input = v_k_input[tid];
|
||
#pragma unroll
|
||
for (int64_t k = 0; k < DataPerTid; ++k) {
|
||
const int64_t idx = k + i * DataPerTid;
|
||
dout_data[idx] = static_cast<U>(v_tmp_dout[k]);
|
||
input_data[idx] = static_cast<U>(v_tmp_input[k]);
|
||
sum_loss1 += dout_data[idx];
|
||
sum_loss2 += dout_data[idx] * (input_data[idx] - c_mean) * c_invvar;
|
||
}
|
||
}
|
||
}
|
||
|
||
// intra-warp reductions
|
||
#pragma unroll
|
||
for (int mask = WarpSize / 2; mask > 0; mask /= 2) {
|
||
#ifdef PADDLE_WITH_HIP
|
||
// WARP_SHFL_XOR(sum_loss, mask);
|
||
sum_loss1 += __shfl_xor(sum_loss1, mask, warpSize);
|
||
sum_loss2 += __shfl_xor(sum_loss2, mask, warpSize);
|
||
#else
|
||
// WARP_SHFL_XOR(sum_loss, mask);
|
||
sum_loss1 += __shfl_xor_sync(0xffffffff, sum_loss1, mask, WarpSize);
|
||
sum_loss2 += __shfl_xor_sync(0xffffffff, sum_loss2, mask, WarpSize);
|
||
#endif
|
||
}
|
||
|
||
// inter-warp reductions
|
||
if (blockDim.y > 1) {
|
||
__shared__ U buf[512];
|
||
for (int offset = blockDim.y / 2; offset > 0; offset /= 2) {
|
||
// upper half of warps write to shared
|
||
if (threadIdx.y >= offset && threadIdx.y < 2 * offset) {
|
||
const int wrt_i = (threadIdx.y - offset) * WarpSize + threadIdx.x;
|
||
buf[2 * wrt_i] = sum_loss1;
|
||
buf[2 * wrt_i + 1] = sum_loss2;
|
||
}
|
||
__syncthreads();
|
||
// lower half merges
|
||
if (threadIdx.y < offset) {
|
||
const int read_i = threadIdx.y * blockDim.x + threadIdx.x;
|
||
sum_loss1 += buf[2 * read_i];
|
||
sum_loss2 += buf[2 * read_i + 1];
|
||
}
|
||
__syncthreads();
|
||
}
|
||
if (threadIdx.y == 0) {
|
||
buf[2 * threadIdx.x] = sum_loss1;
|
||
buf[2 * threadIdx.x + 1] = sum_loss2;
|
||
}
|
||
__syncthreads();
|
||
if (threadIdx.y != 0) {
|
||
sum_loss1 = buf[2 * threadIdx.x];
|
||
sum_loss2 = buf[2 * threadIdx.x + 1];
|
||
}
|
||
}
|
||
|
||
U fH = static_cast<U>(n2);
|
||
U ratio_term = (static_cast<U>(1) / fH) * c_invvar;
|
||
if (gamma != NULL) {
|
||
int64_t tid = thrx;
|
||
for (int64_t i = 0; tid < main_vec_n2; tid += tid_num, ++i) {
|
||
VecT temp_grad;
|
||
#pragma unroll
|
||
for (int k = 0; k < DataPerTid; ++k) {
|
||
const int idx = i * DataPerTid + k;
|
||
const U c_h = input_data[idx];
|
||
const U c_loss = dout_data[idx];
|
||
U f_grad_input = fH * c_loss * gamma_data[idx] - sum_loss1;
|
||
f_grad_input -= (c_h - c_mean) * c_invvar * sum_loss2;
|
||
temp_grad[k] = static_cast<T>(f_grad_input * ratio_term);
|
||
}
|
||
v_grad[tid] = temp_grad;
|
||
}
|
||
} else {
|
||
int64_t tid = thrx;
|
||
for (int64_t i = 0; tid < main_vec_n2; tid += tid_num, ++i) {
|
||
VecT temp_grad;
|
||
#pragma unroll
|
||
for (int k = 0; k < DataPerTid; ++k) {
|
||
const int idx = i * DataPerTid + k;
|
||
const U c_h = input_data[idx];
|
||
const U c_loss = dout_data[idx];
|
||
U f_grad_input = fH * c_loss - sum_loss1;
|
||
f_grad_input -= (c_h - c_mean) * c_invvar * sum_loss2;
|
||
temp_grad[k] = static_cast<T>(f_grad_input * ratio_term);
|
||
}
|
||
v_grad[tid] = temp_grad;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Make sure that d_scale != nullptr && d_bias != nullptr
|
||
// Since d_scale != nullptr, scale would not be nullptr
|
||
template <typename T,
|
||
typename U,
|
||
int BlockDim,
|
||
bool HasDx,
|
||
bool ScaleBiasWithSameTypeX>
|
||
__global__ void LayerNormBackwardGradientAll(
|
||
const T *x,
|
||
const T *d_y,
|
||
LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *d_scale,
|
||
LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *d_bias,
|
||
T *d_x,
|
||
const U *mean,
|
||
const U *var,
|
||
const LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *scale,
|
||
float epsilon,
|
||
int64_t batch_size,
|
||
int64_t feature_size,
|
||
int64_t col_offset) {
|
||
using ScaleBiasT = LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX>;
|
||
int64_t beg_idx = static_cast<int64_t>(threadIdx.x) * feature_size +
|
||
(static_cast<int64_t>(blockIdx.x) + col_offset);
|
||
int64_t end_idx = batch_size * feature_size +
|
||
(static_cast<int64_t>(blockIdx.x) + col_offset);
|
||
int64_t stride = BlockDim * feature_size;
|
||
|
||
U d_scale_partial = static_cast<U>(0), d_bias_partial = static_cast<U>(0);
|
||
|
||
for (int64_t i = beg_idx; i < end_idx; i += stride) {
|
||
int row_idx = i / feature_size;
|
||
auto var_val = rsqrt_(static_cast<U>(var[row_idx]) + epsilon);
|
||
d_scale_partial += static_cast<U>(d_y[i]) *
|
||
(static_cast<U>(x[i]) - mean[row_idx]) * var_val;
|
||
d_bias_partial += static_cast<U>(d_y[i]);
|
||
if (HasDx) {
|
||
d_x[i] = static_cast<T>(static_cast<U>(d_y[i]) *
|
||
static_cast<U>(scale[blockIdx.x + col_offset]) *
|
||
var_val);
|
||
}
|
||
}
|
||
|
||
__shared__ U shared_scale[32]; // threadIdx.x / warpSize <= kMaxBlockDim /
|
||
// warpSize <= 1024/32 = 32;
|
||
__shared__ U shared_bias[32];
|
||
d_scale_partial = BlockReduceSum<U>(d_scale_partial, shared_scale);
|
||
d_bias_partial = BlockReduceSum<U>(d_bias_partial, shared_bias);
|
||
|
||
if (threadIdx.x == 0) {
|
||
d_scale[blockIdx.x + col_offset] = static_cast<ScaleBiasT>(d_scale_partial);
|
||
d_bias[blockIdx.x + col_offset] = static_cast<ScaleBiasT>(d_bias_partial);
|
||
}
|
||
}
|
||
|
||
// Make sure that there is only one true expression: d_scale != nullptr
|
||
// or d_bias != nullptr
|
||
// Notice: scale may be nullptr
|
||
template <typename T,
|
||
typename U,
|
||
int BlockDim,
|
||
bool HasDx,
|
||
bool HasDScale,
|
||
bool ScaleBiasWithSameTypeX>
|
||
__global__ void LayerNormBackwardGradientScaleOrBias(
|
||
const T *x,
|
||
const T *d_y,
|
||
LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *d_scale,
|
||
LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *d_bias,
|
||
T *d_x,
|
||
const U *mean,
|
||
const U *var,
|
||
const LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *scale,
|
||
float epsilon,
|
||
int64_t batch_size,
|
||
int64_t feature_size,
|
||
int col_offset) {
|
||
using ScaleBiasT = LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX>;
|
||
using BlockReduce = cub::BlockReduce<U, BlockDim>;
|
||
__shared__ typename BlockReduce::TempStorage temp_storage;
|
||
int64_t beg_idx = static_cast<int64_t>(threadIdx.x) * feature_size +
|
||
static_cast<int64_t>(blockIdx.x) + col_offset;
|
||
int64_t end_idx =
|
||
batch_size * feature_size + static_cast<int64_t>(blockIdx.x) + col_offset;
|
||
int64_t stride = BlockDim * feature_size;
|
||
U d_scale_or_d_bias_partial = static_cast<U>(0);
|
||
|
||
for (int64_t i = beg_idx; i < end_idx; i += stride) {
|
||
int row_idx = i / feature_size;
|
||
auto var_val =
|
||
static_cast<U>(rsqrt_(static_cast<float>(var[row_idx]) + epsilon));
|
||
if (HasDScale) {
|
||
d_scale_or_d_bias_partial += static_cast<U>(d_y[i]) *
|
||
(static_cast<U>(x[i]) - mean[row_idx]) *
|
||
var_val;
|
||
} else { // d_bias != nullptr
|
||
d_scale_or_d_bias_partial += static_cast<U>(d_y[i]);
|
||
}
|
||
|
||
if (HasDx) {
|
||
if (scale != nullptr) {
|
||
d_x[i] = static_cast<T>(static_cast<U>(d_y[i]) *
|
||
static_cast<U>(scale[blockIdx.x + col_offset]) *
|
||
var_val);
|
||
} else {
|
||
d_x[i] = static_cast<T>(static_cast<U>(d_y[i]) * var_val);
|
||
}
|
||
}
|
||
}
|
||
|
||
d_scale_or_d_bias_partial =
|
||
BlockReduce(temp_storage).Reduce(d_scale_or_d_bias_partial, cub::Sum());
|
||
|
||
if (threadIdx.x == 0) {
|
||
if (HasDScale) {
|
||
d_scale[blockIdx.x + col_offset] =
|
||
static_cast<ScaleBiasT>(d_scale_or_d_bias_partial);
|
||
} else {
|
||
d_bias[blockIdx.x + col_offset] =
|
||
static_cast<ScaleBiasT>(d_scale_or_d_bias_partial);
|
||
}
|
||
}
|
||
}
|
||
|
||
template <typename T, typename U, int BlockDim>
|
||
__global__ void LayerNormBackwardPostProcessToCalculateDX(
|
||
const T *x,
|
||
T *d_x,
|
||
const U *mean,
|
||
const U *var,
|
||
float epsilon,
|
||
int64_t feature_size) {
|
||
using BlockReduce = cub::BlockReduce<PairForLayerNorm<U>, BlockDim>;
|
||
__shared__ typename BlockReduce::TempStorage temp_storage;
|
||
__shared__ U d_x_reduce_tmp[2];
|
||
|
||
int64_t beg_idx = static_cast<int64_t>(blockIdx.x) * feature_size +
|
||
static_cast<int64_t>(threadIdx.x);
|
||
int64_t end_idx = (static_cast<int64_t>(blockIdx.x) + 1) * feature_size;
|
||
|
||
U block_mean = mean[blockIdx.x];
|
||
U block_var = var[blockIdx.x];
|
||
U d_x_mean_partial = static_cast<U>(0), d_x_var_partial = static_cast<U>(0);
|
||
for (int64_t i = beg_idx; i < end_idx; i += BlockDim) {
|
||
d_x_mean_partial += static_cast<U>(d_x[i]);
|
||
d_x_var_partial +=
|
||
static_cast<U>(d_x[i]) * (static_cast<U>(x[i]) - block_mean);
|
||
}
|
||
|
||
auto pair =
|
||
BlockReduce(temp_storage)
|
||
.Reduce(PairForLayerNorm<U>(d_x_mean_partial, d_x_var_partial),
|
||
PairForLayerNormAddFunctor<U>());
|
||
|
||
if (threadIdx.x == 0) {
|
||
d_x_reduce_tmp[0] = static_cast<float>(pair.first_) / feature_size;
|
||
d_x_reduce_tmp[1] =
|
||
static_cast<float>(pair.second_) /
|
||
(feature_size * (static_cast<float>(block_var) + epsilon));
|
||
}
|
||
__syncthreads();
|
||
|
||
d_x_mean_partial = d_x_reduce_tmp[0];
|
||
d_x_var_partial = d_x_reduce_tmp[1];
|
||
for (int64_t i = beg_idx; i < end_idx; i += BlockDim) {
|
||
d_x[i] -= static_cast<T>(d_x_mean_partial);
|
||
d_x[i] -=
|
||
static_cast<T>((static_cast<U>(x[i]) - block_mean) * d_x_var_partial);
|
||
}
|
||
}
|
||
|
||
// Here, we only calculate d_x
|
||
template <typename T, typename U, int BlockDim, bool ScaleBiasWithSameTypeX>
|
||
__global__ void LayerNormBackwardGradientOnlyDX(
|
||
const T *x,
|
||
const T *d_y,
|
||
T *d_x,
|
||
const U *mean,
|
||
const U *var,
|
||
const LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *scale,
|
||
float epsilon,
|
||
int64_t feature_size) {
|
||
using ScaleBiasT = LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX>;
|
||
using BlockReduce = cub::BlockReduce<PairForLayerNorm<U>, BlockDim>;
|
||
__shared__ typename BlockReduce::TempStorage temp_storage;
|
||
__shared__ U d_x_reduce_tmp[2];
|
||
|
||
int64_t beg_idx = static_cast<int64_t>(blockIdx.x) * feature_size +
|
||
static_cast<int64_t>(threadIdx.x);
|
||
int64_t end_idx = (static_cast<int64_t>(blockIdx.x) + 1) * feature_size;
|
||
|
||
U block_mean = mean[blockIdx.x], block_var = var[blockIdx.x];
|
||
U d_x_mean_partial = static_cast<U>(0), d_x_var_partial = static_cast<U>(0);
|
||
for (int64_t i = beg_idx; i < end_idx; i += BlockDim) {
|
||
auto var_val =
|
||
static_cast<U>(rsqrt_(static_cast<float>(block_var) + epsilon));
|
||
if (scale != nullptr) {
|
||
int col_idx = i % feature_size;
|
||
d_x[i] = static_cast<T>(static_cast<U>(d_y[i]) *
|
||
static_cast<U>(scale[col_idx]) * var_val);
|
||
} else {
|
||
d_x[i] = static_cast<T>(static_cast<U>(d_y[i]) * var_val);
|
||
}
|
||
d_x_mean_partial += static_cast<U>(d_x[i]);
|
||
d_x_var_partial +=
|
||
static_cast<U>(d_x[i]) * (static_cast<U>(x[i]) - block_mean);
|
||
}
|
||
|
||
auto pair =
|
||
BlockReduce(temp_storage)
|
||
.Reduce(PairForLayerNorm<U>(d_x_mean_partial, d_x_var_partial),
|
||
PairForLayerNormAddFunctor<U>());
|
||
|
||
if (threadIdx.x == 0) {
|
||
d_x_reduce_tmp[0] = static_cast<float>(pair.first_) / feature_size;
|
||
d_x_reduce_tmp[1] =
|
||
static_cast<float>(pair.second_) /
|
||
(feature_size * (static_cast<float>(block_var) + epsilon));
|
||
}
|
||
__syncthreads();
|
||
|
||
d_x_mean_partial = d_x_reduce_tmp[0];
|
||
d_x_var_partial = d_x_reduce_tmp[1];
|
||
for (int64_t i = beg_idx; i < end_idx; i += BlockDim) {
|
||
d_x[i] -= static_cast<T>(d_x_mean_partial);
|
||
d_x[i] -=
|
||
static_cast<T>((static_cast<U>(x[i]) - block_mean) * d_x_var_partial);
|
||
}
|
||
}
|
||
|
||
template <typename T, typename U, bool ScaleBiasWithSameTypeX>
|
||
__global__ void LayerNormBackwardWhenBatchSizeIsOne(
|
||
const T *x,
|
||
const T *d_y,
|
||
T *d_x,
|
||
LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *d_scale,
|
||
LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *d_bias,
|
||
const U *mean,
|
||
const U *var,
|
||
const LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *scale,
|
||
float epsilon,
|
||
int64_t feature_size) {
|
||
int64_t idx =
|
||
static_cast<int64_t>(threadIdx.x) +
|
||
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x);
|
||
using ScaleBiasT = LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX>;
|
||
if (idx < feature_size) {
|
||
auto var_val = static_cast<U>(rsqrt_(static_cast<float>(var[0]) + epsilon));
|
||
if (d_x != nullptr) {
|
||
if (d_scale == nullptr) {
|
||
d_x[idx] = static_cast<T>(static_cast<U>(d_y[idx]) * var_val);
|
||
} else {
|
||
d_x[idx] = static_cast<T>(static_cast<U>(d_y[idx]) *
|
||
static_cast<U>(scale[idx]) * var_val);
|
||
}
|
||
}
|
||
|
||
if (d_scale != nullptr) {
|
||
d_scale[idx] =
|
||
static_cast<ScaleBiasT>(static_cast<U>(d_y[idx]) *
|
||
(static_cast<U>(x[idx]) - mean[0]) * var_val);
|
||
}
|
||
|
||
if (d_bias != nullptr) {
|
||
d_bias[idx] = static_cast<ScaleBiasT>(d_y[idx]);
|
||
}
|
||
}
|
||
}
|
||
|
||
inline int VecSizeJudgeForeGradInput(const int feature_size,
|
||
const int vec_size) {
|
||
if (!(feature_size & (vec_size - 1))) {
|
||
return vec_size;
|
||
} else if (vec_size == 4) {
|
||
if (!(feature_size & 1)) {
|
||
return 2;
|
||
}
|
||
}
|
||
return 1;
|
||
}
|
||
|
||
template <typename T, typename U, bool ScaleBiasWithSameTypeX = false>
|
||
static void LayerNormBackward(
|
||
const T *x,
|
||
const T *d_y,
|
||
const LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *scale,
|
||
const U *mean,
|
||
const U *var,
|
||
T *d_x,
|
||
LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *d_scale,
|
||
LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX> *d_bias,
|
||
float epsilon,
|
||
int64_t batch_size,
|
||
int64_t feature_size,
|
||
const GPUContext &dev_ctx) {
|
||
auto stream = dev_ctx.stream();
|
||
const int kMaxBlockDim = 512;
|
||
const int kMaxBlockNum = 128;
|
||
// TODO(large-tensor): generic backward kernel launch uses int32 grid dim
|
||
PADDLE_ENFORCE_LE_INT_MAX(batch_size, "batch_size");
|
||
int gradient_flag = ((d_x != nullptr ? 1 : 0) << 2) |
|
||
((d_scale != nullptr ? 1 : 0) << 1) |
|
||
((d_bias != nullptr ? 1 : 0));
|
||
if (gradient_flag == 0) return;
|
||
if (batch_size == 1) {
|
||
// TODO(large-tensor): batch_size==1 path uses int32 grid dim
|
||
PADDLE_ENFORCE_LE_INT_MAX(
|
||
(feature_size + kMaxBlockDim - 1) / kMaxBlockDim,
|
||
"(feature_size + kMaxBlockDim - 1) / kMaxBlockDim");
|
||
LayerNormBackwardWhenBatchSizeIsOne<T, U, ScaleBiasWithSameTypeX>
|
||
<<<(feature_size + kMaxBlockDim - 1) / kMaxBlockDim,
|
||
kMaxBlockDim,
|
||
0,
|
||
stream>>>(x,
|
||
d_y,
|
||
d_x,
|
||
d_scale,
|
||
d_bias,
|
||
mean,
|
||
var,
|
||
scale,
|
||
epsilon,
|
||
feature_size);
|
||
|
||
if (d_x != nullptr) {
|
||
switch (GetDesiredBlockDim(feature_size)) {
|
||
FIXED_BLOCK_DIM_CASE(
|
||
LayerNormBackwardPostProcessToCalculateDX<T, U, kBlockDim>
|
||
<<<1, kBlockDim, 0, stream>>>(
|
||
x, d_x, mean, var, epsilon, feature_size));
|
||
}
|
||
}
|
||
return;
|
||
}
|
||
|
||
auto block_dim = GetDesiredBlockDim(batch_size);
|
||
switch (gradient_flag) {
|
||
case 1: // d_x == nulptr, d_scale == nullptr, d_bias != nullptr
|
||
switch (block_dim) {
|
||
FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE(
|
||
feature_size,
|
||
kMaxBlockNum,
|
||
LayerNormBackwardGradientScaleOrBias<T,
|
||
U,
|
||
kBlockDim,
|
||
false,
|
||
false,
|
||
ScaleBiasWithSameTypeX>
|
||
<<<block_num, kBlockDim, 0, stream>>>(x,
|
||
d_y,
|
||
d_scale,
|
||
d_bias,
|
||
d_x,
|
||
mean,
|
||
var,
|
||
scale,
|
||
epsilon,
|
||
batch_size,
|
||
feature_size,
|
||
col_offset));
|
||
}
|
||
break;
|
||
case 2: // d_x == nullptr, d_scale != nullptr, d_bias == nullptr
|
||
switch (block_dim) {
|
||
FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE(
|
||
feature_size,
|
||
kMaxBlockNum,
|
||
LayerNormBackwardGradientScaleOrBias<T,
|
||
U,
|
||
kBlockDim,
|
||
false,
|
||
true,
|
||
ScaleBiasWithSameTypeX>
|
||
<<<block_num, kBlockDim, 0, stream>>>(x,
|
||
d_y,
|
||
d_scale,
|
||
d_bias,
|
||
d_x,
|
||
mean,
|
||
var,
|
||
scale,
|
||
epsilon,
|
||
batch_size,
|
||
feature_size,
|
||
col_offset));
|
||
}
|
||
break;
|
||
case 3: // d_x == nullptr, d_scale != nulptr, d_bias != nullptr
|
||
switch (block_dim) {
|
||
FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE(
|
||
feature_size,
|
||
kMaxBlockNum,
|
||
LayerNormBackwardGradientAll<T,
|
||
U,
|
||
kBlockDim,
|
||
false,
|
||
ScaleBiasWithSameTypeX>
|
||
<<<block_num, kBlockDim, 0, stream>>>(x,
|
||
d_y,
|
||
d_scale,
|
||
d_bias,
|
||
d_x,
|
||
mean,
|
||
var,
|
||
scale,
|
||
epsilon,
|
||
batch_size,
|
||
feature_size,
|
||
col_offset));
|
||
}
|
||
break;
|
||
case 4: // d_x != nullptr, d_scale == nullptr, d_bias == nullptr
|
||
switch (GetDesiredBlockDim(feature_size)) {
|
||
FIXED_BLOCK_DIM_CASE(
|
||
LayerNormBackwardGradientOnlyDX<T,
|
||
U,
|
||
kBlockDim,
|
||
ScaleBiasWithSameTypeX>
|
||
<<<batch_size, kBlockDim, 0, stream>>>(
|
||
x, d_y, d_x, mean, var, scale, epsilon, feature_size));
|
||
}
|
||
break;
|
||
case 5: // d_x != nulptr, d_scale == nullptr, d_bias != nullptr
|
||
switch (block_dim) {
|
||
FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE(
|
||
feature_size,
|
||
kMaxBlockNum,
|
||
LayerNormBackwardGradientScaleOrBias<T,
|
||
U,
|
||
kBlockDim,
|
||
true,
|
||
false,
|
||
ScaleBiasWithSameTypeX>
|
||
<<<block_num, kBlockDim, 0, stream>>>(x,
|
||
d_y,
|
||
d_scale,
|
||
d_bias,
|
||
d_x,
|
||
mean,
|
||
var,
|
||
scale,
|
||
epsilon,
|
||
batch_size,
|
||
feature_size,
|
||
col_offset));
|
||
}
|
||
switch (GetDesiredBlockDim(feature_size)) {
|
||
FIXED_BLOCK_DIM_CASE(
|
||
LayerNormBackwardPostProcessToCalculateDX<T, U, kBlockDim>
|
||
<<<batch_size, kBlockDim, 0, stream>>>(
|
||
x, d_x, mean, var, epsilon, feature_size));
|
||
}
|
||
break;
|
||
case 6: // d_x != nullptr, d_scale != nullptr, d_bias == nullptr
|
||
switch (block_dim) {
|
||
FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE(
|
||
feature_size,
|
||
kMaxBlockNum,
|
||
LayerNormBackwardGradientScaleOrBias<T,
|
||
U,
|
||
kBlockDim,
|
||
true,
|
||
true,
|
||
ScaleBiasWithSameTypeX>
|
||
<<<block_num, kBlockDim, 0, stream>>>(x,
|
||
d_y,
|
||
d_scale,
|
||
d_bias,
|
||
d_x,
|
||
mean,
|
||
var,
|
||
scale,
|
||
epsilon,
|
||
batch_size,
|
||
feature_size,
|
||
col_offset));
|
||
}
|
||
switch (GetDesiredBlockDim(feature_size)) {
|
||
FIXED_BLOCK_DIM_CASE(
|
||
LayerNormBackwardPostProcessToCalculateDX<T, U, kBlockDim>
|
||
<<<batch_size, kBlockDim, 0, stream>>>(
|
||
x, d_x, mean, var, epsilon, feature_size));
|
||
}
|
||
break;
|
||
case 7: // d_x != nullptr, d_scale != nullptr, d_bias != nullptr
|
||
{
|
||
#ifdef PADDLE_WITH_CUDA
|
||
bool can_call_fast_kernel = false;
|
||
// todo: rule out double type.
|
||
if ((feature_size == 1024 || feature_size == 384 ||
|
||
feature_size == 256) &&
|
||
sizeof(T) <= 4) {
|
||
can_call_fast_kernel = true;
|
||
}
|
||
|
||
VLOG(6) << "can_call_fast_kernel = " << can_call_fast_kernel;
|
||
if (can_call_fast_kernel) {
|
||
ln_bwd_fast_kernel_driver<
|
||
T,
|
||
U,
|
||
LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX>>(dev_ctx,
|
||
batch_size,
|
||
feature_size,
|
||
epsilon,
|
||
x,
|
||
scale,
|
||
mean,
|
||
var,
|
||
d_y,
|
||
d_x,
|
||
d_scale,
|
||
d_bias);
|
||
} else {
|
||
#endif
|
||
using ScaleT = LayerNormScaleBiasT<T, U, ScaleBiasWithSameTypeX>;
|
||
constexpr int BDIMX = 32;
|
||
|
||
constexpr int VPT = 4;
|
||
constexpr int BDIMY1 = 4;
|
||
constexpr int PartSize = BDIMY1 * VPT;
|
||
dim3 threads2(BDIMX, BDIMY1, 1);
|
||
dim3 blocks2((feature_size + BDIMX - 1) / BDIMX, PartSize, 1);
|
||
|
||
int64_t param_num = PartSize * feature_size;
|
||
auto part_grad_param_ptr = phi::memory_utils::Alloc(
|
||
dev_ctx.GetPlace(),
|
||
param_num * sizeof(U) * 2, // for both gamma and beta
|
||
phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
|
||
|
||
U *part_grad_gamma = reinterpret_cast<U *>(part_grad_param_ptr->ptr());
|
||
U *part_grad_beta = reinterpret_cast<U *>(part_grad_gamma + param_num);
|
||
|
||
LayerNormBackwardPartGradGammaBeta<T, U, BDIMX, BDIMY1, VPT>
|
||
<<<blocks2, threads2, 0, stream>>>(d_y,
|
||
x,
|
||
batch_size,
|
||
feature_size,
|
||
mean,
|
||
var,
|
||
epsilon,
|
||
part_grad_gamma,
|
||
part_grad_beta);
|
||
|
||
constexpr int BDIMY2 = 8;
|
||
dim3 threads3(BDIMX, BDIMY2, 1);
|
||
const dim3 blocks3((feature_size + BDIMX - 1) / BDIMX, 1, 1);
|
||
LayerNormBackwardSumGradGammaBeta<T, U, BDIMX, BDIMY2, ScaleT>
|
||
<<<blocks3, threads3, 0, stream>>>(part_grad_gamma,
|
||
part_grad_beta,
|
||
PartSize,
|
||
batch_size,
|
||
feature_size,
|
||
d_scale,
|
||
d_bias);
|
||
|
||
uint64_t addr = reinterpret_cast<uint64_t>(d_y) |
|
||
reinterpret_cast<uint64_t>(x) |
|
||
reinterpret_cast<uint64_t>(d_x);
|
||
int vec_size =
|
||
std::min(4, phi::GetVectorizedSize<T>(reinterpret_cast<T *>(addr)));
|
||
int real_vec = VecSizeJudgeForeGradInput(feature_size, vec_size);
|
||
|
||
if (feature_size <= 2048) {
|
||
// One thread must work with at least real_vec quantity data, at most
|
||
// 8 data.
|
||
int data_per_warp = BDIMX * real_vec;
|
||
uint32_t warp_num =
|
||
feature_size < data_per_warp ? 1 : (feature_size / data_per_warp);
|
||
#if defined(__clang__) || defined(__GNUC__)
|
||
int block_dim_y = std::min(8, 1 << (31 - __builtin_clz(warp_num)));
|
||
#else
|
||
int block_dim_y = 1;
|
||
while (warp_num != 0) {
|
||
warp_num = warp_num >> 1;
|
||
block_dim_y <<= 1;
|
||
}
|
||
block_dim_y = std::min(8, (block_dim_y / 2));
|
||
#endif // __GNUCC__
|
||
|
||
dim3 threads1(BDIMX, block_dim_y, 1);
|
||
#define IMPL_BACKWARD_FOR_INPUT(num) \
|
||
LayerNormBackwardComputeGradInputWithSmallFeatureSize<T, U, ScaleT, num> \
|
||
<<<batch_size, threads1, 0, stream>>>( \
|
||
d_y, x, batch_size, feature_size, mean, var, epsilon, scale, d_x);
|
||
|
||
switch (real_vec) {
|
||
case 4: {
|
||
IMPL_BACKWARD_FOR_INPUT(4);
|
||
} break;
|
||
case 2: {
|
||
IMPL_BACKWARD_FOR_INPUT(2);
|
||
} break;
|
||
default: {
|
||
IMPL_BACKWARD_FOR_INPUT(1);
|
||
}
|
||
}
|
||
#undef IMPL_BACKWARD_FOR_INPUT
|
||
|
||
} else {
|
||
constexpr int BDIMY3 = 4;
|
||
dim3 threads1(BDIMX, BDIMY3, 1);
|
||
LayerNormBackwardComputeGradInput<T, U, BDIMX, BDIMY3, ScaleT>
|
||
<<<batch_size, threads1, 0, stream>>>(d_y,
|
||
x,
|
||
batch_size,
|
||
feature_size,
|
||
mean,
|
||
var,
|
||
epsilon,
|
||
scale,
|
||
d_x);
|
||
}
|
||
#ifdef PADDLE_WITH_CUDA
|
||
}
|
||
#endif
|
||
|
||
break;
|
||
}
|
||
default:
|
||
break;
|
||
}
|
||
}
|
||
|
||
} // namespace funcs
|
||
} // namespace phi
|