fix aten view op (#10050)

* fix aten view op

* add test case

* fix signature

* fix the build

Co-authored-by: Cheng Tang <chenta@microsoft.com@orttrainingdev9.d32nl1ml4oruzj4qz3bqlggovf.px.internal.cloudapp.net>

orttraining/orttraining/eager/ort_aten.cpp

@@ -89,7 +89,7 @@ OrtValue create_ort_value(
     {},
     &ort_val);
   auto* ort_tensor = ort_val.GetMutable<onnxruntime::Tensor>();
-  CopyVectorToTensor<float>(invoker, {val}, *ort_tensor);
+  CopyVectorToTensor<float>(invoker, &val, 1, *ort_tensor);
   return ort_val;
 }
 
@@ -369,7 +369,8 @@ at::Tensor& zero_(at::Tensor& self){
   CreateMLValue(invoker.GetCurrentExecutionProvider().GetAllocator(0, OrtMemTypeDefault),
                 element_type, {}, &flag_val);
   auto* ort_flag_tensor = flag_val.GetMutable<onnxruntime::Tensor>();
-  CopyVectorToTensor<int64_t>(invoker, {1}, *ort_flag_tensor);
+  int64_t one = 1;
+  CopyVectorToTensor<int64_t>(invoker, &one, 1, *ort_flag_tensor);
 
   std::vector<OrtValue> ort_out = {ort_in_self};
   

orttraining/orttraining/eager/ort_aten.h

@@ -46,7 +46,7 @@ OrtValue create_ort_value(
     {1,},
     &ort_val);
   auto* ort_tensor = ort_val.GetMutable<onnxruntime::Tensor>();
-  CopyVectorToTensor<int64_t>(invoker, {val}, *ort_tensor);
+  CopyVectorToTensor<int64_t>(invoker, &val, 1, *ort_tensor);
   return ort_val;
 }
 
@@ -69,7 +69,8 @@ OrtValue create_ort_value(
     &ort_value);
   CopyVectorToTensor<T>(
     invoker,
-    values,
+    values.data(),
+    values.size(),
     *ort_value.GetMutable<onnxruntime::Tensor>());
   return ort_value;
 }

orttraining/orttraining/eager/ort_ops.cpp

@@ -19,5 +19,26 @@ void copy(onnxruntime::ORTInvoker& invoker,
   ORT_THROW_IF_ERROR(ort_ep.GetDataTransfer()->CopyTensor(src_tensor, *dst_tensor));
 }
 
+template <template<class> class V>
+void createInplaceOutputValue(OrtValue& input, V<int64_t> shape, OrtValue* p_mlvalue){
+  auto* input_ort_tensor = input.GetMutable<onnxruntime::Tensor>();
+  auto element_type = onnxruntime::DataTypeImpl::GetType<int64_t>();
+  // the ort TensorShape class only accept std::vector, so have to conversion.
+  std::vector<int64_t> new_shape;
+  new_shape.assign(shape.begin(), shape.end());
+  CreateMLValue(input_ort_tensor->MutableDataRaw(),
+                element_type, new_shape, p_mlvalue);
+}
+
+template <>
+void createInplaceOutputValue<std::vector>(OrtValue& input, std::vector<int64_t> shape, OrtValue* p_mlvalue){
+  auto* input_ort_tensor = input.GetMutable<onnxruntime::Tensor>();
+  auto element_type = onnxruntime::DataTypeImpl::GetType<int64_t>();
+  CreateMLValue(input_ort_tensor->MutableDataRaw(),
+                element_type, shape, p_mlvalue);
+}
+
+template void createInplaceOutputValue<c10::ArrayRef>(OrtValue& input, c10::ArrayRef<int64_t> shape, OrtValue* p_mlvalue);
+
 } // namespace eager
 } // namespace torch_ort

orttraining/orttraining/eager/ort_ops.h

@@ -10,27 +10,27 @@
 namespace torch_ort {
 namespace eager {
 
+template <template<class> class V>
+void createInplaceOutputValue(OrtValue& input, V<int64_t> shape, OrtValue* p_mlvalue);
+
 template <template<class> class V>
 OrtValue reshape_invoke(
   onnxruntime::ORTInvoker& invoker,
   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());
+  // the ort reshape kernel already handle the -1 in target shape
+  // don't need to invoke at::infer_size here.
   OrtValue shape_tensor;
   //todo: avoid the copy on this small shape vector;
   auto element_type = onnxruntime::DataTypeImpl::GetType<int64_t>();
   CreateMLValue(invoker.GetCurrentExecutionProvider().GetAllocator(0, OrtMemTypeDefault),
-                element_type, {(int64_t)new_shape.size(),}, &shape_tensor);
+                element_type, {(int64_t)shape.size(),}, &shape_tensor);
   auto* ort_shape_tensor = shape_tensor.GetMutable<onnxruntime::Tensor>();
-  CopyVectorToTensor<int64_t>(invoker, new_shape, *ort_shape_tensor);
+  CopyVectorToTensor<int64_t>(invoker, shape.data(), shape.size(), *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]);
+    createInplaceOutputValue(input, shape, &result[0]);
   }
   ORT_THROW_IF_ERROR(invoker.Invoke("Reshape", {input, shape_tensor}, result, nullptr));
   return result[0];

orttraining/orttraining/eager/ort_util.h

@@ -19,19 +19,19 @@ void CreateMLValue(void* data_ptr, onnxruntime::MLDataType element_type, const s
 
 template <typename T>
 inline void CopyVectorToTensor(onnxruntime::ORTInvoker& invoker,
-                               const std::vector<T>& value,
+                               const T* value_ptr,
+                               int64_t size,
                                onnxruntime::Tensor& tensor) {
   const auto& execution_provider = invoker.GetCurrentExecutionProvider();
 
   OrtValue* ort_value;
-  int64_t shape = value.size();
   OrtMemoryInfo cpuMemoryInfo;
 
   Ort::ThrowOnError(Ort::GetApi().CreateTensorWithDataAsOrtValue(
     &cpuMemoryInfo,
-    const_cast<void*>(reinterpret_cast<const void*>(value.data())),
-    value.size() * sizeof(T),
-    &shape,
+    const_cast<void*>(reinterpret_cast<const void*>(value_ptr)),
+    size * sizeof(T),
+    &size,
     1,
     Ort::TypeToTensorType<T>::type,
     &ort_value));
@@ -44,11 +44,12 @@ inline void CopyVectorToTensor(onnxruntime::ORTInvoker& invoker,
 // vector<bool> is specialized so we need to handle it separately
 template <>
 inline void CopyVectorToTensor<bool>(onnxruntime::ORTInvoker& /*invoker*/,
-                                     const std::vector<bool>& value,
+                                     const bool* value_ptr,
+                                     int64_t size,
                                      onnxruntime::Tensor& tensor) {
   auto output_span = tensor.MutableDataAsSpan<bool>();
-  for (size_t i = 0, end = value.size(); i < end; ++i) {
-    output_span[i] = value[i];
+  for (size_t i = 0, end = size; i < end; ++i) {
+    output_span[i] = value_ptr[i];
   }
 }
 

orttraining/orttraining/eager/test/ort_tensor.py

@@ -26,6 +26,12 @@ class OrtTensorTests(unittest.TestCase):
     y = ort_ones.reshape(-1)
     assert len(y.size()) == 1
     assert y.size()[0] == 100
+  
+  def test_view(self):
+    cpu_ones = torch.ones(2048)
+    ort_ones = cpu_ones.to('ort')
+    y = ort_ones.view(4, 512)
+    assert y.size() == (4, 512)
 
 if __name__ == '__main__':
   unittest.main()