diff --git a/kernels/portable/cpu/op_div.cpp b/kernels/portable/cpu/op_div.cpp index 700cefcd55c..f896e4151fc 100644 --- a/kernels/portable/cpu/op_div.cpp +++ b/kernels/portable/cpu/op_div.cpp @@ -150,7 +150,11 @@ Tensor& div_out_mode( } CTYPE_COMPUTE value = val_a / val_b; if (mode_is_trunc) { - value = std::trunc(value); + // Integer division already truncates, and std::trunc would round + // int64 values through double. + if constexpr (!std::is_integral_v) { + value = std::trunc(value); + } } else { // We established above that the mode is either trunc or floor, so // it must be floor. @@ -245,15 +249,6 @@ Tensor& div_scalar_mode_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, - "Div mode operation encountered integer division by zero"); - // Check Dim Order ET_KERNEL_CHECK( ctx, tensors_have_same_dim_order(a, out), InvalidArgument, out); @@ -270,13 +265,25 @@ Tensor& div_scalar_mode_out( // @lint-ignore CLANGTIDY facebook-hte-CArray ET_DEFINE_OPERATOR_NAME(op_name, "div.Scalar_mode_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) { + div_by_zero_error = true; + return; + } + } utils::apply_unitensor_elementwise_fn( [val_b, mode_is_trunc](const CTYPE_COMPUTE val_a) { CTYPE_COMPUTE value = val_a / val_b; if (mode_is_trunc) { - value = std::trunc(value); + if constexpr (!std::is_integral_v) { + value = std::trunc(value); + } } else { value = utils::floor_divide(val_a, val_b); } @@ -289,6 +296,13 @@ Tensor& div_scalar_mode_out( utils::SupportedTensorDtypes::REALHBF16); }); + ET_KERNEL_CHECK_MSG( + ctx, + !div_by_zero_error, + InvalidArgument, + out, + "Div mode operation encountered integer division by zero"); + return out; } diff --git a/kernels/portable/cpu/op_fmod.cpp b/kernels/portable/cpu/op_fmod.cpp index de0d44fd2f1..57b2bb1e0c1 100644 --- a/kernels/portable/cpu/op_fmod.cpp +++ b/kernels/portable/cpu/op_fmod.cpp @@ -10,6 +10,7 @@ #include #include +#include #include namespace torch { @@ -45,32 +46,26 @@ Tensor& fmod_Tensor_out( // Compute Dtype ScalarType compute_type = utils::get_compute_type(common_type); - if (compute_type != ScalarType::Float) { - compute_type = ScalarType::Double; - } // @lint-ignore CLANGTIDY facebook-hte-CArray ET_DEFINE_OPERATOR_NAME(op_name, "fmod.Tensor_out"); bool div_by_zero_error = false; - ET_SWITCH_FLOAT_TYPES(compute_type, ctx, op_name, CTYPE_COMPUTE, [&]() { + ET_SWITCH_REAL_TYPES(compute_type, ctx, op_name, CTYPE_COMPUTE, [&]() { utils::apply_bitensor_elementwise_fn< CTYPE_COMPUTE, op_name, utils::SupportedTensorDtypes::REALHBF16>( [&div_by_zero_error]( const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b) { - // TODO: rewrite this to be vectorization-capable? - CTYPE_COMPUTE value = 0; if (is_integral_type::value) { if (val_b == 0) { div_by_zero_error = true; - return value; + return static_cast(0); } } - value = std::fmod(val_a, val_b); - return value; + return utils::fmod_override(val_a, val_b); }, ctx, a, @@ -106,15 +101,6 @@ Tensor& fmod_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, - "Fmod operation encountered integer division by zero"); - // Check Dim Order ET_KERNEL_CHECK( ctx, tensors_have_same_dim_order(a, out), InvalidArgument, out); @@ -125,13 +111,42 @@ Tensor& fmod_Scalar_out( // Compute Dtype ScalarType compute_type = utils::get_compute_type(common_type); - if (compute_type != ScalarType::Float) { - compute_type = ScalarType::Double; - } // @lint-ignore CLANGTIDY facebook-hte-CArray ET_DEFINE_OPERATOR_NAME(op_name, "fmod.Scalar_out"); + if (executorch::runtime::isIntegralType( + compute_type, /*includeBool=*/false)) { + bool div_by_zero_error = false; + ET_SWITCH_INT_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 (val_b == 0) { + div_by_zero_error = true; + return; + } + utils::apply_unitensor_elementwise_fn< + CTYPE_COMPUTE, + op_name, + utils::SupportedTensorDtypes::REALHBF16>( + [val_b](const CTYPE_COMPUTE val_a) { + return utils::fmod_override(val_a, val_b); + }, + ctx, + a, + utils::SupportedTensorDtypes::REALHBBF16, + out); + }); + ET_KERNEL_CHECK_MSG( + ctx, + !div_by_zero_error, + InvalidArgument, + out, + "Fmod operation encountered integer division by zero"); + return out; + } + ET_SWITCH_FLOAT_TYPES(compute_type, ctx, op_name, CTYPE_COMPUTE, [&]() { const CTYPE_COMPUTE val_b = utils::scalar_to(b); utils::apply_unitensor_elementwise_fn< diff --git a/kernels/portable/cpu/util/math_util.h b/kernels/portable/cpu/util/math_util.h index 9499eec671e..41ad2f82898 100644 --- a/kernels/portable/cpu/util/math_util.h +++ b/kernels/portable/cpu/util/math_util.h @@ -230,6 +230,32 @@ CTYPE remainder_override(CTYPE a, CTYPE b) { return rem; } +/** + * std::fmod, with integers computed as integers so that int64 values above + * 2^53 are not rounded through double. Like std::fmod, % keeps the sign of a. + * a % -1 overflows for the most negative signed value, and the result is 0 + * for every a. + */ +template < + typename CTYPE, + typename std::enable_if::value, int>::type = + 0> +CTYPE fmod_override(CTYPE a, CTYPE b) { + return std::fmod(a, b); +} + +template < + typename CTYPE, + typename std::enable_if::value, int>::type = 0> +CTYPE fmod_override(CTYPE a, CTYPE b) { + if constexpr (std::is_signed::value) { + if (b == -1) { + return 0; + } + } + return static_cast(a % b); +} + } // namespace utils } // namespace native } // namespace executor diff --git a/kernels/test/op_div_test.cpp b/kernels/test/op_div_test.cpp index d313e85a338..028757fa790 100644 --- a/kernels/test/op_div_test.cpp +++ b/kernels/test/op_div_test.cpp @@ -562,6 +562,22 @@ TEST_F(OpDivOutTest, DISABLED_DynamicShapeUnbound) { EXPECT_TENSOR_CLOSE(out, expected_result); } +TEST_F(OpDivOutTest, TruncModeKeepsInt64Precision) { + TensorFactory tf; + // Values above 2^53, which a double cannot hold exactly. + Tensor a = tf.make( + {3}, + {(int64_t{1} << 62) + 3, -((int64_t{1} << 62) + 3), 9007199254740993}); + Tensor b = tf.make({3}, {3, 3, 1}); + Tensor out = tf.zeros({3}); + + torch::executor::aten::div_outf(context_, a, b, "trunc", out); + EXPECT_TENSOR_EQ( + out, + tf.make( + {3}, {1537228672809129302, -1537228672809129302, 9007199254740993})); +} + TEST_F(OpDivScalarOutTest, SanityCheckIntScalar) { TensorFactory tf_a; TensorFactory tf_out; @@ -590,6 +606,39 @@ TEST_F(OpDivScalarOutTest, SanityCheckFloatScalar) { EXPECT_TENSOR_EQ(out, tf_out.make(sizes, {0.5, 1.0, 2.0, -4.5})); } +TEST_F(OpDivScalarOutTest, TruncModeKeepsInt64Precision) { + TensorFactory tf; + // Values above 2^53, which a double cannot hold exactly. + Tensor a = tf.make( + {3}, + {(int64_t{1} << 62) + 3, -((int64_t{1} << 62) + 3), 9007199254740993}); + Tensor out = tf.zeros({3}); + + torch::executor::aten::div_outf(context_, a, Scalar(3), "trunc", out); + EXPECT_TENSOR_EQ( + out, + tf.make( + {3}, {1537228672809129302, -1537228672809129302, 3002399751580331})); +} + +TEST_F(OpDivScalarOutTest, ModeIntegerScalarIsCastToTensorDtype) { + TensorFactory tf; + Tensor a = tf.make({3}, {5, -7, 100}); + Tensor out = tf.zeros({3}); + + // Like ATen, the scalar takes the tensor's dtype: 255 is -1 as int8. + torch::executor::aten::div_outf(context_, a, Scalar(255), "trunc", out); + EXPECT_TENSOR_EQ(out, tf.make({3}, {-5, 7, -100})); + + // 256 is 0 as int8. + ET_EXPECT_KERNEL_FAILURE( + context_, + torch::executor::aten::div_outf(context_, a, Scalar(256), "trunc", out)); + ET_EXPECT_KERNEL_FAILURE( + context_, + torch::executor::aten::div_outf(context_, a, Scalar(256), "floor", out)); +} + TEST_F(OpDivScalarOutTest, OptimizedSanityCheck) { TensorFactory tf; diff --git a/kernels/test/op_fmod_test.cpp b/kernels/test/op_fmod_test.cpp index 3227a01a17a..9b1fd8547ee 100644 --- a/kernels/test/op_fmod_test.cpp +++ b/kernels/test/op_fmod_test.cpp @@ -13,6 +13,7 @@ #include #include +#include using namespace ::testing; using executorch::aten::Scalar; @@ -31,6 +32,19 @@ class OpFmodTest : public OperatorTest { op_fmod_scalar_out(const Tensor& self, const Scalar& other, Tensor& out) { return torch::executor::aten::fmod_outf(context_, self, other, out); } + + template + void test_integer_fmod() { + TensorFactory tf; + Tensor self = tf.make({4}, {0, 7, 100, 127}); + Tensor out = tf.zeros({4}); + + op_fmod_tensor_out(self, tf.make({4}, {3, 3, 7, 10}), out); + EXPECT_TENSOR_EQ(out, tf.make({4}, {0, 1, 2, 7})); + + op_fmod_scalar_out(self, 3, out); + EXPECT_TENSOR_EQ(out, tf.make({4}, {0, 1, 1, 1})); + } }; TEST_F(OpFmodTest, SmokeTest) { @@ -46,6 +60,78 @@ TEST_F(OpFmodTest, SmokeTest) { EXPECT_TENSOR_CLOSE(out, out_expected); } +TEST_F(OpFmodTest, Int64KeepsFullPrecision) { + TensorFactory tf; + // Values above 2^53, which a double cannot hold exactly. + Tensor self = tf.make( + {4}, + {(int64_t{1} << 62) + 3, + -((int64_t{1} << 62) + 3), + 9007199254740993, + std::numeric_limits::min()}); + Tensor out = tf.zeros({4}); + + op_fmod_tensor_out(self, tf.make({4}, {10, 10, 2, -1}), out); + EXPECT_TENSOR_EQ(out, tf.make({4}, {7, -7, 1, 0})); + + op_fmod_scalar_out(self, 10, out); + EXPECT_TENSOR_EQ(out, tf.make({4}, {7, -7, 3, -8})); + + op_fmod_scalar_out(self, -1, out); + EXPECT_TENSOR_EQ(out, tf.zeros({4})); +} + +TEST_F(OpFmodTest, AllIntegerDtypes) { +#define TEST_ENTRY(ctype, dtype) test_integer_fmod(); + ET_FORALL_INT_TYPES(TEST_ENTRY); +#undef TEST_ENTRY +} + +TEST_F(OpFmodTest, NegativeIntegers) { + TensorFactory tf; + Tensor self = tf.make({4}, {-7, 7, -7, -128}); + Tensor out = tf.zeros({4}); + + op_fmod_tensor_out(self, tf.make({4}, {3, -3, -3, -1}), out); + EXPECT_TENSOR_EQ(out, tf.make({4}, {-1, 1, -1, 0})); +} + +TEST_F(OpFmodTest, IntegerDivisionByZeroDies) { + TensorFactory tf; + Tensor self = tf.make({2}, {5, 6}); + Tensor out = tf.zeros({2}); + + ET_EXPECT_KERNEL_FAILURE( + context_, op_fmod_tensor_out(self, tf.make({2}, {2, 0}), out)); + ET_EXPECT_KERNEL_FAILURE(context_, op_fmod_scalar_out(self, 0, out)); +} + +TEST_F(OpFmodTest, IntegerScalarIsCastToTensorDtype) { + TensorFactory tf; + Tensor self = tf.make({3}, {5, -7, 100}); + Tensor out = tf.zeros({3}); + + // Like ATen, the scalar takes the tensor's dtype: 255 is -1 and 257 is 1 as + // int8. + op_fmod_scalar_out(self, 255, out); + EXPECT_TENSOR_EQ(out, tf.zeros({3})); + op_fmod_scalar_out(self, 257, out); + EXPECT_TENSOR_EQ(out, tf.zeros({3})); + + // 256 is 0 as int8, and 2^32 is 0 as int32. + ET_EXPECT_KERNEL_FAILURE(context_, op_fmod_scalar_out(self, 256, out)); + TensorFactory tf_byte; + Tensor out_byte = tf_byte.zeros({2}); + ET_EXPECT_KERNEL_FAILURE( + context_, op_fmod_scalar_out(tf_byte.make({2}, {5, 200}), 256, out_byte)); + TensorFactory tf_int; + Tensor out_int = tf_int.zeros({2}); + ET_EXPECT_KERNEL_FAILURE( + context_, + op_fmod_scalar_out( + tf_int.make({2}, {5, 6}), Scalar(int64_t{1} << 32), out_int)); +} + TEST_F(OpFmodTest, ScalarSmokeTest) { TensorFactory tfFloat; std::vector a(18); diff --git a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl index 1eee6f13a15..4ab7f461770 100644 --- a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl +++ b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl @@ -635,6 +635,7 @@ ATEN_OPS = ( "//executorch/kernels/portable/cpu/util:broadcast_util", "//executorch/kernels/portable/cpu/util:dtype_util", "//executorch/kernels/portable/cpu/util:elementwise_util", + "//executorch/kernels/portable/cpu/util:math_util", ], ), op_target(