wehub/paddlepaddle--paddle
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
e25996e7dbc7c36a2d3f3feb318547b89f7fa1d7
paddlepaddle--paddle/paddle/phi/kernels/gpu/group_norm_kernel.cu
T
wehub-resource-sync e25996e7db chore: import upstream snapshot with attribution
2026-07-13 12:40:42 +08:00

1910 lines
70 KiB
Plaintext
Raw Blame History

// 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/group_norm_kernel.h"
#include "paddle/common/layout.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/backends/gpu/gpu_device_function.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/math_function.h"
#include "paddle/phi/kernels/gpu/group_norm_utils.h"
#include "paddle/phi/common/data_type.h"
#include "paddle/phi/core/device_context.h"
#include "paddle/phi/kernels/full_kernel.h"
#include "paddle/common/flags.h"
COMMON_DECLARE_bool(use_accuracy_compatible_kernel);
namespace phi {
template <typename T>
static inline T divUp(T m, T n) {
return (m + n - 1) / n;
}
static inline __device__ __host__ float sigmoid(float x) {
return 1.F / (1.F + expf(-x));
}
#ifdef PADDLE_CUDA_BF16
__host__ __device__ inline float2 bfloat1622float2(const __nv_bfloat162 a) {
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))
return __bfloat1622float2(a);
#else
float hi_float;
float lo_float;
lo_float = __internal_bfloat162float(((__nv_bfloat162_raw)a).x);
hi_float = __internal_bfloat162float(((__nv_bfloat162_raw)a).y);
return make_float2(lo_float, hi_float);
#endif
}
__host__ __device__ inline __nv_bfloat162 float22bfloat162_rn(const float2 a) {
__nv_bfloat162 val;
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))
val = __float22bfloat162_rn(a);
#else
val.x = __float2bfloat16_rn(a.x);
val.y = __float2bfloat16_rn(a.y);
#endif
return val;
}
#endif
template <typename T>
__host__ __device__ inline float __2float(const T a) {
return static_cast<float>(a);
}
template <>
__host__ __device__ inline float __2float<__half>(const __half a) {
return __half2float(a);
}
template <typename T>
__host__ __device__ inline T __2dst(const float a) {
return static_cast<T>(a);
}
template <>
__host__ __device__ inline __half __2dst<__half>(const float a) {
return __float2half(a);
}
struct GroupSums {
// Is it the 1st element of the group?
int32_t flag;
// The sum.
float sum;
// The sum of squares.
float sumSq;
};
struct GroupSumsOp {
inline __device__ GroupSums operator()(GroupSums const& a,
GroupSums const& b) {
GroupSums dst;
dst.sum = b.flag ? b.sum : (a.sum + b.sum);
dst.sumSq = b.flag ? b.sumSq : (a.sumSq + b.sumSq);
dst.flag = a.flag + b.flag;
return dst;
}
};
static int64_t findMaxDivisor(int64_t n, int64_t maxAllowedDivisor) {
int64_t maxDivisor = -1;
for (int64_t i = 1; i <= std::sqrt(n); i++) {
if (n % i == 0) {
int64_t divisor1 = n / i;
int64_t divisor2 = i;
if (divisor1 > maxDivisor && divisor1 < maxAllowedDivisor) {
maxDivisor = divisor1;
}
if (divisor2 > maxDivisor && divisor2 < maxAllowedDivisor) {
maxDivisor = divisor2;
}
}
}
return maxDivisor;
}
template <typename T, int THREADS_PER_CHANNEL>
inline __device__ void UpdateSum(const T* srcX, float* sum, float* sumSq) {
float src_data = __2float<T>(*srcX);
*sum += src_data;
*sumSq += src_data * src_data;
}
template <typename T, int THREADS_PER_CHANNEL>
inline __device__ void UpdateSum(const T* srcX,
const T* srcR,
float* sum,
float* sumSq) {
float src_data = __2float<T>(*srcX);
float srcy_data = __2float<T>(*srcR);
*sum += src_data + srcy_data;
*sumSq += (src_data + srcy_data) * (src_data + srcy_data);
}
template <>
inline __device__ void UpdateSum<__half, 2>(const __half* srcX,
float* sum,
float* sumSq) {
__half2 h2 = *reinterpret_cast<__half2 const*>(srcX);
float2 f2 = __half22float2(h2);
*sum += f2.x + f2.y;
*sumSq += f2.x * f2.x + f2.y * f2.y;
}
template <>
inline __device__ void UpdateSum<__half, 2>(const __half* srcX,
const __half* srcR,
float* sum,
float* sumSq) {
__half2 h2 = *reinterpret_cast<__half2 const*>(srcX);
__half2 h2_r = *reinterpret_cast<__half2 const*>(srcR);
float2 f2 = __half22float2(h2);
float2 f2_r = __half22float2(h2_r);
*sum += f2.x + f2_r.x + f2.y + f2_r.y;
*sumSq +=
(f2.x + f2_r.x) * (f2.x + f2_r.x) + (f2.y + f2_r.y) * (f2.y + f2_r.y);
}
template <>
inline __device__ void UpdateSum<float16, 2>(const float16* srcX,
float* sum,
float* sumSq) {
__half2 h2 = *reinterpret_cast<__half2 const*>(srcX);
float2 f2 = __half22float2(h2);
*sum += f2.x + f2.y;
*sumSq += f2.x * f2.x + f2.y * f2.y;
}
template <>
inline __device__ void UpdateSum<float16, 2>(const float16* srcX,
const float16* srcR,
float* sum,
float* sumSq) {
__half2 h2 = *reinterpret_cast<__half2 const*>(srcX);
__half2 h2_r = *reinterpret_cast<__half2 const*>(srcR);
float2 f2 = __half22float2(h2);
float2 f2_r = __half22float2(h2_r);
*sum += f2.x + f2_r.x + f2.y + f2_r.y;
*sumSq +=
(f2.x + f2_r.x) * (f2.x + f2_r.x) + (f2.y + f2_r.y) * (f2.y + f2_r.y);
}
#ifdef PADDLE_CUDA_BF16
template <>
inline __device__ void UpdateSum<bfloat16, 2>(const bfloat16* srcX,
float* sum,
float* sumSq) {
__nv_bfloat162 h2 = *reinterpret_cast<__nv_bfloat162 const*>(srcX);
float2 f2 = bfloat1622float2(h2);
*sum += f2.x + f2.y;
*sumSq += f2.x * f2.x + f2.y * f2.y;
}
template <>
inline __device__ void UpdateSum<bfloat16, 2>(const bfloat16* srcX,
const bfloat16* srcR,
float* sum,
float* sumSq) {
__nv_bfloat162 h2 = *reinterpret_cast<__nv_bfloat162 const*>(srcX);
__nv_bfloat162 h2_r = *reinterpret_cast<__nv_bfloat162 const*>(srcR);
float2 f2 = bfloat1622float2(h2);
float2 f2_r = bfloat1622float2(h2_r);
*sum += f2.x + f2_r.x + f2.y + f2_r.y;
*sumSq +=
(f2.x + f2_r.x) * (f2.x + f2_r.x) + (f2.y + f2_r.y) * (f2.y + f2_r.y);
}
#endif
template <typename T, int THREADS_PER_BLOCK>
__global__ void groupNormNDHWCSumSingerChannelKernel(
const GroupNormNDHWCParams<T> params) {
// The instance in the batch.
__shared__ float2 smem[THREADS_PER_BLOCK];
int32_t ni = blockIdx.z;
int64_t ci = static_cast<int64_t>(blockIdx.x) *
static_cast<int64_t>(params.cPerBlock) +
static_cast<int64_t>(threadIdx.x);
if (ci >= params.c) {
return;
}
// The first activation loaded by that block.
int64_t dhwBegin = static_cast<int64_t>(blockIdx.y) * params.dhwPerBlock;
// The last activation loaded by that block.
int64_t dhwEnd = min(dhwBegin + params.dhwPerBlock, params.dhw);
// The sums.
float sum = 0.F;
float sumSq = 0.F;
for (int64_t dhwi = dhwBegin; dhwi < dhwEnd; ++dhwi) {
// The offset.
int64_t offset = ni * params.dhwc + dhwi * params.c + ci;
float src_data = *reinterpret_cast<float const*>(&params.srcX[offset]);
if (params.srcR != nullptr) {
int64_t g_offset = params.y_same_with_x ? offset : ci;
UpdateSum<T, 1>(
&params.srcX[offset], &params.srcR[g_offset], &sum, &sumSq);
} else {
UpdateSum<T, 1>(&params.srcX[offset], &sum, &sumSq);
}
}
smem[threadIdx.x] = make_float2(sum, sumSq);
__syncthreads();
float2 sums = smem[threadIdx.x];
atomicAdd(&params.redBuffer[(2 * ni + 0) * params.groups + ci],
sums.x * params.invDHWC);
atomicAdd(&params.redBuffer[(2 * ni + 1) * params.groups + ci], sums.y);
}
template <typename T, int THREADS_PER_BLOCK, int THREADS_PER_CHANNEL>
__global__ void groupNormNDHWCSumKernel(const GroupNormNDHWCParams<T> params) {
// The object in charge of doing the sums for the different blocks.
typedef cub::BlockScan<GroupSums, THREADS_PER_BLOCK> BlockScan;
__shared__ typename BlockScan::TempStorage tempStorage;
// Allocate shared memory for BlockScan.
// Allocate shared memory for the groups. We could reduce the amount of shared
// memory reserved.
__shared__ float2 smem[THREADS_PER_BLOCK];
// The instance in the batch.
int32_t ni = blockIdx.z;
// The channel loaded by that thread (2 channels per thread for F16x2).
int64_t ci = static_cast<int64_t>(blockIdx.x) *
static_cast<int64_t>(params.cPerBlock) +
static_cast<int64_t>(threadIdx.x) * THREADS_PER_CHANNEL;
if (ci >= params.c || threadIdx.x * THREADS_PER_CHANNEL >= params.cPerBlock) {
return;
}
int32_t gj = ci / params.cPerGroup;
int32_t cj = ci % params.cPerGroup;
int64_t dhwBegin = static_cast<int64_t>(blockIdx.y) * params.dhwPerBlock;
// The last activation loaded by that block.
int64_t dhwEnd = min(dhwBegin + params.dhwPerBlock, params.dhw);
// The sums.
float sum = 0.F;
float sumSq = 0.F;
for (int64_t dhwi = dhwBegin; dhwi < dhwEnd; ++dhwi) {
// The offset.
int64_t offset = ni * params.dhwc + dhwi * params.c + ci;
float src_data = *reinterpret_cast<float const*>(&params.srcX[offset]);
if (params.srcR != nullptr) {
int64_t g_offset =
params.y_same_with_x ? offset : gj * params.cPerGroup + cj;
UpdateSum<T, THREADS_PER_CHANNEL>(
&params.srcX[offset], &params.srcR[g_offset], &sum, &sumSq);
} else {
UpdateSum<T, THREADS_PER_CHANNEL>(&params.srcX[offset], &sum, &sumSq);
}
}
// The group that thread works on and the channel in the group (modulus).
int32_t gi =
ci / params.cPerGroup - blockIdx.x * params.cPerBlock / params.cPerGroup;
int flag = (cj == 0 || threadIdx.x == 0) ? 1 : 0;
GroupSums inp{flag, sum, sumSq};
GroupSums out;
BlockScan(tempStorage).InclusiveScan(inp, out, GroupSumsOp());
if (cj == params.cPerGroup - THREADS_PER_CHANNEL ||
threadIdx.x * THREADS_PER_CHANNEL ==
params.cPerBlock - THREADS_PER_CHANNEL) {
smem[gi] = make_float2(out.sum, out.sumSq);
}
__syncthreads();
if (cj == params.cPerGroup - THREADS_PER_CHANNEL ||
threadIdx.x * THREADS_PER_CHANNEL ==
params.cPerBlock - THREADS_PER_CHANNEL) {
float2 sums = smem[gi];
atomicAdd(&params.redBuffer[(2 * ni + 0) * params.groups + gj],
sums.x * params.invDHWC);
atomicAdd(&params.redBuffer[(2 * ni + 1) * params.groups + gj], sums.y);
}
}
template <typename T>
void groupNormNDHWCSum<T>::operator()(GroupNormNDHWCParams<T>* params,
gpuStream_t stream) {
dim3 grid;
grid.x = divUp(params->c, params->cPerBlock);
grid.y = divUp(params->dhw, params->dhwPerBlock);
grid.z = params->n;
if (params->cPerGroup % 2 == 0) {
switch (params->cPerBlock) {
case 512:
case 480:
groupNormNDHWCSumKernel<T, 256, 2><<<grid, 256, 0, stream>>>(*params);
break;
case 320:
groupNormNDHWCSumKernel<T, 160, 2><<<grid, 160, 0, stream>>>(*params);
break;
case 256:
groupNormNDHWCSumKernel<T, 128, 2><<<grid, 128, 0, stream>>>(*params);
break;
case 128:
groupNormNDHWCSumKernel<T, 64, 2><<<grid, 64, 0, stream>>>(*params);
break;
default:
grid.x = divUp(params->c, 128);
params->cPerBlock = 128;
groupNormNDHWCSumKernel<T, 64, 2><<<grid, 64, 0, stream>>>(*params);
}
} else {
if (params->cPerGroup != 1) {
switch (params->cPerBlock) {
case 512:
groupNormNDHWCSumKernel<T, 512, 1><<<grid, 512, 0, stream>>>(*params);
break;
case 480:
groupNormNDHWCSumKernel<T, 480, 1><<<grid, 480, 0, stream>>>(*params);
break;
case 320:
groupNormNDHWCSumKernel<T, 320, 1><<<grid, 320, 0, stream>>>(*params);
break;
case 256:
groupNormNDHWCSumKernel<T, 256, 1><<<grid, 256, 0, stream>>>(*params);
break;
case 128:
groupNormNDHWCSumKernel<T, 128, 1><<<grid, 128, 0, stream>>>(*params);
break;
default:
grid.x = divUp(params->c, 128);
params->cPerBlock = 128;
groupNormNDHWCSumKernel<T, 128, 1><<<grid, 128, 0, stream>>>(*params);
}
} else {
switch (params->cPerBlock) {
case 512:
groupNormNDHWCSumSingerChannelKernel<T, 512>
<<<grid, 512, 0, stream>>>(*params);
break;
case 480:
groupNormNDHWCSumSingerChannelKernel<T, 480>
<<<grid, 480, 0, stream>>>(*params);
break;
case 320:
groupNormNDHWCSumSingerChannelKernel<T, 320>
<<<grid, 320, 0, stream>>>(*params);
break;
case 256:
groupNormNDHWCSumSingerChannelKernel<T, 256>
<<<grid, 256, 0, stream>>>(*params);
break;
case 128:
groupNormNDHWCSumSingerChannelKernel<T, 128>
<<<grid, 128, 0, stream>>>(*params);
break;
default:
grid.x = divUp(params->c, 128);
params->cPerBlock = 128;
groupNormNDHWCSumSingerChannelKernel<T, 128>
<<<grid, 128, 0, stream>>>(*params);
}
}
}
}
template class PADDLE_API groupNormNDHWCSum<half>;
template <typename T, int THREADS_PER_CHANNEL>
inline __device__ void GroupNormCompute(int64_t dhwBegin,
int64_t dhwEnd,
int32_t ci,
const GroupNormNDHWCParams<T>& params,
float mean,
float invStdDev) {
float gamma = __2float<T>(*(reinterpret_cast<T const*>(params.gamma) + ci));
float beta = __2float<T>(*(reinterpret_cast<T const*>(params.beta) + ci));
for (int64_t dhwi = dhwBegin; dhwi < dhwEnd; ++dhwi) {
// The src/dst offset.
int64_t offset =
static_cast<int64_t>(blockIdx.z) * params.dhwc + dhwi * params.c + ci;
float src_data = __2float<T>(params.srcX[offset]);
if (params.srcR != nullptr) {
auto gi = ci / params.cPerGroup;
auto gj = ci % params.cPerGroup;
int64_t g_offset =
params.y_same_with_x ? offset : gi * params.cPerGroup + gj;
src_data += __2float<T>(params.srcR[g_offset]);
*reinterpret_cast<T*>(&params.eleOut[offset]) = __2dst<T>(src_data);
}
// Normalize the channels.
float dst_data = (src_data - mean) * invStdDev;
// Scale by gamma and add beta.
dst_data = gamma * dst_data + beta;
// Apply Silu if needed.
if (params.withSilu) {
dst_data = dst_data * sigmoid(dst_data);
}
// Store the scaled values.
*reinterpret_cast<T*>(&params.dst[offset]) = __2dst<T>(dst_data);
}
}
template <>
inline __device__ void GroupNormCompute<float16, 2>(
int64_t dhwBegin,
int64_t dhwEnd,
int32_t ci,
const GroupNormNDHWCParams<float16>& params,
float mean,
float invStdDev) {
float2 gammaF2, betaF2;
gammaF2 = __half22float2(*reinterpret_cast<__half2 const*>(
reinterpret_cast<half const*>(params.gamma) + ci));
betaF2 = __half22float2(*reinterpret_cast<__half2 const*>(
reinterpret_cast<half const*>(params.beta) + ci));
// Iterate over the activations to compute the sums.
for (int64_t dhwi = dhwBegin; dhwi < dhwEnd; ++dhwi) {
// The src/dst offset.
int64_t offset =
static_cast<int64_t>(blockIdx.z) * params.dhwc + dhwi * params.c + ci;
// Fetch two channels per thread.
__half2 h2 = *reinterpret_cast<__half2 const*>(&params.srcX[offset]);
// Extract the two half values.
float2 f2 = __half22float2(h2);
if (params.srcR != nullptr) {
auto gi = ci / params.cPerGroup;
auto gj = ci % params.cPerGroup;
int64_t g_offset =
params.y_same_with_x ? offset : gi * params.cPerGroup + gj;
__half2 r2 = *reinterpret_cast<__half2 const*>(&params.srcR[g_offset]);
float2 r_f2 = __half22float2(r2);
f2.x += r_f2.x;
f2.y += r_f2.y;
*reinterpret_cast<__half2*>(&params.eleOut[offset]) =
__float22half2_rn(f2);
}
// Normalize the channels.
f2.x = (f2.x - mean) * invStdDev;
f2.y = (f2.y - mean) * invStdDev;
// Scale by gamma and add beta.
f2.x = gammaF2.x * f2.x + betaF2.x;
f2.y = gammaF2.y * f2.y + betaF2.y;
// Apply Silu if needed.
if (params.withSilu) {
f2.x = f2.x * sigmoid(f2.x);
f2.y = f2.y * sigmoid(f2.y);
}
// Store the scaled values.
*reinterpret_cast<__half2*>(&params.dst[offset]) = __float22half2_rn(f2);
}
}
template <>
inline __device__ void GroupNormCompute<__half, 2>(
int64_t dhwBegin,
int64_t dhwEnd,
int32_t ci,
const GroupNormNDHWCParams<__half>& params,
float mean,
float invStdDev) {
float2 gammaF2, betaF2;
gammaF2 = __half22float2(*reinterpret_cast<__half2 const*>(
reinterpret_cast<half const*>(params.gamma) + ci));
betaF2 = __half22float2(*reinterpret_cast<__half2 const*>(
reinterpret_cast<half const*>(params.beta) + ci));
// Iterate over the activations to compute the sums.
for (int64_t dhwi = dhwBegin; dhwi < dhwEnd; ++dhwi) {
// The src/dst offset.
int64_t offset =
static_cast<int64_t>(blockIdx.z) * params.dhwc + dhwi * params.c + ci;
// Fetch two channels per thread.
__half2 h2 = *reinterpret_cast<__half2 const*>(&params.srcX[offset]);
// Extract the two half values.
float2 f2 = __half22float2(h2);
if (params.srcR != nullptr) {
auto gi = ci / params.cPerGroup;
auto gj = ci % params.cPerGroup;
int64_t g_offset =
params.y_same_with_x ? offset : gi * params.cPerGroup + gj;
__half2 r2 = *reinterpret_cast<__half2 const*>(&params.srcR[g_offset]);
float2 r_f2 = __half22float2(r2);
f2.x += r_f2.x;
f2.y += r_f2.y;
*reinterpret_cast<__half2*>(&params.eleOut[offset]) =
__float22half2_rn(f2);
}
// Normalize the channels.
f2.x = (f2.x - mean) * invStdDev;
f2.y = (f2.y - mean) * invStdDev;
// Scale by gamma and add beta.
f2.x = gammaF2.x * f2.x + betaF2.x;
f2.y = gammaF2.y * f2.y + betaF2.y;
// Apply Silu if needed.
if (params.withSilu) {
f2.x = f2.x * sigmoid(f2.x);
f2.y = f2.y * sigmoid(f2.y);
}
// Store the scaled values.
*reinterpret_cast<__half2*>(&params.dst[offset]) = __float22half2_rn(f2);
}
}
#ifdef PADDLE_CUDA_BF16
template <>
inline __device__ void GroupNormCompute<bfloat16, 2>(
int64_t dhwBegin,
int64_t dhwEnd,
int32_t ci,
const GroupNormNDHWCParams<bfloat16>& params,
float mean,
float invStdDev) {
float2 gammaF2, betaF2;
gammaF2 = bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(
reinterpret_cast<__nv_bfloat16 const*>(params.gamma) + ci));
betaF2 = bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(
reinterpret_cast<__nv_bfloat16 const*>(params.beta) + ci));
// Iterate over the activations to compute the sums.
for (int64_t dhwi = dhwBegin; dhwi < dhwEnd; ++dhwi) {
// The src/dst offset.
int64_t offset =
static_cast<int64_t>(blockIdx.z) * params.dhwc + dhwi * params.c + ci;
// Fetch two channels per thread.
__nv_bfloat162 h2 =
*reinterpret_cast<__nv_bfloat162 const*>(&params.srcX[offset]);
// Extract the two half values.
float2 f2 = bfloat1622float2(h2);
if (params.srcR != nullptr) {
auto gi = ci / params.cPerGroup;
auto gj = ci % params.cPerGroup;
int64_t g_offset =
params.y_same_with_x ? offset : gi * params.cPerGroup + gj;
__nv_bfloat162 r2 =
*reinterpret_cast<__nv_bfloat162 const*>(&params.srcR[g_offset]);
float2 r_f2 = bfloat1622float2(r2);
f2.x += r_f2.x;
f2.y += r_f2.y;
*reinterpret_cast<__nv_bfloat162*>(&params.eleOut[offset]) =
float22bfloat162_rn(f2);
}
// Normalize the channels.
f2.x = (f2.x - mean) * invStdDev;
f2.y = (f2.y - mean) * invStdDev;
// Scale by gamma and add beta.
f2.x = gammaF2.x * f2.x + betaF2.x;
f2.y = gammaF2.y * f2.y + betaF2.y;
// Apply Silu if needed.
if (params.withSilu) {
f2.x = f2.x * sigmoid(f2.x);
f2.y = f2.y * sigmoid(f2.y);
}
// Store the scaled values.
*reinterpret_cast<__nv_bfloat162*>(&params.dst[offset]) =
float22bfloat162_rn(f2);
}
}
#endif
template <typename T, int THREADS_PER_CHANNEL>
__global__ void groupNormNDHWCScaleKernel(
const GroupNormNDHWCParams<T> params) {
// The instance in the batch.
int32_t ni = blockIdx.z;
// The channel loaded by that thread (2 channels per thread for F16x2).
int64_t ci = static_cast<int64_t>(blockIdx.x) *
static_cast<int64_t>(params.cPerBlock) +
static_cast<int64_t>(threadIdx.x) * THREADS_PER_CHANNEL;
// The group that thread works on and the channel in the group (modulus).
int32_t gi = ci / params.cPerGroup;
int32_t gj = ci % params.cPerGroup;
if (ci >= params.c || gi >= params.groups) {
return;
}
// Load the sum and sum of squares for the group.
float mean = params.redBuffer[(2 * ni + 0) * params.groups + gi];
float sumSq = params.redBuffer[(2 * ni + 1) * params.groups + gi];
// Compute the variance.
float var = sumSq * params.invDHWC - (mean * mean);
if (params.var_data != nullptr) {
params.var_data[ni * params.groups + gi] = var;
}
// Compute the inverse of the stddev.
float invStdDev = rsqrtf(var + params.eps);
// The first activation loaded by that block.
int64_t dhwBegin = static_cast<int64_t>(blockIdx.y) * params.dhwPerBlock;
// The last activation loaded by that block.
int64_t dhwEnd = min(dhwBegin + params.dhwPerBlock, params.dhw);
GroupNormCompute<T, THREADS_PER_CHANNEL>(
dhwBegin, dhwEnd, ci, params, mean, invStdDev);
}
template <typename T>
void groupNormNDHWCScale<T>::operator()(const GroupNormNDHWCParams<T>& params,
gpuStream_t stream) {
dim3 grid;
// The number of blocks to compute all the channels.
grid.x = divUp(params.c, params.cPerBlock);
// The number of blocks to compute all the activations in a given instance.
grid.y = divUp(params.dhw, params.dhwPerBlock);
// The number of instances.
grid.z = params.n;
if (params.cPerGroup % 2 == 0) {
switch (params.cPerBlock) {
case 512:
case 480:
groupNormNDHWCScaleKernel<T, 2><<<grid, 256, 0, stream>>>(params);
break;
case 320:
groupNormNDHWCScaleKernel<T, 2><<<grid, 160, 0, stream>>>(params);
break;
case 256:
groupNormNDHWCScaleKernel<T, 2><<<grid, 128, 0, stream>>>(params);
break;
case 128:
groupNormNDHWCScaleKernel<T, 2><<<grid, 64, 0, stream>>>(params);
break;
default:
grid.x = divUp(params.c, 128);
groupNormNDHWCScaleKernel<T, 2><<<grid, 64, 0, stream>>>(params);
}
} else {
switch (params.cPerBlock) {
case 512:
groupNormNDHWCScaleKernel<T, 1><<<grid, 512, 0, stream>>>(params);
break;
case 480:
groupNormNDHWCScaleKernel<T, 1><<<grid, 480, 0, stream>>>(params);
break;
case 320:
groupNormNDHWCScaleKernel<T, 1><<<grid, 320, 0, stream>>>(params);
break;
case 256:
groupNormNDHWCScaleKernel<T, 1><<<grid, 256, 0, stream>>>(params);
break;
case 128:
groupNormNDHWCScaleKernel<T, 1><<<grid, 128, 0, stream>>>(params);
break;
default:
grid.x = divUp(params.c, 128);
groupNormNDHWCScaleKernel<T, 1><<<grid, 128, 0, stream>>>(params);
}
}
}
template class PADDLE_API groupNormNDHWCScale<half>;
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) {
const DataLayout data_layout = StringToDataLayout(data_layout_str);
if (data_layout != DataLayout::NHWC) {
PD_THROW("data_layout only supports NHWC and NDHWC");
}
using AccT = typename MPTypeTrait<T>::Type;
GroupNormNDHWCParams<T> params_;
params_.withSilu = activation == "silu" ? true : false;
const auto x_dims = x.dims();
dev_ctx.template Alloc<T>(y);
const T* x_data = x.data<T>();
T* y_data = y->data<T>();
const auto scale_ptr = scale.get_ptr();
const auto bias_ptr = bias.get_ptr();
const T* scale_data = nullptr;
if (scale_ptr) scale_data = scale_ptr->data<T>();
const T* bias_data = nullptr;
if (bias_ptr) bias_data = bias_ptr->data<T>();
const auto d_dim = x_dims.size();
params_.n = x_dims[0];
if (d_dim == 3) {
params_.c = x_dims[2];
params_.d = 1;
params_.h = 1;
params_.w = x_dims[1];
} else if (d_dim == 4) {
params_.c = x_dims[3];
params_.d = 1;
params_.h = x_dims[1];
params_.w = x_dims[2];
} else {
// d_dim == 5
params_.c = x_dims[4];
params_.d = x_dims[1];
params_.h = x_dims[2];
params_.w = x_dims[3];
}
const T* residual_data = nullptr;
const auto residual_ptr = residual.get_ptr();
T* residual_out_data = nullptr;
if (residual_ptr) {
dev_ctx.template Alloc<T>(residual_out);
residual_data = residual_ptr->data<T>();
residual_out_data = residual_out->data<T>();
const auto r_dims = residual_ptr->dims();
int64_t r_dim = 1;
for (size_t i = 0; i < r_dims.size(); i++) {
r_dim *= r_dims[i];
}
params_.y_same_with_x = r_dim == static_cast<int64_t>(params_.n) *
params_.c * params_.d * params_.h *
params_.w
? true
: false;
}
dev_ctx.template Alloc<AccT>(mean);
dev_ctx.template Alloc<AccT>(var);
auto* mean_data = mean->data<AccT>();
auto* var_data = var->data<AccT>();
params_.var_data = var_data;
int32_t cPerBlock = 320;
int32_t maxBlocksPerDHW = 1024;
switch (params_.c) {
case 2048:
case 1024:
cPerBlock = 512;
break;
case 960:
case 1920:
cPerBlock = 480;
break;
case 512:
case 256:
cPerBlock = 256;
break;
case 128:
cPerBlock = 128;
break;
default:
cPerBlock = 320;
}
params_.groups = groups;
params_.cPerGroup = params_.c / params_.groups;
if (cPerBlock % params_.cPerGroup != 0) {
cPerBlock = params_.cPerGroup;
}
params_.srcX = reinterpret_cast<const T*>(x_data);
params_.dst = reinterpret_cast<T*>(y_data);
if (residual_ptr) {
params_.srcR = reinterpret_cast<const T*>(residual_data);
params_.eleOut = reinterpret_cast<T*>(residual_out_data);
}
params_.gamma = scale_data;
params_.beta = bias_data;
params_.dhw = static_cast<int64_t>(params_.d) * params_.h * params_.w;
const int64_t blocksPerDHW =
findMaxDivisor(params_.dhw, static_cast<int64_t>(maxBlocksPerDHW));
params_.dhwPerBlock = divUp(params_.dhw, blocksPerDHW);
params_.cPerBlock = cPerBlock;
params_.dhwc = params_.dhw * params_.c;
params_.invDHWC = 1.F / static_cast<float>(params_.dhw * params_.cPerGroup);
params_.eps = epsilon;
auto stream = dev_ctx.stream();
DenseTensor redBuffer;
int64_t buffer_sizes = 2 * static_cast<int64_t>(params_.n) * groups;
redBuffer.Resize({1, buffer_sizes});
params_.redBuffer = dev_ctx.template Alloc<float>(&redBuffer);
int64_t max_grid_x = dev_ctx.GetCUDAMaxGridDimSize()[0];
int64_t max_grid_y = dev_ctx.GetCUDAMaxGridDimSize()[1];
int64_t max_grid_z = dev_ctx.GetCUDAMaxGridDimSize()[2];
if (params_.n > max_grid_z) {
PADDLE_THROW(common::errors::Unimplemented(
"GroupNorm kernel launch failed: 'batch_size' (%ld) exceeds the "
"maximum supported limit (%ld). "
"Please reduce the batch size or modify the kernel configuration.",
params_.n,
max_grid_z));
}
if (divUp(params_.dhw, params_.dhwPerBlock) > max_grid_y) {
PADDLE_THROW(common::errors::Unimplemented(
"GroupNorm kernel launch failed: computed gridDim.y (%ld) exceeds the "
"hardware limit (%ld). "
"This may be due to excessively large 'dhw' (%ld). Consider reducing "
"input spatial dimensions "
"or adjusting 'dhwPerBlock'.",
divUp(params_.dhw, params_.dhwPerBlock),
max_grid_y,
params_.dhw));
}
if (divUp(params_.c, std::max(params_.cPerBlock, 128)) > max_grid_x) {
PADDLE_THROW(common::errors::Unimplemented(
"GroupNorm kernel launch failed: computed gridDim.x (%ld) exceeds the "
"hardware limit (%ld). "
"This may be due to excessively large channel count 'c' (%ld). "
"Consider reducing the number of channels "
"or adjusting 'cPerBlock'.",
divUp(params_.c, std::max(params_.cPerBlock, 128)),
max_grid_x,
params_.c));
}
#ifdef PADDLE_WITH_HIP
hipMemset(params_.redBuffer, 0, buffer_sizes * sizeof(float));
#else
cudaMemset(params_.redBuffer, 0, buffer_sizes * sizeof(float));
#endif
groupNormNDHWCSum<T> ndhwc_sum;
ndhwc_sum(&params_, stream);
groupNormNDHWCScale<T> ndhwc_scale;
ndhwc_scale(params_, stream);
#ifdef PADDLE_WITH_HIP
backends::gpu::GpuMemcpyAsync(mean_data,
params_.redBuffer,
params_.n * groups * sizeof(float),
hipMemcpyDeviceToHost,
stream);
#else
backends::gpu::GpuMemcpyAsync(mean_data,
params_.redBuffer,
params_.n * groups * sizeof(float),
cudaMemcpyDeviceToHost,
stream);
#endif
}
// ============================================================================
// PyTorch-aligned NCHW forward implementation using Welford algorithm
// ============================================================================
// Welford data structure for online mean/variance computation
template <typename AccT>
struct WelfordData {
AccT mean;
AccT m2;
int64_t n;
AccT nf;
__host__ __device__ WelfordData() : mean(0), m2(0), n(0), nf(0) {}
__host__ __device__ WelfordData(AccT mean, AccT m2, int64_t n, AccT nf)
: mean(mean), m2(m2), n(n), nf(nf) {}
};
// Welford online update: incorporate a single data point
template <typename AccT>
__device__ __forceinline__ WelfordData<AccT> WelfordReduce(
WelfordData<AccT> acc, AccT data) {
int64_t new_n = acc.n + 1;
AccT new_nf = static_cast<AccT>(new_n);
AccT delta = data - acc.mean;
AccT new_mean = acc.mean + delta / new_nf;
AccT new_delta = data - new_mean;
return {new_mean, acc.m2 + delta * new_delta, new_n, new_nf};
}
// Welford combine: merge two partial aggregates
template <typename AccT>
__device__ __forceinline__ WelfordData<AccT> WelfordCombine(
WelfordData<AccT> a, WelfordData<AccT> b) {
if (a.nf == 0) return b;
if (b.nf == 0) return a;
AccT delta = b.mean - a.mean;
AccT new_count = a.nf + b.nf;
AccT nb_over_n = b.nf / new_count;
return {a.mean + delta * nb_over_n,
a.m2 + b.m2 + delta * delta * a.nf * nb_over_n,
-1,
new_count};
}
// Warp-level Welford reduce using shuffle-down
template <typename AccT>
__device__ __forceinline__ WelfordData<AccT> WelfordWarpReduce(
WelfordData<AccT> val) {
#pragma unroll
for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
WelfordData<AccT> other;
other.mean =
backends::gpu::CudaShuffleDownSync(0xffffffff, val.mean, offset);
other.m2 = backends::gpu::CudaShuffleDownSync(0xffffffff, val.m2, offset);
other.n = backends::gpu::CudaShuffleDownSync(0xffffffff, val.n, offset);
other.nf = backends::gpu::CudaShuffleDownSync(0xffffffff, val.nf, offset);
val = WelfordCombine(val, other);
}
return val;
}
// Block-level Welford reduce (for > warpSize threads)
template <typename AccT>
__device__ __forceinline__ WelfordData<AccT> WelfordBlockReduce(
WelfordData<AccT> val) {
// Shared memory for warp results (max 32 warps = 1024/32 or 1024/64)
constexpr int kMaxWarps = 32;
__shared__ AccT s_mean[kMaxWarps];
__shared__ AccT s_m2[kMaxWarps];
__shared__ int64_t s_n[kMaxWarps];
__shared__ AccT s_nf[kMaxWarps];
int tid = threadIdx.x;
int lid = tid % warpSize;
int wid = tid / warpSize;
int num_warps = blockDim.x / warpSize;
// First reduce within each warp
val = WelfordWarpReduce(val);
__syncthreads();
// Warp leaders write to shared memory
if (lid == 0) {
s_mean[wid] = val.mean;
s_m2[wid] = val.m2;
s_n[wid] = val.n;
s_nf[wid] = val.nf;
}
__syncthreads();
// First warp reads and reduces across warps
WelfordData<AccT> zero_val(0, 0, 0, 0);
val = (tid < num_warps)
? WelfordData<AccT>{s_mean[lid], s_m2[lid], s_n[lid], s_nf[lid]}
: zero_val;
if (wid == 0) {
val = WelfordWarpReduce(val);
}
return val;
}
// Constant matching PyTorch's kCUDABlockReduceNumThreads
constexpr int kBlockReduceNumThreads = 512;
// Phase 1: Welford moments kernel - computes mean and rstd per group
// Grid: N*G blocks. Each block processes D*HxW elements for one (n,g) pair.
// Matches PyTorch's RowwiseMomentsCUDAKernel exactly.
template <typename T, typename AccT>
__global__ void WelfordMomentsCUDAKernel(
int64_t D_times_HxW, AccT eps, const T* X, AccT* mean_out, AccT* var_out) {
const int64_t i = blockIdx.x;
WelfordData<AccT> val(0, 0, 0, 0);
for (int64_t j = threadIdx.x; j < D_times_HxW; j += blockDim.x) {
const int64_t index = i * D_times_HxW + j;
val = WelfordReduce(val, static_cast<AccT>(X[index]));
}
if (blockDim.x <= static_cast<unsigned>(warpSize)) {
val = WelfordWarpReduce(val);
} else {
val = WelfordBlockReduce(val);
}
if (threadIdx.x == 0) {
AccT final_mean = val.mean;
AccT final_var = val.m2 / val.nf;
mean_out[i] = final_mean;
var_out[i] = final_var;
}
}
// Phase 2a: Compute fused parameters a and b
// a[n*C + c] = rstd[n*G + c/D] * gamma[c]
// b[n*C + c] = -a[n*C + c] * mean[n*G + c/D] + beta[c]
// Matches PyTorch's ComputeFusedParamsCUDAKernel.
// Key: mean and rstd are first rounded through T (input dtype) to match
// PyTorch's behavior where mean/rstd are stored in T.
template <typename T, typename AccT>
__global__ void ComputeFusedParamsCUDAKernel(int64_t N,
int64_t C,
int64_t G,
AccT eps,
const AccT* mean,
const AccT* var,
const T* gamma,
const T* beta,
AccT* a,
AccT* b) {
const int64_t index =
static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
if (index < N * C) {
const int64_t D = C / G;
const int64_t ng = index / D;
const int64_t c = index % C;
// Round mean and rstd through T to match PyTorch's behavior
// (PyTorch stores mean/rstd in input dtype T, then reads them back)
AccT mean_val = static_cast<AccT>(static_cast<T>(mean[ng]));
AccT rstd_val = static_cast<AccT>(static_cast<T>(rsqrt(var[ng] + eps)));
const AccT scale =
(gamma == nullptr) ? rstd_val : rstd_val * static_cast<AccT>(gamma[c]);
a[index] = scale;
b[index] =
-scale * mean_val +
((beta == nullptr) ? static_cast<AccT>(0) : static_cast<AccT>(beta[c]));
}
}
// Phase 2b: Element-wise normalization: Y = a * X + b
// Matches PyTorch's approach of Y = a[n*C+c] * float(X[idx]) + b[n*C+c].
template <typename T, typename AccT>
__global__ void GroupNormForwardElementwiseCUDAKernel(
int64_t N_C, int64_t HxW, const T* X, const AccT* a, const AccT* b, T* Y) {
const int64_t total = N_C * HxW;
for (int64_t idx =
static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
idx < total;
idx += static_cast<int64_t>(gridDim.x) * blockDim.x) {
const int64_t nc = idx / HxW;
Y[idx] = static_cast<T>(a[nc] * static_cast<AccT>(X[idx]) + b[nc]);
}
}
// Forward without gamma/beta - simple (X - mean) * rstd
// Rounds mean/rstd through T to match PyTorch's precision behavior.
template <typename T, typename AccT>
__global__ void GroupNormForwardNoScaleBiasCUDAKernel(int64_t N_G,
int64_t D_HxW,
const T* X,
const AccT* mean,
const AccT* var,
AccT eps,
T* Y) {
const int64_t total = N_G * D_HxW;
for (int64_t idx =
static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
idx < total;
idx += static_cast<int64_t>(gridDim.x) * blockDim.x) {
const int64_t ng = idx / D_HxW;
// Round through T to match PyTorch behavior
AccT mean_val = static_cast<AccT>(static_cast<T>(mean[ng]));
AccT rstd_val = static_cast<AccT>(static_cast<T>(rsqrt(var[ng] + eps)));
Y[idx] = static_cast<T>((static_cast<AccT>(X[idx]) - mean_val) * rstd_val);
}
}
// Phase 2b vectorized: float4 version for T=float, HxW divisible by 4
template <typename T, typename AccT>
__global__ void GroupNormForwardElementwiseVec4CUDAKernel(
int64_t N_C,
int64_t HxW_vec,
const T* __restrict__ X,
const AccT* __restrict__ a,
const AccT* __restrict__ b,
T* __restrict__ Y) {
const int64_t total = N_C * HxW_vec;
for (int64_t idx =
static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
idx < total;
idx += static_cast<int64_t>(gridDim.x) * blockDim.x) {
const int64_t nc = idx / HxW_vec;
const AccT a_val = a[nc];
const AccT b_val = b[nc];
float4 x4 = reinterpret_cast<const float4*>(X)[idx];
float4 y4;
y4.x = static_cast<float>(a_val * static_cast<AccT>(x4.x) + b_val);
y4.y = static_cast<float>(a_val * static_cast<AccT>(x4.y) + b_val);
y4.z = static_cast<float>(a_val * static_cast<AccT>(x4.z) + b_val);
y4.w = static_cast<float>(a_val * static_cast<AccT>(x4.w) + b_val);
reinterpret_cast<float4*>(Y)[idx] = y4;
}
}
// Phase 2b vectorized: double2 version for T=double, HxW divisible by 2
template <typename T, typename AccT>
__global__ void GroupNormForwardElementwiseVec2DoubleCUDAKernel(
int64_t N_C,
int64_t HxW_vec,
const T* __restrict__ X,
const AccT* __restrict__ a,
const AccT* __restrict__ b,
T* __restrict__ Y) {
const int64_t total = N_C * HxW_vec;
for (int64_t idx =
static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
idx < total;
idx += static_cast<int64_t>(gridDim.x) * blockDim.x) {
const int64_t nc = idx / HxW_vec;
const AccT a_val = a[nc];
const AccT b_val = b[nc];
double2 x2 = reinterpret_cast<const double2*>(X)[idx];
double2 y2;
y2.x = static_cast<double>(a_val * static_cast<AccT>(x2.x) + b_val);
y2.y = static_cast<double>(a_val * static_cast<AccT>(x2.y) + b_val);
reinterpret_cast<double2*>(Y)[idx] = y2;
}
}
// Forward without gamma/beta vectorized: float4 version
template <typename T, typename AccT>
__global__ void GroupNormForwardNoScaleBiasVec4CUDAKernel(
int64_t N_G,
int64_t D_HxW_vec,
const T* __restrict__ X,
const AccT* __restrict__ mean,
const AccT* __restrict__ var,
AccT eps,
T* __restrict__ Y) {
const int64_t total = N_G * D_HxW_vec;
for (int64_t idx =
static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
idx < total;
idx += static_cast<int64_t>(gridDim.x) * blockDim.x) {
const int64_t ng = idx / D_HxW_vec;
// Round through T to match PyTorch behavior
AccT mean_val = static_cast<AccT>(static_cast<T>(mean[ng]));
AccT rstd_val = static_cast<AccT>(static_cast<T>(rsqrt(var[ng] + eps)));
float4 x4 = reinterpret_cast<const float4*>(X)[idx];
float4 y4;
y4.x = static_cast<float>((static_cast<AccT>(x4.x) - mean_val) * rstd_val);
y4.y = static_cast<float>((static_cast<AccT>(x4.y) - mean_val) * rstd_val);
y4.z = static_cast<float>((static_cast<AccT>(x4.z) - mean_val) * rstd_val);
y4.w = static_cast<float>((static_cast<AccT>(x4.w) - mean_val) * rstd_val);
reinterpret_cast<float4*>(Y)[idx] = y4;
}
}
// Forward without gamma/beta vectorized: double2 version
template <typename T, typename AccT>
__global__ void GroupNormForwardNoScaleBiasVec2DoubleCUDAKernel(
int64_t N_G,
int64_t D_HxW_vec,
const T* __restrict__ X,
const AccT* __restrict__ mean,
const AccT* __restrict__ var,
AccT eps,
T* __restrict__ Y) {
const int64_t total = N_G * D_HxW_vec;
for (int64_t idx =
static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
idx < total;
idx += static_cast<int64_t>(gridDim.x) * blockDim.x) {
const int64_t ng = idx / D_HxW_vec;
AccT mean_val = static_cast<AccT>(static_cast<T>(mean[ng]));
AccT rstd_val = static_cast<AccT>(static_cast<T>(rsqrt(var[ng] + eps)));
double2 x2 = reinterpret_cast<const double2*>(X)[idx];
double2 y2;
y2.x = static_cast<double>((static_cast<AccT>(x2.x) - mean_val) * rstd_val);
y2.y = static_cast<double>((static_cast<AccT>(x2.y) - mean_val) * rstd_val);
reinterpret_cast<double2*>(Y)[idx] = y2;
}
}
// Fused normalization kernel that computes Y on the fly without temp buffers
// Rounds mean/rstd through T to match PyTorch's precision behavior.
template <typename T, typename AccT>
__global__ void GroupNormForwardFusedNCHWKernel(int64_t N_C,
int64_t HxW,
int64_t C,
int64_t G,
const T* X,
const AccT* mean,
const AccT* var,
const T* gamma,
const T* beta,
AccT eps,
T* Y) {
const int64_t D = C / G;
const int64_t total = N_C * HxW;
for (int64_t idx =
static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
idx < total;
idx += static_cast<int64_t>(gridDim.x) * blockDim.x) {
const int64_t nc = idx / HxW;
const int64_t c = nc % C;
const int64_t ng = nc / D;
// Round through T to match PyTorch behavior
AccT mean_val = static_cast<AccT>(static_cast<T>(mean[ng]));
AccT rstd_val = static_cast<AccT>(static_cast<T>(rsqrt(var[ng] + eps)));
AccT scale_val =
(gamma == nullptr) ? rstd_val : rstd_val * static_cast<AccT>(gamma[c]);
AccT bias_val =
(beta == nullptr) ? static_cast<AccT>(0) : static_cast<AccT>(beta[c]);
Y[idx] = static_cast<T>(scale_val * static_cast<AccT>(X[idx]) +
(-scale_val * mean_val + bias_val));
}
}
// Keep the old NHWC kernels for GroupNormForwardGetMeanAndVar and
// GroupNormForward (used by NHWC float32/64 path and
// GroupNormDirectCUDAFunctor)
template <typename T, typename AccT>
__global__ void GroupNormForwardGetMeanAndVar(const T* x,
int64_t N,
int64_t C,
int64_t W,
int64_t imsize,
int groups,
int64_t group_size,
AccT* mean,
AccT* var) {
int64_t gid = blockIdx.y;
for (int64_t cid = blockIdx.x; cid < group_size; cid += gridDim.x) {
for (int64_t bid = blockIdx.z; bid < N; bid += gridDim.z) {
int64_t H = imsize / W;
int64_t number =
min(group_size, static_cast<int64_t>(C - gid * group_size));
int64_t ccid = gid * group_size + cid;
if (ccid >= C) return;
AccT x_mean = static_cast<AccT>(0);
AccT x_var = static_cast<AccT>(0);
for (int64_t imid = threadIdx.x; imid < imsize; imid += blockDim.x) {
AccT val;
int64_t hid = imid / W;
int64_t wid = imid % W;
val = static_cast<AccT>(x[(bid * H + hid) * W * C + wid * C + ccid]);
x_mean += val;
x_var += val * val;
}
x_mean /= number * imsize;
x_var /= number * imsize;
#ifdef __NVCC__
CudaAtomicAddWithWarp(&mean[bid * groups + gid], x_mean);
CudaAtomicAddWithWarp(&var[bid * groups + gid], x_var);
#endif
#ifdef __HIPCC__
if (blockDim.x < kps::details::kWarpSize) {
CudaAtomicAdd(&mean[bid * groups + gid], x_mean);
CudaAtomicAdd(&var[bid * groups + gid], x_var);
} else {
CudaAtomicAddWithWarp(&mean[bid * groups + gid], x_mean);
CudaAtomicAddWithWarp(&var[bid * groups + gid], x_var);
}
#endif
}
}
}
template <typename T, typename AccT, int flags>
__global__ void GroupNormForward(const T* x,
const AccT* mean,
const AccT* var,
const T* scale,
const T* bias,
int64_t N,
int64_t C,
int64_t W,
int64_t imsize,
int groups,
int64_t group_size,
AccT epsilon,
T* y,
AccT* real_var,
const DataLayout data_layout) {
int64_t gid = blockIdx.y;
for (int64_t cid = blockIdx.x; cid < group_size; cid += gridDim.x) {
for (int64_t bid = blockIdx.z; bid < N; bid += gridDim.z) {
int64_t H = imsize / W;
int64_t ccid = gid * group_size + cid;
if (ccid >= C) return;
auto ng = bid * groups + gid;
AccT x_mean = mean[ng];
AccT x_var = var[ng];
x_var = x_var - x_mean * x_mean;
AccT var_inv = rsqrt(x_var + epsilon);
if (cid == 0 && threadIdx.x == 0) {
real_var[ng] = x_var;
}
for (int64_t imid = threadIdx.x; imid < imsize; imid += blockDim.x) {
AccT val;
int64_t hid, wid;
int64_t index = (bid * C + ccid) * imsize + imid;
if (data_layout == DataLayout::NCHW) {
val = static_cast<AccT>(x[index]);
} else {
hid = imid / W;
wid = imid % W;
val = static_cast<AccT>(x[(bid * H + hid) * W * C + wid * C + ccid]);
}
val = (val - x_mean) * var_inv;
if (flags & kHasScale) {
val *= static_cast<AccT>(scale[ccid]);
}
if (flags & kHasBias) {
val += static_cast<AccT>(bias[ccid]);
}
if (data_layout == DataLayout::NCHW) {
y[index] = static_cast<T>(val);
} else {
y[(bid * H + hid) * W * C + wid * C + ccid] = static_cast<T>(val);
}
}
}
}
}
template <typename T, typename AccT>
void GroupNormDirectCUDAFunctor<T, AccT>::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) {
const auto input_ddim = make_ddim(input_shape);
const int64_t C =
(data_layout == DataLayout::NCHW ? input_ddim[1]
: input_ddim[input_ddim.size() - 1]);
const int64_t group_size = C / groups;
const int64_t W =
(data_layout == DataLayout::NCHW ? input_ddim[input_ddim.size() - 1]
: input_ddim[input_ddim.size() - 2]);
int64_t image_size = 1;
if (data_layout == DataLayout::NCHW) {
for (int i = 2; i < input_ddim.size(); ++i) {
image_size *= input_ddim[i];
}
} else {
for (int i = 1; i < input_ddim.size() - 1; ++i) {
image_size *= input_ddim[i];
}
}
int block_size = std::min(static_cast<int64_t>(1024), image_size);
int64_t max_grid_x = 65535;
dim3 grid(std::min(group_size, max_grid_x),
groups,
std::min(input_ddim[0], max_grid_x));
dim3 threads(block_size, 1, 1);
if (data_layout == DataLayout::NCHW) {
if (FLAGS_use_accuracy_compatible_kernel) {
// =========================================================
// PyTorch-compatible NCHW path using Welford algorithm
// =========================================================
const int64_t N = input_ddim[0];
const int64_t G = groups;
const int64_t D = C / G;
const int64_t D_times_HxW = D * image_size;
const int64_t num_threads =
D_times_HxW < kBlockReduceNumThreads ? 32 : kBlockReduceNumThreads;
// Phase 1: Welford moments -> mean and centered variance in temp_variance
WelfordMomentsCUDAKernel<T, AccT><<<N * G, num_threads, 0, stream>>>(
D_times_HxW, static_cast<AccT>(eps), input, mean, temp_variance);
// Phase 2: Fused normalization using mean/var on the fly
if (scale != nullptr || bias != nullptr) {
const int64_t total = N * C * image_size;
const int64_t elem_threads = 256;
const int64_t elem_blocks =
std::min((total + elem_threads - 1) / elem_threads, max_grid_x);
GroupNormForwardFusedNCHWKernel<T, AccT>
<<<elem_blocks, elem_threads, 0, stream>>>(N * C,
image_size,
C,
G,
input,
mean,
temp_variance,
scale,
bias,
static_cast<AccT>(eps),
output);
} else {
const int64_t total = N * G * D_times_HxW;
const int64_t elem_threads = 256;
const int64_t elem_blocks =
std::min((total + elem_threads - 1) / elem_threads, max_grid_x);
GroupNormForwardNoScaleBiasCUDAKernel<T, AccT>
<<<elem_blocks, elem_threads, 0, stream>>>(N * G,
D_times_HxW,
input,
mean,
temp_variance,
static_cast<AccT>(eps),
output);
}
// Copy centered variance to variance output
#ifdef PADDLE_WITH_HIP
hipMemcpyAsync(variance,
temp_variance,
N * G * sizeof(AccT),
hipMemcpyDeviceToDevice,
stream);
#else
cudaMemcpyAsync(variance,
temp_variance,
N * G * sizeof(AccT),
cudaMemcpyDeviceToDevice,
stream);
#endif
} else {
// =========================================================
// Original high-performance NCHW path
// =========================================================
constexpr int vec_size = sizeof(float4) / sizeof(T);
int64_t size = group_size * image_size;
const int max_num_threads = 1024;
int max_block_size =
std::min(static_cast<int>(size / vec_size), max_num_threads);
int block_size_nchw = 1;
while (block_size_nchw < max_block_size) {
block_size_nchw *= 2;
}
block_size_nchw = std::max(block_size_nchw, kps::details::kWarpSize);
int64_t n_groups = input_ddim[0] * static_cast<int64_t>(groups);
dim3 grids(std::min(max_grid_x, n_groups));
dim3 blocks(block_size_nchw);
if (size < vec_size * block_size_nchw) {
ScalarGetMeanAndVarNCHW<T, AccT><<<grids, blocks, 0, stream>>>(
input, mean, temp_variance, size, n_groups);
} else {
VectorizedGetMeanAndVarNCHW<T, AccT, vec_size>
<<<grids, blocks, 0, stream>>>(
input, mean, temp_variance, size, n_groups);
}
GroupNormForward<T, AccT, 3>
<<<grid, threads, 0, stream>>>(input,
mean,
temp_variance,
scale,
bias,
input_ddim[0],
C,
W,
image_size,
groups,
group_size,
static_cast<AccT>(eps),
output,
variance,
data_layout);
}
} else {
#ifdef PADDLE_WITH_HIP
hipMemset(mean, 0, sizeof(AccT) * input_ddim[0] * groups);
hipMemset(temp_variance, 0, sizeof(AccT) * input_ddim[0] * groups);
#else
cudaMemset(mean, 0, sizeof(AccT) * input_ddim[0] * groups);
cudaMemset(temp_variance, 0, sizeof(AccT) * input_ddim[0] * groups);
#endif
GroupNormForwardGetMeanAndVar<T, AccT>
<<<grid, threads, 0, stream>>>(input,
input_ddim[0],
C,
W,
image_size,
groups,
group_size,
mean,
temp_variance);
GroupNormForward<T, AccT, 3>
<<<grid, threads, 0, stream>>>(input,
mean,
temp_variance,
scale,
bias,
input_ddim[0],
C,
W,
image_size,
groups,
group_size,
static_cast<AccT>(eps),
output,
variance,
data_layout);
}
}
template class PADDLE_API GroupNormDirectCUDAFunctor<float, float>;
#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP)
template class PADDLE_API GroupNormDirectCUDAFunctor<half, float>;
#endif
template <typename T, typename Context>
void GroupNormGeneralCaseKernel(const Context& dev_ctx,
const DenseTensor& x,
const optional<DenseTensor>& scale,
const optional<DenseTensor>& bias,
double epsilon,
int groups,
const std::string& data_layout_str,
DenseTensor* y,
DenseTensor* mean,
DenseTensor* var) {
using AccT = typename MPTypeTrait<T>::Type;
const DataLayout data_layout = StringToDataLayout(data_layout_str);
const auto scale_ptr = scale.get_ptr();
const auto bias_ptr = bias.get_ptr();
const auto x_dims = x.dims();
const int64_t N = x_dims[0];
const int64_t C =
(data_layout == DataLayout::NCHW ? x_dims[1] : x_dims[x_dims.size() - 1]);
const int64_t group_size = C / groups;
const int64_t W =
(data_layout == DataLayout::NCHW ? x_dims[x_dims.size() - 1]
: x_dims[x_dims.size() - 2]);
const int64_t G = groups;
const int64_t D = C / G;
dev_ctx.template Alloc<T>(y);
dev_ctx.template Alloc<AccT>(mean);
dev_ctx.template Alloc<AccT>(var);
auto* x_data = x.data<T>();
auto* y_data = y->data<T>();
auto* mean_data = mean->data<AccT>();
auto* var_data = var->data<AccT>();
const T* scale_data = nullptr;
if (scale_ptr) scale_data = scale_ptr->data<T>();
const T* bias_data = nullptr;
if (bias_ptr) bias_data = bias_ptr->data<T>();
int64_t imsize = 1;
if (data_layout == DataLayout::NCHW) {
for (int i = 2; i < x_dims.size(); ++i) {
imsize *= x_dims[i];
}
} else {
for (int i = 1; i < x_dims.size() - 1; ++i) {
imsize *= x_dims[i];
}
}
if (data_layout == DataLayout::NCHW) {
if (FLAGS_use_accuracy_compatible_kernel) {
// =========================================================
// PyTorch-compatible NCHW path using Welford algorithm
// =========================================================
const int64_t D_times_HxW = D * imsize;
// Phase 1: Compute moments using Welford algorithm
const int64_t num_threads =
D_times_HxW < kBlockReduceNumThreads ? 32 : kBlockReduceNumThreads;
WelfordMomentsCUDAKernel<T, AccT>
<<<N * G, num_threads, 0, dev_ctx.stream()>>>(
D_times_HxW,
static_cast<AccT>(epsilon),
x_data,
mean_data,
var_data);
// Phase 2: Normalization
if (scale_data != nullptr || bias_data != nullptr) {
// Phase 2a: Compute fused params a and b
DenseTensor a_tensor, b_tensor;
a_tensor.Resize({N, C});
b_tensor.Resize({N, C});
AccT* a_data = dev_ctx.template Alloc<AccT>(&a_tensor);
AccT* b_data = dev_ctx.template Alloc<AccT>(&b_tensor);
constexpr int64_t kNumThreads = 256;
const int64_t B = (N * C + kNumThreads - 1) / kNumThreads;
ComputeFusedParamsCUDAKernel<T, AccT>
<<<B, kNumThreads, 0, dev_ctx.stream()>>>(
N,
C,
G,
static_cast<AccT>(epsilon),
mean_data,
var_data,
scale_data,
bias_data,
a_data,
b_data);
// Phase 2b: Element-wise Y = a * X + b
constexpr int64_t kElemThreads = 256;
int64_t max_grid_x = dev_ctx.GetCUDAMaxGridDimSize()[0];
if (std::is_same<T, float>::value && (imsize % 4 == 0)) {
const int64_t HxW_vec = imsize / 4;
const int64_t total_vec = N * C * HxW_vec;
const int64_t elem_blocks = std::min(
(total_vec + kElemThreads - 1) / kElemThreads, max_grid_x);
GroupNormForwardElementwiseVec4CUDAKernel<T, AccT>
<<<elem_blocks, kElemThreads, 0, dev_ctx.stream()>>>(
N * C, HxW_vec, x_data, a_data, b_data, y_data);
} else if (std::is_same<T, double>::value && (imsize % 2 == 0)) {
const int64_t HxW_vec = imsize / 2;
const int64_t total_vec = N * C * HxW_vec;
const int64_t elem_blocks = std::min(
(total_vec + kElemThreads - 1) / kElemThreads, max_grid_x);
GroupNormForwardElementwiseVec2DoubleCUDAKernel<T, AccT>
<<<elem_blocks, kElemThreads, 0, dev_ctx.stream()>>>(
N * C, HxW_vec, x_data, a_data, b_data, y_data);
} else {
const int64_t total = N * C * imsize;
const int64_t elem_blocks =
std::min((total + kElemThreads - 1) / kElemThreads, max_grid_x);
GroupNormForwardElementwiseCUDAKernel<T, AccT>
<<<elem_blocks, kElemThreads, 0, dev_ctx.stream()>>>(
N * C, imsize, x_data, a_data, b_data, y_data);
}
} else {
// No scale/bias: Y = (X - mean) * rstd
constexpr int64_t kElemThreads = 256;
int64_t max_grid_x = dev_ctx.GetCUDAMaxGridDimSize()[0];
if (std::is_same<T, float>::value && (D_times_HxW % 4 == 0)) {
const int64_t D_HxW_vec = D_times_HxW / 4;
const int64_t total_vec = N * G * D_HxW_vec;
const int64_t elem_blocks = std::min(
(total_vec + kElemThreads - 1) / kElemThreads, max_grid_x);
GroupNormForwardNoScaleBiasVec4CUDAKernel<T, AccT>
<<<elem_blocks, kElemThreads, 0, dev_ctx.stream()>>>(
N * G,
D_HxW_vec,
x_data,
mean_data,
var_data,
static_cast<AccT>(epsilon),
y_data);
} else if (std::is_same<T, double>::value && (D_times_HxW % 2 == 0)) {
const int64_t D_HxW_vec = D_times_HxW / 2;
const int64_t total_vec = N * G * D_HxW_vec;
const int64_t elem_blocks = std::min(
(total_vec + kElemThreads - 1) / kElemThreads, max_grid_x);
GroupNormForwardNoScaleBiasVec2DoubleCUDAKernel<T, AccT>
<<<elem_blocks, kElemThreads, 0, dev_ctx.stream()>>>(
N * G,
D_HxW_vec,
x_data,
mean_data,
var_data,
static_cast<AccT>(epsilon),
y_data);
} else {
const int64_t D_total = N * G * D_times_HxW;
const int64_t elem_blocks =
std::min((D_total + kElemThreads - 1) / kElemThreads, max_grid_x);
GroupNormForwardNoScaleBiasCUDAKernel<T, AccT>
<<<elem_blocks, kElemThreads, 0, dev_ctx.stream()>>>(
N * G,
D_times_HxW,
x_data,
mean_data,
var_data,
static_cast<AccT>(epsilon),
y_data);
}
}
} else {
// =========================================================
// Original high-performance NCHW path
// =========================================================
DenseTensor temp_var;
temp_var.Resize(var->dims());
dev_ctx.template Alloc<AccT>(&temp_var);
auto* temp_var_data = temp_var.data<AccT>();
constexpr int vec_size = sizeof(float4) / sizeof(T);
int64_t size = group_size * imsize;
const int max_num_threads = 1024;
int max_block_size =
std::min(static_cast<int>(size / vec_size), max_num_threads);
int block_size_nchw = 1;
while (block_size_nchw < max_block_size) {
block_size_nchw *= 2;
}
block_size_nchw = std::max(block_size_nchw, kps::details::kWarpSize);
int64_t n_groups = N * static_cast<int64_t>(groups);
int64_t max_grid_x = dev_ctx.GetCUDAMaxGridDimSize()[0];
dim3 grids(std::min(max_grid_x, n_groups));
dim3 blocks(block_size_nchw);
if (size < vec_size * block_size_nchw) {
ScalarGetMeanAndVarNCHW<T, AccT>
<<<grids, blocks, 0, dev_ctx.stream()>>>(
x_data, mean_data, temp_var_data, size, n_groups);
} else {
VectorizedGetMeanAndVarNCHW<T, AccT, vec_size>
<<<grids, blocks, 0, dev_ctx.stream()>>>(
x_data, mean_data, temp_var_data, size, n_groups);
}
int64_t max_grid_z = dev_ctx.GetCUDAMaxGridDimSize()[2];
int block_size_orig = std::min(static_cast<int64_t>(1024), imsize);
dim3 grid(
std::min(max_grid_x, group_size), groups, std::min(max_grid_z, N));
dim3 threads(block_size_orig, 1, 1);
int flags = (scale_data != nullptr) * kHasScale +
(bias_data != nullptr) * kHasBias;
UNROLL_ALL_CASES(flags,
GroupNormForward,
x_data,
mean_data,
temp_var_data,
scale_data,
bias_data,
N,
C,
W,
imsize,
groups,
group_size,
static_cast<AccT>(epsilon),
y_data,
var_data,
data_layout);
}
} else {
// =========================================================
// NHWC float32/float64 path (legacy, kept as-is)
// =========================================================
DenseTensor temp_var;
temp_var.Resize(var->dims());
dev_ctx.template Alloc<AccT>(&temp_var);
auto* temp_var_data = temp_var.data<AccT>();
funcs::SetConstant<GPUContext, AccT> set_zero_AccT;
set_zero_AccT(dev_ctx, mean, static_cast<AccT>(0));
set_zero_AccT(dev_ctx, &temp_var, static_cast<AccT>(0));
int block_size = std::min(static_cast<int64_t>(1024), imsize);
int64_t max_grid_x = dev_ctx.GetCUDAMaxGridDimSize()[0];
int64_t max_grid_z = dev_ctx.GetCUDAMaxGridDimSize()[2];
dim3 grid(
std::min(max_grid_x, group_size), groups, std::min(max_grid_z, N));
dim3 threads(block_size, 1, 1);
GroupNormForwardGetMeanAndVar<T,
AccT><<<grid, threads, 0, dev_ctx.stream()>>>(
x_data, N, C, W, imsize, groups, group_size, mean_data, temp_var_data);
int flags =
(scale_data != nullptr) * kHasScale + (bias_data != nullptr) * kHasBias;
UNROLL_ALL_CASES(flags,
GroupNormForward,
x_data,
mean_data,
temp_var_data,
scale_data,
bias_data,
N,
C,
W,
imsize,
groups,
group_size,
static_cast<AccT>(epsilon),
y_data,
var_data,
data_layout);
}
}
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_str,
DenseTensor* y,
DenseTensor* mean,
DenseTensor* var) {
if (y && y->numel() == 0) {
dev_ctx.template Alloc<T>(y);
if (mean) {
Full<T, Context>(dev_ctx, mean->dims(), 0, mean);
}
if (var) {
Full<T, Context>(dev_ctx, var->dims(), 0, var);
}
return;
}
using std::is_same;
if (is_same<T, float16>::value && data_layout_str == "NHWC") {
const optional<DenseTensor>& residual = optional<DenseTensor>(paddle::none);
DenseTensor empty_tensor;
GroupNormNDHWCKernel<float16, Context>(dev_ctx,
x,
residual,
scale,
bias,
epsilon,
groups,
data_layout_str,
"",
y,
&empty_tensor,
mean,
var);
return;
}
#ifdef PADDLE_CUDA_BF16
if (is_same<T, bfloat16>::value && data_layout_str == "NHWC") {
const optional<DenseTensor>& residual = optional<DenseTensor>(paddle::none);
DenseTensor empty_tensor;
GroupNormNDHWCKernel<bfloat16, Context>(dev_ctx,
x,
residual,
scale,
bias,
epsilon,
groups,
data_layout_str,
"",
y,
&empty_tensor,
mean,
var);
return;
}
#endif
GroupNormGeneralCaseKernel<T, Context>(
dev_ctx, x, scale, bias, epsilon, groups, data_layout_str, y, mean, var);
}
} // namespace phi
PD_REGISTER_KERNEL(group_norm,
GPU,
ALL_LAYOUT,
phi::GroupNormKernel,
float,
double,
phi::bfloat16,
phi::float16) {
if (kernel_key.dtype() == phi::DataType::BFLOAT16 ||
kernel_key.dtype() == phi::DataType::FLOAT16) {
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
}
}
PD_REGISTER_KERNEL(add_group_norm_silu,
GPU,
ALL_LAYOUT,
phi::GroupNormNDHWCKernel,
phi::bfloat16,
phi::float16) {
kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
}
Reference in New Issue View Git Blame Copy Permalink