Add aten::complex.out support for export and CPU execution

Author: mshr-h

exir/dialects/edge/dtype/supported.py

@@ -27,6 +27,7 @@
     torch.float: "Float",
     torch.double: "Double",
     torch.uint16: "UInt16",
+    torch.complex32: "ComplexHalf",
     torch.complex64: "ComplexFloat",
     torch.complex128: "ComplexDouble",
 }

exir/verification/verifier.py

@@ -143,6 +143,9 @@ def _get_core_aten_ops_exception_list(self) -> List[torch._ops.OpOverload]:
                     torch.ops.aten.max.default,  # TODO(T188268054)
                     torch.ops.aten.min.default,  # TODO(T188268054)
                     torch.ops.aten.full_like.default,  # TODO(T183507359)
+                    # TODO: Add complex.default to PyTorch's Core ATen opset.
+                    # Allow it here until the upstream change is available in supported versions.
+                    torch.ops.aten.complex.default,
                 ]
                 + list(_EXECUTORCH_SYM_OPS)
                 + DISALLOW_LIST

extension/pybindings/test/test_pybindings.py

@@ -874,6 +874,24 @@ def test_non_dense_numpy_input_is_rejected(self):
         with self.assertRaisesRegex(BufferError, "dense, non-overlapping"):
             executorch_module(non_dense)
 
+    def test_complex_construction(self):
+        class Complex(torch.nn.Module):
+            def forward(self, real, imag):
+                return torch.complex(real, imag)
+
+        for dtype in (torch.float16, torch.float32, torch.float64):
+            with self.subTest(dtype=dtype):
+                real = torch.tensor([[1.0], [-2.0]], dtype=dtype)
+                imag = torch.tensor([[3.0, 0.0, -4.0]], dtype=dtype)
+                program = to_edge(export(Complex(), (real, imag))).to_executorch()
+                module = self.load_fn(program.buffer)
+                output = torch.from_dlpack(module.forward([real, imag])[0])
+                expected = torch.complex(real, imag)
+                self.assertEqual(output.dtype, expected.dtype)
+                self.assertEqual(output.shape, expected.shape)
+                torch.testing.assert_close(output.real, expected.real, rtol=0, atol=0)
+                torch.testing.assert_close(output.imag, expected.imag, rtol=0, atol=0)
+
     def test_complex_numpy_input(self):
         class Identity(torch.nn.Module):
             def forward(self, x):

kernels/aten/functions.yaml

@@ -119,6 +119,8 @@
 
 - op: clone.out
 
+- op: complex.out
+
 - op: constant_pad_nd.out
 
 - op: conv1d.out

kernels/portable/cpu/op_complex.cpp

@@ -0,0 +1,68 @@
+/*
+ * 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/broadcast_util.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+
+namespace torch {
+namespace executor {
+namespace native {
+
+Tensor& complex_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& real,
+    const Tensor& imag,
+    Tensor& out) {
+  // Check Input Dtypes
+  const ScalarType dtype = real.scalar_type();
+  ET_KERNEL_CHECK(
+      ctx,
+      dtype == ScalarType::Half || dtype == ScalarType::Float ||
+          dtype == ScalarType::Double,
+      InvalidArgument,
+      out);
+  ET_KERNEL_CHECK(ctx, imag.scalar_type() == dtype, InvalidArgument, out);
+
+  // Check Output Dtype
+  ET_KERNEL_CHECK(
+      ctx,
+      out.scalar_type() == executorch::runtime::toComplexType(dtype),
+      InvalidArgument,
+      out);
+
+  // Check Dim Order
+  ET_KERNEL_CHECK(
+      ctx,
+      tensor_is_default_dim_order(real) && tensor_is_default_dim_order(imag) &&
+          tensor_is_default_dim_order(out),
+      InvalidArgument,
+      out);
+
+  // Resize
+  ET_KERNEL_CHECK(
+      ctx,
+      resize_to_broadcast_target_size(real, imag, out) == Error::Ok,
+      InvalidArgument,
+      out);
+
+  ET_SWITCH_FLOATH_TYPES(dtype, ctx, "complex.out", CTYPE, [&] {
+    using Complex = executorch::aten::complex<CTYPE>;
+    apply_binary_elementwise_fn<CTYPE, CTYPE, Complex>(
+        [](CTYPE real_value, CTYPE imag_value) {
+          return Complex(real_value, imag_value);
+        },
+        real,
+        imag,
+        out);
+  });
+  return out;
+}
+
+} // namespace native
+} // namespace executor
+} // namespace torch

kernels/portable/functions.yaml

@@ -288,6 +288,11 @@
     - arg_meta: null
       kernel_name: torch::executor::clone_out
 
+- op: complex.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::complex_out
+
 - op: constant_pad_nd.out
   kernels:
     - arg_meta: null

kernels/test/CMakeLists.txt

@@ -191,6 +191,7 @@ set(all_test_sources
     "op_ceil_test.cpp"
     "op_clamp_test.cpp"
     "op_clone_test.cpp"
+    "op_complex_test.cpp"
     "op_constant_pad_nd_test.cpp"
     "op_convolution_backward_test.cpp"
     "op_convolution_test.cpp"

kernels/test/op_complex_test.cpp

@@ -0,0 +1,239 @@
+/*
+ * 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>
+#include <executorch/kernels/test/TestUtil.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
+
+#include <cmath>
+#include <limits>
+
+using executorch::aten::ScalarType;
+using executorch::aten::Tensor;
+using executorch::aten::TensorShapeDynamism;
+using torch::executor::testing::TensorFactory;
+
+class OpComplexOutTest : public OperatorTest {
+ protected:
+  Tensor& op_complex_out(const Tensor& real, const Tensor& imag, Tensor& out) {
+    return torch::executor::aten::complex_outf(context_, real, imag, out);
+  }
+
+  template <ScalarType REAL, ScalarType COMPLEX>
+  void test_dtype() {
+    TensorFactory<REAL> tf;
+    TensorFactory<COMPLEX> cf;
+    using C = typename TensorFactory<COMPLEX>::true_ctype;
+    Tensor real = tf.make({2, 2}, {1, -2, 0, 4});
+    Tensor imag = tf.make({2, 2}, {-5, 0, 7, -8});
+    Tensor out = cf.full({2, 2}, C(0, 0));
+    EXPECT_EQ(&op_complex_out(real, imag, out), &out);
+    EXPECT_TENSOR_EQ(
+        out, cf.make({2, 2}, {C(1, -5), C(-2, 0), C(0, 7), C(4, -8)}));
+  }
+
+  template <ScalarType DTYPE>
+  void test_unsupported_input() {
+    TensorFactory<DTYPE> tf;
+    TensorFactory<ScalarType::ComplexFloat> cf;
+    Tensor input = tf.zeros({2});
+    Tensor out = cf.zeros({2});
+    ET_EXPECT_KERNEL_FAILURE(context_, op_complex_out(input, input, out));
+  }
+};
+
+TEST_F(OpComplexOutTest, Half) {
+  test_dtype<ScalarType::Half, ScalarType::ComplexHalf>();
+}
+
+TEST_F(OpComplexOutTest, Float) {
+  test_dtype<ScalarType::Float, ScalarType::ComplexFloat>();
+}
+
+TEST_F(OpComplexOutTest, Double) {
+  test_dtype<ScalarType::Double, ScalarType::ComplexDouble>();
+}
+
+TEST_F(OpComplexOutTest, SpecialValues) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::ComplexFloat> cf;
+  const float inf = std::numeric_limits<float>::infinity();
+  const float nan = std::numeric_limits<float>::quiet_NaN();
+  Tensor real = tf.make({4}, {0.0f, -0.0f, inf, nan});
+  Tensor imag = tf.make({4}, {-0.0f, 0.0f, nan, -inf});
+  Tensor out = cf.zeros({4});
+  op_complex_out(real, imag, out);
+  const auto* data = out.const_data_ptr<executorch::aten::complex<float>>();
+  EXPECT_EQ(data[0].real(), 0.0f);
+  EXPECT_FALSE(std::signbit(data[0].real()));
+  EXPECT_EQ(data[0].imag(), 0.0f);
+  EXPECT_TRUE(std::signbit(data[0].imag()));
+  EXPECT_EQ(data[1].real(), 0.0f);
+  EXPECT_TRUE(std::signbit(data[1].real()));
+  EXPECT_EQ(data[1].imag(), 0.0f);
+  EXPECT_FALSE(std::signbit(data[1].imag()));
+  EXPECT_EQ(data[2].real(), inf);
+  EXPECT_TRUE(std::isnan(data[2].imag()));
+  EXPECT_TRUE(std::isnan(data[3].real()));
+  EXPECT_EQ(data[3].imag(), -inf);
+}
+
+TEST_F(OpComplexOutTest, BroadcastBothInputs) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::ComplexFloat> cf;
+  using C = executorch::aten::complex<float>;
+  Tensor real = tf.make({2, 1}, {1, 2});
+  Tensor imag = tf.make({1, 3}, {3, 4, 5});
+  Tensor out = cf.zeros({2, 3});
+  op_complex_out(real, imag, out);
+  EXPECT_TENSOR_EQ(
+      out,
+      cf.make({2, 3}, {C(1, 3), C(1, 4), C(1, 5), C(2, 3), C(2, 4), C(2, 5)}));
+}
+
+TEST_F(OpComplexOutTest, BroadcastDifferentRanks) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::ComplexFloat> cf;
+  using C = executorch::aten::complex<float>;
+  Tensor real = tf.make({2}, {1, 2});
+  Tensor imag = tf.make({2, 1}, {3, 4});
+  Tensor out = cf.zeros({2, 2});
+  op_complex_out(real, imag, out);
+  EXPECT_TENSOR_EQ(out, cf.make({2, 2}, {C(1, 3), C(2, 3), C(1, 4), C(2, 4)}));
+}
+
+TEST_F(OpComplexOutTest, Scalars) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::ComplexFloat> cf;
+  Tensor real = tf.make({}, {1});
+  Tensor imag = tf.make({}, {-2});
+  Tensor out = cf.zeros({});
+  op_complex_out(real, imag, out);
+  EXPECT_TENSOR_EQ(out, cf.make({}, {executorch::aten::complex<float>(1, -2)}));
+}
+
+TEST_F(OpComplexOutTest, BroadcastScalar) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::ComplexFloat> cf;
+  using C = executorch::aten::complex<float>;
+  Tensor real = tf.make({}, {1});
+  Tensor imag = tf.make({2}, {2, 3});
+  Tensor out = cf.zeros({2});
+  op_complex_out(real, imag, out);
+  EXPECT_TENSOR_EQ(out, cf.make({2}, {C(1, 2), C(1, 3)}));
+}
+
+TEST_F(OpComplexOutTest, EmptyBroadcast) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::ComplexFloat> cf;
+  Tensor real = tf.zeros({0, 3});
+  Tensor imag = tf.ones({1, 3});
+  Tensor out = cf.zeros({0, 3});
+  op_complex_out(real, imag, out);
+  EXPECT_TENSOR_EQ(out, cf.zeros({0, 3}));
+}
+
+TEST_F(OpComplexOutTest, DynamicOutput) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::ComplexFloat> cf;
+  using C = executorch::aten::complex<float>;
+  Tensor real = tf.make({2, 1}, {1, 2});
+  Tensor imag = tf.make({2}, {3, 4});
+  Tensor out = cf.zeros({4, 4}, TensorShapeDynamism::DYNAMIC_BOUND);
+  op_complex_out(real, imag, out);
+  EXPECT_TENSOR_EQ(out, cf.make({2, 2}, {C(1, 3), C(1, 4), C(2, 3), C(2, 4)}));
+}
+
+TEST_F(OpComplexOutTest, MismatchedInputs) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Double> df;
+  TensorFactory<ScalarType::ComplexFloat> cf;
+  Tensor real = tf.ones({2});
+  Tensor imag = df.ones({2});
+  Tensor out = cf.zeros({2});
+  ET_EXPECT_KERNEL_FAILURE(context_, op_complex_out(real, imag, out));
+}
+
+TEST_F(OpComplexOutTest, WrongComplexOutputType) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::ComplexDouble> cf;
+  Tensor input = tf.ones({2});
+  Tensor out = cf.zeros({2});
+  ET_EXPECT_KERNEL_FAILURE(context_, op_complex_out(input, input, out));
+}
+
+TEST_F(OpComplexOutTest, RealOutputType) {
+  TensorFactory<ScalarType::Float> tf;
+  Tensor input = tf.ones({2});
+  Tensor out = tf.zeros({2});
+  ET_EXPECT_KERNEL_FAILURE(context_, op_complex_out(input, input, out));
+}
+
+TEST_F(OpComplexOutTest, BFloat16Input) {
+  test_unsupported_input<ScalarType::BFloat16>();
+}
+
+TEST_F(OpComplexOutTest, IntegerInput) {
+  test_unsupported_input<ScalarType::Int>();
+}
+
+TEST_F(OpComplexOutTest, BoolInput) {
+  test_unsupported_input<ScalarType::Bool>();
+}
+
+TEST_F(OpComplexOutTest, ComplexInput) {
+  test_unsupported_input<ScalarType::ComplexFloat>();
+}
+
+TEST_F(OpComplexOutTest, IncompatibleShapes) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::ComplexFloat> cf;
+  Tensor real = tf.ones({2});
+  Tensor imag = tf.ones({3});
+  Tensor out = cf.zeros({3});
+  ET_EXPECT_KERNEL_FAILURE(context_, op_complex_out(real, imag, out));
+}
+
+#ifndef USE_ATEN_LIB
+TEST_F(OpComplexOutTest, OutputCapacityTooSmall) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::ComplexFloat> cf;
+  Tensor input = tf.ones({2, 3});
+  Tensor out = cf.zeros({1, 1}, TensorShapeDynamism::DYNAMIC_BOUND);
+  ET_EXPECT_KERNEL_FAILURE(context_, op_complex_out(input, input, out));
+}
+
+TEST_F(OpComplexOutTest, StaticOutputWrongShape) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::ComplexFloat> cf;
+  Tensor input = tf.ones({2, 3});
+  Tensor out = cf.zeros({3, 2});
+  ET_EXPECT_KERNEL_FAILURE(context_, op_complex_out(input, input, out));
+}
+
+TEST_F(OpComplexOutTest, NonDefaultDimOrder) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::ComplexFloat> cf;
+  for (int index = 0; index < 3; ++index) {
+    context_ = {};
+    Tensor real = tf.ones({2, 3});
+    Tensor imag = tf.ones({2, 3});
+    Tensor out = cf.zeros({2, 3});
+    if (index == 0) {
+      real = tf.make_with_dimorder({2, 3}, {1, 2, 3, 4, 5, 6}, {1, 0});
+    } else if (index == 1) {
+      imag = tf.make_with_dimorder({2, 3}, {1, 2, 3, 4, 5, 6}, {1, 0});
+    } else {
+      out = cf.make_with_dimorder(
+          {2, 3}, std::vector<executorch::aten::complex<float>>(6), {1, 0});
+    }
+    ET_EXPECT_KERNEL_FAILURE(context_, op_complex_out(real, imag, out));
+  }
+}
+#endif

kernels/test/targets.bzl

@@ -242,6 +242,7 @@ def define_common_targets():
     _common_op_test("op_ceil_test", ["aten", "portable"])
     _common_op_test("op_clamp_test", ["aten", "portable"])
     _common_op_test("op_clone_test", ["aten", "portable"])
+    _common_op_test("op_complex_test", ["aten", "portable"])
     _common_op_test("op_constant_pad_nd_test", ["aten", "portable"])
     _common_op_test("op_convolution_test", ["aten", "portable"])
     _common_op_test("op_convolution_backward_test", ["aten", "portable"])