   }
 }
 
+/* Use the Chan's Parallel Algorithm to merge all (mean, sigma2, counts) within a warp of threads.
+ * After calling the function, threadIdx.x == 0 will store the result of the
+ * aggregated (mean, sigma2, counts).
+ *
+ *
+ */
+template<typename DType>
+__device__ __forceinline__ void warp_merge_mean_sigma2(DType mean, DType sigma2, DType count) {
+  for (int l = 0; l <= 4; ++l) {
+    int src_lane = (threadIdx.x + (1<<l)) & 31;
+    DType meanB = WARP_SHFL(mean, src_lane);
+    DType sigma2B = WARP_SHFL(sigma2, src_lane);
+    DType countB = WARP_SHFL(count, src_lane);
+    chan_merge_partition(meanB, sigma2B, countB, mean, sigma2, count);
+  }
+}
+
 /* Fused CUDA kernel for layer normalization. It computes the LayerNorm when axis=-1.
  * Shape of the input tensors:
  *      in_data = (nbatch, nchannel)

 *  TODO(sxjscience) Explore the possibility of int lhs_count and rhs_count
 */
template<typename DType>
__device__ void chan_merge_partition(const DType lhs_mean,
                                     const DType lhs_sigma2,
                                     const DType lhs_count,
                                     DType& rhs_mean,

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