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 41ad2f82898..be29bb366f0 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..765d84a9447 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; @@ -35,6 +37,22 @@ class OpRemainderOutTest : public OperatorTest { Tensor& out) { return torch::executor::aten::remainder_outf(context_, self, other, out); } + + template + void test_most_negative_by_minus_one() { + TensorFactory tf; + using CTYPE = typename TensorFactory::ctype; + + const CTYPE min = std::numeric_limits::min(); + Tensor self = tf.make({3}, {min, 7, -7}); + Tensor out = tf.full({3}, 1); + op_remainder_tensor_out(self, tf.full({3}, -1), out); + EXPECT_TENSOR_EQ(out, tf.zeros({3})); + + out = tf.full({3}, 1); + op_remainder_scalar_out(self, -1, out); + EXPECT_TENSOR_EQ(out, tf.zeros({3})); + } }; TEST_F(OpRemainderOutTest, SmokeTest) { @@ -75,3 +93,61 @@ TEST_F(OpRemainderOutTest, DoubleKeepsDoublePrecision) { EXPECT_TENSOR_CLOSE_WITH_TOL( out, tf.make({2}, {789.123456789, 210.876543211}), 0, 1e-6); } + +TEST_F(OpRemainderOutTest, MostNegativeByMinusOneIsZero) { + test_most_negative_by_minus_one(); + test_most_negative_by_minus_one(); +} + +TEST_F(OpRemainderOutTest, TensorIntegerDivisionByZeroFails) { + 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)); +} + +TEST_F(OpRemainderOutTest, ScalarIntegerDivisionByZeroFails) { + TensorFactory tf; + + Tensor self = tf.make({2}, {5, -3}); + Tensor out = tf.zeros({2}); + 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, ScalarIsCastToUnsignedDtype) { + TensorFactory tf; + + Tensor self = tf.make({3}, {5, 255, 0}); + Tensor out = tf.full({3}, 1); + // -1 becomes 255 as uint8. + op_remainder_scalar_out(self, -1, out); + EXPECT_TENSOR_EQ(out, tf.make({3}, {5, 0, 0})); + // 256 becomes 0 as uint8. + 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, 0, out); + EXPECT_TENSOR_EQ(out, tf.make({0}, {})); +}