From 324dc1623e2f91892038fb1b151450a7c6529dd9 Mon Sep 17 00:00:00 2001 From: Cloud Han Date: Sun, 10 May 2020 23:13:33 -0700 Subject: [PATCH] add dtype checking for gather and scatter (#38025) Summary: Fixed https://github.com/pytorch/pytorch/issues/37996 in the `cpu_scatter_gather_base_kernel`, it interpret a pointer as `int64_t` regardless the actual dtype. https://github.com/pytorch/pytorch/blob/2b41b9bceb01851df83d40c1280b3d3b09e1395b/aten/src/ATen/native/cpu/ScatterGatherKernel.cpp#L106 add a index dtype checking will avoid the nasty index out of bound error. As using `int64_t` is convention in ATen code (a.k.a, a limitation), no further fix is needed at the moment. Pull Request resolved: https://github.com/pytorch/pytorch/pull/38025 Differential Revision: D21498146 Pulled By: ezyang fbshipit-source-id: b1f96f394a460c4bc63d21ec8d4a2cfbf3e97b03 --- aten/src/ATen/native/TensorAdvancedIndexing.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/aten/src/ATen/native/TensorAdvancedIndexing.cpp b/aten/src/ATen/native/TensorAdvancedIndexing.cpp index 6288490ca6e6e..edf694e28b460 100644 --- a/aten/src/ATen/native/TensorAdvancedIndexing.cpp +++ b/aten/src/ATen/native/TensorAdvancedIndexing.cpp @@ -518,40 +518,48 @@ Tensor index_fill(const Tensor & self, int64_t dim, const Tensor & index, const } Tensor & gather_out_cpu(Tensor & result, const Tensor & self, int64_t dim, const Tensor & index, bool sparse_grad) { + TORCH_CHECK_INDEX(index.scalar_type() == ScalarType::Long, "gather_out(): Expected dtype int64 for index"); result.resize_(index.sizes()); gather_stub(result.device().type(), result, self, dim, index); return result; } Tensor gather_cpu(const Tensor & self, int64_t dim, const Tensor & index, bool sparse_grad) { + TORCH_CHECK_INDEX(index.scalar_type() == ScalarType::Long, "gather(): Expected dtype int64 for index"); Tensor result = at::empty({0}, self.options()); return gather_out_cpu(result, self, dim, index, sparse_grad); } Tensor & scatter_(Tensor & self, int64_t dim, const Tensor & index, const Tensor & src) { + TORCH_CHECK_INDEX(index.scalar_type() == ScalarType::Long, "scatter_(): Expected dtype int64 for index"); scatter_stub(self.device().type(), self, dim, index, src); return self; } Tensor & scatter_fill_(Tensor & self, int64_t dim, const Tensor & index, Scalar src) { + TORCH_CHECK_INDEX(index.scalar_type() == ScalarType::Long, "scatter_fill_(): Expected dtype int64 for index"); scatter_fill_stub(self.device().type(), self, dim, index, src); return self; } Tensor scatter(const Tensor & self, int64_t dim, const Tensor & index, const Tensor & source) { + TORCH_CHECK_INDEX(index.scalar_type() == ScalarType::Long, "scatter(): Expected dtype int64 for index"); return self.clone(at::MemoryFormat::Preserve).scatter_(dim, index, source); } Tensor scatter(const Tensor & self, int64_t dim, const Tensor & index, Scalar source) { + TORCH_CHECK_INDEX(index.scalar_type() == ScalarType::Long, "scatter(): Expected dtype int64 for index"); return self.clone(at::MemoryFormat::Preserve).scatter_(dim, index, source); } Tensor & scatter_add_cpu_(Tensor & self, int64_t dim, const Tensor & index, const Tensor & src) { + TORCH_CHECK_INDEX(index.scalar_type() == ScalarType::Long, "scatter_add_(): Expected dtype int64 for index"); scatter_add_stub(self.device().type(), self, dim, index, src); return self; } Tensor scatter_add(const Tensor & self, int64_t dim, const Tensor & index, const Tensor & source) { + TORCH_CHECK_INDEX(index.scalar_type() == ScalarType::Long, "scatter_add(): Expected dtype int64 for index"); return self.clone(at::MemoryFormat::Preserve).scatter_add_(dim, index, source); }