fix register spill for thread_reduce

Author: coderfeli

aten/src/ATen/native/cuda/Reduce.cuh

@@ -569,14 +569,11 @@ struct ReduceOp {
     using load_t = at::native::memory::aligned_vector<scalar_t, output_vec_size>;
 
     // Multiple accumulators to remove dependency between unrolled loops.
-    arg_vec_t value_list[vt0];
+    arg_vec_t value_list;
 
     #pragma unroll
-    for (int i = 0; i < vt0; i++) {
-      #pragma unroll
-      for (int j = 0; j < output_vec_size; j++) {
-        value_list[i][j] = ident;
-      }
+    for (int j = 0; j < output_vec_size; j++) {
+      value_list[j] = ident;
     }
 
     load_t values[vt0];
@@ -591,7 +588,7 @@ struct ReduceOp {
       for (index_t i = 0; i < vt0; i++) {
         #pragma unroll
         for (index_t j = 0; j < output_vec_size; j++) {
-          value_list[i][j] = ops.reduce(value_list[i][j], values[i].val[j], idx + i * stride);
+          value_list[j] = ops.reduce(value_list[j], values[i].val[j], idx + i * stride);
         }
       }
       idx += stride * vt0;
@@ -616,20 +613,12 @@ struct ReduceOp {
       }
       #pragma unroll
       for (index_t j = 0; j < output_vec_size; j++) {
-        value_list[i][j] = ops.reduce(value_list[i][j], values[i].val[j], idx);
+        value_list[j] = ops.reduce(value_list[j], values[i].val[j], idx);
       }
       idx += stride;
     }
 
-    // combine accumulators
-    #pragma unroll
-    for (int i = 1; i < vt0; i++) {
-      #pragma unroll
-      for (index_t j = 0; j < output_vec_size; j++) {
-        value_list[0][j] = ops.combine(value_list[0][j], value_list[i][j]);
-      }
-    }
-    return value_list[0];
+    return value_list;
   }
 
   template <int output_vec_size>