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/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(),

orttraining/orttraining/eager/ort_aten.cpp

@@ -2,7 +2,6 @@
 // Licensed under the MIT License.
 
 #include "ort_aten.h"
-#include "ort_tensor.h"
 #include <c10/core/TensorImpl.h>
 #include <ATen/native/CPUFallback.h>
 #include <ATen/InferSize.h>
@@ -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<ORTTensorImpl*>(out.unsafeGetTensorImpl()),
+                self.sizes());
+
+  auto ort_out = create_ort_value(invoker, out);
+  std::vector<OrtValue> 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

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<at::ScalarType> 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
+} // namespace torch_ort

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()