wehub-resource-sync
226 lines
9.0 KiB
Plaintext
226 lines
9.0 KiB
Plaintext
// Copyright (c) 2024 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.
|
|
|
|
#include "paddle/phi/kernels/sync_batch_norm_kernel.h"
|
|
#include "paddle/common/enforce.h"
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/funcs/sync_batch_norm_utils.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
void SyncBatchNormKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& mean,
|
|
const DenseTensor& variance,
|
|
const DenseTensor& scale,
|
|
const DenseTensor& bias,
|
|
bool is_test,
|
|
float momentum,
|
|
float epsilon_f,
|
|
const std::string& data_layout_str,
|
|
bool use_global_stats,
|
|
bool trainable_statistics,
|
|
DenseTensor* y,
|
|
DenseTensor* mean_out,
|
|
DenseTensor* variance_out,
|
|
DenseTensor* saved_mean,
|
|
DenseTensor* saved_variance,
|
|
DenseTensor* reserve_space) {
|
|
PADDLE_ENFORCE_EQ(use_global_stats,
|
|
false,
|
|
common::errors::InvalidArgument(
|
|
"sync_batch_norm doesn't support "
|
|
"to set use_global_stats True. Please use batch_norm "
|
|
"in this case."));
|
|
|
|
double epsilon = epsilon_f;
|
|
const bool trainable_stats = trainable_statistics;
|
|
const DataLayout layout = StringToDataLayout(data_layout_str);
|
|
bool test_mode = is_test && (!trainable_statistics);
|
|
const auto& x_dims = x.dims();
|
|
PADDLE_ENFORCE_GE(x_dims.size(),
|
|
2,
|
|
common::errors::InvalidArgument(
|
|
"The Input dim size should be larger than 1."));
|
|
PADDLE_ENFORCE_LE(x_dims.size(),
|
|
5,
|
|
common::errors::InvalidArgument(
|
|
"The Input dim size should be less than 6."));
|
|
int N, C, H, W, D;
|
|
funcs::ExtractNCWHD(x_dims, layout, &N, &C, &H, &W, &D);
|
|
int64_t x_numel = x.numel();
|
|
const int64_t fsize = static_cast<int64_t>(H) * W * D;
|
|
|
|
const T* x_d = x.template data<T>();
|
|
const auto* s_d = scale.template data<BatchNormParamType<T>>();
|
|
const auto* b_d = bias.template data<BatchNormParamType<T>>();
|
|
|
|
T* y_d = dev_ctx.template Alloc<T>(y);
|
|
|
|
const BatchNormParamType<T>* mean_data = nullptr;
|
|
const BatchNormParamType<T>* var_data = nullptr;
|
|
|
|
auto stream = dev_ctx.stream();
|
|
const int block = 512;
|
|
int max_threads = dev_ctx.GetMaxPhysicalThreadCount();
|
|
|
|
Allocator::AllocationPtr alloc_ptr{nullptr};
|
|
|
|
if (test_mode) {
|
|
mean_data = mean.template data<BatchNormParamType<T>>();
|
|
var_data = variance.template data<BatchNormParamType<T>>();
|
|
} else {
|
|
// x, x^2, 1, here 1 is used to calc device num
|
|
// device num also can be got from DeviceContextPool
|
|
const int64_t bytes_64 =
|
|
(static_cast<int64_t>(C) * 2 + 1) * sizeof(BatchNormParamType<T>);
|
|
DenseTensor stats_tensor;
|
|
stats_tensor.Resize({bytes_64});
|
|
dev_ctx.template Alloc<BatchNormParamType<T>>(&stats_tensor);
|
|
auto* stats_data = stats_tensor.data<BatchNormParamType<T>>();
|
|
auto* stats = reinterpret_cast<BatchNormParamType<T>*>(stats_data);
|
|
const int threads = 512;
|
|
int grid = std::min(C, (max_threads + threads - 1) / threads);
|
|
if (layout == DataLayout::NCHW) {
|
|
KeLocalStats<T, threads, DataLayout::NCHW>
|
|
<<<grid, threads, 0, stream>>>(x_d, N, fsize, C, stats);
|
|
} else {
|
|
KeLocalStats<T, threads, DataLayout::NHWC>
|
|
<<<grid, threads, 0, stream>>>(x_d, N, fsize, C, stats);
|
|
}
|
|
|
|
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
|
auto comm_ctx =
|
|
static_cast<distributed::NCCLCommContext*>(dev_ctx.GetCommContext());
|
|
if (comm_ctx) {
|
|
comm_ctx->AllReduce(&stats_tensor, stats_tensor, ncclSum, stream);
|
|
}
|
|
#endif
|
|
|
|
auto* est_mean_data =
|
|
dev_ctx.template Alloc<BatchNormParamType<T>>(mean_out);
|
|
auto* est_var_data =
|
|
dev_ctx.template Alloc<BatchNormParamType<T>>(variance_out);
|
|
|
|
auto* sv_mean_data =
|
|
dev_ctx.template Alloc<BatchNormParamType<T>>(saved_mean);
|
|
auto* sv_inv_var_data =
|
|
dev_ctx.template Alloc<BatchNormParamType<T>>(saved_variance);
|
|
|
|
int64_t reserve_space_size = 0;
|
|
DenseTensor tmp_reserve_space;
|
|
if (reserve_space == nullptr) {
|
|
reserve_space = &tmp_reserve_space;
|
|
}
|
|
reserve_space->Resize({reserve_space_size});
|
|
dev_ctx.template Alloc<T>(reserve_space);
|
|
|
|
// Note, Input('Mean')/Input('Variance') share variable with
|
|
// Output('MeanOut')/Output('VarianceOut')
|
|
KeSyncAndMovingStats<T>
|
|
<<<(C + block - 1) / block, block, 0, stream>>>(stats,
|
|
stats + C,
|
|
stats + 2 * C,
|
|
C,
|
|
momentum,
|
|
epsilon,
|
|
sv_mean_data,
|
|
sv_inv_var_data,
|
|
est_mean_data,
|
|
est_var_data);
|
|
|
|
mean_data = sv_mean_data;
|
|
var_data = stats + C;
|
|
}
|
|
|
|
const int64_t grid2_64 =
|
|
(std::min(x_numel, static_cast<int64_t>(max_threads)) + block - 1) /
|
|
block;
|
|
uint32_t grid2 = static_cast<uint32_t>(grid2_64);
|
|
if (layout == DataLayout::NCHW) {
|
|
KeNormAffine<T, DataLayout::NCHW><<<grid2, block, 0, stream>>>(
|
|
x_d, s_d, b_d, mean_data, var_data, epsilon, C, fsize, x_numel, y_d);
|
|
} else {
|
|
KeNormAffine<T, DataLayout::NHWC><<<grid2, block, 0, stream>>>(
|
|
x_d, s_d, b_d, mean_data, var_data, epsilon, C, fsize, x_numel, y_d);
|
|
}
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
PD_REGISTER_KERNEL(sync_batch_norm,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::SyncBatchNormKernel,
|
|
float,
|
|
phi::float16) {
|
|
if (kernel_key.dtype() == phi::DataType::FLOAT16) {
|
|
kernel->InputAt(1).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->InputAt(2).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->InputAt(3).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->InputAt(4).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
|
|
}
|
|
}
|
|
#else
|
|
#if CUDNN_VERSION_MIN(8, 1, 0)
|
|
PD_REGISTER_KERNEL(sync_batch_norm,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::SyncBatchNormKernel,
|
|
float,
|
|
double,
|
|
phi::float16,
|
|
phi::bfloat16) {
|
|
if (kernel_key.dtype() == phi::DataType::FLOAT16 ||
|
|
kernel_key.dtype() == phi::DataType::BFLOAT16) {
|
|
kernel->InputAt(1).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->InputAt(2).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->InputAt(3).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->InputAt(4).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
|
|
}
|
|
}
|
|
#else
|
|
PD_REGISTER_KERNEL(sync_batch_norm,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::SyncBatchNormKernel,
|
|
float,
|
|
double,
|
|
phi::float16) {
|
|
if (kernel_key.dtype() == phi::DataType::FLOAT16) {
|
|
kernel->InputAt(1).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->InputAt(2).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->InputAt(3).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->InputAt(4).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
|
|
}
|
|
}
|
|
#endif
|
|
#endif
|