chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,291 @@
|
||||
// Copyright (c) 2022 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/instance_norm_kernel.h"
|
||||
#include "glog/logging.h"
|
||||
#include "paddle/common/layout.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/full_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/funcs/norm_utils.h"
|
||||
#include "paddle/phi/kernels/gpu/instance_norm_utils.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void InstanceNormKernel(const Context &dev_ctx,
|
||||
const DenseTensor &x,
|
||||
const optional<DenseTensor> &scale,
|
||||
const optional<DenseTensor> &bias,
|
||||
float epsilon_f,
|
||||
DenseTensor *y,
|
||||
DenseTensor *saved_mean,
|
||||
DenseTensor *saved_variance) {
|
||||
using AccT = typename MPTypeTrait<T>::Type;
|
||||
double epsilon = static_cast<double>(epsilon_f);
|
||||
auto &x_dims = x.dims();
|
||||
PADDLE_ENFORCE_GE(x_dims.size(),
|
||||
2,
|
||||
common::errors::InvalidArgument(
|
||||
"The `shape` in InstanceNormOp is invalid: "
|
||||
"the size of X's dimensions must greater than "
|
||||
"or equal to 2. But received: "
|
||||
"the size of X's dimensions is [%d]",
|
||||
x_dims.size()));
|
||||
PADDLE_ENFORCE_LE(x_dims.size(),
|
||||
5,
|
||||
common::errors::InvalidArgument(
|
||||
"The `shape` in InstanceNormOp is invalid: "
|
||||
"the size of X's dimensions must smaller than "
|
||||
"or equal to 5. But received: "
|
||||
"the size of X's dimensions is [%d]",
|
||||
x_dims.size()));
|
||||
int N, C, H, W, D;
|
||||
funcs::ExtractNCWHD(x_dims, DataLayout::NCHW, &N, &C, &H, &W, &D);
|
||||
const int64_t NxC64 = static_cast<int64_t>(N) * C;
|
||||
PADDLE_ENFORCE_LE_INT_MAX(NxC64, "NxC");
|
||||
const int NxC = static_cast<int>(NxC64);
|
||||
DenseTensor x_tmp;
|
||||
x_tmp.ShareDataWith(x).Resize({1, NxC, H, W, D});
|
||||
dev_ctx.template Alloc<T>(y);
|
||||
funcs::SetConstant<GPUContext, BatchNormParamType<T>> functor;
|
||||
funcs::SetConstant<GPUContext, T> functor_y;
|
||||
if (x.numel() == 0) {
|
||||
functor_y(dev_ctx, y, static_cast<T>(0));
|
||||
if (saved_mean) {
|
||||
dev_ctx.template Alloc<BatchNormParamType<T>>(saved_mean);
|
||||
functor(dev_ctx, saved_mean, static_cast<BatchNormParamType<T>>(0));
|
||||
}
|
||||
if (saved_variance) {
|
||||
dev_ctx.template Alloc<BatchNormParamType<T>>(saved_variance);
|
||||
functor(dev_ctx, saved_variance, static_cast<BatchNormParamType<T>>(0));
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
miopenTensorDescriptor_t data_desc_;
|
||||
miopenTensorDescriptor_t in_param_desc_;
|
||||
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
dynload::miopenCreateTensorDescriptor(&data_desc_));
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
dynload::miopenCreateTensorDescriptor(&in_param_desc_));
|
||||
#else
|
||||
cudnnTensorDescriptor_t data_desc_;
|
||||
cudnnTensorDescriptor_t in_param_desc_;
|
||||
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cudnnCreateTensorDescriptor(&data_desc_));
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
dynload::cudnnCreateTensorDescriptor(&in_param_desc_));
|
||||
#endif
|
||||
if (epsilon <= CUDNN_BN_MIN_EPSILON - FLT_EPSILON) {
|
||||
LOG(ERROR) << "Provided epsilon is smaller than "
|
||||
<< "CUDNN_BN_MIN_EPSILON. Setting it to "
|
||||
<< "CUDNN_BN_MIN_EPSILON instead.";
|
||||
}
|
||||
epsilon = std::max(epsilon, CUDNN_BN_MIN_EPSILON);
|
||||
|
||||
VLOG(3) << "Setting descriptors.";
|
||||
std::vector<int> dims;
|
||||
std::vector<int> strides;
|
||||
const int64_t stride0 = NxC64 * H * W * D;
|
||||
const int64_t stride1 = static_cast<int64_t>(H) * W * D;
|
||||
const int64_t stride2 = static_cast<int64_t>(W) * D;
|
||||
PADDLE_ENFORCE_LE_INT_MAX(stride0, "cudnn tensor descriptor stride0");
|
||||
PADDLE_ENFORCE_LE_INT_MAX(stride1, "cudnn tensor descriptor stride1");
|
||||
PADDLE_ENFORCE_LE_INT_MAX(stride2, "cudnn tensor descriptor stride2");
|
||||
dims = {1, NxC, H, W, D};
|
||||
strides = {static_cast<int>(stride0),
|
||||
static_cast<int>(stride1),
|
||||
static_cast<int>(stride2),
|
||||
D,
|
||||
1};
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
dynload::miopenSetTensorDescriptor(data_desc_,
|
||||
CudnnDataType<T>::type,
|
||||
x_dims.size() > 3 ? x_dims.size() : 4,
|
||||
const_cast<int *>(dims.data()),
|
||||
const_cast<int *>(strides.data())));
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::miopenDeriveBNTensorDescriptor(
|
||||
in_param_desc_, data_desc_, miopenBNSpatial));
|
||||
#else
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
dynload::cudnnSetTensorNdDescriptor(data_desc_,
|
||||
CudnnDataType<T>::type,
|
||||
x_dims.size() > 3 ? x_dims.size() : 4,
|
||||
dims.data(),
|
||||
strides.data()));
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cudnnDeriveBNTensorDescriptor(
|
||||
in_param_desc_, data_desc_, CUDNN_BATCHNORM_SPATIAL));
|
||||
#endif
|
||||
|
||||
const auto scale_ptr = scale.get_ptr();
|
||||
const auto bias_ptr = bias.get_ptr();
|
||||
|
||||
DenseTensor scale_tmp;
|
||||
scale_tmp.Resize({NxC});
|
||||
dev_ctx.template Alloc<AccT>(&scale_tmp);
|
||||
DenseTensor bias_tmp;
|
||||
bias_tmp.Resize({NxC});
|
||||
dev_ctx.template Alloc<AccT>(&bias_tmp);
|
||||
|
||||
// TODO(large-tensor): downstream functors may still use int; guard until
|
||||
// upgraded.
|
||||
int64_t n = x.numel();
|
||||
|
||||
const int block = 512;
|
||||
int max_threads = dev_ctx.GetMaxPhysicalThreadCount();
|
||||
const int max_blocks = std::max(max_threads / block, 1);
|
||||
const int grid = std::min((NxC + block - 1) / block, max_blocks);
|
||||
|
||||
funcs::SetConstant<GPUContext, AccT> set_constant;
|
||||
if (scale_ptr) {
|
||||
repeat_param<AccT><<<grid, block, 0, dev_ctx.stream()>>>(
|
||||
scale_ptr->data<AccT>(), scale_tmp.data<AccT>(), N, C);
|
||||
} else {
|
||||
set_constant(dev_ctx, &scale_tmp, static_cast<AccT>(1));
|
||||
}
|
||||
if (bias_ptr) {
|
||||
repeat_param<AccT><<<grid, block, 0, dev_ctx.stream()>>>(
|
||||
bias_ptr->data<AccT>(), bias_tmp.data<AccT>(), N, C);
|
||||
} else {
|
||||
set_constant(dev_ctx, &bias_tmp, static_cast<AccT>(0));
|
||||
}
|
||||
|
||||
auto handle = dev_ctx.cudnn_handle();
|
||||
|
||||
DenseTensor saved_mean_tmp, saved_variance_tmp;
|
||||
|
||||
if (saved_mean) {
|
||||
dev_ctx.template Alloc<BatchNormParamType<T>>(saved_mean);
|
||||
functor(dev_ctx, saved_mean, static_cast<BatchNormParamType<T>>(0));
|
||||
} else {
|
||||
saved_mean_tmp = Full<BatchNormParamType<T>>(
|
||||
dev_ctx, {NxC}, static_cast<BatchNormParamType<T>>(0));
|
||||
}
|
||||
if (saved_variance) {
|
||||
dev_ctx.template Alloc<BatchNormParamType<T>>(saved_variance);
|
||||
functor(dev_ctx, saved_variance, static_cast<BatchNormParamType<T>>(0));
|
||||
} else {
|
||||
saved_variance_tmp = Full<BatchNormParamType<T>>(
|
||||
dev_ctx, {NxC}, static_cast<BatchNormParamType<T>>(0));
|
||||
}
|
||||
auto *saved_mean_data = saved_mean
|
||||
? saved_mean->data<BatchNormParamType<T>>()
|
||||
: saved_mean_tmp.data<BatchNormParamType<T>>();
|
||||
auto *saved_variance_data =
|
||||
saved_variance ? saved_variance->data<BatchNormParamType<T>>()
|
||||
: saved_variance_tmp.data<BatchNormParamType<T>>();
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::miopenBatchNormalizationForwardTraining(
|
||||
handle,
|
||||
miopenBNSpatial,
|
||||
const_cast<void *>(static_cast<const void *>(CudnnDataType<T>::kOne())),
|
||||
const_cast<void *>(static_cast<const void *>(CudnnDataType<T>::kZero())),
|
||||
data_desc_,
|
||||
static_cast<const void *>(x_tmp.template data<T>()),
|
||||
data_desc_,
|
||||
static_cast<void *>(y->template data<T>()),
|
||||
in_param_desc_,
|
||||
const_cast<void *>(static_cast<const void *>(
|
||||
scale_tmp.template data<BatchNormParamType<T>>())),
|
||||
const_cast<void *>(static_cast<const void *>(
|
||||
bias_tmp.template data<BatchNormParamType<T>>())),
|
||||
0,
|
||||
nullptr,
|
||||
nullptr,
|
||||
epsilon,
|
||||
static_cast<void *>(saved_mean_data),
|
||||
static_cast<void *>(saved_variance_data)));
|
||||
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
dynload::miopenDestroyTensorDescriptor(data_desc_));
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
dynload::miopenDestroyTensorDescriptor(in_param_desc_));
|
||||
#else
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cudnnBatchNormalizationForwardTraining(
|
||||
handle,
|
||||
CUDNN_BATCHNORM_SPATIAL,
|
||||
CudnnDataType<T>::kOne(),
|
||||
CudnnDataType<T>::kZero(),
|
||||
data_desc_,
|
||||
x_tmp.template data<T>(),
|
||||
data_desc_,
|
||||
y->template data<T>(),
|
||||
in_param_desc_,
|
||||
scale_tmp.template data<BatchNormParamType<T>>(),
|
||||
bias_tmp.template data<BatchNormParamType<T>>(),
|
||||
0,
|
||||
nullptr,
|
||||
nullptr,
|
||||
epsilon,
|
||||
saved_mean_data,
|
||||
saved_variance_data));
|
||||
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cudnnDestroyTensorDescriptor(data_desc_));
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
dynload::cudnnDestroyTensorDescriptor(in_param_desc_));
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
// MIOPEN do not support double
|
||||
PD_REGISTER_KERNEL(instance_norm,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::InstanceNormKernel,
|
||||
float,
|
||||
phi::float16) {
|
||||
if (kernel_key.dtype() == phi::DataType::FLOAT16) {
|
||||
kernel->InputAt(1).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->InputAt(2).SetDataType(phi::DataType::FLOAT32);
|
||||
}
|
||||
}
|
||||
#elif CUDNN_VERSION_MIN(8, 1, 0)
|
||||
PD_REGISTER_KERNEL(instance_norm,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::InstanceNormKernel,
|
||||
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);
|
||||
}
|
||||
}
|
||||
#else
|
||||
PD_REGISTER_KERNEL(instance_norm,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::InstanceNormKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16) {
|
||||
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);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
Reference in New Issue
Block a user