From 97659495d982982a06eee6a4e3560923b988ff7e Mon Sep 17 00:00:00 2001 From: "Tang, Cheng" Date: Tue, 4 Jan 2022 08:29:30 -0800 Subject: [PATCH] fix aten view op (#10050) * fix aten view op * add test case * fix signature * fix the build Co-authored-by: Cheng Tang --- orttraining/orttraining/eager/ort_aten.cpp | 5 +++-- orttraining/orttraining/eager/ort_aten.h | 5 +++-- orttraining/orttraining/eager/ort_ops.cpp | 21 +++++++++++++++++++ orttraining/orttraining/eager/ort_ops.h | 16 +++++++------- orttraining/orttraining/eager/ort_util.h | 17 ++++++++------- .../orttraining/eager/test/ort_tensor.py | 6 ++++++ 6 files changed, 50 insertions(+), 20 deletions(-) diff --git a/orttraining/orttraining/eager/ort_aten.cpp b/orttraining/orttraining/eager/ort_aten.cpp index ca33b09d12..d62e783805 100644 --- a/orttraining/orttraining/eager/ort_aten.cpp +++ b/orttraining/orttraining/eager/ort_aten.cpp @@ -89,7 +89,7 @@ OrtValue create_ort_value( {}, &ort_val); auto* ort_tensor = ort_val.GetMutable(); - CopyVectorToTensor(invoker, {val}, *ort_tensor); + CopyVectorToTensor(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(); - CopyVectorToTensor(invoker, {1}, *ort_flag_tensor); + int64_t one = 1; + CopyVectorToTensor(invoker, &one, 1, *ort_flag_tensor); std::vector ort_out = {ort_in_self}; diff --git a/orttraining/orttraining/eager/ort_aten.h b/orttraining/orttraining/eager/ort_aten.h index a63f914b30..fc560914e2 100644 --- a/orttraining/orttraining/eager/ort_aten.h +++ b/orttraining/orttraining/eager/ort_aten.h @@ -46,7 +46,7 @@ OrtValue create_ort_value( {1,}, &ort_val); auto* ort_tensor = ort_val.GetMutable(); - CopyVectorToTensor(invoker, {val}, *ort_tensor); + CopyVectorToTensor(invoker, &val, 1, *ort_tensor); return ort_val; } @@ -69,7 +69,8 @@ OrtValue create_ort_value( &ort_value); CopyVectorToTensor( invoker, - values, + values.data(), + values.size(), *ort_value.GetMutable()); return ort_value; } diff --git a/orttraining/orttraining/eager/ort_ops.cpp b/orttraining/orttraining/eager/ort_ops.cpp index 2cc98b356b..9b5f2e8f50 100644 --- a/orttraining/orttraining/eager/ort_ops.cpp +++ b/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 class V> +void createInplaceOutputValue(OrtValue& input, V shape, OrtValue* p_mlvalue){ + auto* input_ort_tensor = input.GetMutable(); + auto element_type = onnxruntime::DataTypeImpl::GetType(); + // the ort TensorShape class only accept std::vector, so have to conversion. + std::vector new_shape; + new_shape.assign(shape.begin(), shape.end()); + CreateMLValue(input_ort_tensor->MutableDataRaw(), + element_type, new_shape, p_mlvalue); +} + +template <> +void createInplaceOutputValue(OrtValue& input, std::vector shape, OrtValue* p_mlvalue){ + auto* input_ort_tensor = input.GetMutable(); + auto element_type = onnxruntime::DataTypeImpl::GetType(); + CreateMLValue(input_ort_tensor->MutableDataRaw(), + element_type, shape, p_mlvalue); +} + +template void createInplaceOutputValue(OrtValue& input, c10::ArrayRef shape, OrtValue* p_mlvalue); + } // namespace eager } // namespace torch_ort \ No newline at end of file diff --git a/orttraining/orttraining/eager/ort_ops.h b/orttraining/orttraining/eager/ort_ops.h index b949c329a2..f8c5564a20 100644 --- a/orttraining/orttraining/eager/ort_ops.h +++ b/orttraining/orttraining/eager/ort_ops.h @@ -10,27 +10,27 @@ namespace torch_ort { namespace eager { +template class V> +void createInplaceOutputValue(OrtValue& input, V shape, OrtValue* p_mlvalue); + template class V> OrtValue reshape_invoke( onnxruntime::ORTInvoker& invoker, OrtValue& input, V shape, bool in_place) { - // TODO: actual reshape on buffer - const onnxruntime::Tensor& input_tensor = input.Get(); - 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(); 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(); - CopyVectorToTensor(invoker, new_shape, *ort_shape_tensor); + CopyVectorToTensor(invoker, shape.data(), shape.size(), *ort_shape_tensor); std::vector result(1); if (in_place){ - auto* input_ort_tensor = input.GetMutable(); - 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]; diff --git a/orttraining/orttraining/eager/ort_util.h b/orttraining/orttraining/eager/ort_util.h index 4b467ee000..cf82b47183 100644 --- a/orttraining/orttraining/eager/ort_util.h +++ b/orttraining/orttraining/eager/ort_util.h @@ -19,19 +19,19 @@ void CreateMLValue(void* data_ptr, onnxruntime::MLDataType element_type, const s template inline void CopyVectorToTensor(onnxruntime::ORTInvoker& invoker, - const std::vector& 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(reinterpret_cast(value.data())), - value.size() * sizeof(T), - &shape, + const_cast(reinterpret_cast(value_ptr)), + size * sizeof(T), + &size, 1, Ort::TypeToTensorType::type, &ort_value)); @@ -44,11 +44,12 @@ inline void CopyVectorToTensor(onnxruntime::ORTInvoker& invoker, // vector is specialized so we need to handle it separately template <> inline void CopyVectorToTensor(onnxruntime::ORTInvoker& /*invoker*/, - const std::vector& value, + const bool* value_ptr, + int64_t size, onnxruntime::Tensor& tensor) { auto output_span = tensor.MutableDataAsSpan(); - 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]; } } diff --git a/orttraining/orttraining/eager/test/ort_tensor.py b/orttraining/orttraining/eager/test/ort_tensor.py index 772c26287a..b821a9179a 100644 --- a/orttraining/orttraining/eager/test/ort_tensor.py +++ b/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() \ No newline at end of file