 
 
 template<typename AType, typename DType, typename IType>
-__device__ __forceinline__ void BlockWelfordOnlineSum(const int tid,
-                                                      const int nthread,
-                                                      const DType* __restrict__ col_vals,
+__device__ __forceinline__ void BlockWelfordOnlineSum(const DType* __restrict__ col_vals,
                                                       const int nchannel,
                                                       AType& mean,
                                                       AType& sigma2,
                                                       IType& count) {
+  int tid = threadIdx.x + threadIdx.y * blockDim.x;
+  const int nthread = blockDim.x * blockDim.y;
   // Each thread takes charge of 4 consecutive numbers. This should optimize the loading speed using
   // vectorized types like float4.
   // Also, to minimize branch divergence, we split the for-loop into two parts.

    }
    __syncthreads();
  }
  for(int offset = blockDim.y/2;  offset > 0;  offset /= 2) {
    if(threadIdx.y < offset) {
      int idx1 = threadIdx.y * blockDim.x + threadIdx.x;
      int idx2 = (threadIdx.y + offset) * blockDim.x + threadIdx.x;

    __syncthreads();
    // Begin for inter-warp reduce
    if(npart > 1) {
      for(int offset = blockDim.y/2; offset > 0; offset >>= 1) {
        if(threadIdx.y < offset) {
          int idx1 = tid;
          int idx2 = tid + offset * blockDim.x;
