 __global__ void ElementwiseLinearGradientKernel(const int N, const int D,
   const float* g_o_data, const float* X_data, const float* a_data,
   float* g_X_data, float* g_a_data, float* g_b_data) {
-  CUDA_1D_KERNEL_LOOP(d, D) {
-    for (int n = 0; n < N; ++n) {
-      g_X_data[n * D + d] = g_o_data[n * D + d] * a_data[d];
-      g_a_data[d] += g_o_data[n * D + d] * X_data[n * D + d];
-      g_b_data[d] += g_o_data[n * D + d];
-    }
+  int d = blockIdx.x; // One block per D
+
+  float g_a_sum = 0;
+  float g_b_sum = 0;
+  for (int n = threadIdx.x; n < N; n += blockDim.x) {
+    const float gox = g_o_data[n * D + d];
+    g_X_data[n * D + d] = gox * a_data[d];
+    g_a_sum += gox * X_data[n * D + d];
+    g_b_sum += gox;
+  }
+
+  typedef cub::BlockReduce<float, CAFFE_CUDA_NUM_THREADS> BlockReduce;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
+
+  float g_a_sum_tot = BlockReduce(temp_storage).Sum(g_a_sum);
+  float g_b_sum_tot = BlockReduce(temp_storage).Sum(g_b_sum);
+
+  if (threadIdx.x == 0) {
+    g_a_data[d] = g_a_sum_tot;
+    g_b_data[d] = g_b_sum_tot;
   }
 }
 


  float* g_a_data = g_a->mutable_data<float>();
  float* g_b_data = g_b->mutable_data<float>();
  math::Set<float, CUDAContext>(g_a->size(), 0.f, g_a_data, &context_);
  math::Set<float, CUDAContext>(g_b->size(), 0.f, g_b_data, &context_);

  ElementwiseLinearGradientKernel<<<CAFFE_GET_BLOCKS(D), CAFFE_CUDA_NUM_THREADS,
                                  0, context_.cuda_stream()>>>(
    N, D, g_o.data<float>(), X.data<float>(), a.data<float>(),
    g_X->mutable_data<float>(), g_a_data, g_b_data);
  return true;
}


#include "caffe2/core/context_gpu.h"
#include "caffe2/operators/operator_fallback_gpu.h"

#include <cub/block/block_reduce.cuh>

namespace caffe2 {

