If the input is contiguous, short-circuit infer_size_dv in reshape (p…

…ytorch#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 <[email protected]> Pull Request resolved: pytorch#95216 Approved by: https://github.com/albanD
ROCm · Feb 21, 2023 · e5785f1 · e5785f1
1 parent 7b403c8
commit e5785f1
Showing 1 changed file with 5 additions and 0 deletions.
diff --git 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()) {