From ae88f435505ec3db7b3c378fefa1f006657637e5 Mon Sep 17 00:00:00 2001 From: Jameson Miller Date: Fri, 1 Jul 2022 13:35:12 -0400 Subject: [PATCH] Eager mode: structure for supporting out= operators (#12066) * Add utility methods for resize_output * Eager mode: implement abs.out This is an initial hand written implementation of an out= operator to demonstrate how to structure out= methods using resize_out helper methods. This is meant to be used as a reference when we update the code generator to generate implementations for out= operations. --- .../orttraining/eager/opgen/opgen/atenops.py | 2 +- orttraining/orttraining/eager/ort_aten.cpp | 66 ++++++++++++++++++- orttraining/orttraining/eager/ort_aten.h | 14 +++- orttraining/orttraining/eager/test/ort_ops.py | 35 ++++++++++ 4 files changed, 114 insertions(+), 3 deletions(-) diff --git a/orttraining/orttraining/eager/opgen/opgen/atenops.py b/orttraining/orttraining/eager/opgen/opgen/atenops.py index e3586a0f46..1bd426958a 100644 --- a/orttraining/orttraining/eager/opgen/opgen/atenops.py +++ b/orttraining/orttraining/eager/opgen/opgen/atenops.py @@ -60,7 +60,6 @@ for binary_op, onnx_op in { type_promotion_ops.append(f"aten::{binary_op}{variant}.{dtype}") for unary_op in [ - "abs", "acos", "acosh", "asinh", @@ -100,6 +99,7 @@ for unary_op in [ ops[f"{aten_name}_"] = onnx_op hand_implemented = { + "aten::abs.out": SignatureOnly(), "aten::empty.memory_format": SignatureOnly(), "aten::empty_strided": SignatureOnly(), "aten::zero_": SignatureOnly(), diff --git a/orttraining/orttraining/eager/ort_aten.cpp b/orttraining/orttraining/eager/ort_aten.cpp index b70e158197..d26dfd71ec 100644 --- a/orttraining/orttraining/eager/ort_aten.cpp +++ b/orttraining/orttraining/eager/ort_aten.cpp @@ -2,7 +2,6 @@ // Licensed under the MIT License. #include "ort_aten.h" -#include "ort_tensor.h" #include #include #include @@ -344,6 +343,35 @@ OrtValue CastToType(onnxruntime::ORTInvoker& invoker, const OrtValue& input, at: return output[0]; } +/* + * Utility function for resizing output tensor + * Only resizes if: + * - The shape is different + * - The output tensor is empty + * + * We do not support resizing non-empty output tensors. + * PyToch implementation of resize will warn about resizing + * non-empty and indicate this is deprecated behavior that + * can / will change. + * + * In PyTorch repository see: aten/src/ATen/native/Resize.{h|cpp} + */ +void resize_output( + onnxruntime::ORTInvoker& invoker, + ORTTensorImpl* output, + at::IntArrayRef shape) { + if (output->sizes().equals(shape)) { + return; + } + + if (output->numel() != 0) { + throw std::runtime_error( + "resizing a non-empty output tensor is not supported."); + } + + resize_impl_ort_(invoker, output, shape); +} + //#pragma endregion /* @@ -873,6 +901,42 @@ const at::Tensor& resize_( return self; } +// aten::abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) +at::Tensor& abs_out( + const at::Tensor& self, + // *, + at::Tensor& out) { + ORT_LOG_FN(self, out); + + if ( + !IsSupportedType(self, {at::kHalf,at::kByte,at::kInt,at::kBFloat16,at::kFloat,at::kDouble,at::kShort,at::kLong})) { + return at::native::call_fallback_fn< + &at::native::cpu_fallback, + ATEN_OP(abs_out)>::call(self, out); + } + auto& invoker = GetORTInvoker(self.device()); + + auto ort_input_self = create_ort_value(invoker, self); + + resize_output(invoker, + dynamic_cast(out.unsafeGetTensorImpl()), + self.sizes()); + + auto ort_out = create_ort_value(invoker, out); + std::vector ort_outputs_0_Abs{ort_out}; + + auto status = invoker.Invoke("Abs", { + std::move(ort_input_self), + }, ort_outputs_0_Abs, nullptr); + + if (!status.IsOK()) { + throw std::runtime_error( + "ORT return failure status:" + status.ErrorMessage()); + } + + return out; +} + } // namespace aten //#pragma endregion diff --git a/orttraining/orttraining/eager/ort_aten.h b/orttraining/orttraining/eager/ort_aten.h index 6e86198f31..a29b272ec6 100644 --- a/orttraining/orttraining/eager/ort_aten.h +++ b/orttraining/orttraining/eager/ort_aten.h @@ -9,6 +9,7 @@ #include "ort_util.h" #include "ort_ops.h" #include "ort_log.h" +#include "ort_tensor.h" namespace torch_ort { namespace eager { @@ -121,5 +122,16 @@ c10::optional PromoteScalarTypesWithCategory( ONNX_NAMESPACE::TensorProto_DataType GetONNXTensorProtoDataType(at::ScalarType dtype); OrtValue CastToType(onnxruntime::ORTInvoker& invoker, const OrtValue& input, at::ScalarType type); + +void resize_output( + onnxruntime::ORTInvoker& invoker, + ORTTensorImpl* output, + at::IntArrayRef shape); + +void resize_impl_ort_( + onnxruntime::ORTInvoker& invoker, + ORTTensorImpl* self, + at::IntArrayRef size); + } // namespace eager -} // namespace torch_ort \ No newline at end of file +} // namespace torch_ort diff --git a/orttraining/orttraining/eager/test/ort_ops.py b/orttraining/orttraining/eager/test/ort_ops.py index 04e2c23b30..7f6b412aa9 100644 --- a/orttraining/orttraining/eager/test/ort_ops.py +++ b/orttraining/orttraining/eager/test/ort_ops.py @@ -254,6 +254,41 @@ class OrtOpTests(unittest.TestCase): self.assertEqual(cpu_tensor.size(), ort_tensor.size()) self.assertTrue(torch.allclose(cpu_tensor, ort_tensor.cpu())) + def test_abs(self): + device = self.get_device() + cpu_tensor = torch.tensor([-1, -2, 3, -6, -7]) + ort_tensor = cpu_tensor.to(device) + + cpu_result = torch.abs(cpu_tensor) + ort_result = torch.abs(ort_tensor) + + assert torch.equal(cpu_result, ort_result.cpu()) + + def test_abs_(self): + device = self.get_device() + cpu_tensor = torch.tensor([-1, -2, 3, -6, -7]) + ort_tensor = cpu_tensor.to(device) + + torch.abs_(cpu_tensor) + torch.abs_(ort_tensor) + + assert torch.equal(cpu_tensor, ort_tensor.cpu()) + + def test_abs_out(self): + device = self.get_device() + cpu_tensor = torch.tensor([-1, -2, 3, -6, -7]) + ort_tensor = cpu_tensor.to(device) + + cpu_out_tensor = torch.tensor([], dtype=torch.long) + ort_out_tensor = cpu_out_tensor.to(device) + + cpu_result = torch.abs(cpu_tensor, out=cpu_out_tensor) + ort_result = torch.abs(ort_tensor, out=ort_out_tensor) + + assert torch.equal(cpu_result, ort_result.cpu()) + assert torch.equal(cpu_out_tensor, ort_out_tensor.cpu()) + assert torch.equal(ort_result.cpu(), ort_out_tensor.cpu()) + if __name__ == "__main__": unittest.main()