145 lines
4.4 KiB
C++
145 lines
4.4 KiB
C++
// 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.
|
|
|
|
#pragma once
|
|
|
|
#include <string>
|
|
|
|
#include "paddle/phi/backends/gpu/gpu_decls.h"
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#ifdef PADDLE_WITH_CUDA
|
|
#include <cuda.h>
|
|
#include <cuda_fp16.h>
|
|
#endif
|
|
#include <stdint.h>
|
|
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
void GroupNormKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const optional<DenseTensor>& scale,
|
|
const optional<DenseTensor>& bias,
|
|
double epsilon,
|
|
int groups,
|
|
const std::string& data_layout,
|
|
DenseTensor* y,
|
|
DenseTensor* mean,
|
|
DenseTensor* variance);
|
|
|
|
template <typename T, typename Context>
|
|
void GroupNormNDHWCKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const optional<DenseTensor>& residual,
|
|
const optional<DenseTensor>& scale,
|
|
const optional<DenseTensor>& bias,
|
|
double epsilon,
|
|
int groups,
|
|
const std::string& data_layout_str,
|
|
const std::string& activation,
|
|
DenseTensor* y,
|
|
DenseTensor* residual_out,
|
|
DenseTensor* mean,
|
|
DenseTensor* var);
|
|
|
|
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
|
template <typename T, typename AccT = T>
|
|
class GroupNormDirectCUDAFunctor {
|
|
public:
|
|
void operator()(gpuStream_t stream,
|
|
const T* input,
|
|
std::vector<int> input_shape,
|
|
const T* bias,
|
|
const T* scale,
|
|
AccT* temp_variance,
|
|
int groups,
|
|
float eps,
|
|
T* output,
|
|
AccT* mean,
|
|
AccT* variance,
|
|
const DataLayout data_layout);
|
|
};
|
|
#endif
|
|
|
|
template <typename T>
|
|
struct GroupNormNDHWCParams {
|
|
// The output buffer. Layout NDHWC.
|
|
T* dst;
|
|
// The output buffer. Layout NDHWC.
|
|
T* eleOut;
|
|
// The input buffer. Layout NDHWC.
|
|
T const* srcX;
|
|
// The input buffer. Layout NDHWC.
|
|
T const* srcY;
|
|
// The input buffer. Layout NDHWC.
|
|
T const* srcR = nullptr;
|
|
// The gamma scaling factor.
|
|
void const* gamma;
|
|
// The beta term to add in GN.
|
|
void const* beta;
|
|
// The temporary buffer to do the global parallel reduction. Size:
|
|
// BLOCKS_PER_BATCH x C x 2.
|
|
float* redBuffer;
|
|
|
|
float* var_data;
|
|
|
|
// The number of instances in the batch.
|
|
int32_t n;
|
|
// The depth, height and width of each activation map.
|
|
int32_t d, h, w;
|
|
// The number of channels.
|
|
int32_t c;
|
|
// The number of groups.
|
|
int32_t groups;
|
|
// Do we apply the Silu activation function?
|
|
bool withSilu;
|
|
//
|
|
bool y_same_with_x = false;
|
|
// Precomputed values and parameters to control the execution of the kernels.
|
|
|
|
// The number of activations per instance (d * h * w) and the number of
|
|
// activations per block.
|
|
int64_t dhw, dhwPerBlock;
|
|
// The number of channels per group and blocks per activation in the C
|
|
// dimension.
|
|
int32_t cPerBlock, cPerGroup;
|
|
|
|
// The precomputed stride between instances.
|
|
int64_t dhwc;
|
|
// The inverse of dhwc in floats (to compute mean/var).
|
|
float invDHWC;
|
|
// The precomputed number of groups per block.
|
|
int32_t groupsPerBlock;
|
|
// epsilon, Constant for numerical stability
|
|
float eps;
|
|
// for NCDHW32 int8 use
|
|
float dqScaleIn;
|
|
float inv_qScale;
|
|
};
|
|
|
|
template <typename T>
|
|
class groupNormNDHWCSum {
|
|
public:
|
|
void operator()(GroupNormNDHWCParams<T>* params, const gpuStream_t stream);
|
|
};
|
|
|
|
template <typename T>
|
|
class groupNormNDHWCScale {
|
|
public:
|
|
void operator()(const GroupNormNDHWCParams<T>& params,
|
|
const gpuStream_t stream);
|
|
};
|
|
|
|
} // namespace phi
|