wehub-resource-sync
794 lines
33 KiB
C++
794 lines
33 KiB
C++
/* Copyright (c) 2020 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 <algorithm>
|
|
#include <cfloat>
|
|
#include <string>
|
|
#include <vector>
|
|
#include "paddle/common/enforce.h"
|
|
#include "paddle/common/layout.h"
|
|
#include "paddle/phi/kernels/funcs/cub.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
#include "paddle/phi/kernels/funcs/reduce_function.h"
|
|
|
|
#ifdef __HIPCC__
|
|
#define LAUNCH_BOUNDS(BlockDim) __launch_bounds__(BlockDim)
|
|
#else
|
|
#define LAUNCH_BOUNDS(BlockDim)
|
|
#endif
|
|
|
|
namespace phi {
|
|
namespace funcs {
|
|
|
|
// math: dx = scale * ((x - mean) * inv_var / NxHxW * (np.mean(ddx,
|
|
// axis=(n,h,w)) *
|
|
// np.sum(dy, axis=(n,h,w)) -
|
|
// np.sum(dy * ddx, axis=(n,h,w)) + 3 * np.mean(dy * (x -
|
|
// mean),
|
|
// axis=(n,h,w)) * inv_var.pow(2) *
|
|
// np.sum(ddx * (x - mean), axis=(n,h,w))) + inv_var.pow(3) /
|
|
// NxHxW *
|
|
// np.sum(ddx * (x - mean)) *
|
|
// (np.mean(dy, axis=(n,h,w)) - dy) + inv_var.pow(3) / NxHxW *
|
|
// np.sum(dy,
|
|
// axis=(n,h,w)) * (x - mean) *
|
|
// (np.mean(ddx, axis=(n,h,w)) - ddx)) + ddr * (dy * inv_var -
|
|
// inv_var
|
|
// *
|
|
// np.mean(dy, axis=(n,h,w)) -
|
|
// inv_var.pow(3) * (x - mean) * np.mean(dy * (x - mean),
|
|
// axis=(n,h,w)))
|
|
|
|
template <typename T, int BlockDim, DataLayout layout>
|
|
__global__ LAUNCH_BOUNDS(BlockDim) void DoubleGradComputeDX(
|
|
const T *x,
|
|
const T *mean,
|
|
const T *variance,
|
|
const T *ddx,
|
|
const T *dy,
|
|
const T *scale,
|
|
const T *ddscale,
|
|
const int N,
|
|
const int C,
|
|
const int sample_size,
|
|
const double epsilon,
|
|
T *dx) {
|
|
const int outer_size = C;
|
|
const int64_t inner_size = static_cast<int64_t>(N) * sample_size;
|
|
|
|
typedef cub::BlockReduce<T, BlockDim> BlockReduce;
|
|
__shared__ typename BlockReduce::TempStorage dy_storage;
|
|
__shared__ typename BlockReduce::TempStorage ddx_storage;
|
|
__shared__ typename BlockReduce::TempStorage dy_mul_ddx_storage;
|
|
__shared__ typename BlockReduce::TempStorage dy_mul_x_sub_mean_storage;
|
|
__shared__ typename BlockReduce::TempStorage ddx_mul_x_sub_mean_storage;
|
|
__shared__ T dy_sum_val;
|
|
__shared__ T ddx_sum_val;
|
|
__shared__ T dy_mul_ddx_sum_val;
|
|
__shared__ T dy_mul_x_sub_mean_sum_val;
|
|
__shared__ T ddx_mul_x_sub_mean_sum_val;
|
|
|
|
for (int i = blockIdx.x; i < outer_size; i += gridDim.x) {
|
|
T mean_val = mean[i];
|
|
T var_val = variance[i];
|
|
T dy_sum = 0;
|
|
T ddx_sum = 0;
|
|
T dy_mul_ddx_sum = 0;
|
|
T dy_mul_x_sub_mean_sum = 0;
|
|
T ddx_mul_x_sub_mean_sum = 0;
|
|
for (int64_t j = threadIdx.x; j < inner_size; j += blockDim.x) {
|
|
const int64_t index =
|
|
layout == DataLayout::NCHW
|
|
? (j / sample_size * C + i) * sample_size + j % sample_size
|
|
: j * outer_size + i;
|
|
T ddx_i = ddx[index];
|
|
T dy_i = dy[index];
|
|
T tmp = x[index] - mean_val;
|
|
|
|
dy_sum += dy_i;
|
|
ddx_sum += ddx_i;
|
|
dy_mul_ddx_sum += (ddx_i * dy_i);
|
|
|
|
dy_mul_x_sub_mean_sum += (dy_i * tmp);
|
|
ddx_mul_x_sub_mean_sum += (ddx_i * tmp);
|
|
}
|
|
|
|
dy_sum = BlockReduce(dy_storage).Reduce(dy_sum, cub::Sum());
|
|
ddx_sum = BlockReduce(ddx_storage).Reduce(ddx_sum, cub::Sum());
|
|
dy_mul_ddx_sum =
|
|
BlockReduce(dy_mul_ddx_storage).Reduce(dy_mul_ddx_sum, cub::Sum());
|
|
dy_mul_x_sub_mean_sum = BlockReduce(dy_mul_x_sub_mean_storage)
|
|
.Reduce(dy_mul_x_sub_mean_sum, cub::Sum());
|
|
ddx_mul_x_sub_mean_sum = BlockReduce(ddx_mul_x_sub_mean_storage)
|
|
.Reduce(ddx_mul_x_sub_mean_sum, cub::Sum());
|
|
|
|
if (threadIdx.x == 0) {
|
|
dy_sum_val = dy_sum;
|
|
ddx_sum_val = ddx_sum;
|
|
dy_mul_ddx_sum_val = dy_mul_ddx_sum;
|
|
dy_mul_x_sub_mean_sum_val = dy_mul_x_sub_mean_sum;
|
|
ddx_mul_x_sub_mean_sum_val = ddx_mul_x_sub_mean_sum;
|
|
}
|
|
__syncthreads();
|
|
|
|
if (ddx != nullptr) {
|
|
for (int64_t j = threadIdx.x; j < inner_size; j += blockDim.x) {
|
|
const int64_t index =
|
|
layout == DataLayout::NCHW
|
|
? (j / sample_size * C + i) * sample_size + j % sample_size
|
|
: j * outer_size + i;
|
|
dx[index] +=
|
|
((x[index] - mean_val) * var_val * var_val * var_val / inner_size *
|
|
(ddx_sum_val * dy_sum_val / inner_size - dy_mul_ddx_sum_val +
|
|
3. * dy_mul_x_sub_mean_sum_val * var_val *
|
|
ddx_mul_x_sub_mean_sum_val * var_val / inner_size) +
|
|
ddx_mul_x_sub_mean_sum_val * var_val / inner_size * var_val *
|
|
var_val * (dy_sum_val / inner_size - dy[index]) +
|
|
dy_mul_x_sub_mean_sum_val * var_val / inner_size * var_val *
|
|
var_val * (ddx_sum_val / inner_size - ddx[index])) *
|
|
scale[i];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
if (ddscale != nullptr) {
|
|
for (int64_t j = threadIdx.x; j < inner_size; j += blockDim.x) {
|
|
const int64_t index =
|
|
layout == DataLayout::NCHW
|
|
? (j / sample_size * C + i) * sample_size + j % sample_size
|
|
: j * outer_size + i;
|
|
dx[index] += (dy[index] * var_val - dy_sum_val / inner_size * var_val -
|
|
(x[index] - mean_val) * var_val * var_val *
|
|
dy_mul_x_sub_mean_sum_val * var_val / inner_size) *
|
|
ddscale[i];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// math: ddy = (x - mean) * inv_var * ddscale + ddbias +
|
|
// scale * inv_var * (ddx - (x - mean) * inv_var.pow(2) *
|
|
// np.mean(ddx * (x - mean), axis=(n,h,w)))
|
|
template <typename T, int BlockDim, DataLayout layout>
|
|
__global__ LAUNCH_BOUNDS(BlockDim) void DoubleGradComputeDDY(
|
|
const T *x,
|
|
const T *mean,
|
|
const T *variance,
|
|
const T *ddscale,
|
|
const T *ddbias,
|
|
const T *ddx,
|
|
const T *scale,
|
|
const int N,
|
|
const int C,
|
|
const int sample_size,
|
|
const double epsilon,
|
|
T *ddy) {
|
|
const int outer_size = C;
|
|
const int64_t inner_size = static_cast<int64_t>(N) * sample_size;
|
|
|
|
typedef cub::BlockReduce<T, BlockDim> BlockReduce;
|
|
__shared__ typename BlockReduce::TempStorage ddx_storage;
|
|
__shared__ typename BlockReduce::TempStorage ddx_mul_x_sub_mean_storage;
|
|
__shared__ T ddx_sum_val;
|
|
__shared__ T ddx_mul_x_sub_mean_sum_val;
|
|
|
|
for (int i = blockIdx.x; i < outer_size; i += gridDim.x) {
|
|
T mean_val = mean[i];
|
|
T var_val = variance[i];
|
|
T ddx_sum = 0;
|
|
T ddx_mul_x_sub_mean_sum = 0;
|
|
for (int64_t j = threadIdx.x; j < inner_size; j += blockDim.x) {
|
|
const int64_t index =
|
|
layout == DataLayout::NCHW
|
|
? (j / sample_size * C + i) * sample_size + j % sample_size
|
|
: j * outer_size + i;
|
|
T ddx_i = ddx[index];
|
|
ddx_sum += ddx_i;
|
|
ddx_mul_x_sub_mean_sum += (ddx_i * (x[index] - mean_val));
|
|
}
|
|
ddx_sum = BlockReduce(ddx_storage).Reduce(ddx_sum, cub::Sum());
|
|
ddx_mul_x_sub_mean_sum = BlockReduce(ddx_mul_x_sub_mean_storage)
|
|
.Reduce(ddx_mul_x_sub_mean_sum, cub::Sum());
|
|
|
|
if (threadIdx.x == 0) {
|
|
ddx_sum_val = ddx_sum;
|
|
ddx_mul_x_sub_mean_sum_val = ddx_mul_x_sub_mean_sum;
|
|
}
|
|
__syncthreads();
|
|
|
|
if (ddx != nullptr) {
|
|
for (int64_t j = threadIdx.x; j < inner_size; j += blockDim.x) {
|
|
const int64_t index =
|
|
layout == DataLayout::NCHW
|
|
? (j / sample_size * C + i) * sample_size + j % sample_size
|
|
: j * outer_size + i;
|
|
ddy[index] += scale[i] * var_val *
|
|
(ddx[index] - ddx_sum_val / inner_size -
|
|
(x[index] - mean_val) * var_val *
|
|
ddx_mul_x_sub_mean_sum_val * var_val / inner_size);
|
|
}
|
|
}
|
|
__syncthreads();
|
|
if (ddscale != nullptr) {
|
|
for (int64_t j = threadIdx.x; j < inner_size; j += blockDim.x) {
|
|
const int64_t index =
|
|
layout == DataLayout::NCHW
|
|
? (j / sample_size * C + i) * sample_size + j % sample_size
|
|
: j * outer_size + i;
|
|
ddy[index] += (x[index] - mean_val) * var_val * ddscale[i];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
if (ddbias != nullptr) {
|
|
for (int64_t j = threadIdx.x; j < inner_size; j += blockDim.x) {
|
|
const int64_t index =
|
|
layout == DataLayout::NCHW
|
|
? (j / sample_size * C + i) * sample_size + j % sample_size
|
|
: j * outer_size + i;
|
|
ddy[index] += ddbias[i];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// math: dscale = inv_var * (dy - np.mean(dy, axis=(n,h,w) - (x-mean) *
|
|
// inv_var.pow(2) * np.mean(dy * (x-mean), axis=(n,h,w)))) *
|
|
// ddx
|
|
template <typename T, int BlockDim, DataLayout layout>
|
|
__global__ LAUNCH_BOUNDS(BlockDim) void DoubleGradComputeDScale(
|
|
const T *x,
|
|
const T *mean,
|
|
const T *variance,
|
|
const T *ddx,
|
|
const T *dy,
|
|
const int N,
|
|
const int C,
|
|
const int sample_size,
|
|
const double epsilon,
|
|
T *dscale) {
|
|
const int outer_size = C;
|
|
const int64_t inner_size = static_cast<int64_t>(N) * sample_size;
|
|
|
|
typedef cub::BlockReduce<T, BlockDim> BlockReduce;
|
|
__shared__ typename BlockReduce::TempStorage dy_storage;
|
|
__shared__ typename BlockReduce::TempStorage dy_mul_x_sub_mean_storage;
|
|
__shared__ typename BlockReduce::TempStorage dscale_tmp_storage;
|
|
__shared__ T dy_sum_val;
|
|
__shared__ T dy_mul_x_sub_mean_sum_val;
|
|
|
|
for (int i = blockIdx.x; i < outer_size; i += gridDim.x) {
|
|
T dy_sum = 0;
|
|
T dy_mul_x_sub_mean_sum = 0;
|
|
T mean_val = mean[i];
|
|
T var_val = variance[i];
|
|
for (int64_t j = threadIdx.x; j < inner_size; j += blockDim.x) {
|
|
const int64_t index =
|
|
layout == DataLayout::NCHW
|
|
? (j / sample_size * C + i) * sample_size + j % sample_size
|
|
: j * outer_size + i;
|
|
T dy_i = dy[index];
|
|
dy_sum += dy_i;
|
|
dy_mul_x_sub_mean_sum += (dy_i * (x[index] - mean_val));
|
|
}
|
|
dy_sum = BlockReduce(dy_storage).Reduce(dy_sum, cub::Sum());
|
|
dy_mul_x_sub_mean_sum = BlockReduce(dy_mul_x_sub_mean_storage)
|
|
.Reduce(dy_mul_x_sub_mean_sum, cub::Sum());
|
|
|
|
if (threadIdx.x == 0) {
|
|
dy_sum_val = dy_sum;
|
|
dy_mul_x_sub_mean_sum_val = dy_mul_x_sub_mean_sum;
|
|
}
|
|
__syncthreads();
|
|
|
|
if (ddx != nullptr) {
|
|
T dscale_tmp = 0;
|
|
for (int64_t j = threadIdx.x; j < inner_size; j += blockDim.x) {
|
|
const int64_t index =
|
|
layout == DataLayout::NCHW
|
|
? (j / sample_size * C + i) * sample_size + j % sample_size
|
|
: j * outer_size + i;
|
|
dscale_tmp += ddx[index] * var_val *
|
|
(dy[index] - dy_sum_val / inner_size -
|
|
dy_mul_x_sub_mean_sum_val * (x[index] - mean_val) *
|
|
var_val * var_val / inner_size);
|
|
}
|
|
dscale_tmp =
|
|
BlockReduce(dscale_tmp_storage).Reduce(dscale_tmp, cub::Sum());
|
|
|
|
if (threadIdx.x == 0) {
|
|
dscale[i] += dscale_tmp;
|
|
}
|
|
__syncthreads();
|
|
}
|
|
}
|
|
}
|
|
|
|
// math: dscale = np.sum(ddx * dy, axis=(n,h,w)) * inv_var
|
|
template <typename T, int BlockDim, DataLayout layout>
|
|
__global__ LAUNCH_BOUNDS(BlockDim) void DoubleGradComputeDScaleWithGlobal(
|
|
const T *ddx,
|
|
const T *variance,
|
|
const T *dy,
|
|
const double epsilon,
|
|
const int N,
|
|
const int C,
|
|
const int sample_size,
|
|
T *dscale) {
|
|
int outer_size = C;
|
|
int64_t inner_size = static_cast<int64_t>(N) * sample_size;
|
|
typedef cub::BlockReduce<T, BlockDim> BlockReduce;
|
|
__shared__ typename BlockReduce::TempStorage ddx_mul_dy_storage;
|
|
__shared__ T ddx_mul_dy_sum_val;
|
|
for (int i = blockIdx.x; i < outer_size; i += gridDim.x) {
|
|
T inv_var_i = 1.0 / sqrt(variance[i] + epsilon);
|
|
T ddx_mul_dy_sum = 0;
|
|
for (int64_t j = threadIdx.x; j < inner_size; j += blockDim.x) {
|
|
const int64_t index =
|
|
layout == DataLayout::NCHW
|
|
? (j / sample_size * C + i) * sample_size + j % sample_size
|
|
: j * outer_size + i;
|
|
T ddx_i = ddx[index];
|
|
T dy_i = dy[index];
|
|
ddx_mul_dy_sum += (ddx_i * dy_i);
|
|
}
|
|
ddx_mul_dy_sum =
|
|
BlockReduce(ddx_mul_dy_storage).Reduce(ddx_mul_dy_sum, cub::Sum());
|
|
if (threadIdx.x == 0) {
|
|
ddx_mul_dy_sum_val = ddx_mul_dy_sum;
|
|
}
|
|
__syncthreads();
|
|
|
|
if (ddx != nullptr) {
|
|
dscale[i] = inv_var_i * ddx_mul_dy_sum_val;
|
|
}
|
|
}
|
|
}
|
|
|
|
// math: dx = ddscale * dy * inv_var
|
|
template <typename T, DataLayout layout>
|
|
__global__ void DoubleGradComputeDXWithGlobal(const T *dy,
|
|
const T *ddscale,
|
|
const T *variance,
|
|
const double epsilon,
|
|
const int C,
|
|
const int sample_size,
|
|
const int64_t num,
|
|
T *dx) {
|
|
int64_t gid =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
int stride = blockDim.x * gridDim.x;
|
|
if (ddscale != nullptr) {
|
|
for (int64_t i = gid; i < num; i += stride) {
|
|
const int c = layout == DataLayout::NCHW ? i / sample_size % C : i % C;
|
|
T inv_var = 1.0 / sqrt(variance[c] + epsilon);
|
|
dx[i] = dy[i] * ddscale[c] * inv_var;
|
|
}
|
|
}
|
|
}
|
|
|
|
// math: ddy = scale * ddx * inv_var + ddbias +
|
|
// ddscale * (x - mean) * inv_var
|
|
template <typename T, DataLayout layout>
|
|
__global__ void DoubleGradComputeDDYWithGlobal(const T *ddx,
|
|
const T *scale,
|
|
const T *mean,
|
|
const T *variance,
|
|
const T *x,
|
|
const T *ddbias,
|
|
const T *ddscale,
|
|
const double epsilon,
|
|
const int C,
|
|
const int sample_size,
|
|
const int64_t num,
|
|
T *ddy) {
|
|
int64_t gid =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
int stride = blockDim.x * gridDim.x;
|
|
|
|
if (ddx != nullptr) {
|
|
for (int64_t i = gid; i < num; i += stride) {
|
|
const int c = layout == DataLayout::NCHW ? i / sample_size % C : i % C;
|
|
T inv_var = 1.0 / sqrt(variance[c] + epsilon);
|
|
ddy[i] += ddx[i] * scale[c] * inv_var;
|
|
}
|
|
}
|
|
__syncthreads();
|
|
if (ddscale != nullptr) {
|
|
for (int64_t i = gid; i < num; i += stride) {
|
|
const int c = layout == DataLayout::NCHW ? i / sample_size % C : i % C;
|
|
T inv_var = 1.0 / sqrt(variance[c] + epsilon);
|
|
ddy[i] += (x[i] - mean[c]) * inv_var * ddscale[c];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
if (ddbias != nullptr) {
|
|
for (int64_t i = gid; i < num; i += stride) {
|
|
const int c = layout == DataLayout::NCHW ? i / sample_size % C : i % C;
|
|
ddy[i] += ddbias[c];
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename DeviceContext, typename T>
|
|
void NormDoubleGradFunctor(const DeviceContext &dev_ctx,
|
|
const DataLayout data_layout,
|
|
const DenseTensor *X,
|
|
const DenseTensor *Scale,
|
|
const DenseTensor *dY,
|
|
const DenseTensor *Saved_mean,
|
|
const DenseTensor *Saved_variance,
|
|
const DenseTensor *Mean,
|
|
const DenseTensor *Variance,
|
|
const double epsilon,
|
|
const bool use_global_stats,
|
|
const DenseTensor *ddX,
|
|
const DenseTensor *ddScale,
|
|
const DenseTensor *ddBias,
|
|
DenseTensor *dX,
|
|
DenseTensor *dScale,
|
|
DenseTensor *ddY) {
|
|
const T *x_data = X->data<T>();
|
|
const T *dy_data = dY->data<T>();
|
|
const T *ddx_data = (ddX == nullptr ? nullptr : ddX->data<T>());
|
|
|
|
const T *ddscale_data = (ddScale == nullptr ? nullptr : ddScale->data<T>());
|
|
const T *ddbias_data = (ddBias == nullptr ? nullptr : ddBias->data<T>());
|
|
|
|
funcs::SetConstant<DeviceContext, T> set_constant;
|
|
|
|
auto &x_dims = X->dims();
|
|
const int64_t C_64 =
|
|
(data_layout == DataLayout::NCHW ? x_dims[1] : x_dims[x_dims.size() - 1]);
|
|
const int64_t N_64 = x_dims[0];
|
|
const int64_t num = X->numel();
|
|
const int64_t sample_size_64 = num / N_64 / C_64;
|
|
PADDLE_ENFORCE_LE_INT_MAX(C_64, "norm double grad C");
|
|
PADDLE_ENFORCE_LE_INT_MAX(N_64, "norm double grad N");
|
|
PADDLE_ENFORCE_LE_INT_MAX(sample_size_64, "norm double grad sample_size");
|
|
const int C = static_cast<int>(C_64);
|
|
const int N = static_cast<int>(N_64);
|
|
const int sample_size = static_cast<int>(sample_size_64);
|
|
DenseTensor scale_tmp;
|
|
if (!Scale) {
|
|
scale_tmp.Resize({C});
|
|
dev_ctx.template Alloc<T>(&scale_tmp);
|
|
set_constant(dev_ctx, &scale_tmp, static_cast<T>(1));
|
|
}
|
|
const T *scale_data = Scale ? Scale->data<T>() : scale_tmp.data<T>();
|
|
constexpr uint32_t block = 512;
|
|
const int max_threads = dev_ctx.GetMaxPhysicalThreadCount();
|
|
const int64_t max_blocks =
|
|
std::max(static_cast<int64_t>(max_threads / static_cast<int>(block)),
|
|
static_cast<int64_t>(1));
|
|
const uint32_t channel_grid =
|
|
static_cast<uint32_t>(std::min(C_64, max_blocks));
|
|
const uint32_t element_grid =
|
|
static_cast<uint32_t>(std::min((num + block - 1) / block, max_blocks));
|
|
|
|
const T *mean_data, *variance_data;
|
|
if (use_global_stats) {
|
|
const auto *running_mean = Mean;
|
|
const auto *running_var = Variance;
|
|
const auto *running_mean_data = running_mean->template data<T>();
|
|
const auto *running_var_data = running_var->template data<T>();
|
|
mean_data = running_mean_data;
|
|
variance_data = running_var_data;
|
|
} else {
|
|
const T *smean_data = Saved_mean->data<T>();
|
|
const T *svariance_data = Saved_variance->data<T>();
|
|
|
|
mean_data = smean_data;
|
|
variance_data = svariance_data;
|
|
}
|
|
|
|
if (dX) {
|
|
T *dx_data = dev_ctx.template Alloc<T>(dX);
|
|
set_constant(dev_ctx, dX, static_cast<T>(0));
|
|
if (use_global_stats) {
|
|
if (data_layout == DataLayout::NHWC) {
|
|
DoubleGradComputeDXWithGlobal<T, DataLayout::NHWC>
|
|
<<<element_grid, block, 0, dev_ctx.stream()>>>(dy_data,
|
|
ddscale_data,
|
|
variance_data,
|
|
epsilon,
|
|
C,
|
|
sample_size,
|
|
num,
|
|
dx_data);
|
|
} else {
|
|
DoubleGradComputeDXWithGlobal<T, DataLayout::NCHW>
|
|
<<<element_grid, block, 0, dev_ctx.stream()>>>(dy_data,
|
|
ddscale_data,
|
|
variance_data,
|
|
epsilon,
|
|
C,
|
|
sample_size,
|
|
num,
|
|
dx_data);
|
|
}
|
|
} else {
|
|
if (data_layout == DataLayout::NHWC) {
|
|
DoubleGradComputeDX<T, block, DataLayout::NHWC>
|
|
<<<channel_grid, block, 0, dev_ctx.stream()>>>(x_data,
|
|
mean_data,
|
|
variance_data,
|
|
ddx_data,
|
|
dy_data,
|
|
scale_data,
|
|
ddscale_data,
|
|
N,
|
|
C,
|
|
sample_size,
|
|
epsilon,
|
|
dx_data);
|
|
} else {
|
|
DoubleGradComputeDX<T, block, DataLayout::NCHW>
|
|
<<<channel_grid, block, 0, dev_ctx.stream()>>>(x_data,
|
|
mean_data,
|
|
variance_data,
|
|
ddx_data,
|
|
dy_data,
|
|
scale_data,
|
|
ddscale_data,
|
|
N,
|
|
C,
|
|
sample_size,
|
|
epsilon,
|
|
dx_data);
|
|
}
|
|
}
|
|
}
|
|
if (dScale) {
|
|
T *dscale_data = dev_ctx.template Alloc<T>(dScale);
|
|
set_constant(dev_ctx, dScale, static_cast<T>(0));
|
|
if (use_global_stats) {
|
|
if (data_layout == DataLayout::NHWC) {
|
|
DoubleGradComputeDScaleWithGlobal<T, block, DataLayout::NHWC>
|
|
<<<channel_grid, block, 0, dev_ctx.stream()>>>(ddx_data,
|
|
variance_data,
|
|
dy_data,
|
|
epsilon,
|
|
N,
|
|
C,
|
|
sample_size,
|
|
dscale_data);
|
|
} else {
|
|
DoubleGradComputeDScaleWithGlobal<T, block, DataLayout::NCHW>
|
|
<<<channel_grid, block, 0, dev_ctx.stream()>>>(ddx_data,
|
|
variance_data,
|
|
dy_data,
|
|
epsilon,
|
|
N,
|
|
C,
|
|
sample_size,
|
|
dscale_data);
|
|
}
|
|
} else {
|
|
if (data_layout == DataLayout::NHWC) {
|
|
DoubleGradComputeDScale<T, block, DataLayout::NHWC>
|
|
<<<channel_grid, block, 0, dev_ctx.stream()>>>(x_data,
|
|
mean_data,
|
|
variance_data,
|
|
ddx_data,
|
|
dy_data,
|
|
N,
|
|
C,
|
|
sample_size,
|
|
epsilon,
|
|
dscale_data);
|
|
} else {
|
|
DoubleGradComputeDScale<T, block, DataLayout::NCHW>
|
|
<<<channel_grid, block, 0, dev_ctx.stream()>>>(x_data,
|
|
mean_data,
|
|
variance_data,
|
|
ddx_data,
|
|
dy_data,
|
|
N,
|
|
C,
|
|
sample_size,
|
|
epsilon,
|
|
dscale_data);
|
|
}
|
|
}
|
|
}
|
|
if (ddY) {
|
|
T *ddy_data = dev_ctx.template Alloc<T>(ddY);
|
|
set_constant(dev_ctx, ddY, static_cast<T>(0));
|
|
if (use_global_stats) {
|
|
if (data_layout == DataLayout::NHWC) {
|
|
DoubleGradComputeDDYWithGlobal<T, DataLayout::NHWC>
|
|
<<<element_grid, block, 0, dev_ctx.stream()>>>(ddx_data,
|
|
scale_data,
|
|
mean_data,
|
|
variance_data,
|
|
x_data,
|
|
ddbias_data,
|
|
ddscale_data,
|
|
epsilon,
|
|
C,
|
|
sample_size,
|
|
num,
|
|
ddy_data);
|
|
} else {
|
|
DoubleGradComputeDDYWithGlobal<T, DataLayout::NCHW>
|
|
<<<element_grid, block, 0, dev_ctx.stream()>>>(ddx_data,
|
|
scale_data,
|
|
mean_data,
|
|
variance_data,
|
|
x_data,
|
|
ddbias_data,
|
|
ddscale_data,
|
|
epsilon,
|
|
C,
|
|
sample_size,
|
|
num,
|
|
ddy_data);
|
|
}
|
|
} else {
|
|
if (data_layout == DataLayout::NHWC) {
|
|
DoubleGradComputeDDY<T, block, DataLayout::NHWC>
|
|
<<<channel_grid, block, 0, dev_ctx.stream()>>>(x_data,
|
|
mean_data,
|
|
variance_data,
|
|
ddscale_data,
|
|
ddbias_data,
|
|
ddx_data,
|
|
scale_data,
|
|
N,
|
|
C,
|
|
sample_size,
|
|
epsilon,
|
|
ddy_data);
|
|
} else {
|
|
DoubleGradComputeDDY<T, block, DataLayout::NCHW>
|
|
<<<channel_grid, block, 0, dev_ctx.stream()>>>(x_data,
|
|
mean_data,
|
|
variance_data,
|
|
ddscale_data,
|
|
ddbias_data,
|
|
ddx_data,
|
|
scale_data,
|
|
N,
|
|
C,
|
|
sample_size,
|
|
epsilon,
|
|
ddy_data);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename BnT>
|
|
__device__ __forceinline__ void BlockReduceByVertical(BnT x_sum,
|
|
BnT x_square_sum,
|
|
BnT *smem_sum,
|
|
BnT *smem_square_sum,
|
|
BnT *x_sum_out,
|
|
BnT *x_square_sum_out) {
|
|
int tid = threadIdx.x + threadIdx.y * blockDim.x;
|
|
#pragma unroll
|
|
for (int offset = blockDim.y / 2; offset > 0; offset >>= 1) {
|
|
if (threadIdx.y < offset * 2) {
|
|
smem_sum[tid] = x_sum;
|
|
smem_square_sum[tid] = x_square_sum;
|
|
}
|
|
__syncthreads();
|
|
if (threadIdx.y < offset) {
|
|
int pair_tid = tid + offset * blockDim.x;
|
|
x_sum += smem_sum[pair_tid];
|
|
x_square_sum += smem_square_sum[pair_tid];
|
|
}
|
|
}
|
|
if (threadIdx.y == 0) {
|
|
*x_sum_out = x_sum;
|
|
*x_square_sum_out = x_square_sum;
|
|
}
|
|
}
|
|
|
|
template <typename T, typename BnT>
|
|
__device__ __forceinline__ void ReduceSumPost(const int C, // channels
|
|
const int c, // channel index
|
|
BnT *sum1,
|
|
BnT *sum2,
|
|
bool *is_last_block_done,
|
|
BnT *cache1,
|
|
BnT *cache2,
|
|
BnT *block_data_ptr,
|
|
int *flag_ptr) {
|
|
volatile BnT *staging_sum = block_data_ptr;
|
|
volatile BnT *staging_sum2 = &block_data_ptr[C * gridDim.y];
|
|
// write block data to global memory
|
|
if (threadIdx.y == 0) {
|
|
staging_sum[c + blockIdx.y * C] = *sum1;
|
|
staging_sum2[c + blockIdx.y * C] = *sum2;
|
|
}
|
|
|
|
// make sure write is visible to all blocks
|
|
__threadfence();
|
|
__syncthreads();
|
|
|
|
// mark block done
|
|
if (threadIdx.x == 0 && threadIdx.y == 0) {
|
|
int old = atomicAdd(&flag_ptr[blockIdx.x], 1);
|
|
*is_last_block_done = (old == (gridDim.y - 1));
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
if (*is_last_block_done) {
|
|
*sum1 = static_cast<BnT>(0);
|
|
*sum2 = static_cast<BnT>(0);
|
|
// thread sum
|
|
for (int y = threadIdx.y; y < gridDim.y; y += blockDim.y) {
|
|
*sum1 += staging_sum[c + y * C];
|
|
*sum2 += staging_sum2[c + y * C];
|
|
}
|
|
|
|
// vertical block sum
|
|
funcs::BlockReduceByVertical<T, BnT>(
|
|
*sum1, *sum2, &cache1[0], &cache2[0], sum1, sum2);
|
|
}
|
|
}
|
|
|
|
template <typename T, typename BnT, typename Context>
|
|
void SetLaunchConfigInfoForChannelLast(const Context &dev_ctx,
|
|
DenseTensor *block_data_tensor,
|
|
DenseTensor *flag_tensor,
|
|
BnT **block_data_ptr,
|
|
int **flag_ptr,
|
|
const int N,
|
|
const int H,
|
|
const int W,
|
|
const int D,
|
|
const int C,
|
|
const int block_size,
|
|
dim3 *block,
|
|
dim3 *grid) {
|
|
const int64_t MAX_GRID_SIZE = 128;
|
|
const int64_t WARP_SIZE = 32;
|
|
|
|
int block_x = std::min(funcs::details::GetLastPow2(C), WARP_SIZE);
|
|
int block_y = std::min(
|
|
funcs::details::GetLastPow2(static_cast<int64_t>(N) * H * W * D / 16),
|
|
static_cast<int64_t>(block_size / block_x));
|
|
if (block_x * block_y != block_size) {
|
|
block_x = std::min(funcs::details::GetLastPow2(C),
|
|
static_cast<int64_t>(block_size / block_y));
|
|
}
|
|
int grid_x = (C + block_x - 1) / block_x;
|
|
int grid_y = std::min(
|
|
(static_cast<int64_t>(N) * H * W * D + block_y * 16 - 1) / (block_y * 16),
|
|
MAX_GRID_SIZE);
|
|
|
|
block->x = block_x;
|
|
block->y = block_y;
|
|
grid->x = grid_x;
|
|
grid->y = grid_y;
|
|
|
|
if (grid->y > 1) {
|
|
*block_data_tensor = Empty<BnT, Context>(dev_ctx, {2 * C * grid->y});
|
|
*flag_tensor = Empty<int, Context>(dev_ctx, {grid->x});
|
|
|
|
*block_data_ptr = block_data_tensor->data<BnT>();
|
|
*flag_ptr = flag_tensor->data<int>();
|
|
funcs::SetConstant<Context, int> set_zero;
|
|
set_zero(dev_ctx, flag_tensor, static_cast<int>(0));
|
|
}
|
|
}
|
|
|
|
} // namespace funcs
|
|
} // namespace phi
|