   }
 }
 
+void GetGammaBetaGradKernelParams(const int nbatch, const int nchannel,
+                                  dim3* part_grad_block_dim, dim3* part_grad_grid_dim,
+                                  dim3* gb_block_dim, dim3* gb_grid_dim, int* row_repeat) {
+  const int npart = 16;
+  *part_grad_block_dim = dim3(32, 16);
+  *part_grad_grid_dim = dim3((nchannel + 32 - 1) / 32, npart);
+  *row_repeat = 16;
+  *gb_block_dim = dim3(32, npart);
+  *gb_grid_dim = dim3((nchannel + 32 - 1) / 32);
+  CheckLaunchParam(*part_grad_grid_dim, *part_grad_block_dim);
+  CheckLaunchParam(*gb_grid_dim, *gb_block_dim);
+}
+
 void LayerNormGradGPUContig(const LayerNormParam param,
                             const OpContext& ctx, const std::vector<TBlob>& inputs,
                             const std::vector<OpReqType>& req,

    // within a warp of threads.
    // After calling the function, threadIdx.x == 0 will store the result of
    // the aggregated (mean, sigma2, counts).
    for (int l = 0; l <= 4; ++l) {
      int src_lane = (threadIdx.x + (1<<l)) & 31;
      DType meanB = WARP_SHFL(mean, src_lane);

  CHECK_NE(data_grad_req, kWriteInplace);
  CHECK_NE(gamma_grad_req, kWriteInplace);
  CHECK_NE(beta_grad_req, kWriteInplace);
  Stream<gpu> *s = ctx.get_stream<gpu>();
  cudaStream_t stream = Stream<gpu>::GetStream(s);

  // Calculate the gradient for gamma/beta
  CHECK_EQ(gamma_grad.CheckContiguous(), true);
  CHECK_EQ(beta_grad.CheckContiguous(), true);
  dim3 part_grad_block_dim, part_grad_grid_dim, gb_block_dim, gb_grid_dim;
  int row_repeat;
  GetGammaBetaGradKernelParams(nbatch, nchannel, &part_grad_block_dim, &part_grad_grid_dim,
                               &gb_block_dim, &gb_grid_dim);
  if(gamma_grad_req != kNullOp || beta_grad_req != kNullOp) {
    MSHADOW_REAL_TYPE_SWITCH(in_data.type_flag_, DType, {
      Tensor<gpu, 1, DType> workspace =
        ctx.requested[0].get_space_typed<gpu, 1, DType>(Shape1(2 * npart * nchannel), s);
      DType* part_gamma_grad_ptr = workspace.dptr_;
      DType* part_beta_grad_ptr = workspace.dptr_ + npart * nchannel;
      const int nshared_K1 = (2 * part_grad_block_dim.y * row_repeat
                              + 2 * part_grad_block_dim.x * part_grad_block_dim.y) * sizeof(DType);
      const int nshared_K2 = 2 * part_grad_block_dim.x * part_grad_block_dim.y * sizeof(DType);
      DType* gamma_grad_ptr = (gamma_grad_req != kNullOp) ? gamma_grad.dptr<DType>() : nullptr;
      DType* beta_grad_ptr = (beta_grad_req != kNullOp) ? beta_grad.dptr<DType>() : nullptr;
      LayerNormFusedBackwardKernel_PartGammaBeta
        <<<part_grad_grid_dim, part_grid_block_dim, nshared_K1, stream>>>
        (nbatch, nchannel, row_repeat, in_data.dptr<DType>(), out_grad.dptr<DType>(),
         mean_data.dptr<DType>(), std_data.dptr<DType>(), part_gamma_grad_ptr, part_beta_grad_ptr);
      MSHADOW_CUDA_POST_KERNEL_CHECK(LayerNormFusedBackwardKernel_PartGammaBeta);
      if(gamma_grad_req == kAddTo && beta_grad_req != kAddTo) {
        LayerNormFusedBackwardKernel_GammaBeta<true, false>
          <<<gb_grid_dim, gb_block_dim, nshared_K2, stream>>>
          (nbatch, nchannel, npart, part_gamma_grad_ptr, part_beta_grad_ptr,
           gamma_grad_ptr, beta_grad_ptr);
      } else if(gamma_grad_req != kAddTo && beta_grad_req == kAddTo) {
        LayerNormFusedBackwardKernel_GammaBeta<false, true>
          <<<gb_grid_dim, gb_block_dim, nshared_K2, stream>>>
          (nbatch, nchannel, npart, part_gamma_grad_ptr, part_beta_grad_ptr,
            gamma_grad_ptr, beta_grad_ptr);
      } else if(gamma_grad_req == kAddTo && beta_grad_req == kAddTo) {
        LayerNormFusedBackwardKernel_GammaBeta<true, true>
          <<<gb_grid_dim, gb_block_dim, nshared_K2, stream>>>
          (nbatch, nchannel, npart, part_gamma_grad_ptr, part_beta_grad_ptr,
            gamma_grad_ptr, beta_grad_ptr);
      } else {
        LayerNormFusedBackwardKernel_GammaBeta<false, false>
          <<<gb_grid_dim, gb_block_dim, nshared_K2, stream>>>
          (nbatch, nchannel, npart, part_gamma_grad_ptr, part_beta_grad_ptr,
            gamma_grad_ptr, beta_grad_ptr);
      }
    });
    MSHADOW_CUDA_POST_KERNEL_CHECK(LayerNormFusedBackwardKernel_GammaBeta);
  }

  // Calculate the gradient for data
  CHECK_EQ(data_grad.CheckContiguous(), true);
  int ngrid_x = (nbatch > kMaxGridDim) ? (nbatch + kBaseGridNum - 1) / kBaseGridNum : nbatch;
  int ngrid_y = (nbatch > kMaxGridDim) ? kBaseGridNum : 1;
  const dim3 data_grid_dim(ngrid_x, ngrid_y);
  int nthread_y = 0;
  if(nchannel <= 32) {
    nthread_y = 1;
  } else {
    nthread_y = 2;
  }
  const dim3 data_block_dim(32, nthread_y);
  const int LOAD_UNROLL = 4;
  if(data_grad_req != kNullOp) {
    MSHADOW_REAL_TYPE_SWITCH(in_data.type_flag_, DType, {
      int nshared = data_block_dim.y > 1 ? data_block_dim.y * data_block_dim.x * sizeof(DType) : 0;
      CheckLaunchParam(data_grid_dim, data_block_dim);
      if(data_grad_req == kAddTo) {
        LayerNormFusedBackwardKernel_Data<LOAD_UNROLL, true>
          <<<data_grid_dim, data_block_dim, nshared, stream>>>
          (nbatch, nchannel, in_data.dptr<DType>(), out_grad.dptr<DType>(), mean_data.dptr<DType>(),
           std_data.dptr<DType>(), gamma.dptr<DType>(), data_grad.dptr<DType>());
      } else {
        LayerNormFusedBackwardKernel_Data<LOAD_UNROLL, false>
          <<<data_grid_dim, data_block_dim, nshared, stream>>>
          (nbatch, nchannel, in_data.dptr<DType>(), out_grad.dptr<DType>(), mean_data.dptr<DType>(),
           std_data.dptr<DType>(), gamma.dptr<DType>(), data_grad.dptr<DType>());
      }
    });
    MSHADOW_CUDA_POST_KERNEL_CHECK(LayerNormFusedBackwardKernel_Data);
  }
}

