fix group_norm address misalignment (#40657)

* fix group_norm address misalignment * fix vectorize * fix code * fix vectorize length * optimize code

Author: Zjq9409

paddle/fluid/operators/group_norm_op.cu

@@ -152,6 +152,21 @@ __device__ __forceinline__ void ThreadReduce(phi::Array<const T*, Num> arrs,
  }
 }
 
+template <typename T>
+__device__ __forceinline__ void ReduceMeanAndVar(T* mean, T* var, T x_mean,
+ T x_var, int size) {
+ const int nc = blockIdx.x;
+ x_mean = kps::details::BlockXReduce<T, kps::AddFunctor<T>>(
+ x_mean, kps::AddFunctor<T>());
+ x_var = kps::details::BlockXReduce<T, kps::AddFunctor<T>>(
+ x_var, kps::AddFunctor<T>());
+ __syncthreads();
+ if (threadIdx.x == 0) {
+ mean[nc] = static_cast<T>(x_mean / size);
+ var[nc] = static_cast<T>(x_var / size);
+ }
+}
+
 template <typename T>
 __global__ void ScalarGetMeanAndVarNCHW(const T* x, T* mean, T* var, int size) {
  int i = blockIdx.x;
@@ -162,10 +177,7 @@ __global__ void ScalarGetMeanAndVarNCHW(const T* x, T* mean, T* var, int size) {
  x_mean += val;
  x_var += val * val;
  }
- x_mean /= size;
- x_var /= size;
- CudaAtomicAddWithWarp(&mean[i], x_mean);
- CudaAtomicAddWithWarp(&var[i], x_var);
+ ReduceMeanAndVar<T>(mean, var, x_mean, x_var, size);
 }
 
 template <typename T, typename AccT, int VecSize>
@@ -174,21 +186,12 @@ __global__ void VectorizedGetMeanAndVarNCHW(const T* x, T* mean, T* var,
  int i = blockIdx.x;
  AccT x_mean = static_cast<AccT>(0);
  AccT x_var = static_cast<AccT>(0);
- const int input_offset = ((uint64_t)x) % ALIGN_BYTES / sizeof(T);
  x += i * size;
+ const int input_offset = ((uint64_t)x) % ALIGN_BYTES / sizeof(T);
  phi::Array<const T*, 1> ins;
  ins[0] = x;
  ThreadReduce<T, AccT, VecSize, 1>(ins, size, input_offset, &x_mean, &x_var);
-
- x_mean = kps::details::BlockXReduce<AccT, kps::AddFunctor<AccT>>(
- x_mean, kps::AddFunctor<AccT>());
- x_var = kps::details::BlockXReduce<AccT, kps::AddFunctor<AccT>>(
- x_var, kps::AddFunctor<AccT>());
- __syncthreads();
- if (threadIdx.x == 0) {
- mean[i] = static_cast<T>(x_mean / size);
- var[i] = static_cast<T>(x_var / size);
- }
+ ReduceMeanAndVar<AccT>(mean, var, x_mean, x_var, size);
 }
 
 template <typename T, int flags>
@@ -272,10 +275,6 @@ class GroupNormKernel<platform::CUDADeviceContext, T>
  auto& dev_ctx = ctx.template device_context<platform::CUDADeviceContext>();
  Tensor temp_var;
  temp_var.mutable_data<T>(var->dims(), ctx.GetPlace());
-
- set_zero(dev_ctx, mean, static_cast<T>(0));
- set_zero(dev_ctx, &temp_var, static_cast<T>(0));
-
  auto* x_data = x->data<T>();
  auto* y_data = y->data<T>();
  auto* mean_data = mean->data<T>();
@@ -319,7 +318,7 @@ class GroupNormKernel<platform::CUDADeviceContext, T>
  block_size_nchw = std::max(block_size_nchw, kps::details::kWarpSize);
  dim3 grids(x_dims[0] * groups);
  dim3 blocks(block_size_nchw);
- if (size < vec_size) {
+ if (size < vec_size * block_size_nchw) {
  ScalarGetMeanAndVarNCHW<T><<<grids, blocks, 0, dev_ctx.stream()>>>(
  x_data, mean_data, temp_var_data, size);
  } else {
@@ -328,6 +327,8 @@ class GroupNormKernel<platform::CUDADeviceContext, T>
  x_data, mean_data, temp_var_data, size);
  }
  } else {
+ set_zero(dev_ctx, mean, static_cast<T>(0));
+ set_zero(dev_ctx, &temp_var, static_cast<T>(0));
  GroupNormForwardGetMeanAndVar<T><<<grid, threads, 0, dev_ctx.stream()>>>(
  x_data, x_dims[0], C, W, imsize, groups, group_size, mean_data,
  temp_var_data);
@@ -424,24 +425,15 @@ __global__ void VectorizedGetDsDbCUDAKernel(int imsize, const T* x, const T* dy,
  int i = blockIdx.x;
  AccT ds_sum = static_cast<AccT>(0);
  AccT db_sum = static_cast<AccT>(0);
- const int input_offset = ((uint64_t)x) % ALIGN_BYTES / sizeof(T);
  x += i * imsize;
+ const int input_offset = ((uint64_t)x) % ALIGN_BYTES / sizeof(T);
 
  phi::Array<const T*, 2> ins;
  ins[0] = x;
  ins[1] = dy;
  ThreadReduce<T, AccT, VecSize, 2>(ins, imsize, input_offset, &db_sum,
  &ds_sum);
-
- ds_sum = kps::details::BlockXReduce<AccT, kps::AddFunctor<AccT>>(
- ds_sum, kps::AddFunctor<AccT>());
- db_sum = kps::details::BlockXReduce<AccT, kps::AddFunctor<AccT>>(
- db_sum, kps::AddFunctor<AccT>());
- __syncthreads();
- if (threadIdx.x == 0) {
- ds[i] = ds_sum;
- db[i] = db_sum;
- }
+ ReduceMeanAndVar<AccT>(db, ds, db_sum, ds_sum, 1);
 }
 
 template <typename T>
@@ -455,8 +447,7 @@ __global__ void ScalarGetDsDbCUDAKernel(int imsize, const T* x, const T* dy,
  ds_sum += dy[index] * x[index];
  db_sum += dy[index];
  }
- CudaAtomicAddWithWarp(&ds[nc], ds_sum);
- CudaAtomicAddWithWarp(&db[nc], db_sum);
+ ReduceMeanAndVar<T>(db, ds, db_sum, ds_sum, 1);
 }
 
 template <typename T>
@@ -641,13 +632,7 @@ class GroupNormGradKernel<platform::CUDADeviceContext, T>
  }
  block_size_nchw = std::max(block_size_nchw, kps::details::kWarpSize);
  dim3 blocks(block_size_nchw);
- if (imsize < vec_size) {
- if (d_scale) {
- set_zero(dev_ctx, d_scale, static_cast<T>(0));
- }
- if (d_bias) {
- set_zero(dev_ctx, d_bias, static_cast<T>(0));
- }
+ if (imsize < vec_size * block_size_nchw) {
  ScalarGetDsDbCUDAKernel<
  T><<<x_dims[0] * C, blocks, 0, dev_ctx.stream()>>>(
  imsize, x_data, dy_data, ds_data, db_data);
@@ -687,7 +672,6 @@ class GroupNormGradKernel<platform::CUDADeviceContext, T>
  imsize, C, group_size, groups, p1_data, p2_data, p3_data, x_data,
  dy_data, d_x_data);
  }
-
  } else {
  if (d_scale) {
  set_zero(dev_ctx, d_scale, static_cast<T>(0));