Fix topk work size computation (#5053)

albanD · soumith · commit 17f94d20aa2c · 2018-02-07T09:20:13.000-08:00
* fix grid computation for topk kernel

* backslash alignment, no change in code
diff --git a/torch/lib/THC/generic/THCTensorTopK.cu b/torch/lib/THC/generic/THCTensorTopK.cu
@@ -29,7 +29,7 @@ THC_API void THCTensor_(topk)(THCState* state,
   THLongStorage_free(topKSize);
 
 #define RUN_K(INDEX_T, DIM, DIR)                                        \
-  gatherTopK<real, INDEX_T, DIM, DIR>                                         \
+  gatherTopK<real, INDEX_T, DIM, DIR>                                   \
     <<<grid, block, 0, THCState_getCurrentStream(state)>>>(             \
       inputInfo,                                                        \
       sliceSize,                                                        \
@@ -63,10 +63,10 @@ THC_API void THCTensor_(topk)(THCState* state,
   }
 
 #define RUN_T(INDEX_T)                                                  \
-  TensorInfo<real, INDEX_T> inputInfo =                                \
-    getTensorInfo<THCTensor, INDEX_T>(state, input);                 \
-  TensorInfo<real, INDEX_T> topKInfo =                                 \
-    getTensorInfo<THCTensor, INDEX_T>(state, topK);                  \
+  TensorInfo<real, INDEX_T> inputInfo =                                 \
+    getTensorInfo<THCTensor, INDEX_T>(state, input);                    \
+  TensorInfo<real, INDEX_T> topKInfo =                                  \
+    getTensorInfo<THCTensor, INDEX_T>(state, topK);                     \
   TensorInfo<int64_t, INDEX_T> indicesInfo =                            \
     getTensorInfo<THCudaLongTensor, INDEX_T>(state, indices);           \
                                                                         \
@@ -82,9 +82,11 @@ THC_API void THCTensor_(topk)(THCState* state,
   int collapseIndicesDim = indicesInfo.collapseDims(dim);               \
                                                                         \
   int64_t inputSlices = 1;                                              \
-  int64_t topKSlices = 1;                                               \
-  for (int i = 0; i < numDims; ++i) {                                   \
+  for (int i = 0; i < inputInfo.dims; ++i) {                            \
     inputSlices *= inputInfo.sizes[i];                                  \
+  }                                                                     \
+  int64_t topKSlices = 1;                                               \
+  for (int i = 0; i < topKInfo.dims; ++i) {                             \
     topKSlices *= topKInfo.sizes[i];                                    \
   }                                                                     \
                                                                         \
