chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,225 @@
|
||||
// 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
|
||||
Reference in New Issue
Block a user