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

t-vi · facebook-github-bot · commit 8444e1660b66 · 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: #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());
 }
 
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -3194,6 +3194,14 @@ def test_topk_arguments(self):
         # Make sure True isn't mistakenly taken as the 2nd dimension (interpreted as 1)
         self.assertRaises(TypeError, lambda: q.topk(4, True))
 
+    @unittest.skipIf(not torch.cuda.is_available(), 'no CUDA')
+    def test_topk_noncontiguous_gpu(self):
+        t = torch.randn(20, device="cuda")[::2]
+        top1, idx1 = t.topk(5)
+        top2, idx2 = t.contiguous().topk(5)
+        self.assertEqual(top1, top2)
+        self.assertEqual(idx1, idx2)
+
     def test_kthvalue(self):
         SIZE = 50
         x = torch.rand(SIZE, SIZE, SIZE)
