     const float* labeldata,
     const float* weights,
     float* Ydata) {
-  CUDA_1D_KERNEL_LOOP(i, N) {
+  typedef cub::BlockReduce<float, CAFFE_CUDA_NUM_THREADS> BlockReduce;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
+
+  for (int i = blockIdx.x; i < N; i += gridDim.x) {
     float weight = weights ? weights[i] : 1.0;
+    float sum = 0.0;
     float total_prob = 0.0;
-    Ydata[i] = 0.0;
-    for (int j = 0; j < D; j++) {
+    for (int j = threadIdx.x; j < D; j += blockDim.x) {
       int idx = i * D + j;
+      CUDA_KERNEL_ASSERT(labeldata[idx] >= 0);
       total_prob += labeldata[idx];
-      CUDA_KERNEL_ASSERT(labeldata[idx] >= 0);
-      Ydata[i] += -logf(max(Pdata[idx], FLT_MIN)) * labeldata[idx] * weight;
-    }
-    CUDA_KERNEL_ASSERT(abs(total_prob - 1.0) < 1e-5f);
+      sum += -logf(max(Pdata[idx], FLT_MIN)) * labeldata[idx] * weight;
+    }
+    float tot = BlockReduce(temp_storage).Sum(sum);
+    __syncthreads();
+    float total_prob_sum = BlockReduce(temp_storage).Sum(total_prob);
+    if (threadIdx.x == 0) {
+      Ydata[i] = tot;
+      // Sanity check
+      CUDA_KERNEL_ASSERT(abs(1.0 - total_prob_sum) < 1e-5f);
+    }
+    __syncthreads();
   }
 }
 

      }
    } else {
      ProbCrossEntropyGradientKernel<<<
          CAFFE_GET_BLOCKS(N),
          CAFFE_CUDA_NUM_THREADS,
          0,
          context_.cuda_stream()>>>(

    float* dXdata,
    const float* weights) {
  if (weights == NULL) {
    CUDA_1D_KERNEL_LOOP(idx, N * D) {
      dXdata[idx] = Pdata[idx] - labeldata[idx];
    }
  } else {
    CUDA_1D_KERNEL_LOOP(idx, N * D) {
      dXdata[idx] = (Pdata[idx] - labeldata[idx]) * weights[idx / D];
    }
  }
}
