diff --git a/kernels/portable/cpu/op_remainder.cpp b/kernels/portable/cpu/op_remainder.cpp index 912827cf3a8..d65ca088531 100644 --- a/kernels/portable/cpu/op_remainder.cpp +++ b/kernels/portable/cpu/op_remainder.cpp @@ -104,15 +104,6 @@ Tensor& remainder_Scalar_out( InvalidArgument, out); - // Check for intergral division by zero - ET_KERNEL_CHECK_MSG( - ctx, - !(executorch::runtime::isIntegralType(common_type, true) && - utils::scalar_to(b) == 0), - InvalidArgument, - out, - "Remainder operation encountered integer division by zero"); - // Check Dim Order ET_KERNEL_CHECK( ctx, tensors_have_same_dim_order(a, out), InvalidArgument, out); @@ -127,8 +118,18 @@ Tensor& remainder_Scalar_out( // @lint-ignore CLANGTIDY facebook-hte-CArray ET_DEFINE_OPERATOR_NAME(op_name, "remainder.Scalar_out"); + bool div_by_zero_error = false; + ET_SWITCH_REAL_TYPES(compute_type, ctx, op_name, CTYPE_COMPUTE, [&]() { + // Like ATen, the scalar takes the compute dtype first, so a scalar such as + // 256 for an int8 tensor becomes 0. const CTYPE_COMPUTE val_b = utils::scalar_to(b); + if constexpr (std::is_integral_v) { + if (val_b == 0 && a.numel() > 0) { + div_by_zero_error = true; + return; + } + } utils::apply_unitensor_elementwise_fn< CTYPE_COMPUTE, op_name, @@ -143,6 +144,13 @@ Tensor& remainder_Scalar_out( out); }); + ET_KERNEL_CHECK_MSG( + ctx, + !div_by_zero_error, + InvalidArgument, + out, + "Remainder operation encountered integer division by zero"); + return out; } diff --git a/kernels/portable/cpu/util/math_util.h b/kernels/portable/cpu/util/math_util.h index 9499eec671e..84ad0421143 100644 --- a/kernels/portable/cpu/util/math_util.h +++ b/kernels/portable/cpu/util/math_util.h @@ -221,6 +221,12 @@ template < typename CTYPE, typename std::enable_if::value, int>::type = 0> CTYPE remainder_override(CTYPE a, CTYPE b) { + if constexpr (std::is_signed::value) { + // a % -1 overflows for the most negative value; the result is always 0. + if (b == -1) { + return 0; + } + } CTYPE rem = a % b; if constexpr (std::is_signed::value) { if (((a < 0) ^ (b < 0)) && rem != 0) { diff --git a/kernels/test/op_remainder_test.cpp b/kernels/test/op_remainder_test.cpp index 998743cdee5..29e64e9cc5c 100644 --- a/kernels/test/op_remainder_test.cpp +++ b/kernels/test/op_remainder_test.cpp @@ -14,6 +14,8 @@ #include +#include + using namespace ::testing; using executorch::aten::Scalar; using executorch::aten::ScalarType; @@ -75,3 +77,53 @@ TEST_F(OpRemainderOutTest, DoubleKeepsDoublePrecision) { EXPECT_TENSOR_CLOSE_WITH_TOL( out, tf.make({2}, {789.123456789, 210.876543211}), 0, 1e-6); } + +TEST_F(OpRemainderOutTest, MostNegativeByMinusOneIsZero) { + TensorFactory tf; + + const int64_t min = std::numeric_limits::min(); + Tensor self = tf.make({3}, {min, 7, -7}); + Tensor out = tf.zeros({3}); + op_remainder_tensor_out(self, tf.full({3}, -1), out); + EXPECT_TENSOR_EQ(out, tf.zeros({3})); + + Tensor scalar_out = tf.full({3}, 1); + op_remainder_scalar_out(self, -1, scalar_out); + EXPECT_TENSOR_EQ(scalar_out, tf.zeros({3})); +} + +TEST_F(OpRemainderOutTest, IntegerDivisionByZeroFails) { + TensorFactory tf; + + Tensor self = tf.make({2}, {5, -3}); + Tensor out = tf.zeros({2}); + ET_EXPECT_KERNEL_FAILURE( + context_, op_remainder_tensor_out(self, tf.make({2}, {1, 0}), out)); + ET_EXPECT_KERNEL_FAILURE(context_, op_remainder_scalar_out(self, 0, out)); +} + +TEST_F(OpRemainderOutTest, ScalarIsCastToIntegerDtype) { + TensorFactory tf; + + Tensor self = tf.make({3}, {5, -3, 100}); + Tensor out = tf.full({3}, 1); + // 255 and 257 become -1 and 1 as int8. + op_remainder_scalar_out(self, 255, out); + EXPECT_TENSOR_EQ(out, tf.zeros({3})); + out = tf.full({3}, 1); + op_remainder_scalar_out(self, 257, out); + EXPECT_TENSOR_EQ(out, tf.zeros({3})); + // 256 becomes 0 as int8. + ET_EXPECT_KERNEL_FAILURE(context_, op_remainder_scalar_out(self, 256, out)); +} + +TEST_F(OpRemainderOutTest, EmptyInputWithZeroScalarDivisor) { + TensorFactory tf; + + Tensor self = tf.make({0}, {}); + Tensor out = tf.make({0}, {}); + op_remainder_scalar_out(self, 256, out); + EXPECT_TENSOR_EQ(out, tf.make({0}, {})); + op_remainder_scalar_out(self, 0, out); + EXPECT_TENSOR_EQ(out, tf.make({0}, {})); +}