From 0d9a02e64775c58080641e00fb1139dc3c687dcd Mon Sep 17 00:00:00 2001
From: msftlincoln <107071614+msftlincoln@users.noreply.github.com>
Date: Tue, 9 Aug 2022 17:16:18 -0400
Subject: [PATCH] Eager Mode - Support Concatenation via aten::cat.out (#12527)

* support concatenation via aten::cat.out

* wrap dims

* rename vars in tests, test wrapped dims
---
 .../orttraining/eager/opgen/opgen/atenops.py  |  1 +
 orttraining/orttraining/eager/ort_aten.cpp    | 53 +++++++++++++++++++
 orttraining/orttraining/eager/test/ort_ops.py | 18 +++++++
 3 files changed, 72 insertions(+)

diff --git a/orttraining/orttraining/eager/opgen/opgen/atenops.py b/orttraining/orttraining/eager/opgen/opgen/atenops.py
index 1e49fe86da..1e77d159b5 100644
--- a/orttraining/orttraining/eager/opgen/opgen/atenops.py
+++ b/orttraining/orttraining/eager/opgen/opgen/atenops.py
@@ -153,6 +153,7 @@ hand_implemented = {
     "aten::gelu": Gelu("self"),
     "aten::max": ReduceMax("self", keepdims=0),
     "aten::min": ReduceMin("self", keepdims=0),
+    "aten::cat.out": SignatureOnly(),
     "aten::fill_.Scalar": SignatureOnly(),
     "aten::ne.Scalar_out": Cast(Not(Equal("self", "other")), to="GetONNXTensorProtoDataType(out.scalar_type())"),
     "aten::ne.Tensor_out": Cast(Not(Equal("self", "other")), to="GetONNXTensorProtoDataType(out.scalar_type())"),
diff --git a/orttraining/orttraining/eager/ort_aten.cpp b/orttraining/orttraining/eager/ort_aten.cpp
index a24250cfd8..5afae07f4d 100644
--- a/orttraining/orttraining/eager/ort_aten.cpp
+++ b/orttraining/orttraining/eager/ort_aten.cpp
@@ -920,6 +920,59 @@ const at::Tensor& resize_(
   return self;
 }
 
+// aten::cat.out(Tensor[] tensors, int dim=0, *, Tensor(a!) out) -> Tensor(a!)
+at::Tensor& cat_out(
+  at::TensorList tensors,
+  int64_t dim,
+  // *,
+  at::Tensor& out) {
+  ORT_LOG_FN(tensors, dim, out);
+
+  assert(tensors.size() > 0);
+  if (
+      std::vector<at::ScalarType> supportedTypes =
+        {at::kBFloat16, at::kBool, at::kByte, at::kDouble, at::kFloat, at::kHalf, at::kInt, at::kLong, at::kShort};
+      !IsSupportedType(tensors, supportedTypes)) {
+    return at::native::call_fallback_fn<
+      &at::native::cpu_fallback,
+      ATEN_OP(cat_out)>::call(tensors, dim, out);
+  }
+  int64_t ndim = tensors[0].dim();
+  assert(ndim != 0);
+  dim = at::maybe_wrap_dim(dim, ndim);
+
+  auto& invoker = GetORTInvoker(tensors[0].device());
+
+  // IntArrayRef isn't writeable, convert to vector.
+  std::vector<int64_t> sizes;
+  for (auto s : tensors[0].sizes())
+    sizes.push_back(s);
+
+  // Calculate the new size of the dimension being concatenated.
+  sizes[dim] = 0;
+  for (auto t : tensors)
+    sizes[dim] += t.size(dim);
+
+  resize_output(invoker, dynamic_cast<ORTTensorImpl*>(out.unsafeGetTensorImpl()), at::IntArrayRef(sizes));
+  auto ort_input_out = create_ort_value(invoker, out);
+
+  auto ort_input_0_tensors = create_ort_value(invoker, tensors);
+
+  NodeAttributes attrs_0(1);
+  attrs_0["axis"] = create_ort_attribute(
+    "axis", dim, at::ScalarType::Int);
+
+  std::vector<OrtValue> ort_outputs_0_Concat(1);
+  ort_outputs_0_Concat[0] = ort_input_out;
+
+  auto status = invoker.Invoke("Concat", {
+    std::move(ort_input_0_tensors),
+  }, ort_outputs_0_Concat, &attrs_0);
+  CHECK_STATUS(status);
+
+  return out;
+}
+
 // aten::fill_.Scalar(Tensor(a!) self, Scalar value) -> Tensor(a!)
 at::Tensor& fill__Scalar(
     at::Tensor& self,
diff --git a/orttraining/orttraining/eager/test/ort_ops.py b/orttraining/orttraining/eager/test/ort_ops.py
index 9a9892137b..77e5da6c4c 100644
--- a/orttraining/orttraining/eager/test/ort_ops.py
+++ b/orttraining/orttraining/eager/test/ort_ops.py
@@ -108,6 +108,24 @@ class OrtOpTests(unittest.TestCase):
         assert torch.allclose(cpu_min, ort_min.cpu())
         assert cpu_min.dim() == ort_min.dim()
 
+    def test_cat(self):
+        device = self.get_device()
+        cpu_out_tensor = torch.tensor([])
+        ort_out_tensor = cpu_out_tensor.to(device)
+        cpu_x = torch.randn((128, 64))
+        cpu_y = torch.randn((128, 64))
+        cpu_z = torch.randn((128, 64))
+        cpu_ans_0 = torch.cat((cpu_x, cpu_y, cpu_z), 0)
+        cpu_ans_1 = torch.cat((cpu_x, cpu_y, cpu_z), -1, out=cpu_out_tensor)
+        ort_x = cpu_x.to(device)
+        ort_y = cpu_y.to(device)
+        ort_z = cpu_z.to(device)
+        ort_ans_0 = torch.cat((ort_x, ort_y, ort_z), 0)
+        ort_ans_1 = torch.cat((ort_x, ort_y, ort_z), -1, out=ort_out_tensor)
+        assert torch.allclose(cpu_ans_0, ort_ans_0.cpu())
+        assert torch.allclose(cpu_ans_1, ort_ans_1.cpu())
+        assert torch.allclose(cpu_out_tensor, ort_out_tensor.cpu())
+
     def test_equal(self):
         device = self.get_device()
         cpu_a = torch.Tensor([1.0, 1.5])