From 825068d16640fe90c80e56cbe74a85a37dae0bfd Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Fri, 30 Oct 2026 08:15:51 +0000 Subject: [PATCH] Add `aten::complex.out` support for export and CPU execution --- exir/dialects/edge/dtype/supported.py | 1 + exir/verification/verifier.py | 3 + extension/pybindings/test/test_pybindings.py | 18 ++ kernels/aten/functions.yaml | 2 + kernels/portable/cpu/op_complex.cpp | 68 +++++ kernels/portable/functions.yaml | 5 + kernels/test/CMakeLists.txt | 1 + kernels/test/op_complex_test.cpp | 239 ++++++++++++++++++ kernels/test/targets.bzl | 1 + .../kernels/portable/op_registration_util.bzl | 6 + 10 files changed, 344 insertions(+) create mode 100644 kernels/portable/cpu/op_complex.cpp create mode 100644 kernels/test/op_complex_test.cpp diff --git a/exir/dialects/edge/dtype/supported.py b/exir/dialects/edge/dtype/supported.py index fb37e544de1..01fc6c5060f 100644 --- a/exir/dialects/edge/dtype/supported.py +++ b/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", } diff --git a/exir/verification/verifier.py b/exir/verification/verifier.py index 52bf257241b..5020fed63e8 100644 --- a/exir/verification/verifier.py +++ b/exir/verification/verifier.py @@ -144,6 +144,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 diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index 1b661995501..2d9a2349202 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/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): diff --git a/kernels/aten/functions.yaml b/kernels/aten/functions.yaml index eb8758e30a0..7ca675faa69 100644 --- a/kernels/aten/functions.yaml +++ b/kernels/aten/functions.yaml @@ -121,6 +121,8 @@ - op: clone.out +- op: complex.out + - op: constant_pad_nd.out - op: conv1d.out diff --git a/kernels/portable/cpu/op_complex.cpp b/kernels/portable/cpu/op_complex.cpp new file mode 100644 index 00000000000..37e4c07c20f --- /dev/null +++ b/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 +#include + +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; + apply_binary_elementwise_fn( + [](CTYPE real_value, CTYPE imag_value) { + return Complex(real_value, imag_value); + }, + real, + imag, + out); + }); + return out; +} + +} // namespace native +} // namespace executor +} // namespace torch diff --git a/kernels/portable/functions.yaml b/kernels/portable/functions.yaml index c2b3a47397b..d6bbcbc9cb7 100644 --- a/kernels/portable/functions.yaml +++ b/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 diff --git a/kernels/test/CMakeLists.txt b/kernels/test/CMakeLists.txt index de96692da7c..8dce0a91dcc 100644 --- a/kernels/test/CMakeLists.txt +++ b/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" diff --git a/kernels/test/op_complex_test.cpp b/kernels/test/op_complex_test.cpp new file mode 100644 index 00000000000..6e5b93d0af0 --- /dev/null +++ b/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 +#include +#include +#include + +#include +#include + +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 + void test_dtype() { + TensorFactory tf; + TensorFactory cf; + using C = typename TensorFactory::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 + void test_unsupported_input() { + TensorFactory tf; + TensorFactory 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(); +} + +TEST_F(OpComplexOutTest, Float) { + test_dtype(); +} + +TEST_F(OpComplexOutTest, Double) { + test_dtype(); +} + +TEST_F(OpComplexOutTest, SpecialValues) { + TensorFactory tf; + TensorFactory cf; + const float inf = std::numeric_limits::infinity(); + const float nan = std::numeric_limits::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>(); + 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 tf; + TensorFactory cf; + using C = executorch::aten::complex; + 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 tf; + TensorFactory cf; + using C = executorch::aten::complex; + 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 tf; + TensorFactory 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(1, -2)})); +} + +TEST_F(OpComplexOutTest, BroadcastScalar) { + TensorFactory tf; + TensorFactory cf; + using C = executorch::aten::complex; + 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 tf; + TensorFactory 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 tf; + TensorFactory cf; + using C = executorch::aten::complex; + 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 tf; + TensorFactory df; + TensorFactory 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 tf; + TensorFactory 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 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(); +} + +TEST_F(OpComplexOutTest, IntegerInput) { + test_unsupported_input(); +} + +TEST_F(OpComplexOutTest, BoolInput) { + test_unsupported_input(); +} + +TEST_F(OpComplexOutTest, ComplexInput) { + test_unsupported_input(); +} + +TEST_F(OpComplexOutTest, IncompatibleShapes) { + TensorFactory tf; + TensorFactory 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 tf; + TensorFactory 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 tf; + TensorFactory 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 tf; + TensorFactory 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>(6), {1, 0}); + } + ET_EXPECT_KERNEL_FAILURE(context_, op_complex_out(real, imag, out)); + } +} +#endif diff --git a/kernels/test/targets.bzl b/kernels/test/targets.bzl index 4da702cf6be..babb5cdc7b9 100644 --- a/kernels/test/targets.bzl +++ b/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"]) diff --git a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl index 1eee6f13a15..d0da4571be3 100644 --- a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl +++ b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl @@ -456,6 +456,12 @@ ATEN_OPS = ( op_target( name = "op_clone", ), + op_target( + name = "op_complex", + deps = [ + "//executorch/kernels/portable/cpu/util:broadcast_util", + ], + ), op_target( name = "op_constant_pad_nd", deps = [