 #include "THCUNN.h"
 #include "common.h"
 
-#define LOGSOFTMAX_THREADS 128
-
-__global__ void cunn_SpatialLogSoftMax_updateOutput_kernel(float *output, float *input, int nframe, int dim, int stride)
-{
-  __shared__ float buffer[LOGSOFTMAX_THREADS+1];
-  int k = blockIdx.x;
-  float *input_k = input + k*dim*stride + blockIdx.y;
-  float *output_k = output + k*dim*stride + blockIdx.y;
-  int tx = threadIdx.x;
-
-  int i_start = threadIdx.x;
-  int i_end = dim;
-  int i_step = blockDim.x;
-
-  // max?
-  buffer[tx] = -FLT_MAX;
-  for (int i=i_start; i<i_end; i+=i_step)
-  {
-    float z = input_k[i*stride];
-    if(buffer[tx] < z)
-      buffer[tx] = z;
-  }
-
-  // reduce
-  for (unsigned int stride_ = blockDim.x >> 1; stride_ > 0; stride_ >>= 1)
-  {
-    __syncthreads();
-    if ((tx < stride_) && (buffer[tx] < buffer[tx+stride_]))
-      buffer[tx] = buffer[tx+stride_];
-  }
-  if (tx == 0)
-  {
-    float max_k = -FLT_MAX;
-    if(max_k < buffer[0])
-      max_k = buffer[0];
-    buffer[LOGSOFTMAX_THREADS] = max_k;
-  }
-
-  __syncthreads();
-
-  // logadd?
-  float max_k = buffer[LOGSOFTMAX_THREADS];
-  buffer[tx] = 0;
-  for (int i=i_start; i<i_end; i+=i_step)
-    buffer[tx] += expf(input_k[i*stride]-max_k);
-
-  // reduce
-  for (unsigned int stride_ = blockDim.x >> 1; stride_ > 0; stride_ >>= 1)
-  {
-    __syncthreads();
-    if (tx < stride_)
-      buffer[tx] += buffer[tx+stride_];
-  }
-  if (tx == 0)
-    buffer[LOGSOFTMAX_THREADS] = max_k + logf(buffer[0]);
-
-  __syncthreads();
-
-  // logsoftmax
-  float logsum_k = buffer[LOGSOFTMAX_THREADS];
-  for (int i=i_start; i<i_end; i+=i_step)
-    output_k[i*stride] = input_k[i*stride] - logsum_k;
-}
-
-
-__global__ void cunn_SpatialLogSoftMax_updateGradInput_kernel(float *gradInput, float *output, float *gradOutput, int nframe, int dim, int stride)
-{
-  __shared__ float buffer[LOGSOFTMAX_THREADS];
-  int k = blockIdx.x;
-  float *gradInput_k = gradInput + k*dim*stride + blockIdx.y;
-  float *output_k = output + k*dim*stride + blockIdx.y;
-  float *gradOutput_k = gradOutput + k*dim*stride + blockIdx.y;
-  int tx = threadIdx.x;
-
-  int i_end = dim;
-  int i_step = blockDim.x;
-
-  // sum?
-  buffer[tx] = 0;
-  for (int i=tx; i<i_end; i+=i_step)
-    buffer[tx] += gradOutput_k[i*stride];
-
-  // reduce
-  for (unsigned int stride_ = blockDim.x >> 1; stride_ > 0; stride_ >>= 1)
-  {
-    __syncthreads();
-    if (tx < stride_)
-      buffer[tx] += buffer[tx+stride_];
-  }
-
-  __syncthreads();
-
-  float sum_k = buffer[0];
-  for (int i=tx; i<i_end; i+=i_step)
-    gradInput_k[i*stride] = gradOutput_k[i*stride] - __expf(output_k[i*stride])*sum_k;
-}
-
-// here starts the 1D/2D implementation
+__global__ void cunn_SpatialLogSoftMax_updateOutput_kernel(float *output, float *input, int classSize, int height, int width)
+{
+  int batchIndex = blockIdx.x;
+  int index = threadIdx.x;
+
+  while (index < height*width) {
+    int y = index / width;
+    int x = index % width;
+    if (y >= height)
+      break;
+
+    // calculate input starting index in cuda layout (B x H x W x P)
+    int inputStartIndex = 0
+    + (height*width*classSize)*batchIndex
+    + (width*classSize)*y
+    + (classSize)*x;
+
+    // compute a
+    float a = 0;
+    for (int i = 0; i < classSize; i++) {
+      a += __expf(input[inputStartIndex + i]);
+    }
+
+    // invert a
+    a = 1.0f / a;
+
+    for (int i = 0; i < classSize; i++) {
+      // calculate output index in torch layout (B x P x H x W)
+      int outputIndex = 0
+      + (classSize*height*width)*batchIndex
+      + (height*width)*i
+      + (width)*y
+      + x;
+      output[outputIndex] = logf(a * __expf(input[inputStartIndex + i]));
+    }
+    index += 1024;
+  }
+}
+
+__global__ void cunn_SpatialLogSoftMax_updateGradInput_kernel(float *gradInput, float *output, float *gradOutput, int classSize, int height, int width)
+{
+  int batchIndex = blockIdx.x;
+  int index = threadIdx.x;
+
+  while (index < height*width) {
+    int y = index / width;
+    int x = index % width;
+    if (y >= height)
+      break;
+
+    // calculate output starting index in cuda layout (B x H x W x P)
+    int outputStartIndex = 0
+    + (height*width*classSize)*batchIndex
+    + (width*classSize)*y
+    + (classSize)*x;
+
+    // compute sum
+    float a = 0;
+    for (int i = 0; i < classSize; i++) {
+      a += gradOutput[outputStartIndex + i];
+    }
+
+    for (int i = 0; i < classSize; i++) {
+      // calculate input index in torch layout (B x P x H x W)
+      int inputIndex = 0
+      + (classSize*height*width)*batchIndex
+      + (height*width)*i
+      + (width)*y
+      + x;
+      gradInput[inputIndex] = gradOutput[outputStartIndex + i] - __expf(output[outputStartIndex + i]) * a;
+    }
+    index += 1024;
+  }
+}
 
 struct MaxFloat
 {

  THCudaTensor_free(state, gradOutput);
  THCudaTensor_free(state, output);
}

#undef LOGSOFTMAX_THREADS

{
  THCUNN_assertSameGPU(state, 3, output, gradOutput, gradInput);

  THCudaTensor_resizeAs(state, gradInput, output);

  bool spatial  = false;
  int batchSize = 1;
  int classSize = 0;
  int height = 0;
  int width = 0;

  int ndims = THCudaTensor_nDimension(state, input);

  if (ndims == 1)
  {
    classSize = THCudaTensor_size(state, gradInput, 0);
    output = THCudaTensor_newContiguous(state, output);
    gradOutput = THCudaTensor_newContiguous(state, gradOutput);
  }
  else if (ndims == 2)
  {
    batchSize = THCudaTensor_size(state, gradInput, 0);
    classSize = THCudaTensor_size(state, gradInput, 1);
    output = THCudaTensor_newContiguous(state, output);
    gradOutput = THCudaTensor_newContiguous(state, gradOutput);
  }
  else if (ndims == 3)
  {
    spatial = true;
    classSize = THCudaTensor_size(state, input, 0);
    height = THCudaTensor_size(state, input, 1);
    width = THCudaTensor_size(state, input, 2);

    THCudaTensor *transposedOutput = THCudaTensor_new(state);
    THCudaTensor_resizeAs(state, transposedOutput, output);
    THCudaTensor_copy(state, transposedOutput, output);
    output = transposedOutput;
    // convert tensor from torch layout to cuda layout
    // P x H x W -> W x H x P
    THCudaTensor_transpose(state, output, output, 0, 2);
    // W x H x P -> H x W x P
    THCudaTensor_transpose(state, output, output, 0, 1);
    output = THCudaTensor_newContiguous(state, output);

    THCudaTensor *transposedGradOutput = THCudaTensor_new(state);
    THCudaTensor_resizeAs(state, transposedGradOutput, gradOutput);
    THCudaTensor_copy(state, transposedGradOutput, gradOutput);
    gradOutput = transposedGradOutput;
    // convert tensor from torch layout to cuda layout
    // P x H x W -> W x H x P
    THCudaTensor_transpose(state, gradOutput, gradOutput, 0, 2);
    // W x H x P -> H x W x P
    THCudaTensor_transpose(state, gradOutput, gradOutput, 0, 1);
    gradOutput = THCudaTensor_newContiguous(state, gradOutput);
  }
  else if (ndims == 4)
  {
    spatial = true;
    batchSize = THCudaTensor_size(state, gradInput, 0);
    classSize = THCudaTensor_size(state, input, 1);
    height = THCudaTensor_size(state, input, 2);
    width = THCudaTensor_size(state, input, 3);

    THCudaTensor *transposedOutput = THCudaTensor_new(state);
    THCudaTensor_resizeAs(state, transposedOutput, output);
    THCudaTensor_copy(state, transposedOutput, output);
    output = transposedOutput;
    // convert tensor from torch layout to cuda layout
    // B x P x H x W -> B x W x H x P
    THCudaTensor_transpose(state, output, output, 1, 3);
    // B x W x H x P -> B x H x W x P
    THCudaTensor_transpose(state, output, output, 1, 2);
    output = THCudaTensor_newContiguous(state, output);

    THCudaTensor *transposedGradOutput = THCudaTensor_new(state);
    THCudaTensor_resizeAs(state, transposedGradOutput, gradOutput);
    THCudaTensor_copy(state, transposedGradOutput, gradOutput);
    gradOutput = transposedGradOutput;
    // convert tensor from torch layout to cuda layout
    // B x P x H x W -> B x W x H x P
    THCudaTensor_transpose(state, gradOutput, gradOutput, 1, 3);
    // B x W x H x P -> B x H x W x P
    THCudaTensor_transpose(state, gradOutput, gradOutput, 1, 2);
    gradOutput = THCudaTensor_newContiguous(state, gradOutput);
  }
  else
  {
    THError("1D, 2D, 3D or 4D Tensor expected");
  }

  if (!spatial)
  {
    dim3 grid(batchSize);
    dim3 block(1024);
