// 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_grad_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/full_kernel.h"
#include "paddle/phi/kernels/funcs/math_function.h"
#include "paddle/phi/kernels/gpu/group_norm_utils.h"

namespace phi {

// ============================================================================
// NHWC backward kernels (kept unchanged for NHWC data layout)
// ============================================================================

template <typename T, typename AccT, int flags>
__global__ void GroupNormBackwardGetMeanAndVar(const T* x,
                                               const T* scale,
                                               const T* bias,
                                               const T* d_y,
                                               int64_t N,
                                               int64_t C,
                                               int64_t W,
                                               int64_t imsize,
                                               int groups,
                                               int64_t group_size,
                                               double epsilon,
                                               AccT* d_mean,
                                               AccT* d_var,
                                               T* d_scale,
                                               T* d_bias) {
  int 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, C - gid * group_size);
      int64_t ccid = gid * group_size + cid;
      if (ccid >= C) return;
      T x_scale = (flags & kHasScale) ? scale[ccid] : static_cast<T>(1);
      T x_bias = (flags & kHasBias) ? bias[ccid] : static_cast<T>(0);
      T x_scale_inv = static_cast<T>(0);
      if (x_scale != static_cast<T>(0))
        x_scale_inv = static_cast<T>(1.0) / x_scale;
      AccT d_mean_data = static_cast<AccT>(0);
      AccT d_var_data = static_cast<AccT>(0);
      AccT d_scale_data = static_cast<AccT>(0);
      AccT d_bias_data = static_cast<AccT>(0);

      for (int64_t imid = threadIdx.x; imid < imsize; imid += blockDim.x) {
        AccT val, dval;

        int64_t hid = imid / W;
        int64_t wid = imid % W;
        val = static_cast<AccT>(x[(bid * H + hid) * W * C + wid * C + ccid]) -
              static_cast<AccT>(x_bias);
        dval = static_cast<AccT>(d_y[(bid * H + hid) * W * C + wid * C + ccid]);

        d_var_data += val * dval;
        d_mean_data += dval * static_cast<AccT>(x_scale);

        val = val * static_cast<AccT>(x_scale_inv);
        d_bias_data += dval;
        d_scale_data += val * dval;
      }
      CudaAtomicAddWithWarp(&(d_mean[bid * groups + gid]),
                            static_cast<AccT>(d_mean_data));
      CudaAtomicAddWithWarp(&(d_var[bid * groups + gid]),
                            static_cast<AccT>(d_var_data));

      if (flags & kHasScale) {
        CudaAtomicAdd(&(d_scale[ccid]), static_cast<T>(d_scale_data));
      }
      if (flags & kHasBias) {
        CudaAtomicAdd(&(d_bias[ccid]), static_cast<T>(d_bias_data));
      }
    }
  }
}

template <typename T, typename AccT, int flags>
__global__ void GroupNormBackward(const T* x,
                                  const T* d_y,
                                  const T* scale,
                                  const T* bias,
                                  const AccT* var,
                                  const AccT* d_mean,
                                  const AccT* d_var,
                                  int64_t N,
                                  int64_t C,
                                  int64_t W,
                                  int64_t imsize,
                                  int groups,
                                  int64_t group_size,
                                  double epsilon,
                                  T* d_x) {
  int 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, C - gid * group_size);
      int64_t ccid = gid * group_size + cid;
      if (ccid >= C) return;
      AccT x_var = var[bid * groups + gid];
      AccT d_x_mean = static_cast<AccT>(d_mean[bid * groups + gid]);
      AccT d_x_var = static_cast<AccT>(d_var[bid * groups + gid]);

      AccT x_var_inv = static_cast<AccT>(1.0) / sqrt((x_var) + epsilon);
      AccT number_inv =
          static_cast<AccT>(1.0) / static_cast<AccT>((number * imsize));

      AccT x_scale = (flags & kHasScale) ? static_cast<AccT>(scale[ccid])
                                         : static_cast<AccT>(1);
      AccT x_bias = (flags & kHasBias) ? static_cast<AccT>(bias[ccid])
                                       : static_cast<AccT>(0);
      AccT x_scale_inv = static_cast<AccT>(0);
      if (x_scale != static_cast<AccT>(0))
        x_scale_inv = static_cast<AccT>(1.0) / x_scale;

      for (int64_t imid = threadIdx.x; imid < imsize; imid += blockDim.x) {
        int64_t hid = imid / W;
        int64_t wid = imid % W;
        AccT tmp =
            static_cast<AccT>(x[(bid * H + hid) * W * C + wid * C + ccid]);
        AccT v_y = (tmp - x_bias) * x_scale_inv;
        AccT dly =
            static_cast<AccT>(d_y[(bid * H + hid) * W * C + wid * C + ccid]);
        d_x[(bid * H + hid) * W * C + wid * C + ccid] =
            static_cast<T>(x_var_inv * ((dly) * (x_scale)-number_inv * d_x_var *
                                        (v_y)-number_inv * d_x_mean));
      }
    }
  }
}

// ============================================================================
// PyTorch-aligned NCHW backward implementation
// ============================================================================

// Constant matching PyTorch's kCUDABlockReduceNumThreads
constexpr int kGradBlockReduceNumThreads = 512;
constexpr int kGradReduceTileSize = 32;

// Warp reduce sum using shuffle-down (matching PyTorch's WarpReduceSum)
template <typename AccT>
__device__ __forceinline__ AccT GradWarpReduceSum(AccT val) {
#pragma unroll
  for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
    val += backends::gpu::CudaShuffleDownSync(0xffffffff, val, offset);
  }
  return val;
}

// Block reduce sum (matching PyTorch's BlockReduceSum)
template <typename AccT>
__device__ __forceinline__ AccT GradBlockReduceSum(AccT val, AccT* shared) {
  const int tid = threadIdx.x;
  const int lid = tid % warpSize;
  const int wid = tid / warpSize;
  const int num_warps = blockDim.x / warpSize;
  val = GradWarpReduceSum(val);
  __syncthreads();
  if (lid == 0) {
    shared[wid] = val;
  }
  __syncthreads();
  val = (tid < num_warps) ? shared[lid] : AccT(0);
  if (wid == 0) {
    val = GradWarpReduceSum(val);
  }
  return val;
}

// Stage 1: Compute internal gradients ds[n,c] and db[n,c]
// ds[nc] = sum_hw(dY * X), db[nc] = sum_hw(dY)
// Grid: N*C blocks. Matches PyTorch's ComputeInternalGradientsCUDAKernel.
template <typename T, typename AccT>
__global__ void ComputeInternalGradientsCUDAKernel(
    int64_t HxW, const T* dY, const T* X, AccT* ds, AccT* db) {
  const int64_t nc = blockIdx.x;
  AccT sum1 = 0;
  AccT sum2 = 0;
  for (int64_t hw = threadIdx.x; hw < HxW; hw += blockDim.x) {
    const int64_t index = nc * HxW + hw;
    sum1 += static_cast<AccT>(dY[index]) * static_cast<AccT>(X[index]);
    sum2 += static_cast<AccT>(dY[index]);
  }
  if (blockDim.x <= static_cast<unsigned>(warpSize)) {
    sum1 = GradWarpReduceSum(sum1);
    sum2 = GradWarpReduceSum(sum2);
  } else {
    constexpr int kMaxWarps = 32;
    __shared__ AccT ds_shared[kMaxWarps];
    __shared__ AccT db_shared[kMaxWarps];
    sum1 = GradBlockReduceSum(sum1, ds_shared);
    sum2 = GradBlockReduceSum(sum2, db_shared);
  }
  if (threadIdx.x == 0) {
    ds[nc] = sum1;
    db[nc] = sum2;
  }
}

// Stage 1 vectorized: float4 version for T=float, HxW divisible by 4
// Processes 4 elements per thread per iteration for better memory throughput.
template <typename T, typename AccT>
__global__ void ComputeInternalGradientsVec4CUDAKernel(
    int64_t HxW_vec, const T* dY, const T* X, AccT* ds, AccT* db) {
  const int64_t nc = blockIdx.x;
  AccT sum1 = 0;
  AccT sum2 = 0;
  const float4* dY_vec = reinterpret_cast<const float4*>(dY + nc * HxW_vec * 4);
  const float4* X_vec = reinterpret_cast<const float4*>(X + nc * HxW_vec * 4);
  for (int64_t i = threadIdx.x; i < HxW_vec; i += blockDim.x) {
    float4 dy4 = dY_vec[i];
    float4 x4 = X_vec[i];
    sum1 += static_cast<AccT>(dy4.x) * static_cast<AccT>(x4.x);
    sum1 += static_cast<AccT>(dy4.y) * static_cast<AccT>(x4.y);
    sum1 += static_cast<AccT>(dy4.z) * static_cast<AccT>(x4.z);
    sum1 += static_cast<AccT>(dy4.w) * static_cast<AccT>(x4.w);
    sum2 += static_cast<AccT>(dy4.x);
    sum2 += static_cast<AccT>(dy4.y);
    sum2 += static_cast<AccT>(dy4.z);
    sum2 += static_cast<AccT>(dy4.w);
  }
  if (blockDim.x <= static_cast<unsigned>(warpSize)) {
    sum1 = GradWarpReduceSum(sum1);
    sum2 = GradWarpReduceSum(sum2);
  } else {
    constexpr int kMaxWarps = 32;
    __shared__ AccT ds_shared[kMaxWarps];
    __shared__ AccT db_shared[kMaxWarps];
    sum1 = GradBlockReduceSum(sum1, ds_shared);
    sum2 = GradBlockReduceSum(sum2, db_shared);
  }
  if (threadIdx.x == 0) {
    ds[nc] = sum1;
    db[nc] = sum2;
  }
}

// Stage 1 vectorized: double2 version for T=double, HxW divisible by 2
template <typename T, typename AccT>
__global__ void ComputeInternalGradientsVec2DoubleCUDAKernel(
    int64_t HxW_vec, const T* dY, const T* X, AccT* ds, AccT* db) {
  const int64_t nc = blockIdx.x;
  AccT sum1 = 0;
  AccT sum2 = 0;
  const double2* dY_vec =
      reinterpret_cast<const double2*>(dY + nc * HxW_vec * 2);
  const double2* X_vec = reinterpret_cast<const double2*>(X + nc * HxW_vec * 2);
  for (int64_t i = threadIdx.x; i < HxW_vec; i += blockDim.x) {
    double2 dy2 = dY_vec[i];
    double2 x2 = X_vec[i];
    sum1 += static_cast<AccT>(dy2.x) * static_cast<AccT>(x2.x);
    sum1 += static_cast<AccT>(dy2.y) * static_cast<AccT>(x2.y);
    sum2 += static_cast<AccT>(dy2.x);
    sum2 += static_cast<AccT>(dy2.y);
  }
  if (blockDim.x <= static_cast<unsigned>(warpSize)) {
    sum1 = GradWarpReduceSum(sum1);
    sum2 = GradWarpReduceSum(sum2);
  } else {
    constexpr int kMaxWarps = 32;
    __shared__ AccT ds_shared[kMaxWarps];
    __shared__ AccT db_shared[kMaxWarps];
    sum1 = GradBlockReduceSum(sum1, ds_shared);
    sum2 = GradBlockReduceSum(sum2, db_shared);
  }
  if (threadIdx.x == 0) {
    ds[nc] = sum1;
    db[nc] = sum2;
  }
}

// Stage 2: Compute backward fused params c2[n,g] and c3[n,g]
// Reduces ds*gamma and db*gamma across channels in group.
// Grid: dim3(N, G). Matches PyTorch's ComputeBackwardFusedParamsCUDAKernel.
// Rounds mean/rstd through T to match PyTorch's precision behavior.
template <typename T, typename AccT>
__global__ void ComputeBackwardFusedParamsCUDAKernel(int64_t C,
                                                     int64_t HxW,
                                                     int64_t G,
                                                     const AccT* mean,
                                                     const AccT* var,
                                                     const T* gamma,
                                                     const AccT* ds,
                                                     const AccT* db,
                                                     AccT eps,
                                                     AccT* c2,
                                                     AccT* c3) {
  const int64_t D = C / G;
  const int64_t n = blockIdx.x;
  const int64_t g = blockIdx.y;
  const int64_t ng = n * G + g;
  AccT sum1 = 0;
  AccT sum2 = 0;
  for (int64_t i = threadIdx.x; i < D; i += blockDim.x) {
    const int64_t index = ng * D + i;
    const int64_t c_idx = g * D + i;
    const AccT gamma_v =
        gamma == nullptr ? AccT(1) : static_cast<AccT>(gamma[c_idx]);
    sum1 += ds[index] * gamma_v;
    sum2 += db[index] * gamma_v;
  }
  if (blockDim.x <= static_cast<unsigned>(warpSize)) {
    sum1 = GradWarpReduceSum(sum1);
    sum2 = GradWarpReduceSum(sum2);
  } else {
    constexpr int kMaxWarps = 32;
    __shared__ AccT ds_shared[kMaxWarps];
    __shared__ AccT db_shared[kMaxWarps];
    sum1 = GradBlockReduceSum(sum1, ds_shared);
    sum2 = GradBlockReduceSum(sum2, db_shared);
  }
  if (threadIdx.x == 0) {
    // Round mean and rstd through T to match PyTorch
    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 s = AccT(1) / static_cast<AccT>(D * HxW);
    const AccT x =
        (sum2 * mean_val - sum1) * rstd_val * rstd_val * rstd_val * s;
    c2[ng] = x;
    c3[ng] = -x * mean_val - sum2 * rstd_val * s;
  }
}

// ============================================================================
// Optimized dX kernels for NCHW backward path
// ============================================================================

// Stage 3: Pre-compute c1[n,c] = rstd[n,g] * gamma[c]
// This removes per-element rsqrt and gamma lookup from the inner dX loop.
// Grid: ceil(N*C / 256) blocks.
// Rounds rstd through T to match PyTorch's precision behavior.
template <typename T, typename AccT>
__global__ void ComputeC1CUDAKernel(int64_t N_C,
                                    int64_t C,
                                    int64_t G,
                                    AccT eps,
                                    const AccT* var,
                                    const T* gamma,
                                    AccT* c1) {
  const int64_t D = C / G;
  for (int64_t index =
           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
       index < N_C;
       index += static_cast<int64_t>(gridDim.x) * blockDim.x) {
    const int64_t ng = index / D;
    const int64_t c = index % C;
    // Round rstd through T to match PyTorch
    AccT rstd_val = static_cast<AccT>(static_cast<T>(rsqrt(var[ng] + eps)));
    AccT gamma_val = (gamma != nullptr) ? static_cast<AccT>(gamma[c]) : AccT(1);
    c1[index] = rstd_val * gamma_val;
  }
}

// Stage 4 optimized: Vectorized dX kernel using float4
// dX[idx] = c1[nc] * dY[idx] + c2[ng] * X[idx] + c3[ng]
// Processes 4 float elements at a time within each (n,c) spatial plane.
template <typename T, typename AccT>
__global__ void GroupNormBackwardDxVec4CUDAKernel(int64_t N_C,
                                                  int64_t HxW_vec,
                                                  int64_t D,
                                                  const T* __restrict__ dY,
                                                  const T* __restrict__ X,
                                                  const AccT* __restrict__ c1,
                                                  const AccT* __restrict__ c2,
                                                  const AccT* __restrict__ c3,
                                                  T* __restrict__ dX) {
  const int64_t total_vec = N_C * HxW_vec;
  for (int64_t vidx =
           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
       vidx < total_vec;
       vidx += static_cast<int64_t>(gridDim.x) * blockDim.x) {
    const int64_t nc = vidx / HxW_vec;
    const int64_t ng = nc / D;

    const AccT c1v = c1[nc];
    const AccT c2v = c2[ng];
    const AccT c3v = c3[ng];

    // Vectorized load: 4 consecutive float elements
    float4 dy4 = reinterpret_cast<const float4*>(dY)[vidx];
    float4 x4 = reinterpret_cast<const float4*>(X)[vidx];

    float4 dx4;
    dx4.x = static_cast<T>(c1v * static_cast<AccT>(dy4.x) +
                           c2v * static_cast<AccT>(x4.x) + c3v);
    dx4.y = static_cast<T>(c1v * static_cast<AccT>(dy4.y) +
                           c2v * static_cast<AccT>(x4.y) + c3v);
    dx4.z = static_cast<T>(c1v * static_cast<AccT>(dy4.z) +
                           c2v * static_cast<AccT>(x4.z) + c3v);
    dx4.w = static_cast<T>(c1v * static_cast<AccT>(dy4.w) +
                           c2v * static_cast<AccT>(x4.w) + c3v);

    reinterpret_cast<float4*>(dX)[vidx] = dx4;
  }
}

// Stage 4 optimized: Vectorized dX kernel using double2 for float64
template <typename T, typename AccT>
__global__ void GroupNormBackwardDxVec2DoubleCUDAKernel(
    int64_t N_C,
    int64_t HxW_vec,
    int64_t D,
    const T* __restrict__ dY,
    const T* __restrict__ X,
    const AccT* __restrict__ c1,
    const AccT* __restrict__ c2,
    const AccT* __restrict__ c3,
    T* __restrict__ dX) {
  const int64_t total_vec = N_C * HxW_vec;
  for (int64_t vidx =
           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
       vidx < total_vec;
       vidx += static_cast<int64_t>(gridDim.x) * blockDim.x) {
    const int64_t nc = vidx / HxW_vec;
    const int64_t ng = nc / D;

    const AccT c1v = c1[nc];
    const AccT c2v = c2[ng];
    const AccT c3v = c3[ng];

    // Vectorized load: 2 consecutive double elements
    double2 dy2 = reinterpret_cast<const double2*>(dY)[vidx];
    double2 x2 = reinterpret_cast<const double2*>(X)[vidx];

    double2 dx2;
    dx2.x = static_cast<T>(c1v * static_cast<AccT>(dy2.x) +
                           c2v * static_cast<AccT>(x2.x) + c3v);
    dx2.y = static_cast<T>(c1v * static_cast<AccT>(dy2.y) +
                           c2v * static_cast<AccT>(x2.y) + c3v);

    reinterpret_cast<double2*>(dX)[vidx] = dx2;
  }
}

// Stage 4 optimized: Scalar fallback for when HxW is not divisible by vec_size
// Uses pre-computed c1[n,c].
template <typename T, typename AccT>
__global__ void GroupNormBackwardDxOptCUDAKernel(int64_t N_C,
                                                 int64_t HxW,
                                                 int64_t D,
                                                 const T* __restrict__ dY,
                                                 const T* __restrict__ X,
                                                 const AccT* __restrict__ c1,
                                                 const AccT* __restrict__ c2,
                                                 const AccT* __restrict__ c3,
                                                 T* __restrict__ dX) {
  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 ng = nc / D;
    dX[idx] = static_cast<T>(c1[nc] * static_cast<AccT>(dY[idx]) +
                             c2[ng] * static_cast<AccT>(X[idx]) + c3[ng]);
  }
}

// Stage 4 optimized: Vectorized no-gamma dX kernel using float4
// When gamma is null, c1 = rstd[ng], constant across D*HxW per (n,g) group.
template <typename T, typename AccT>
__global__ void GroupNormBackwardDxNoGammaVec4CUDAKernel(
    int64_t N_G,
    int64_t D_HxW_vec,
    const T* __restrict__ dY,
    const T* __restrict__ X,
    const AccT* __restrict__ c1_ng,
    const AccT* __restrict__ c2,
    const AccT* __restrict__ c3,
    T* __restrict__ dX) {
  const int64_t total_vec = N_G * D_HxW_vec;
  for (int64_t vidx =
           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
       vidx < total_vec;
       vidx += static_cast<int64_t>(gridDim.x) * blockDim.x) {
    const int64_t ng = vidx / D_HxW_vec;

    const AccT rstd_val = c1_ng[ng];
    const AccT c2v = c2[ng];
    const AccT c3v = c3[ng];

    float4 dy4 = reinterpret_cast<const float4*>(dY)[vidx];
    float4 x4 = reinterpret_cast<const float4*>(X)[vidx];

    float4 dx4;
    dx4.x = static_cast<T>(rstd_val * static_cast<AccT>(dy4.x) +
                           c2v * static_cast<AccT>(x4.x) + c3v);
    dx4.y = static_cast<T>(rstd_val * static_cast<AccT>(dy4.y) +
                           c2v * static_cast<AccT>(x4.y) + c3v);
    dx4.z = static_cast<T>(rstd_val * static_cast<AccT>(dy4.z) +
                           c2v * static_cast<AccT>(x4.z) + c3v);
    dx4.w = static_cast<T>(rstd_val * static_cast<AccT>(dy4.w) +
                           c2v * static_cast<AccT>(x4.w) + c3v);

    reinterpret_cast<float4*>(dX)[vidx] = dx4;
  }
}

// Stage 4 optimized: Vectorized no-gamma dX kernel using double2
template <typename T, typename AccT>
__global__ void GroupNormBackwardDxNoGammaVec2DoubleCUDAKernel(
    int64_t N_G,
    int64_t D_HxW_vec,
    const T* __restrict__ dY,
    const T* __restrict__ X,
    const AccT* __restrict__ c1_ng,
    const AccT* __restrict__ c2,
    const AccT* __restrict__ c3,
    T* __restrict__ dX) {
  const int64_t total_vec = N_G * D_HxW_vec;
  for (int64_t vidx =
           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
       vidx < total_vec;
       vidx += static_cast<int64_t>(gridDim.x) * blockDim.x) {
    const int64_t ng = vidx / D_HxW_vec;

    const AccT rstd_val = c1_ng[ng];
    const AccT c2v = c2[ng];
    const AccT c3v = c3[ng];

    double2 dy2 = reinterpret_cast<const double2*>(dY)[vidx];
    double2 x2 = reinterpret_cast<const double2*>(X)[vidx];

    double2 dx2;
    dx2.x = static_cast<T>(rstd_val * static_cast<AccT>(dy2.x) +
                           c2v * static_cast<AccT>(x2.x) + c3v);
    dx2.y = static_cast<T>(rstd_val * static_cast<AccT>(dy2.y) +
                           c2v * static_cast<AccT>(x2.y) + c3v);

    reinterpret_cast<double2*>(dX)[vidx] = dx2;
  }
}

// Stage 4 optimized: Scalar no-gamma dX using pre-computed rstd
template <typename T, typename AccT>
__global__ void GroupNormBackwardDxNoGammaOptCUDAKernel(
    int64_t N_G,
    int64_t D_HxW,
    const T* __restrict__ dY,
    const T* __restrict__ X,
    const AccT* __restrict__ c1_ng,
    const AccT* __restrict__ c2,
    const AccT* __restrict__ c3,
    T* __restrict__ dX) {
  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;
    const AccT rstd_val = c1_ng[ng];
    dX[idx] = static_cast<T>(rstd_val * static_cast<AccT>(dY[idx]) +
                             c2[ng] * static_cast<AccT>(X[idx]) + c3[ng]);
  }
}

// Pre-compute rstd values rounded through T for the no-gamma path.
// c1_ng[ng] = round_through_T(rsqrt(var[ng] + eps))
template <typename T, typename AccT>
__global__ void ComputeRstdCUDAKernel(int64_t N_G,
                                      AccT eps,
                                      const AccT* var,
                                      AccT* c1_ng) {
  for (int64_t index =
           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
       index < N_G;
       index += static_cast<int64_t>(gridDim.x) * blockDim.x) {
    c1_ng[index] = static_cast<AccT>(static_cast<T>(rsqrt(var[index] + eps)));
  }
}

// ============================================================================
// End of optimized dX kernels
// ============================================================================

// dgamma/dbeta variant 1: small batch (N <= 128)
// Per-channel thread, loop over batch dimension.
// Matches PyTorch's GammaBetaBackwardCUDAKernel1.
// Rounds mean/rstd through T to match PyTorch.
template <typename T, typename AccT>
__global__ void GammaBetaBackwardCUDAKernel1(int64_t N,
                                             int64_t C,
                                             int64_t G,
                                             const AccT* mean,
                                             const AccT* var,
                                             const AccT* ds,
                                             const AccT* db,
                                             AccT eps,
                                             T* dgamma,
                                             T* dbeta) {
  const int64_t c = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
  if (c < C) {
    const int64_t D = C / G;
    AccT sum1 = 0;
    AccT sum2 = 0;
    for (int64_t n = 0; n < N; ++n) {
      const int64_t nc = n * C + c;
      const int64_t ng = n * G + c / D;
      // Round through T to match PyTorch
      AccT mean_val = static_cast<AccT>(static_cast<T>(mean[ng]));
      AccT rstd_val = static_cast<AccT>(static_cast<T>(rsqrt(var[ng] + eps)));
      sum1 += (dgamma == nullptr) ? AccT(0)
                                  : ((ds[nc] - db[nc] * mean_val) * rstd_val);
      sum2 += (dbeta == nullptr) ? AccT(0) : db[nc];
    }
    if (dgamma != nullptr) {
      dgamma[c] = static_cast<T>(sum1);
    }
    if (dbeta != nullptr) {
      dbeta[c] = static_cast<T>(sum2);
    }
  }
}

// dgamma/dbeta variant 2: large batch (N > 128)
// 32x32 tile with shared memory.
// Matches PyTorch's GammaBetaBackwardCUDAKernel2.
// Rounds mean/rstd through T to match PyTorch.
template <typename T, typename AccT>
__global__ void GammaBetaBackwardCUDAKernel2(int64_t N,
                                             int64_t C,
                                             int64_t G,
                                             const AccT* mean,
                                             const AccT* var,
                                             const AccT* ds,
                                             const AccT* db,
                                             AccT eps,
                                             T* dgamma,
                                             T* dbeta) {
  __shared__ AccT g_shared[kGradReduceTileSize][kGradReduceTileSize + 1];
  __shared__ AccT b_shared[kGradReduceTileSize][kGradReduceTileSize + 1];
  const int64_t c = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
  AccT dg_sum1 = 0;
  AccT dg_sum2 = 0;
  AccT db_sum1 = 0;
  AccT db_sum2 = 0;
  if (c < C) {
    const int64_t D = C / G;
    for (int64_t n = threadIdx.y; n < N; n += blockDim.y * 2) {
      const int64_t n1 = n;
      const int64_t n2 = n + blockDim.y;
      const int64_t nc1 = n1 * C + c;
      const int64_t nc2 = n2 * C + c;
      const int64_t ng1 = n1 * G + c / D;
      const int64_t ng2 = n2 * G + c / D;
      // Round through T to match PyTorch
      AccT mean_val1 = static_cast<AccT>(static_cast<T>(mean[ng1]));
      AccT rstd1 = static_cast<AccT>(static_cast<T>(rsqrt(var[ng1] + eps)));
      dg_sum1 += dgamma == nullptr ? AccT(0)
                                   : ((ds[nc1] - db[nc1] * mean_val1) * rstd1);
      db_sum1 += dbeta == nullptr ? AccT(0) : db[nc1];
      if (n2 < N) {
        AccT mean_val2 = static_cast<AccT>(static_cast<T>(mean[ng2]));
        AccT rstd2 = static_cast<AccT>(static_cast<T>(rsqrt(var[ng2] + eps)));
        dg_sum2 += dgamma == nullptr
                       ? AccT(0)
                       : ((ds[nc2] - db[nc2] * mean_val2) * rstd2);
        db_sum2 += dbeta == nullptr ? AccT(0) : db[nc2];
      }
    }
  }

  g_shared[threadIdx.y][threadIdx.x] = dg_sum1;
  g_shared[threadIdx.y + blockDim.y][threadIdx.x] = dg_sum2;
  b_shared[threadIdx.y][threadIdx.x] = db_sum1;
  b_shared[threadIdx.y + blockDim.y][threadIdx.x] = db_sum2;
  __syncthreads();

  // Warp reduce for 1st 16 cols in tile
  AccT sum1 = g_shared[threadIdx.x][threadIdx.y];
  AccT sum2 = b_shared[threadIdx.x][threadIdx.y];
  sum1 = GradWarpReduceSum(sum1);
  sum2 = GradWarpReduceSum(sum2);
  if (threadIdx.x == 0) {
    const int64_t c_out = blockIdx.x * blockDim.x + threadIdx.y;
    if (c_out < C) {
      if (dgamma != nullptr) {
        dgamma[c_out] = static_cast<T>(sum1);
      }
      if (dbeta != nullptr) {
        dbeta[c_out] = static_cast<T>(sum2);
      }
    }
  }

  // Warp reduce for 2nd 16 cols in tile
  sum1 = g_shared[threadIdx.x][threadIdx.y + blockDim.y];
  sum2 = b_shared[threadIdx.x][threadIdx.y + blockDim.y];
  sum1 = GradWarpReduceSum(sum1);
  sum2 = GradWarpReduceSum(sum2);
  if (threadIdx.x == 0) {
    const int64_t c_out = blockIdx.x * blockDim.x + threadIdx.y + blockDim.y;
    if (c_out < C) {
      if (dgamma != nullptr) {
        dgamma[c_out] = static_cast<T>(sum1);
      }
      if (dbeta != nullptr) {
        dbeta[c_out] = static_cast<T>(sum2);
      }
    }
  }
}

// ============================================================================
// Main backward kernel
// ============================================================================

template <typename T, typename Context>
void GroupNormGradKernel(const Context& dev_ctx,
                         const DenseTensor& x,
                         const optional<DenseTensor>& scale,
                         const optional<DenseTensor>& bias,
                         const DenseTensor& y,
                         const DenseTensor& mean,
                         const DenseTensor& var,
                         const DenseTensor& d_y,
                         double epsilon,
                         int groups,
                         const std::string& data_layout_str,
                         DenseTensor* d_x,
                         DenseTensor* d_scale,
                         DenseTensor* d_bias) {
  if (x.numel() == 0) {
    dev_ctx.template Alloc<T>(d_x);
    if (d_scale) {
      if (x.dims().size() > 0 && x.dims()[0] == 0) {
        Full<T, Context>(dev_ctx, d_scale->dims(), 0, d_scale);
      } else {
        Full<T, Context>(dev_ctx, d_scale->dims(), NAN, d_scale);
      }
    }
    if (d_bias) {
      Full<T, Context>(dev_ctx, d_bias->dims(), 0, d_bias);
    }
    return;
  }
  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;

  if (d_x) {
    dev_ctx.template Alloc<T>(d_x);
  }
  funcs::SetConstant<GPUContext, T> set_zero;
  funcs::SetConstant<GPUContext, AccT> set_zero_AccT;

  auto* x_data = x.data<T>();
  T* d_x_data = nullptr;
  if (d_x) d_x_data = d_x->data<T>();
  auto* dy_data = d_y.data<T>();
  auto* var_data = var.data<AccT>();
  auto* mean_data = mean.data<AccT>();
  T* d_scale_data = nullptr;
  if (d_scale) {
    dev_ctx.template Alloc<T>(d_scale);
    d_scale_data = d_scale->data<T>();
  }
  T* d_bias_data = nullptr;
  if (d_bias) {
    dev_ctx.template Alloc<T>(d_bias);
    d_bias_data = d_bias->data<T>();
  }

  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) {
    // =========================================================
    // PyTorch-aligned NCHW backward path (optimized)
    // =========================================================
    const int64_t HxW = imsize;

    // Stage 1: Compute internal gradients ds[n,c] = sum_hw(dY*X),
    //          db[n,c] = sum_hw(dY)
    DenseTensor ds_tensor, db_tensor;
    ds_tensor.Resize({N, C});
    db_tensor.Resize({N, C});
    AccT* ds_data = dev_ctx.template Alloc<AccT>(&ds_tensor);
    AccT* db_data = dev_ctx.template Alloc<AccT>(&db_tensor);

    {
      // Use scalar kernel to match PyTorch's exact accumulation order
      // (Vec4 changes thread-to-element mapping, causing 1-2 ULP differences)
      {
        int64_t num_threads =
            HxW < kGradBlockReduceNumThreads ? 32 : kGradBlockReduceNumThreads;
        ComputeInternalGradientsCUDAKernel<T, AccT>
            <<<N * C, num_threads, 0, dev_ctx.stream()>>>(
                HxW, dy_data, x_data, ds_data, db_data);
      }
    }

    // Stage 2+3+4: Compute dX
    if (d_x_data != nullptr) {
      DenseTensor c2_tensor, c3_tensor;
      c2_tensor.Resize({N, G});
      c3_tensor.Resize({N, G});
      AccT* c2_data = dev_ctx.template Alloc<AccT>(&c2_tensor);
      AccT* c3_data = dev_ctx.template Alloc<AccT>(&c3_tensor);

      // Stage 2: Compute backward fused params c2, c3
      {
        int64_t num_threads =
            D < kGradBlockReduceNumThreads ? 32 : kGradBlockReduceNumThreads;
        ComputeBackwardFusedParamsCUDAKernel<T, AccT>
            <<<dim3(N, G), num_threads, 0, dev_ctx.stream()>>>(
                C,
                HxW,
                G,
                mean_data,
                var_data,
                scale_data,
                ds_data,
                db_data,
                static_cast<AccT>(epsilon),
                c2_data,
                c3_data);
      }

      // Stage 4: dX = c1 * dY + c2 * X + c3 (optimized)
      int64_t max_grid_x = dev_ctx.GetCUDAMaxGridDimSize()[0];
      constexpr int64_t kElemThreads = 256;

      if (scale_data != nullptr) {
        // Pre-compute c1[n,c] = rstd[n,g] * gamma[c]
        // This eliminates per-element rsqrt and gamma lookup.
        DenseTensor c1_tensor;
        c1_tensor.Resize({N, C});
        AccT* c1_data = dev_ctx.template Alloc<AccT>(&c1_tensor);

        {
          const int64_t N_C = N * C;
          const int64_t c1_blocks =
              std::min((N_C + kElemThreads - 1) / kElemThreads, max_grid_x);
          ComputeC1CUDAKernel<T, AccT>
              <<<c1_blocks, kElemThreads, 0, dev_ctx.stream()>>>(
                  N_C,
                  C,
                  G,
                  static_cast<AccT>(epsilon),
                  var_data,
                  scale_data,
                  c1_data);
        }

        // Dispatch vectorized or scalar dX kernel
        if (std::is_same<T, float>::value && (HxW % 4 == 0)) {
          const int64_t HxW_vec = HxW / 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);
          GroupNormBackwardDxVec4CUDAKernel<T, AccT>
              <<<elem_blocks, kElemThreads, 0, dev_ctx.stream()>>>(N * C,
                                                                   HxW_vec,
                                                                   D,
                                                                   dy_data,
                                                                   x_data,
                                                                   c1_data,
                                                                   c2_data,
                                                                   c3_data,
                                                                   d_x_data);
        } else if (std::is_same<T, double>::value && (HxW % 2 == 0)) {
          const int64_t HxW_vec = HxW / 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);
          GroupNormBackwardDxVec2DoubleCUDAKernel<T, AccT>
              <<<elem_blocks, kElemThreads, 0, dev_ctx.stream()>>>(N * C,
                                                                   HxW_vec,
                                                                   D,
                                                                   dy_data,
                                                                   x_data,
                                                                   c1_data,
                                                                   c2_data,
                                                                   c3_data,
                                                                   d_x_data);
        } else {
          // Scalar fallback with pre-computed c1
          const int64_t total = N * C * HxW;
          const int64_t elem_blocks =
              std::min((total + kElemThreads - 1) / kElemThreads, max_grid_x);
          GroupNormBackwardDxOptCUDAKernel<T, AccT>
              <<<elem_blocks, kElemThreads, 0, dev_ctx.stream()>>>(N * C,
                                                                   HxW,
                                                                   D,
                                                                   dy_data,
                                                                   x_data,
                                                                   c1_data,
                                                                   c2_data,
                                                                   c3_data,
                                                                   d_x_data);
        }
      } else {
        // No gamma path: c1 = rstd[ng], pre-compute rstd rounded through T
        DenseTensor c1_ng_tensor;
        c1_ng_tensor.Resize({N, G});
        AccT* c1_ng_data = dev_ctx.template Alloc<AccT>(&c1_ng_tensor);

        {
          const int64_t N_G = N * G;
          const int64_t rstd_blocks =
              std::min((N_G + kElemThreads - 1) / kElemThreads, max_grid_x);
          ComputeRstdCUDAKernel<T, AccT>
              <<<rstd_blocks, kElemThreads, 0, dev_ctx.stream()>>>(
                  N_G, static_cast<AccT>(epsilon), var_data, c1_ng_data);
        }

        const int64_t D_HxW = D * HxW;

        if (std::is_same<T, float>::value && (D_HxW % 4 == 0)) {
          const int64_t D_HxW_vec = D_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);
          GroupNormBackwardDxNoGammaVec4CUDAKernel<T, AccT>
              <<<elem_blocks, kElemThreads, 0, dev_ctx.stream()>>>(N * G,
                                                                   D_HxW_vec,
                                                                   dy_data,
                                                                   x_data,
                                                                   c1_ng_data,
                                                                   c2_data,
                                                                   c3_data,
                                                                   d_x_data);
        } else if (std::is_same<T, double>::value && (D_HxW % 2 == 0)) {
          const int64_t D_HxW_vec = D_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);
          GroupNormBackwardDxNoGammaVec2DoubleCUDAKernel<T, AccT>
              <<<elem_blocks, kElemThreads, 0, dev_ctx.stream()>>>(N * G,
                                                                   D_HxW_vec,
                                                                   dy_data,
                                                                   x_data,
                                                                   c1_ng_data,
                                                                   c2_data,
                                                                   c3_data,
                                                                   d_x_data);
        } else {
          // Scalar fallback with pre-computed rstd
          const int64_t total = N * G * D_HxW;
          const int64_t elem_blocks =
              std::min((total + kElemThreads - 1) / kElemThreads, max_grid_x);
          GroupNormBackwardDxNoGammaOptCUDAKernel<T, AccT>
              <<<elem_blocks, kElemThreads, 0, dev_ctx.stream()>>>(N * G,
                                                                   D_HxW,
                                                                   dy_data,
                                                                   x_data,
                                                                   c1_ng_data,
                                                                   c2_data,
                                                                   c3_data,
                                                                   d_x_data);
        }
      }
    }

    // Stage 5: dgamma/dbeta
    if (d_scale || d_bias) {
      if (N <= 128) {
        constexpr int kNumThreads = 256;
        const int64_t B = (C + kNumThreads - 1) / kNumThreads;
        GammaBetaBackwardCUDAKernel1<T, AccT>
            <<<B, kNumThreads, 0, dev_ctx.stream()>>>(
                N,
                C,
                G,
                mean_data,
                var_data,
                ds_data,
                db_data,
                static_cast<AccT>(epsilon),
                d_scale_data,
                d_bias_data);
      } else {
        const int64_t B = (C + kGradReduceTileSize - 1) / kGradReduceTileSize;
        constexpr int kThreadX = kGradReduceTileSize;
        constexpr int kThreadY = kGradReduceTileSize / 2;
        GammaBetaBackwardCUDAKernel2<T, AccT>
            <<<B, dim3(kThreadX, kThreadY), 0, dev_ctx.stream()>>>(
                N,
                C,
                G,
                mean_data,
                var_data,
                ds_data,
                db_data,
                static_cast<AccT>(epsilon),
                d_scale_data,
                d_bias_data);
      }
    }
  } else {
    // =========================================================
    // NHWC backward path (kept unchanged)
    // =========================================================
    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);

    auto* y_data = y.data<T>();

    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));
    }

    DenseTensor temp_var;
    temp_var.Resize(var.dims());
    dev_ctx.template Alloc<AccT>(&temp_var);
    set_zero_AccT(dev_ctx, &temp_var, static_cast<AccT>(0));
    auto* temp_var_data = temp_var.data<AccT>();

    DenseTensor temp_mean;
    temp_mean.Resize(var.dims());
    dev_ctx.template Alloc<AccT>(&temp_mean);
    set_zero_AccT(dev_ctx, &temp_mean, static_cast<AccT>(0));
    auto* temp_mean_data = temp_mean.data<AccT>();

    int flags =
        (scale_data != nullptr) * kHasScale + (bias_data != nullptr) * kHasBias;
    UNROLL_ALL_CASES(flags,
                     GroupNormBackwardGetMeanAndVar,
                     y_data,
                     scale_data,
                     bias_data,
                     dy_data,
                     N,
                     C,
                     W,
                     imsize,
                     groups,
                     group_size,
                     epsilon,
                     temp_mean_data,
                     temp_var_data,
                     d_scale_data,
                     d_bias_data);
    if (d_x_data != nullptr) {
      UNROLL_ALL_CASES(flags,
                       GroupNormBackward,
                       y_data,
                       dy_data,
                       scale_data,
                       bias_data,
                       var_data,
                       temp_mean_data,
                       temp_var_data,
                       N,
                       C,
                       W,
                       imsize,
                       groups,
                       group_size,
                       epsilon,
                       d_x_data);
    }
  }
}

}  // namespace phi

PD_REGISTER_KERNEL(group_norm_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::GroupNormGradKernel,
                   float,
                   double,
                   phi::bfloat16,
                   phi::float16) {}