From 4e70d090a30c9fe65c19a1d5d76d714607518082 Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Sun, 15 Nov 2026 04:44:12 +0000 Subject: [PATCH 1/4] 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 --- kernels/optimized/cpu/op_div.cpp | 19 +- kernels/portable/cpu/op_div.cpp | 19 +- kernels/portable/cpu/util/broadcast_util.h | 24 +- .../portable/cpu/util/test/broadcast_test.cpp | 126 ++++++++++ kernels/test/op_div_test.cpp | 235 ++++++++++++++++++ 5 files changed, 403 insertions(+), 20 deletions(-) diff --git a/kernels/optimized/cpu/op_div.cpp b/kernels/optimized/cpu/op_div.cpp index 11deb7d9f5b..a49c0b1cae2 100644 --- a/kernels/optimized/cpu/op_div.cpp +++ b/kernels/optimized/cpu/op_div.cpp @@ -21,9 +21,6 @@ namespace native { namespace { ScalarType get_common_type(ScalarType a_type, ScalarType b_type) { - if (isComplexType(a_type) || isComplexType(b_type)) { - return promoteTypes(a_type, b_type); - } ET_CHECK(!isQIntType(a_type) && !isBitsType(a_type)); ET_CHECK(!isQIntType(b_type) && !isBitsType(b_type)); @@ -64,14 +61,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 f896e4151fc..0b1b25ffb3d 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..5a6babc4c62 100644 --- a/kernels/portable/cpu/util/broadcast_util.h +++ b/kernels/portable/cpu/util/broadcast_util.h @@ -247,6 +247,15 @@ 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. `out` must + * already be resized to the broadcast target shape, and `a` and `b` must each + * have at most `out.dim()` dimensions and be broadcastable to `out`. The + * strides of `a`, `b`, and `out` are all respected. If all three tensors have + * identical sizes and strides, `compute_fun` is called in ascending physical + * offset order; otherwise, it is called in the logical row-major coordinate + * order of `out`. The call order is not guaranteed to stay the same across + * layouts or shapes. */ template inline void apply_binary_elementwise_fn( @@ -258,8 +267,19 @@ 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 (const auto i : c10::irange(out.numel())) { + data_out[i] = compute_fun(data_a[i], data_b[i]); + } + return; + } + + // The first index is a logical linear index that ignores out's strides. As + // documented in BroadcastIndexesRange, pass out as an extra input to get its + // physical offset instead. + for (const auto [unused_linear_index, 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..cfec3bf2a54 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,131 @@ 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 a_channels_last = tf.channels_last_like(a); + Tensor b_channels_last = tf.channels_last_like(b); + 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; }; + + Tensor out_channels_last_1 = tf.channels_last_like(tf.zeros({1, 2, 1, 2})); + apply_binary_elementwise_fn( + subtract, a, b_channels_last, out_channels_last_1); + EXPECT_TENSOR_EQ(out_channels_last_1, expected_channels_last); + + Tensor out_channels_last_2 = tf.channels_last_like(tf.zeros({1, 2, 1, 2})); + apply_binary_elementwise_fn( + subtract, a_channels_last, b, out_channels_last_2); + EXPECT_TENSOR_EQ(out_channels_last_2, expected_channels_last); + + Tensor out_contiguous = tf.zeros({1, 2, 1, 2}); + apply_binary_elementwise_fn( + subtract, a_channels_last, b_channels_last, out_contiguous); + EXPECT_TENSOR_EQ(out_contiguous, 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..c8dded56f57 100644 --- a/kernels/test/op_div_test.cpp +++ b/kernels/test/op_div_test.cpp @@ -31,6 +31,143 @@ 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() { + 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 +744,104 @@ 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) { + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen supports mixed complex dtypes"); + test_mismatched_complex_dtypes_dies< + ScalarType::ComplexFloat, + ScalarType::ComplexDouble, + ScalarType::ComplexDouble>(); +} + +TEST_F(OpDivOutTest, ComplexMismatchedOutputDtypeDies) { + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen supports mixed complex dtypes"); + 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) { + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen supports mixed complex dtypes"); + test_mismatched_complex_dtypes_dies< + ScalarType::ComplexFloat, + ScalarType::Float, + ScalarType::ComplexFloat>(); +} + +TEST_F(OpDivOutTest, RealAndComplexInputsDies) { + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen supports mixed complex dtypes"); + test_mismatched_complex_dtypes_dies< + ScalarType::Float, + ScalarType::ComplexFloat, + ScalarType::ComplexFloat>(); +} + TEST_F(OpDivOutTest, ComplexFloatBasic) { TensorFactory tf; From 06aa91ee2a7f2eb626f0a6ddd303b2e59912b70a Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Sun, 15 Nov 2026 08:00:11 +0000 Subject: [PATCH 2/4] Fix channels-last masked_scatter traversal order Consume source values in logical output order for channels-last tensors while preserving the default-layout fast path. Add regression tests for matching shapes, mixed ranks, mask broadcasting, and insufficient sources. Include masked_scatter tests in the CMake suite. --- kernels/portable/cpu/op_masked_scatter.cpp | 33 +++++---- kernels/test/CMakeLists.txt | 1 + kernels/test/op_masked_scatter_test.cpp | 82 ++++++++++++++++++++++ 3 files changed, 104 insertions(+), 12 deletions(-) diff --git a/kernels/portable/cpu/op_masked_scatter.cpp b/kernels/portable/cpu/op_masked_scatter.cpp index e91a8d0a08d..14c2ed833d9 100644 --- a/kernels/portable/cpu/op_masked_scatter.cpp +++ b/kernels/portable/cpu/op_masked_scatter.cpp @@ -49,18 +49,27 @@ Tensor& masked_scatter_out( ET_SWITCH_REALHBBF16_TYPES(in_type, ctx, name, CTYPE, [&]() { const CTYPE* const src_data = src.const_data_ptr(); - apply_binary_elementwise_fn( - [src_data, &idx, &src_numel, &src_numel_check]( - const CTYPE val_in, const bool val_mask) { - if (val_mask && idx >= src_numel) { - src_numel_check = false; - return val_in; - } - return val_mask ? src_data[idx++] : val_in; - }, - in, - mask, - out); + const auto compute = [src_data, &idx, &src_numel, &src_numel_check]( + const CTYPE val_in, const bool val_mask) { + if (val_mask && idx >= src_numel) { + src_numel_check = false; + return val_in; + } + return val_mask ? src_data[idx++] : val_in; + }; + if (is_contiguous_dim_order(out.dim_order().data(), out.dim())) { + apply_binary_elementwise_fn(compute, in, mask, out); + } else { + const CTYPE* const in_data = in.const_data_ptr(); + const bool* const mask_data = mask.const_data_ptr(); + CTYPE* const out_data = out.mutable_data_ptr(); + // Source consumption must follow logical output order, even when all + // tensors have matching channels-last layouts. + for (const auto [unused_linear_index, in_index, mask_index, out_index] : + BroadcastIndexesRange<3, true>(out, in, mask, out)) { + out_data[out_index] = compute(in_data[in_index], mask_data[mask_index]); + } + } }); ET_KERNEL_CHECK_MSG( diff --git a/kernels/test/CMakeLists.txt b/kernels/test/CMakeLists.txt index 1b00ebc9fcc..072a1ee91c3 100644 --- a/kernels/test/CMakeLists.txt +++ b/kernels/test/CMakeLists.txt @@ -246,6 +246,7 @@ set(all_test_sources "op_logit_test.cpp" "op_lt_test.cpp" "op_masked_fill_test.cpp" + "op_masked_scatter_test.cpp" "op_masked_select_test.cpp" "op_max_test.cpp" "op_max_pool2d_with_indices_test.cpp" diff --git a/kernels/test/op_masked_scatter_test.cpp b/kernels/test/op_masked_scatter_test.cpp index 99c463464bf..da58d6632ab 100644 --- a/kernels/test/op_masked_scatter_test.cpp +++ b/kernels/test/op_masked_scatter_test.cpp @@ -14,6 +14,7 @@ #include #include +#include using namespace ::testing; using executorch::aten::ScalarType; @@ -74,6 +75,87 @@ TEST_F(OpMaskedScatterOutTest, BroadcastMask) { EXPECT_TENSOR_EQ(out, tf.make({2, 3}, {1, 10, 3, 4, 20, 6})); } +TEST_F(OpMaskedScatterOutTest, ChannelsLastSameShape) { + TensorFactory tf; + TensorFactory tfBool; + + Tensor in = tf.channels_last_like(tf.make({1, 2, 1, 2}, {1, 2, 3, 4})); + Tensor mask = + tfBool.make_channels_last({1, 2, 1, 2}, {true, true, true, false}); + Tensor src = tf.make({3}, {10, 20, 30}); + Tensor out = tf.channels_last_like(tf.zeros({1, 2, 1, 2})); + + op_masked_scatter_out(in, mask, src, out); + EXPECT_TENSOR_EQ( + out, tf.channels_last_like(tf.make({1, 2, 1, 2}, {10, 20, 30, 4}))); +} + +TEST_F(OpMaskedScatterOutTest, ChannelsLastMixedRank) { + TensorFactory tf; + TensorFactory tfBool; + const std::vector lower_sizes = {2, 3, 1, 1}; + const std::vector upper_sizes = {1, 2, 3, 1, 1}; + + for (bool broadcast_input : {false, true}) { + SCOPED_TRACE(broadcast_input); + Tensor in = tf.channels_last_like(tf.make( + broadcast_input ? lower_sizes : upper_sizes, {1, 2, 3, 4, 5, 6})); + Tensor mask = tfBool.make_channels_last( + broadcast_input ? upper_sizes : lower_sizes, + broadcast_input + ? std::vector{true, true, true, false, false, true} + : std::vector{true, true, false, true, false, true}); + Tensor src = tf.make({4}, {10, 20, 30, 40}); + Tensor out = tf.channels_last_like(tf.zeros({1, 2, 3, 1, 1})); + + op_masked_scatter_out(in, mask, src, out); + EXPECT_TENSOR_EQ( + out, + tf.channels_last_like( + tf.make({1, 2, 3, 1, 1}, {10, 20, 3, 30, 5, 40}))); + } +} + +TEST_F(OpMaskedScatterOutTest, ChannelsLastBroadcastMask) { + TensorFactory tf; + TensorFactory tfBool; + + Tensor in = tf.channels_last_like(tf.make({1, 2, 1, 2}, {1, 2, 3, 4})); + Tensor mask = tfBool.make_channels_last({1, 2, 1, 1}, {true, false}); + Tensor src = tf.make({2}, {10, 20}); + Tensor out = tf.channels_last_like(tf.zeros({1, 2, 1, 2})); + + op_masked_scatter_out(in, mask, src, out); + EXPECT_TENSOR_EQ( + out, tf.channels_last_like(tf.make({1, 2, 1, 2}, {10, 20, 3, 4}))); +} + +TEST_F(OpMaskedScatterOutTest, ChannelsLastSrcWithLessElementsFails) { + TensorFactory tf; + TensorFactory tfBool; + + Tensor in = tf.channels_last_like(tf.make({1, 2, 1, 2}, {1, 2, 3, 4})); + Tensor mask = + tfBool.make_channels_last({1, 2, 1, 2}, {true, true, true, false}); + Tensor src = tf.make({2}, {10, 20}); + Tensor out = tf.channels_last_like(tf.zeros({1, 2, 1, 2})); + + ET_EXPECT_KERNEL_FAILURE(context_, op_masked_scatter_out(in, mask, src, out)); +} + +TEST_F(OpMaskedScatterOutTest, ChannelsLastMixedRankSrcWithLessElementsFails) { + TensorFactory tf; + TensorFactory tfBool; + + Tensor in = tf.channels_last_like(tf.make({2, 3, 1, 1}, {1, 2, 3, 4, 5, 6})); + Tensor mask = tfBool.make_channels_last( + {1, 2, 3, 1, 1}, {true, true, true, false, false, true}); + Tensor src = tf.make({3}, {10, 20, 30}); + Tensor out = tf.channels_last_like(tf.zeros({1, 2, 3, 1, 1})); + + ET_EXPECT_KERNEL_FAILURE(context_, op_masked_scatter_out(in, mask, src, out)); +} + TEST_F(OpMaskedScatterOutTest, SrcWithMoreElements) { TensorFactory tf; TensorFactory tfBool; From 58c7746ff13854ccee4ab414795e012f77f52b47 Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Sun, 15 Nov 2026 17:32:35 +0900 Subject: [PATCH 3/4] Apply batched suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- kernels/portable/cpu/op_masked_scatter.cpp | 28 +++++++++++++++------- kernels/test/op_div_test.cpp | 1 + 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/kernels/portable/cpu/op_masked_scatter.cpp b/kernels/portable/cpu/op_masked_scatter.cpp index 14c2ed833d9..23ea43f690b 100644 --- a/kernels/portable/cpu/op_masked_scatter.cpp +++ b/kernels/portable/cpu/op_masked_scatter.cpp @@ -49,14 +49,26 @@ Tensor& masked_scatter_out( ET_SWITCH_REALHBBF16_TYPES(in_type, ctx, name, CTYPE, [&]() { const CTYPE* const src_data = src.const_data_ptr(); - const auto compute = [src_data, &idx, &src_numel, &src_numel_check]( - const CTYPE val_in, const bool val_mask) { - if (val_mask && idx >= src_numel) { - src_numel_check = false; - return val_in; - } - return val_mask ? src_data[idx++] : val_in; - }; + const bool src_is_contiguous = + is_contiguous_dim_order(src.dim_order().data(), src.dim()); + const auto compute = + [src_data, &src, src_is_contiguous, &idx, &src_numel, &src_numel_check]( + const CTYPE val_in, const bool val_mask) { + if (!val_mask) { + return val_in; + } + if (idx >= src_numel) { + src_numel_check = false; + return val_in; + } + const size_t src_index = idx++; + if (src_is_contiguous) { + return src_data[src_index]; + } + size_t src_coord[kTensorDimensionLimit]; + delinearize_index(src_index, src, src_coord, kTensorDimensionLimit); + return src_data[linearize_access_indexes(src_coord, src.dim(), src)]; + }; if (is_contiguous_dim_order(out.dim_order().data(), out.dim())) { apply_binary_elementwise_fn(compute, in, mask, out); } else { diff --git a/kernels/test/op_div_test.cpp b/kernels/test/op_div_test.cpp index c8dded56f57..383aa9ac7a3 100644 --- a/kernels/test/op_div_test.cpp +++ b/kernels/test/op_div_test.cpp @@ -152,6 +152,7 @@ class OpDivOutTest : public OperatorTest { Tensor a = tf.ones({2}); Tensor b = tf.ones({3}); Tensor out = tf.zeros({3}); + context_.fail(torch::executor::Error::Ok); ET_EXPECT_KERNEL_FAILURE(context_, op_div_out(a, b, out)); } From 135c5b65a39e5cad4a42ef2c04b15ccde4543907 Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Sun, 15 Nov 2026 17:48:01 +0900 Subject: [PATCH 4/4] Potential fix for pull request finding 'Add channels-last source coverage for masked scatter stride remapping' Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- kernels/test/op_masked_scatter_test.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/kernels/test/op_masked_scatter_test.cpp b/kernels/test/op_masked_scatter_test.cpp index da58d6632ab..6c1042df70f 100644 --- a/kernels/test/op_masked_scatter_test.cpp +++ b/kernels/test/op_masked_scatter_test.cpp @@ -82,7 +82,8 @@ TEST_F(OpMaskedScatterOutTest, ChannelsLastSameShape) { Tensor in = tf.channels_last_like(tf.make({1, 2, 1, 2}, {1, 2, 3, 4})); Tensor mask = tfBool.make_channels_last({1, 2, 1, 2}, {true, true, true, false}); - Tensor src = tf.make({3}, {10, 20, 30}); + Tensor src = + tf.channels_last_like(tf.make({1, 2, 1, 2}, {10, 20, 30, 40})); Tensor out = tf.channels_last_like(tf.zeros({1, 2, 1, 2})); op_masked_scatter_out(in, mask, src, out);