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/op_masked_scatter.cpp b/kernels/portable/cpu/op_masked_scatter.cpp index e91a8d0a08d..23ea43f690b 100644 --- a/kernels/portable/cpu/op_masked_scatter.cpp +++ b/kernels/portable/cpu/op_masked_scatter.cpp @@ -49,18 +49,39 @@ 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 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 && idx >= src_numel) { + if (!val_mask) { + return val_in; + } + if (idx >= src_numel) { src_numel_check = false; return val_in; } - return val_mask ? src_data[idx++] : val_in; - }, - in, - mask, - out); + 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 { + 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/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/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_div_test.cpp b/kernels/test/op_div_test.cpp index d313e85a338..383aa9ac7a3 100644 --- a/kernels/test/op_div_test.cpp +++ b/kernels/test/op_div_test.cpp @@ -31,6 +31,144 @@ 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}); + context_.fail(torch::executor::Error::Ok); + 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 +745,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; diff --git a/kernels/test/op_masked_scatter_test.cpp b/kernels/test/op_masked_scatter_test.cpp index 99c463464bf..6c1042df70f 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,88 @@ 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.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); + 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;