 
 using DataLayout = framework::DataLayout;
 enum GroupNormKernelFlags { kHasScale = 1, kHasBias = 2 };
+#define ALIGN_BYTES 16
 
 #define CHECK_CASE(i, flags, kernel_name, ...)                              \
   if (i == flags) {                                                         \

        imsize *= x_dims[i];
      }
    }
#ifdef __HIPCC__
    int block_size = std::max(std::min(256, imsize), 64);
#else
    int block_size = std::min(1024, imsize);
#endif
    dim3 grid(group_size, groups, x_dims[0]);
    dim3 threads(block_size, 1, 1);
    GroupNormForwardGetMeanAndVar<T><<<grid, threads, 0, dev_ctx.stream()>>>(
        x_data, x_dims[0], C, W, imsize, groups, group_size, mean_data,
        temp_var_data, data_layout);
    int flags =
        (scale_data != nullptr) * kHasScale + (bias_data != nullptr) * kHasBias;
    UNROLL_ALL_CASES(flags, GroupNormForward, x_data, mean_data, temp_var_data,

  x_var /= number * imsize;
  CudaAtomicAddWithWarp(&mean[bid * groups + gid], x_mean);
  CudaAtomicAddWithWarp(&var[bid * groups + gid], x_var);
}

template <typename T, typename AccT, int VecSize>
__device__ __forceinline__ void ThreadReduce(const T* input, int size,
                                             const int offset, AccT* mean,
                                             AccT* var) {
  using VecT = kps::details::VectorType<T, VecSize>;
  int tid = threadIdx.x;
  if (offset > 0) {
    input -= offset;
    size += offset;
    if (tid >= offset) {
      AccT temp = input[tid];
      *mean += temp;
      *var += temp * temp;
    }
    size -= blockDim.x;
    input += blockDim.x;
  }
  int remain = size % (VecSize * blockDim.x);

  T ins[VecSize];
  VecT* ins_vec = reinterpret_cast<VecT*>(&ins);

  // vector part
  for (; VecSize * tid < (size - remain); tid += blockDim.x) {
    *ins_vec = reinterpret_cast<const VecT*>(input)[tid];

#pragma unroll
    for (int i = 0; i < VecSize; ++i) {
      AccT temp = ins[i];
      *mean += temp;
      *var += temp * temp;
    }
  }

  // scalar part
  tid = size - remain + threadIdx.x;
  for (; tid < size; tid += blockDim.x) {
    AccT temp = input[tid];
    *mean += temp;
    *var += temp * temp;
  }
}

template <typename T>
__global__ void ScalarGetMeanAndVarNCHW(const T* x, T* mean, T* var, int size) {
  int i = blockIdx.x;
  T x_mean = 0, x_var = 0;
  for (int j = threadIdx.x; j < size; j += blockDim.x) {
    T val;
    val = x[i * size + j];
    x_mean += val;
    x_var += val * val;
  }
  x_mean /= size;
  x_var /= size;
  CudaAtomicAddWithWarp(&mean[i], x_mean);
  CudaAtomicAddWithWarp(&var[i], x_var);
}

template <typename T, typename AccT, int VecSize>
__global__ void VectorizedGetMeanAndVarNCHW(const T* x, T* mean, T* var,
                                            int size) {
  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;
  ThreadReduce<T, AccT, VecSize>(x, 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);
  }
}

template <typename T, int flags>
