From e5785f1e3496750ecc683ba8dbd670c51f2904b8 Mon Sep 17 00:00:00 2001 From: "Edward Z. Yang" Date: Tue, 21 Feb 2023 11:13:38 -0500 Subject: [PATCH] If the input is contiguous, short-circuit infer_size_dv in reshape (#95216) The main improvement is that this avoids guards from infer_size_dv, although this also counts as a minor perf improvement too. Signed-off-by: Edward Z. Yang Pull Request resolved: https://github.com/pytorch/pytorch/pull/95216 Approved by: https://github.com/albanD --- aten/src/ATen/native/TensorShape.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/aten/src/ATen/native/TensorShape.cpp b/aten/src/ATen/native/TensorShape.cpp index 7192ef85e3d05..26b03289494cf 100644 --- a/aten/src/ATen/native/TensorShape.cpp +++ b/aten/src/ATen/native/TensorShape.cpp @@ -1567,6 +1567,11 @@ Tensor reshape_symint(const Tensor& self, c10::SymIntArrayRef proposed_shape) { if (self.is_sparse()) { AT_ERROR("reshape is not implemented for sparse tensors"); } + + if (self.is_contiguous() && !self.is_mkldnn()) { + return self.view_symint(proposed_shape); + } + c10::SymDimVector shape = infer_size_dv(proposed_shape, self.sym_numel()); if (self.is_mkldnn()) {