saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-23 19:32:23 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

use inplace reshape (#9991)

Browse source 
Co-authored-by: Cheng Tang <chenta@microsoft.com@orttrainingdev9.d32nl1ml4oruzj4qz3bqlggovf.px.internal.cloudapp.net>
This commit is contained in:
Tang, Cheng 2021-12-15 21:17:29 -08:00 committed by GitHub
parent 7922a8c22f
commit 6357c12977
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 21 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
orttraining/orttraining/eager/ort_aten.cpp
View file

@ -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());
}

12
orttraining/orttraining/eager/ort_ops.h
View file

@ -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];
}