Fix integer divisor handling in portable remainder

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<double>(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<CTYPE_COMPUTE>(b);
+    if constexpr (std::is_integral_v<CTYPE_COMPUTE>) {
+      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;
 }
 

kernels/portable/cpu/util/math_util.h

@@ -221,6 +221,12 @@ template <
     typename CTYPE,
     typename std::enable_if<std::is_integral<CTYPE>::value, int>::type = 0>
 CTYPE remainder_override(CTYPE a, CTYPE b) {
+  if constexpr (std::is_signed<CTYPE>::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<CTYPE>::value) {
     if (((a < 0) ^ (b < 0)) && rem != 0) {

kernels/test/op_remainder_test.cpp

@@ -14,6 +14,8 @@
 
 #include <gtest/gtest.h>
 
+#include <limits>
+
 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 <ScalarType DTYPE>
+  void test_most_negative_by_minus_one() {
+    TensorFactory<DTYPE> tf;
+    using CTYPE = typename TensorFactory<DTYPE>::ctype;
+
+    const CTYPE min = std::numeric_limits<CTYPE>::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<ScalarType::Int>();
+  test_most_negative_by_minus_one<ScalarType::Long>();
+}
+
+TEST_F(OpRemainderOutTest, TensorIntegerDivisionByZeroFails) {
+  TensorFactory<ScalarType::Int> 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<ScalarType::Int> 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<ScalarType::Char> 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<ScalarType::Byte> 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<ScalarType::Char> 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}, {}));
+}