   }
 }
 
+/**
+ * Calculate RowwiseSparseAdagrad
+ * M: gradients.dims[0]
+ * N: gradients.size_from_dim(1)
+ * grad: pointer to the gradients
+ * param: pointer to weights
+ * param_mom: pointer to the momentum
+ * indices: keys
+ */
+template <typename SIndex>
+__global__ void RowWiseSparseAdagradKernel(
+    const int M,
+    const int N,
+    const float epsilon,
+    float* param,
+    float* param_mom,
+    const SIndex* indices,
+    const float* grad,
+    const float* lr) {
+  typedef cub::BlockReduce<float, CAFFE_CUDA_NUM_THREADS> BlockReduce;
+  __shared__ BlockReduce::TempStorage temp_storage;
+  // in case gridDim is smaller than M
+  for (int i = blockIdx.x; i < M; i += gridDim.x) {
+    const SIndex index = indices[i];
+    float sum_squares = 0.0;
+    __shared__ float row_sum_squares_avg;
+
+    // in case N is bigger than block size which is 512 by default
+    for (int j = threadIdx.x; j < N; j += blockDim.x) {
+      const float x_ij = grad[i * N + j];
+      sum_squares += x_ij * x_ij;
+    }
+    float reduce_result = BlockReduce(temp_storage).Sum(sum_squares);
+    if (threadIdx.x == 0) {
+      row_sum_squares_avg = reduce_result / (float)N;
+      param_mom[index] += row_sum_squares_avg;
+    }
+    __syncthreads();
+    // update param
+    float step = lr[0] / (std::sqrt(param_mom[index]) + epsilon);
+    for (int j = threadIdx.x; j < N; j += blockDim.x) {
+      param[index * N + j] = param[index * N + j] + grad[i * N + j] * step;
+    }
+  }
+}
+
 template<>
 template<typename SIndex>
 bool SparseAdagradOp<float, CUDAContext>::DoRunWithType()

  return true;
}

REGISTER_CUDA_OPERATOR(Adagrad, AdagradOp<float, CUDAContext>);
REGISTER_CUDA_OPERATOR(SparseAdagrad, SparseAdagradOp<float, CUDAContext>);
}

 * limitations under the License.
 */

#include <cub/block/block_reduce.cuh>
#include "adagrad_op.h"
#include "caffe2/core/common_gpu.h"
#include "caffe2/core/context_gpu.h"
namespace caffe2 {

__global__ void AdagradUpdate(
