     // within a warp of threads.
     // After calling the function, threadIdx.x == 0 will store the result of
     // the aggregated (mean, sigma2, counts).
+#pragma unroll
     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);
  cudaStream_t stream = Stream<gpu>::GetStream(ctx.get_stream<gpu>());

  // Calculate the gradient for gamma/beta
  CHECK_EQ(gamma_grad.CheckContiguous(), true);
  CHECK_EQ(beta_grad.CheckContiguous(), true);
  const dim3 dimBlock(32, 4);
  const int LOAD_UNROLL = 4;
  if(gamma_grad_req != kNullOp || beta_grad_req != kNullOp) {
    MSHADOW_REAL_TYPE_SWITCH(in_data.type_flag_, 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;
      if(gamma_grad_req == kAddTo && beta_grad_req != kAddTo) {
        LayerNormFusedBackwardKernel_GammaBeta<LOAD_UNROLL, true, false>
          <<<1, dimBlock, 0, stream>>>
          (nbatch, nchannel, in_data.dptr<DType>(), out_grad.dptr<DType>(), mean_data.dptr<DType>(),
           std_data.dptr<DType>(), gamma_grad_ptr, beta_grad_ptr);
      } else if(gamma_grad_req != kAddTo && beta_grad_req == kAddTo) {
        LayerNormFusedBackwardKernel_GammaBeta<LOAD_UNROLL, false, true>
          <<<1, dimBlock, 0, stream>>>
          (nbatch, nchannel, in_data.dptr<DType>(), out_grad.dptr<DType>(), mean_data.dptr<DType>(),
           std_data.dptr<DType>(), gamma_grad_ptr, beta_grad_ptr);
      } else if(gamma_grad_req == kAddTo && beta_grad_req == kAddTo) {
        LayerNormFusedBackwardKernel_GammaBeta<LOAD_UNROLL, true, true>
          <<<1, dimBlock, 0, stream>>>
          (nbatch, nchannel, in_data.dptr<DType>(), out_grad.dptr<DType>(), mean_data.dptr<DType>(),
           std_data.dptr<DType>(), gamma_grad_ptr, beta_grad_ptr);
      } else {
        LayerNormFusedBackwardKernel_GammaBeta<LOAD_UNROLL, false, false>
          <<<1, dimBlock, 0, stream>>>
          (nbatch, nchannel, in_data.dptr<DType>(), out_grad.dptr<DType>(), mean_data.dptr<DType>(),
            std_data.dptr<DType>(), gamma_grad_ptr, beta_grad_ptr);
      }
    });
  }

  // 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 dimGrid(ngrid_x, ngrid_y);
  int nthread_y = 0;
  if(nchannel <= 32) {
    nthread_y = 1;
  } else {
    nthread_y = 2;
  }
  const dim3 dimBlockData(32, nthread_y);
  if(data_grad_req != kNullOp) {
    MSHADOW_REAL_TYPE_SWITCH(in_data.type_flag_, DType, {
      int nshared = dimBlock.y > 1 ? dimBlock.y * dimBlock.x * sizeof(DType) : 0;
      CheckLaunchParam(dimGrid, dimBlockData);
      if(data_grad_req == kAddTo) {
        LayerNormFusedBackwardKernel_Data<LOAD_UNROLL, true>
          <<<dimGrid, dimBlockData, 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>
          <<<dimGrid, dimBlockData, 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>());
      }
    });
  }
}


  }
}

/* Fused CUDA kernel for the forward pass of layer normalization.
 * It computes the LayerNorm when axis=-1, i.e., contiguous reduction scenario.
 * Shape of the input tensors:
 *      in_data = (nbatch, nchannel)
 *      gamma = (nchannel,)
