From 6357c1297798705f173a0b990e7a118233b8cfa0 Mon Sep 17 00:00:00 2001
From: "Tang, Cheng" <souptc@gmail.com>
Date: Wed, 15 Dec 2021 21:17:29 -0800
Subject: [PATCH] use inplace reshape (#9991)

Co-authored-by: Cheng Tang <chenta@microsoft.com@orttrainingdev9.d32nl1ml4oruzj4qz3bqlggovf.px.internal.cloudapp.net>
---
 orttraining/orttraining/eager/ort_aten.cpp | 18 ++++++++++++------
 orttraining/orttraining/eager/ort_ops.h    | 12 +++++++++---
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/orttraining/orttraining/eager/ort_aten.cpp b/orttraining/orttraining/eager/ort_aten.cpp
index 9493c18f12..ca33b09d12 100644
--- a/orttraining/orttraining/eager/ort_aten.cpp
+++ b/orttraining/orttraining/eager/ort_aten.cpp
@@ -255,22 +255,28 @@ at::Tensor _reshape_alias(
   ORT_LOG_FN(self, size, stride);
   // TODO: support stride
   auto& invoker = GetORTInvoker(self.device());
+  auto ort_input = create_ort_value(invoker, self);
   return aten_tensor_from_ort(
-    reshape_copy(
+    reshape_invoke(
       invoker,
-      create_ort_value(invoker, self),
-      size),
+      ort_input,
+      size,
+      // invoke reshape kernel inplace
+      true),
     self.options());
 }
 
 at::Tensor view(const at::Tensor& self, at::IntArrayRef size) {
   ORT_LOG_FN(self, size);
   auto& invoker = GetORTInvoker(self.device());
+  auto ort_input = create_ort_value(invoker, self);
   return aten_tensor_from_ort(
-    reshape_copy(
+    reshape_invoke(
       invoker,
-      create_ort_value(invoker, self),
-      size),
+      ort_input,
+      size,
+      // invoke reshape kernel inplace
+      true),
     self.options());
 }
 
diff --git a/orttraining/orttraining/eager/ort_ops.h b/orttraining/orttraining/eager/ort_ops.h
index 13f7c30be4..b949c329a2 100644
--- a/orttraining/orttraining/eager/ort_ops.h
+++ b/orttraining/orttraining/eager/ort_ops.h
@@ -11,10 +11,11 @@ namespace torch_ort {
 namespace eager {
 
 template <template<class> class V>
-OrtValue reshape_copy(
+OrtValue reshape_invoke(
   onnxruntime::ORTInvoker& invoker,
-  const OrtValue& input,
-  V<int64_t> shape) {
+  OrtValue& input,
+  V<int64_t> shape,
+  bool in_place) {
   // TODO: actual reshape on buffer
   const onnxruntime::Tensor& input_tensor = input.Get<onnxruntime::Tensor>();
   auto new_shape = at::infer_size(shape, input_tensor.Shape().Size());
@@ -26,6 +27,11 @@ OrtValue reshape_copy(
   auto* ort_shape_tensor = shape_tensor.GetMutable<onnxruntime::Tensor>();
   CopyVectorToTensor<int64_t>(invoker, new_shape, *ort_shape_tensor);
   std::vector<OrtValue> result(1);
+  if (in_place){
+    auto* input_ort_tensor = input.GetMutable<onnxruntime::Tensor>();
+    CreateMLValue(input_ort_tensor->MutableDataRaw(),
+                element_type, new_shape, &result[0]);
+  }
   ORT_THROW_IF_ERROR(invoker.Invoke("Reshape", {input, shape_tensor}, result, nullptr));
   return result[0];
 }