From fb256ec78557d88b1ee3526d44e3f8b7ecd946e4 Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Mon, 2 Nov 2026 01:29:15 +0000 Subject: [PATCH] Fix complex div broadcasting in CPU kernels Use broadcast-aware indexing for complex div in optimized and portable CPU kernels, and require matching complex input and output dtypes. Respect output strides while preserving linear iteration for matching dense layouts. Cover scalar and channels-last broadcasting, mixed ranks, empty tensors, bounded output resizing, and invalid inputs. Refs #24408 AI assistance: GPT 6 Astra Max and Copilot Autofix. Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- kernels/optimized/cpu/op_div.cpp | 16 +- kernels/portable/cpu/op_div.cpp | 19 +- kernels/portable/cpu/util/broadcast_util.h | 14 +- .../portable/cpu/util/test/broadcast_test.cpp | 123 ++++++++++ kernels/test/op_div_test.cpp | 226 ++++++++++++++++++ 5 files changed, 381 insertions(+), 17 deletions(-) diff --git a/kernels/optimized/cpu/op_div.cpp b/kernels/optimized/cpu/op_div.cpp index 855b2bb23f4..82a0b1f4688 100644 --- a/kernels/optimized/cpu/op_div.cpp +++ b/kernels/optimized/cpu/op_div.cpp @@ -64,14 +64,14 @@ Tensor& opt_div_out( // Handle complex types if (isComplexType(a_type) || isComplexType(b_type)) { - ScalarType common_type = get_common_type(a_type, b_type); - ET_SWITCH_COMPLEX_TYPES(common_type, ctx, op_name, CTYPE, [&]() { - const CTYPE* a_data = a.const_data_ptr(); - const CTYPE* b_data = b.const_data_ptr(); - CTYPE* out_data = out.mutable_data_ptr(); - for (size_t i = 0; i < out.numel(); ++i) { - out_data[i] = a_data[i] / b_data[i]; - } + ET_KERNEL_CHECK( + ctx, a_type == b_type && a_type == out_type, InvalidArgument, out); + ET_SWITCH_COMPLEX_TYPES(out_type, ctx, op_name, CTYPE, [&]() { + apply_binary_elementwise_fn( + [](const CTYPE val_a, const CTYPE val_b) { return val_a / val_b; }, + a, + b, + out); }); return out; } diff --git a/kernels/portable/cpu/op_div.cpp b/kernels/portable/cpu/op_div.cpp index 700cefcd55c..49e9038e68a 100644 --- a/kernels/portable/cpu/op_div.cpp +++ b/kernels/portable/cpu/op_div.cpp @@ -58,13 +58,18 @@ Tensor& div_out( ET_DEFINE_OPERATOR_NAME(op_name, "div.out"); if (executorch::runtime::isComplexType(common_type)) { - ET_SWITCH_COMPLEX_TYPES(common_type, ctx, op_name, CTYPE, [&]() { - const CTYPE* a_data = a.const_data_ptr(); - const CTYPE* b_data = b.const_data_ptr(); - CTYPE* out_data = out.mutable_data_ptr(); - for (ssize_t i = 0; i < out.numel(); ++i) { - out_data[i] = a_data[i] / b_data[i]; - } + ET_KERNEL_CHECK( + ctx, + a.scalar_type() == b.scalar_type() && + a.scalar_type() == out.scalar_type(), + InvalidArgument, + out); + ET_SWITCH_COMPLEX_TYPES(out.scalar_type(), ctx, op_name, CTYPE, [&]() { + apply_binary_elementwise_fn( + [](const CTYPE val_a, const CTYPE val_b) { return val_a / val_b; }, + a, + b, + out); }); } else { // Compute Dtype for real types diff --git a/kernels/portable/cpu/util/broadcast_util.h b/kernels/portable/cpu/util/broadcast_util.h index 2b10ee24411..48f7d13fc43 100644 --- a/kernels/portable/cpu/util/broadcast_util.h +++ b/kernels/portable/cpu/util/broadcast_util.h @@ -247,6 +247,7 @@ size_t linearize_access_indexes( * Useful for binary elementwise operators. For each element of the inputs, * perform a computation and write to the corresponding element of the output. * Tensor broadcasting is applied wherever it is required. + * The tensors must be dense in their respective dimension orders. */ template inline void apply_binary_elementwise_fn( @@ -258,8 +259,17 @@ inline void apply_binary_elementwise_fn( const CTYPE_B* const data_b = b.const_data_ptr(); CTYPE_OUT* const data_out = out.mutable_data_ptr(); - for (const auto [out_index, a_index, b_index] : - BroadcastIndexesRange<2>(out, a, b)) { + if (a.sizes() == out.sizes() && b.sizes() == out.sizes() && + a.strides() == out.strides() && b.strides() == out.strides()) { + for (size_t i = 0; i < out.numel(); ++i) { + data_out[i] = compute_fun(data_a[i], data_b[i]); + } + return; + } + + // Use an input index to respect output strides. + for (const auto [_, a_index, b_index, out_index] : + BroadcastIndexesRange<3, true>(out, a, b, out)) { data_out[out_index] = compute_fun(data_a[a_index], data_b[b_index]); } } diff --git a/kernels/portable/cpu/util/test/broadcast_test.cpp b/kernels/portable/cpu/util/test/broadcast_test.cpp index 7ffd95b6c52..a585ce061b3 100644 --- a/kernels/portable/cpu/util/test/broadcast_test.cpp +++ b/kernels/portable/cpu/util/test/broadcast_test.cpp @@ -23,6 +23,7 @@ using executorch::aten::ScalarType; using executorch::aten::Tensor; using executorch::runtime::ArrayRef; using executorch::runtime::testing::TensorFactory; +using torch::executor::apply_binary_elementwise_fn; using torch::executor::broadcast_tensor; using torch::executor::delinearize_index; using torch::executor::get_broadcast_target_size; @@ -30,6 +31,128 @@ using torch::executor::linearize_access_indexes; using torch::executor::tensor_is_broadcastable_to; using torch::executor::tensors_are_broadcastable_between; +TEST(BroadcastUtilTest, BinaryElementwiseDefaultPhysicalOrder) { + TensorFactory tf; + Tensor a = tf.make({2, 2}, {1, 2, 3, 4}); + Tensor b = tf.make({2, 2}, {11, 13, 17, 19}); + Tensor out = tf.zeros({2, 2}); + std::vector visited; + apply_binary_elementwise_fn( + [&visited](int32_t x, int32_t y) { + visited.push_back(x); + return x - y; + }, + a, + b, + out); + EXPECT_EQ(visited, (std::vector{1, 2, 3, 4})); + EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {-10, -11, -14, -15})); +} + +TEST(BroadcastUtilTest, BinaryElementwiseChannelsLastPhysicalOrder) { + TensorFactory tf; + Tensor a = tf.channels_last_like(tf.make({1, 2, 1, 2}, {1, 2, 3, 4})); + Tensor b = tf.channels_last_like(tf.make({1, 2, 1, 2}, {11, 13, 17, 19})); + Tensor out = tf.channels_last_like(tf.zeros({1, 2, 1, 2})); + std::vector visited; + apply_binary_elementwise_fn( + [&visited](int32_t x, int32_t y) { + visited.push_back(x); + return x - y; + }, + a, + b, + out); + EXPECT_EQ(visited, (std::vector{1, 3, 2, 4})); + EXPECT_TENSOR_EQ( + out, tf.channels_last_like(tf.make({1, 2, 1, 2}, {-10, -11, -14, -15}))); +} + +TEST(BroadcastUtilTest, BinaryElementwiseDifferentStrides) { + TensorFactory tf; + Tensor a = tf.make({1, 2, 1, 2}, {1, 2, 3, 4}); + Tensor b = tf.make({1, 2, 1, 2}, {11, 13, 17, 19}); + Tensor out = tf.zeros({1, 2, 1, 2}); + Tensor a_channels_last = tf.channels_last_like(a); + Tensor b_channels_last = tf.channels_last_like(b); + Tensor out_channels_last = tf.channels_last_like(out); + Tensor expected = tf.make({1, 2, 1, 2}, {-10, -11, -14, -15}); + Tensor expected_channels_last = tf.channels_last_like(expected); + const auto subtract = [](int32_t x, int32_t y) { return x - y; }; + + apply_binary_elementwise_fn( + subtract, a, b_channels_last, out_channels_last); + EXPECT_TENSOR_EQ(out_channels_last, expected_channels_last); + apply_binary_elementwise_fn( + subtract, a_channels_last, b, out_channels_last); + EXPECT_TENSOR_EQ(out_channels_last, expected_channels_last); + apply_binary_elementwise_fn( + subtract, a_channels_last, b_channels_last, out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST(BroadcastUtilTest, BinaryElementwiseChannelsLastBroadcast) { + TensorFactory tf; + Tensor a = tf.channels_last_like(tf.make({1, 2, 1, 2}, {1, 2, 3, 4})); + Tensor b = tf.channels_last_like(tf.make({1, 2, 1, 1}, {10, 20})); + Tensor out = tf.channels_last_like(tf.zeros({1, 2, 1, 2})); + const auto subtract = [](int32_t x, int32_t y) { return x - y; }; + + apply_binary_elementwise_fn(subtract, a, b, out); + EXPECT_TENSOR_EQ( + out, tf.channels_last_like(tf.make({1, 2, 1, 2}, {-9, -8, -17, -16}))); + apply_binary_elementwise_fn(subtract, b, a, out); + EXPECT_TENSOR_EQ( + out, tf.channels_last_like(tf.make({1, 2, 1, 2}, {9, 8, 17, 16}))); +} + +TEST(BroadcastUtilTest, BinaryElementwiseMixedRank) { + TensorFactory tf; + for (bool channels_last : {false, true}) { + Tensor a = tf.make({2, 3, 1, 1}, {1, 2, 3, 4, 5, 6}); + Tensor b = tf.make({1, 2, 3, 1, 1}, {10, 20, 30, 40, 50, 60}); + Tensor out = tf.zeros({1, 2, 3, 1, 1}); + Tensor expected = tf.make({1, 2, 3, 1, 1}, {-9, -18, -27, -36, -45, -54}); + if (channels_last) { + a = tf.channels_last_like(a); + b = tf.channels_last_like(b); + out = tf.channels_last_like(out); + expected = tf.channels_last_like(expected); + } + apply_binary_elementwise_fn( + [](int32_t x, int32_t y) { return x - y; }, a, b, out); + EXPECT_TENSOR_EQ(out, expected); + } +} + +TEST(BroadcastUtilTest, BinaryElementwiseBroadcastBothInputs) { + TensorFactory tf; + Tensor a = tf.make({2, 1}, {1, 4}); + Tensor b = tf.make({1, 2}, {10, 20}); + Tensor out = tf.zeros({2, 2}); + apply_binary_elementwise_fn( + [](int32_t x, int32_t y) { return x - y; }, a, b, out); + EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {-9, -19, -6, -16})); +} + +TEST(BroadcastUtilTest, BinaryElementwiseEmpty) { + TensorFactory tf; + Tensor a = tf.zeros({0, 2}); + Tensor b = tf.ones({1, 2}); + Tensor out = tf.zeros({0, 2}); + bool called = false; + apply_binary_elementwise_fn( + [&called](int32_t x, int32_t y) { + called = true; + return x - y; + }, + a, + b, + out); + EXPECT_FALSE(called); + EXPECT_TENSOR_EQ(out, a); +} + TEST(BroadcastUtilTest, BroadcastTensor) { TensorFactory tf; diff --git a/kernels/test/op_div_test.cpp b/kernels/test/op_div_test.cpp index d313e85a338..c1d2fca9131 100644 --- a/kernels/test/op_div_test.cpp +++ b/kernels/test/op_div_test.cpp @@ -31,6 +31,146 @@ class OpDivOutTest : public OperatorTest { return torch::executor::aten::div_outf(context_, a, b, out); } + template + void test_complex_scalar_broadcast() { + TensorFactory tf; + Tensor a = + tf.make({2, 2}, {CTYPE(1, 1), CTYPE(2, -2), CTYPE(4, 4), CTYPE(8, -8)}); + Tensor right_expected = tf.make( + {2, 2}, {CTYPE(0.5, 0), CTYPE(0, -1), CTYPE(2, 0), CTYPE(0, -4)}); + Tensor left_expected = tf.make( + {2, 2}, {CTYPE(2, 0), CTYPE(0, 1), CTYPE(0.5, 0), CTYPE(0, 0.25)}); + const std::vector> scalar_sizes = {{}, {1}, {1, 1}}; + for (const auto& sizes : scalar_sizes) { + Tensor b = tf.make(sizes, {CTYPE(2, 2)}); + Tensor out = tf.full({2, 2}, CTYPE(0, 0)); + EXPECT_TENSOR_CLOSE(op_div_out(a, b, out), right_expected); + EXPECT_TENSOR_CLOSE(op_div_out(b, a, out), left_expected); + } + + Tensor scalar_a = tf.make({}, {CTYPE(1, 1)}); + Tensor scalar_b = tf.make({}, {CTYPE(2, 2)}); + Tensor out = tf.full({}, CTYPE(0, 0)); + EXPECT_TENSOR_CLOSE( + op_div_out(scalar_a, scalar_b, out), tf.make({}, {CTYPE(0.5, 0)})); + } + + template + void test_complex_broadcast() { + TensorFactory tf; + Tensor a = tf.make({2, 1}, {CTYPE(2, 2), CTYPE(4, -4)}); + Tensor b = tf.make({1, 2}, {CTYPE(1, 1), CTYPE(2, -2)}); + Tensor out = tf.full({2, 2}, CTYPE(0, 0)); + EXPECT_TENSOR_CLOSE( + op_div_out(a, b, out), + tf.make({2, 2}, {CTYPE(2, 0), CTYPE(0, 1), CTYPE(0, -4), CTYPE(2, 0)})); + + Tensor empty_a = tf.full({0, 2}, CTYPE(0, 0)); + Tensor empty_out = tf.full({0, 2}, CTYPE(0, 0)); + EXPECT_TENSOR_CLOSE(op_div_out(empty_a, b, empty_out), empty_a); + EXPECT_TENSOR_CLOSE(op_div_out(b, empty_a, empty_out), empty_a); + } + + template + void test_complex_channels_last_broadcast() { + TensorFactory tf; + Tensor a = tf.channels_last_like(tf.make( + {1, 2, 1, 2}, + {CTYPE(2, 2), CTYPE(4, -4), CTYPE(8, 8), CTYPE(16, -16)})); + Tensor b = tf.channels_last_like( + tf.make({1, 2, 1, 1}, {CTYPE(1, 1), CTYPE(2, 2)})); + Tensor out = tf.channels_last_like(tf.full({1, 2, 1, 2}, CTYPE(0, 0))); + Tensor right_expected = tf.channels_last_like(tf.make( + {1, 2, 1, 2}, {CTYPE(2, 0), CTYPE(0, -4), CTYPE(4, 0), CTYPE(0, -8)})); + Tensor left_expected = tf.channels_last_like(tf.make( + {1, 2, 1, 2}, + {CTYPE(0.5, 0), CTYPE(0, 0.25), CTYPE(0.25, 0), CTYPE(0, 0.125)})); + + EXPECT_TENSOR_CLOSE(op_div_out(a, b, out), right_expected); + EXPECT_TENSOR_CLOSE(op_div_out(b, a, out), left_expected); + } + + template + void test_complex_mixed_rank_channels_last_broadcast(bool varied) { + TensorFactory tf; + const std::vector a_values = varied + ? std::vector< + CTYPE>{CTYPE(1, 1), CTYPE(2, -2), CTYPE(3, 3), CTYPE(4, -4), CTYPE(5, 5), CTYPE(6, -6)} + : std::vector{ + CTYPE(1, 0), + CTYPE(2, 0), + CTYPE(3, 0), + CTYPE(4, 0), + CTYPE(5, 0), + CTYPE(6, 0)}; + const std::vector b_values = varied + ? std::vector< + CTYPE>{CTYPE(1, -1), CTYPE(2, 2), CTYPE(3, -3), CTYPE(4, 4), CTYPE(5, -5), CTYPE(6, 6)} + : std::vector(6, CTYPE(1, 0)); + const std::vector right_values = varied + ? std::vector< + CTYPE>{CTYPE(0, 1), CTYPE(0, -1), CTYPE(0, 1), CTYPE(0, -1), CTYPE(0, 1), CTYPE(0, -1)} + : a_values; + const std::vector left_values = varied + ? std::vector< + CTYPE>{CTYPE(0, -1), CTYPE(0, 1), CTYPE(0, -1), CTYPE(0, 1), CTYPE(0, -1), CTYPE(0, 1)} + : std::vector{ + CTYPE(1, 0), + CTYPE(0.5, 0), + CTYPE(1.0 / 3, 0), + CTYPE(0.25, 0), + CTYPE(0.2, 0), + CTYPE(1.0 / 6, 0)}; + + Tensor a = tf.channels_last_like(tf.make({2, 3, 1, 1}, a_values)); + Tensor b = tf.channels_last_like(tf.make({1, 2, 3, 1, 1}, b_values)); + Tensor out = tf.channels_last_like(tf.zeros({1, 2, 3, 1, 1})); + Tensor right_expected = + tf.channels_last_like(tf.make({1, 2, 3, 1, 1}, right_values)); + Tensor left_expected = + tf.channels_last_like(tf.make({1, 2, 3, 1, 1}, left_values)); + + EXPECT_TENSOR_CLOSE(op_div_out(a, b, out), right_expected); + EXPECT_TENSOR_CLOSE(op_div_out(b, a, out), left_expected); + } + + template + void test_complex_dynamic_shape() { + TensorFactory tf; + Tensor a = tf.make({2, 1}, {CTYPE(2, 2), CTYPE(4, -4)}); + Tensor b = tf.make({1, 2}, {CTYPE(1, 1), CTYPE(2, -2)}); + Tensor out = + tf.zeros({3, 4}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND); + Tensor expected = + tf.make({2, 2}, {CTYPE(2, 0), CTYPE(0, 1), CTYPE(0, -4), CTYPE(2, 0)}); + EXPECT_TENSOR_CLOSE(op_div_out(a, b, out), expected); + } + + template + void test_complex_mismatched_shapes_dies() { + TensorFactory tf; + Tensor a = tf.ones({2}); + Tensor b = tf.ones({3}); + Tensor out = tf.zeros({3}); + ET_EXPECT_KERNEL_FAILURE(context_, op_div_out(a, b, out)); + } + + template + void test_mismatched_complex_dtypes_dies() { + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen supports mixed complex dtypes"); + TensorFactory tf_a; + TensorFactory tf_b; + TensorFactory tf_out; + + Tensor a = tf_a.ones({2}); + Tensor b = tf_b.ones({2}); + Tensor out = tf_out.zeros({2}); + + ET_EXPECT_KERNEL_FAILURE(context_, op_div_out(a, b, out)); + } + template void test_div() { TensorFactory tf_a; @@ -607,6 +747,92 @@ TEST_F(OpDivScalarOutTest, OptimizedSanityCheck) { // Complex Type Tests // +TEST_F(OpDivOutTest, ComplexScalarBroadcast) { +#define TEST_ENTRY(ctype, dtype) \ + test_complex_scalar_broadcast(); + ET_FORALL_COMPLEX_TYPES(TEST_ENTRY); +#undef TEST_ENTRY +} + +TEST_F(OpDivOutTest, ComplexBroadcast) { +#define TEST_ENTRY(ctype, dtype) \ + test_complex_broadcast(); + ET_FORALL_COMPLEX_TYPES(TEST_ENTRY); +#undef TEST_ENTRY +} + +TEST_F(OpDivOutTest, ComplexChannelsLastBroadcast) { +#define TEST_ENTRY(ctype, dtype) \ + test_complex_channels_last_broadcast(); + ET_FORALL_COMPLEX_TYPES(TEST_ENTRY); +#undef TEST_ENTRY +} + +TEST_F(OpDivOutTest, ComplexMixedRankChannelsLastBroadcastOnes) { +#define TEST_ENTRY(ctype, dtype) \ + test_complex_mixed_rank_channels_last_broadcast( \ + false); + ET_FORALL_COMPLEX_TYPES(TEST_ENTRY); +#undef TEST_ENTRY +} + +TEST_F(OpDivOutTest, ComplexMixedRankChannelsLastBroadcastVaried) { +#define TEST_ENTRY(ctype, dtype) \ + test_complex_mixed_rank_channels_last_broadcast( \ + true); + ET_FORALL_COMPLEX_TYPES(TEST_ENTRY); +#undef TEST_ENTRY +} + +TEST_F(OpDivOutTest, ComplexDynamicShapeUpperBoundLargerThanExpected) { +#define TEST_ENTRY(ctype, dtype) \ + test_complex_dynamic_shape(); + ET_FORALL_COMPLEX_TYPES(TEST_ENTRY); +#undef TEST_ENTRY +} + +TEST_F(OpDivOutTest, ComplexMismatchedShapesDies) { +#define TEST_ENTRY(ctype, dtype) \ + test_complex_mismatched_shapes_dies(); + ET_FORALL_COMPLEX_TYPES(TEST_ENTRY); +#undef TEST_ENTRY +} + +TEST_F(OpDivOutTest, ComplexMismatchedInputDtypesDies) { + test_mismatched_complex_dtypes_dies< + ScalarType::ComplexFloat, + ScalarType::ComplexDouble, + ScalarType::ComplexDouble>(); +} + +TEST_F(OpDivOutTest, ComplexMismatchedOutputDtypeDies) { + test_mismatched_complex_dtypes_dies< + ScalarType::ComplexFloat, + ScalarType::ComplexFloat, + ScalarType::ComplexDouble>(); +} + +TEST_F(OpDivOutTest, ComplexInputsRealOutputDies) { + test_mismatched_complex_dtypes_dies< + ScalarType::ComplexFloat, + ScalarType::ComplexFloat, + ScalarType::Float>(); +} + +TEST_F(OpDivOutTest, ComplexAndRealInputsDies) { + test_mismatched_complex_dtypes_dies< + ScalarType::ComplexFloat, + ScalarType::Float, + ScalarType::ComplexFloat>(); +} + +TEST_F(OpDivOutTest, RealAndComplexInputsDies) { + test_mismatched_complex_dtypes_dies< + ScalarType::Float, + ScalarType::ComplexFloat, + ScalarType::ComplexFloat>(); +} + TEST_F(OpDivOutTest, ComplexFloatBasic) { TensorFactory tf;