Add view_as_real_copy.out (pytorch#10207)

Summary:

As title. Needed for multichannel ASR.

Differential Revision: D72294238

Author: pssrawat

kernels/aten/functions.yaml

@@ -423,6 +423,8 @@
 
 - op: var.out
 
+- op: view_as_real_copy.out
+
 - op: view_copy.out
 
 - op: where.self_out

kernels/portable/cpu/op_view_as_real_copy.cpp

@@ -0,0 +1,77 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/kernels/portable/cpu/util/copy_ops_util.h>
+#include <executorch/kernels/portable/cpu/util/functional_util.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+#include <executorch/runtime/platform/assert.h>
+
+namespace torch {
+namespace executor {
+namespace native {
+
+using Tensor = executorch::aten::Tensor;
+
+template <typename SELF_CTYPE, typename OUT_CTYPE>
+inline void _to_impl(const Tensor& self, Tensor& out) {
+  auto self_data = self.mutable_data_ptr<SELF_CTYPE>();
+  auto out_data = out.mutable_data_ptr<OUT_CTYPE>();
+
+  for (size_t i = 0, e = self.numel(); i < e; i++) {
+    auto val_in = self_data[i];
+    out_data[2 * i] = static_cast<OUT_CTYPE>(val_in.real_);
+    out_data[2 * i + 1] = static_cast<OUT_CTYPE>(val_in.imag_);
+  }
+}
+
+// view_as_real_copy(Tensor self) -> Tensor
+Tensor& view_as_real_copy_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& self,
+    Tensor& out) {
+  (void)ctx;
+
+  // Get the output shape
+  Tensor::SizesType expected_output_size[kTensorDimensionLimit];
+  get_view_as_real_copy_out_target_size(self, expected_output_size);
+
+  // Resize for dynamic shape
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      resize_tensor(
+          out, {expected_output_size, static_cast<size_t>(out.dim())}) ==
+          Error::Ok,
+      InvalidArgument,
+      out,
+      "Failed to resize output tensor.");
+
+  // The input tensor must be complex type
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      executorch::runtime::isComplexType(self.scalar_type()),
+      InvalidArgument,
+      out,
+      "Input tensor must be complex type");
+
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(self, out), InvalidArgument, out);
+
+  constexpr auto op_name = "view_as_real_copy.out";
+
+  ET_SWITCH_COMPLEXH_TYPES(self.scalar_type(), ctx, op_name, CTYPE_IN, [&] {
+    ET_SWITCH_FLOATH_TYPES(out.scalar_type(), ctx, op_name, CTYPE_OUT, [&] {
+      _to_impl<CTYPE_IN, CTYPE_OUT>(self, out);
+    });
+  });
+
+  return out;
+}
+
+} // namespace native
+} // namespace executor
+} // namespace torch

kernels/portable/cpu/util/copy_ops_util.cpp

@@ -1018,5 +1018,14 @@ void get_unfold_copy_out_target_size(
   *out_ndim = self.dim() + 1;
 }
 
+void get_view_as_real_copy_out_target_size(
+    const Tensor& self,
+    executorch::aten::SizesType* out_sizes) {
+  for (auto i : c10::irange(self.dim())) {
+    out_sizes[i] = self.size(i);
+  }
+  out_sizes[self.dim()] = 2;
+}
+
 } // namespace executor
 } // namespace torch

kernels/portable/cpu/util/copy_ops_util.h

@@ -247,5 +247,9 @@ void get_unfold_copy_out_target_size(
     executorch::aten::SizesType* out_sizes,
     size_t* out_ndim);
 
+void get_view_as_real_copy_out_target_size(
+    const Tensor& self,
+    executorch::aten::SizesType* out_sizes);
+
 } // namespace executor
 } // namespace torch

kernels/portable/functions.yaml

@@ -957,6 +957,11 @@
     - arg_meta: null
       kernel_name: torch::executor::var_out
 
+- op: view_as_real_copy.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::view_as_real_copy_out
+
 - op: view_copy.out
   kernels:
     - arg_meta: null

kernels/test/CMakeLists.txt

@@ -242,6 +242,7 @@ set(all_test_sources
     "op_upsample_bilinear2d_test.cpp"
     "op_upsample_nearest2d_test.cpp"
     "op_var_test.cpp"
+    "op_view_as_real_copy_test.cpp"
     "op_view_copy_test.cpp"
     "op_where_test.cpp"
     "op_zeros_test.cpp"

kernels/test/op_view_as_real_copy_test.cpp

@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator
+#include <executorch/kernels/test/TestUtil.h>
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
+
+#include <gtest/gtest.h>
+
+using namespace ::testing;
+using executorch::aten::ScalarType;
+using executorch::aten::Tensor;
+using torch::executor::testing::TensorFactory;
+
+class OpViewAsRealTest : public OperatorTest {
+ protected:
+  Tensor& view_as_real_copy_out(const Tensor& self, Tensor& out) {
+    return torch::executor::aten::view_as_real_copy_outf(context_, self, out);
+  }
+
+  template <typename CTYPE, ScalarType DTYPE>
+  void run_complex_smoke_test() {
+    TensorFactory<DTYPE> tf;
+    constexpr auto REAL_DTYPE = executorch::runtime::toRealValueType(DTYPE);
+    TensorFactory<REAL_DTYPE> tf_out;
+    using REAL_CTYPE =
+        typename executorch::runtime::ScalarTypeToCppType<REAL_DTYPE>::type;
+    Tensor in = tf.make(
+        {2, 2},
+        {CTYPE{REAL_CTYPE(3), REAL_CTYPE(4)},
+         CTYPE{REAL_CTYPE(-1.7), REAL_CTYPE(7.4)},
+         CTYPE{REAL_CTYPE(5), REAL_CTYPE(-12)},
+         CTYPE{REAL_CTYPE(8.3), REAL_CTYPE(0.1)}});
+    Tensor out = tf_out.zeros({2, 2, 2});
+    Tensor expected =
+        tf_out.make({2, 2, 2}, {3, 4, -1.7, 7.4, 5, -12, 8.3, 0.1});
+    Tensor ret = view_as_real_copy_out(in, out);
+    EXPECT_TENSOR_EQ(out, ret);
+    EXPECT_TENSOR_EQ(out, expected);
+  }
+};
+
+TEST_F(OpViewAsRealTest, ComplexSmokeTest) {
+#define RUN_SMOKE_TEST(ctype, dtype) \
+  run_complex_smoke_test<ctype, ScalarType::dtype>();
+  ET_FORALL_COMPLEXH_TYPES(RUN_SMOKE_TEST);
+#undef RUN_SMOKE_TEST
+}

kernels/test/targets.bzl

@@ -331,6 +331,7 @@ def define_common_targets():
     _common_op_test("op_upsample_bilinear2d_test", ["aten", "portable"])
     _common_op_test("op_upsample_nearest2d_test", ["aten", "portable"])
     _common_op_test("op_var_test", ["aten", "portable"])
+    _common_op_test("op_view_as_real_copy_test", ["aten", "portable"])
     _common_op_test("op_view_copy_test", ["aten", "portable"])
     _common_op_test("op_where_test", ["aten", "portable"])
     _common_op_test("op_zeros_test", ["aten", "portable"])

shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl

@@ -1269,6 +1269,13 @@ ATEN_OPS = (
             "//executorch/kernels/portable/cpu/util:reduce_util",
         ],
     ),
+    op_target(
+        name = "op_view_as_real_copy",
+        deps = [
+            "//executorch/kernels/portable/cpu/util:functional_util",
+            "//executorch/kernels/portable/cpu/util:copy_ops_util",
+        ],
+    ),
     op_target(
         name = "op_view_copy",
         deps = [