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;