                                                    long rdim, thrust::pair<float,float> init,
                                                    BinaryFunction binary_op)
 {
-  const size_t reduce = 3;
-  dim4 src_stride(0);
-  dim4 tgt1_stride(0);
-  dim4 tgt2_stride(0);
-  dim4 size(1);
-
   unsigned ndim = THCudaTensor_nDimension(state, src);
-  for(unsigned idim=0, o=ndim-2; idim < ndim; idim++) {
-    unsigned odim = idim == rdim ? reduce : o--;
-    src_stride[odim] = THCudaTensor_stride(state, src, idim);
-    tgt1_stride[odim] = THCudaTensor_stride(state, tgt1, idim);
-    tgt2_stride[odim] = THCudaTensor_stride(state, tgt2, idim);
-    size[odim]       = THCudaTensor_size(state, src, idim);
-  }
-
-  const unsigned nThreadPerBlock = 256;
-  unsigned nBlockPerColumn = DIVUP(size[0], nThreadPerBlock);
-  dim3 threads(nThreadPerBlock);
-  unsigned maxGridDim = 1024; // anything < 64k is fine. The choice has no impact on performance.
-  dim3 grid(min(maxGridDim, nBlockPerColumn), min(maxGridDim, size[1]), min(maxGridDim, size[2]));
+  unsigned num_orows = 1;
+  for (unsigned dim = 0; dim < rdim; dim++) {
+    num_orows *= THCudaTensor_size(state, src, dim);
+  }
+  unsigned row_size = THCudaTensor_size(state, src, rdim);
+  unsigned num_irows = 1;
+  for (unsigned dim = rdim + 1; dim < ndim; dim++) {
+    num_irows *= THCudaTensor_size(state, src, dim);
+  }
+
+  dim3 threads(min(512, num_irows));
+  unsigned maxGridDim = 1024;
+  dim3 grid(min(maxGridDim, num_orows), min(maxGridDim, DIVUP(num_irows, threads.x)));
 
   THCudaTensor_kernel_transformReduceOuterDimIndex<<<grid, threads>>>(
     THCudaTensor_data(state, tgt1), THCudaTensor_data(state, tgt2),
-    THCudaTensor_data(state, src), src_stride, tgt1_stride, tgt2_stride, size, init, binary_op);
+    THCudaTensor_data(state, src), num_orows, num_irows, row_size, init, binary_op);
   cudaError errcode = cudaGetLastError();
   if(errcode != cudaSuccess) {
     THError(cudaGetErrorString(errcode));


/* Reduce the innermost dimension of a tensor
 *
 * For an n-d tensor (n <= 4) where the reduction is along the innermost dimension:
 *
 * - block.x is the innermost dimension, i.e. dimension 0;
 * - block.y and grid.y make up dimension 1; and
 * - grid.x and grid z are the remaining two outer dimensions (if any)
 *
 * Reduction along other dimensions is handled in a separate kernel.
 */
template<class UnaryFunction, class BinaryFunction>
__global__ void THCudaTensor_kernel_transformReduceInnermostDim(float *tgt, float *src_,
        dim4 src_stride, dim4 tgt_stride, dim4 size, UnaryFunction unary_op, float init, BinaryFunction binary_op)
{
  __shared__ float sbuf[16][32]; // 8kB

  for(unsigned z = blockIdx.z; z < size[3] ; z += gridDim.z)
  for(unsigned x = blockIdx.x; x < size[2] ; x += gridDim.x)
  for(unsigned bRow = blockIdx.y * blockDim.y; bRow < size[1]; bRow += blockDim.y * gridDim.y) {

    float acc = init;
    unsigned row = bRow + threadIdx.y;
    float *src = src_ + z * src_stride[3] + x * src_stride[2] + row * src_stride[1];
    bool reducing = threadIdx.x < blockDim.y && bRow + threadIdx.x < size[1] && threadIdx.y == 0;

    for(unsigned bCol=0; bCol < size[0]; bCol += blockDim.x) {

      sbuf[threadIdx.y][threadIdx.x] = init;
      unsigned col = bCol + threadIdx.x;
      if(row < size[1] && col < size[0]) {
        sbuf[threadIdx.y][threadIdx.x] = unary_op(src[col]);
      }
      __syncthreads();

      float* line = &sbuf[threadIdx.y][0];
      for(unsigned s = 16; s > 1; s >>= 1) {
        if(row < size[1] && threadIdx.x < s) {
          line[threadIdx.x] = binary_op(line[threadIdx.x], line[threadIdx.x + s]);
        }
        __syncthreads();
      }
      if(reducing) {
        sbuf[threadIdx.x][0] = binary_op(sbuf[threadIdx.x][0], sbuf[threadIdx.x][1]);
        acc = binary_op(acc, sbuf[threadIdx.x][0]);
      }
      __syncthreads();
    }

    if(reducing) {
      unsigned row = bRow + threadIdx.x;
      unsigned tgt_offset = z * tgt_stride[3] + x * tgt_stride[2];
      tgt[tgt_offset + row] = acc;
    }
  }
}


        long dimension, UnaryFunction unary_op, float init, BinaryFunction binary_op)
{
  THArgCheck(dimension >= 0 && dimension < THCudaTensor_nDimension(state, src), 3, "dimension out of range");

  THLongStorage *dim = THCudaTensor_newSizeOf(state, src);
  THLongStorage_set(dim, dimension, 1);
