Only accept continguous tensors in TopK for cuda (#9441)

t-vi · facebook-github-bot · commit c8bd62c20c6b · 2018-07-14T12:24:52.000-07:00
Summary: Fixes: #9421 I don't think it is easy to deal with non-contiguous array in cuda topk, so I'm adding a check. The argument number is a bit confusing when it shows in PyTorch but it is consistent with the other checks. (Not sure whether it would make sense to eliminate argument numbers from the error TH/THC error messages given that they're probably off more than once...) Do we need a test that it indeed refuses non-contiguous? Pull Request resolved: pytorch/pytorch#9441 Reviewed By: soumith Differential Revision: D8850719 Pulled By: ezyang fbshipit-source-id: d50561bb37ed50ab97aeaf54d8e3fc6c765bdc7c
diff --git a/aten/src/THC/generic/THCTensorTopK.cu b/aten/src/THC/generic/THCTensorTopK.cu
@@ -5,21 +5,23 @@
 THC_API void THCTensor_(topk)(THCState* state,
                                THCTensor *topK,
                                THCudaLongTensor *indices,
-                               THCTensor *input,
+                               THCTensor *input_,
                                int64_t k, int dim, int dir, int sorted) {
-  THAssert(topK != NULL && indices != NULL && input != NULL);
-  THCAssertSameGPU(THCTensor_(checkGPU)(state, 3, topK, indices, input));
+  THAssert(topK != NULL && indices != NULL && input_ != NULL);
+  THCAssertSameGPU(THCTensor_(checkGPU)(state, 3, topK, indices, input_));
   THArgCheck(THCTensor_(_nDimension)(state, topK) <= MAX_CUTORCH_DIMS, 2, CUTORCH_DIM_WARNING);
   int64_t dims = THCudaLongTensor__nDimension(state, indices);
   THArgCheck(dims <= MAX_CUTORCH_DIMS, 3, CUTORCH_DIM_WARNING);
-  int numDims = THCTensor_(_nDimension)(state, input);
+  int numDims = THCTensor_(_nDimension)(state, input_);
   THArgCheck(numDims <= MAX_CUTORCH_DIMS, 4, CUTORCH_DIM_WARNING);
 
   THArgCheck(dim >= 0 && dim < numDims, 6, "dim not in range");
 
-  int64_t sliceSize = THCTensor_(size)(state, input, dim);
+  int64_t sliceSize = THCTensor_(size)(state, input_, dim);
   THArgCheck(k > 0 && k <= sliceSize, 5, "k not in range for dimension");
 
+  THCTensor *input = THCTensor_(newContiguous)(state, input_);
+
   // Build the output size, which is the dim being selected set to
   // size k
   THLongStorage* topKSize = THCTensor_(newSizeOf)(state, input);
@@ -155,6 +157,8 @@ THC_API void THCTensor_(topk)(THCState* state,
     }
   }
 
+  THCudaLongTensor_free(state, input);
+
   THCudaCheck(cudaGetLastError());
 }
 
