Skip to content

Commit 186ca41

Browse files
usmanmateenusmanmateen
authored
Keep int64 precision in portable div trunc mode and fmod (pytorch#24170)
### Summary Fixes pytorch#24169. The portable `div.out_mode` and `div.Scalar_mode_out` kernels compute the integer quotient with `/` and then call `std::trunc` on it, which converts it to `double` and back. Integer division already truncates, so `std::trunc` now only runs for floating-point types. For int64 values above 2^53, `div(a, 3, rounding_mode="trunc")` gave `1537228672809129216` instead of `1537228672809129302`. `fmod.Tensor_out` and `fmod.Scalar_out` switched every integer compute type to `Double`. Integer types now use `%` through a new `utils::integer_fmod` in `math_util.h`, which returns 0 for a divisor of -1 so the most negative value can't overflow. Float types keep the existing `std::fmod` path, including the vectorized one for the Scalar variant. ### Test plan Added `OpDivOutTest.TruncModeKeepsInt64Precision` (Tensor and Scalar divisor) and `OpFmodTest.Int64KeepsFullPrecision` (Tensor, Scalar, and a divisor of -1 with `INT64_MIN`). The expected values come from eager PyTorch. I built `portable_kernels_test` with CMake (`EXECUTORCH_BUILD_TESTS=ON`, Release) on Linux x86_64 (WSL2) at 5b14a2f. Before this change both new tests fail; after it, the full `portable_kernels_test` run passes (1696 tests). `clang-format` 18.1.3 reports no changes on the changed files. This PR was authored with AI assistance (Claude Code). cc @nil-is-all @larryliu0820 @manuelcandales @JakeStevens
1 parent 249237c commit 186ca41

6 files changed

Lines changed: 223 additions & 32 deletions

File tree

‎kernels/portable/cpu/op_div.cpp‎

Lines changed: 25 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,11 @@ Tensor& div_out_mode(
150150
}
151151
CTYPE_COMPUTE value = val_a / val_b;
152152
if (mode_is_trunc) {
153-
value = std::trunc(value);
153+
// Integer division already truncates, and std::trunc would round
154+
// int64 values through double.
155+
if constexpr (!std::is_integral_v<CTYPE_COMPUTE>) {
156+
value = std::trunc(value);
157+
}
154158
} else {
155159
// We established above that the mode is either trunc or floor, so
156160
// it must be floor.
@@ -245,15 +249,6 @@ Tensor& div_scalar_mode_out(
245249
InvalidArgument,
246250
out);
247251

248-
// Check for intergral division by zero
249-
ET_KERNEL_CHECK_MSG(
250-
ctx,
251-
!(executorch::runtime::isIntegralType(common_type, true) &&
252-
utils::scalar_to<double>(b) == 0),
253-
InvalidArgument,
254-
out,
255-
"Div mode operation encountered integer division by zero");
256-
257252
// Check Dim Order
258253
ET_KERNEL_CHECK(
259254
ctx, tensors_have_same_dim_order(a, out), InvalidArgument, out);
@@ -270,13 +265,25 @@ Tensor& div_scalar_mode_out(
270265
// @lint-ignore CLANGTIDY facebook-hte-CArray
271266
ET_DEFINE_OPERATOR_NAME(op_name, "div.Scalar_mode_out");
272267

268+
bool div_by_zero_error = false;
269+
273270
ET_SWITCH_REAL_TYPES(compute_type, ctx, op_name, CTYPE_COMPUTE, [&]() {
271+
// Like ATen, the scalar takes the compute dtype first, so a scalar such as
272+
// 256 for an int8 tensor becomes 0.
274273
const CTYPE_COMPUTE val_b = utils::scalar_to<CTYPE_COMPUTE>(b);
274+
if constexpr (std::is_integral_v<CTYPE_COMPUTE>) {
275+
if (val_b == 0) {
276+
div_by_zero_error = true;
277+
return;
278+
}
279+
}
275280
utils::apply_unitensor_elementwise_fn<CTYPE_COMPUTE, op_name>(
276281
[val_b, mode_is_trunc](const CTYPE_COMPUTE val_a) {
277282
CTYPE_COMPUTE value = val_a / val_b;
278283
if (mode_is_trunc) {
279-
value = std::trunc(value);
284+
if constexpr (!std::is_integral_v<CTYPE_COMPUTE>) {
285+
value = std::trunc(value);
286+
}
280287
} else {
281288
value = utils::floor_divide(val_a, val_b);
282289
}
@@ -289,6 +296,13 @@ Tensor& div_scalar_mode_out(
289296
utils::SupportedTensorDtypes::REALHBF16);
290297
});
291298

299+
ET_KERNEL_CHECK_MSG(
300+
ctx,
301+
!div_by_zero_error,
302+
InvalidArgument,
303+
out,
304+
"Div mode operation encountered integer division by zero");
305+
292306
return out;
293307
}
294308

‎kernels/portable/cpu/op_fmod.cpp‎

Lines changed: 36 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
#include <executorch/kernels/portable/cpu/scalar_utils.h>
1212
#include <executorch/kernels/portable/cpu/util/elementwise_util.h>
13+
#include <executorch/kernels/portable/cpu/util/math_util.h>
1314
#include <executorch/runtime/kernel/kernel_includes.h>
1415

1516
namespace torch {
@@ -45,32 +46,26 @@ Tensor& fmod_Tensor_out(
4546

4647
// Compute Dtype
4748
ScalarType compute_type = utils::get_compute_type(common_type);
48-
if (compute_type != ScalarType::Float) {
49-
compute_type = ScalarType::Double;
50-
}
5149

5250
// @lint-ignore CLANGTIDY facebook-hte-CArray
5351
ET_DEFINE_OPERATOR_NAME(op_name, "fmod.Tensor_out");
5452

5553
bool div_by_zero_error = false;
5654

57-
ET_SWITCH_FLOAT_TYPES(compute_type, ctx, op_name, CTYPE_COMPUTE, [&]() {
55+
ET_SWITCH_REAL_TYPES(compute_type, ctx, op_name, CTYPE_COMPUTE, [&]() {
5856
utils::apply_bitensor_elementwise_fn<
5957
CTYPE_COMPUTE,
6058
op_name,
6159
utils::SupportedTensorDtypes::REALHBF16>(
6260
[&div_by_zero_error](
6361
const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b) {
64-
// TODO: rewrite this to be vectorization-capable?
65-
CTYPE_COMPUTE value = 0;
6662
if (is_integral_type<CTYPE_COMPUTE, /*includeBool=*/true>::value) {
6763
if (val_b == 0) {
6864
div_by_zero_error = true;
69-
return value;
65+
return static_cast<CTYPE_COMPUTE>(0);
7066
}
7167
}
72-
value = std::fmod(val_a, val_b);
73-
return value;
68+
return utils::fmod_override(val_a, val_b);
7469
},
7570
ctx,
7671
a,
@@ -106,15 +101,6 @@ Tensor& fmod_Scalar_out(
106101
InvalidArgument,
107102
out);
108103

109-
// Check for intergral division by zero
110-
ET_KERNEL_CHECK_MSG(
111-
ctx,
112-
!(executorch::runtime::isIntegralType(common_type, true) &&
113-
utils::scalar_to<double>(b) == 0),
114-
InvalidArgument,
115-
out,
116-
"Fmod operation encountered integer division by zero");
117-
118104
// Check Dim Order
119105
ET_KERNEL_CHECK(
120106
ctx, tensors_have_same_dim_order(a, out), InvalidArgument, out);
@@ -125,13 +111,42 @@ Tensor& fmod_Scalar_out(
125111

126112
// Compute Dtype
127113
ScalarType compute_type = utils::get_compute_type(common_type);
128-
if (compute_type != ScalarType::Float) {
129-
compute_type = ScalarType::Double;
130-
}
131114

132115
// @lint-ignore CLANGTIDY facebook-hte-CArray
133116
ET_DEFINE_OPERATOR_NAME(op_name, "fmod.Scalar_out");
134117

118+
if (executorch::runtime::isIntegralType(
119+
compute_type, /*includeBool=*/false)) {
120+
bool div_by_zero_error = false;
121+
ET_SWITCH_INT_TYPES(compute_type, ctx, op_name, CTYPE_COMPUTE, [&]() {
122+
// Like ATen, the scalar takes the compute dtype first, so a scalar such
123+
// as 256 for an int8 tensor becomes 0.
124+
const CTYPE_COMPUTE val_b = utils::scalar_to<CTYPE_COMPUTE>(b);
125+
if (val_b == 0) {
126+
div_by_zero_error = true;
127+
return;
128+
}
129+
utils::apply_unitensor_elementwise_fn<
130+
CTYPE_COMPUTE,
131+
op_name,
132+
utils::SupportedTensorDtypes::REALHBF16>(
133+
[val_b](const CTYPE_COMPUTE val_a) {
134+
return utils::fmod_override(val_a, val_b);
135+
},
136+
ctx,
137+
a,
138+
utils::SupportedTensorDtypes::REALHBBF16,
139+
out);
140+
});
141+
ET_KERNEL_CHECK_MSG(
142+
ctx,
143+
!div_by_zero_error,
144+
InvalidArgument,
145+
out,
146+
"Fmod operation encountered integer division by zero");
147+
return out;
148+
}
149+
135150
ET_SWITCH_FLOAT_TYPES(compute_type, ctx, op_name, CTYPE_COMPUTE, [&]() {
136151
const CTYPE_COMPUTE val_b = utils::scalar_to<CTYPE_COMPUTE>(b);
137152
utils::apply_unitensor_elementwise_fn<

‎kernels/portable/cpu/util/math_util.h‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,32 @@ CTYPE remainder_override(CTYPE a, CTYPE b) {
230230
return rem;
231231
}
232232

233+
/**
234+
* std::fmod, with integers computed as integers so that int64 values above
235+
* 2^53 are not rounded through double. Like std::fmod, % keeps the sign of a.
236+
* a % -1 overflows for the most negative signed value, and the result is 0
237+
* for every a.
238+
*/
239+
template <
240+
typename CTYPE,
241+
typename std::enable_if<std::is_floating_point<CTYPE>::value, int>::type =
242+
0>
243+
CTYPE fmod_override(CTYPE a, CTYPE b) {
244+
return std::fmod(a, b);
245+
}
246+
247+
template <
248+
typename CTYPE,
249+
typename std::enable_if<std::is_integral<CTYPE>::value, int>::type = 0>
250+
CTYPE fmod_override(CTYPE a, CTYPE b) {
251+
if constexpr (std::is_signed<CTYPE>::value) {
252+
if (b == -1) {
253+
return 0;
254+
}
255+
}
256+
return static_cast<CTYPE>(a % b);
257+
}
258+
233259
} // namespace utils
234260
} // namespace native
235261
} // namespace executor

‎kernels/test/op_div_test.cpp‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -562,6 +562,22 @@ TEST_F(OpDivOutTest, DISABLED_DynamicShapeUnbound) {
562562
EXPECT_TENSOR_CLOSE(out, expected_result);
563563
}
564564

565+
TEST_F(OpDivOutTest, TruncModeKeepsInt64Precision) {
566+
TensorFactory<ScalarType::Long> tf;
567+
// Values above 2^53, which a double cannot hold exactly.
568+
Tensor a = tf.make(
569+
{3},
570+
{(int64_t{1} << 62) + 3, -((int64_t{1} << 62) + 3), 9007199254740993});
571+
Tensor b = tf.make({3}, {3, 3, 1});
572+
Tensor out = tf.zeros({3});
573+
574+
torch::executor::aten::div_outf(context_, a, b, "trunc", out);
575+
EXPECT_TENSOR_EQ(
576+
out,
577+
tf.make(
578+
{3}, {1537228672809129302, -1537228672809129302, 9007199254740993}));
579+
}
580+
565581
TEST_F(OpDivScalarOutTest, SanityCheckIntScalar) {
566582
TensorFactory<ScalarType::Int> tf_a;
567583
TensorFactory<ScalarType::Float> tf_out;
@@ -590,6 +606,39 @@ TEST_F(OpDivScalarOutTest, SanityCheckFloatScalar) {
590606
EXPECT_TENSOR_EQ(out, tf_out.make(sizes, {0.5, 1.0, 2.0, -4.5}));
591607
}
592608

609+
TEST_F(OpDivScalarOutTest, TruncModeKeepsInt64Precision) {
610+
TensorFactory<ScalarType::Long> tf;
611+
// Values above 2^53, which a double cannot hold exactly.
612+
Tensor a = tf.make(
613+
{3},
614+
{(int64_t{1} << 62) + 3, -((int64_t{1} << 62) + 3), 9007199254740993});
615+
Tensor out = tf.zeros({3});
616+
617+
torch::executor::aten::div_outf(context_, a, Scalar(3), "trunc", out);
618+
EXPECT_TENSOR_EQ(
619+
out,
620+
tf.make(
621+
{3}, {1537228672809129302, -1537228672809129302, 3002399751580331}));
622+
}
623+
624+
TEST_F(OpDivScalarOutTest, ModeIntegerScalarIsCastToTensorDtype) {
625+
TensorFactory<ScalarType::Char> tf;
626+
Tensor a = tf.make({3}, {5, -7, 100});
627+
Tensor out = tf.zeros({3});
628+
629+
// Like ATen, the scalar takes the tensor's dtype: 255 is -1 as int8.
630+
torch::executor::aten::div_outf(context_, a, Scalar(255), "trunc", out);
631+
EXPECT_TENSOR_EQ(out, tf.make({3}, {-5, 7, -100}));
632+
633+
// 256 is 0 as int8.
634+
ET_EXPECT_KERNEL_FAILURE(
635+
context_,
636+
torch::executor::aten::div_outf(context_, a, Scalar(256), "trunc", out));
637+
ET_EXPECT_KERNEL_FAILURE(
638+
context_,
639+
torch::executor::aten::div_outf(context_, a, Scalar(256), "floor", out));
640+
}
641+
593642
TEST_F(OpDivScalarOutTest, OptimizedSanityCheck) {
594643
TensorFactory<ScalarType::Float> tf;
595644

‎kernels/test/op_fmod_test.cpp‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
#include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
1414

1515
#include <gtest/gtest.h>
16+
#include <limits>
1617

1718
using namespace ::testing;
1819
using executorch::aten::Scalar;
@@ -31,6 +32,19 @@ class OpFmodTest : public OperatorTest {
3132
op_fmod_scalar_out(const Tensor& self, const Scalar& other, Tensor& out) {
3233
return torch::executor::aten::fmod_outf(context_, self, other, out);
3334
}
35+
36+
template <ScalarType DTYPE>
37+
void test_integer_fmod() {
38+
TensorFactory<DTYPE> tf;
39+
Tensor self = tf.make({4}, {0, 7, 100, 127});
40+
Tensor out = tf.zeros({4});
41+
42+
op_fmod_tensor_out(self, tf.make({4}, {3, 3, 7, 10}), out);
43+
EXPECT_TENSOR_EQ(out, tf.make({4}, {0, 1, 2, 7}));
44+
45+
op_fmod_scalar_out(self, 3, out);
46+
EXPECT_TENSOR_EQ(out, tf.make({4}, {0, 1, 1, 1}));
47+
}
3448
};
3549

3650
TEST_F(OpFmodTest, SmokeTest) {
@@ -46,6 +60,78 @@ TEST_F(OpFmodTest, SmokeTest) {
4660
EXPECT_TENSOR_CLOSE(out, out_expected);
4761
}
4862

63+
TEST_F(OpFmodTest, Int64KeepsFullPrecision) {
64+
TensorFactory<ScalarType::Long> tf;
65+
// Values above 2^53, which a double cannot hold exactly.
66+
Tensor self = tf.make(
67+
{4},
68+
{(int64_t{1} << 62) + 3,
69+
-((int64_t{1} << 62) + 3),
70+
9007199254740993,
71+
std::numeric_limits<int64_t>::min()});
72+
Tensor out = tf.zeros({4});
73+
74+
op_fmod_tensor_out(self, tf.make({4}, {10, 10, 2, -1}), out);
75+
EXPECT_TENSOR_EQ(out, tf.make({4}, {7, -7, 1, 0}));
76+
77+
op_fmod_scalar_out(self, 10, out);
78+
EXPECT_TENSOR_EQ(out, tf.make({4}, {7, -7, 3, -8}));
79+
80+
op_fmod_scalar_out(self, -1, out);
81+
EXPECT_TENSOR_EQ(out, tf.zeros({4}));
82+
}
83+
84+
TEST_F(OpFmodTest, AllIntegerDtypes) {
85+
#define TEST_ENTRY(ctype, dtype) test_integer_fmod<ScalarType::dtype>();
86+
ET_FORALL_INT_TYPES(TEST_ENTRY);
87+
#undef TEST_ENTRY
88+
}
89+
90+
TEST_F(OpFmodTest, NegativeIntegers) {
91+
TensorFactory<ScalarType::Char> tf;
92+
Tensor self = tf.make({4}, {-7, 7, -7, -128});
93+
Tensor out = tf.zeros({4});
94+
95+
op_fmod_tensor_out(self, tf.make({4}, {3, -3, -3, -1}), out);
96+
EXPECT_TENSOR_EQ(out, tf.make({4}, {-1, 1, -1, 0}));
97+
}
98+
99+
TEST_F(OpFmodTest, IntegerDivisionByZeroDies) {
100+
TensorFactory<ScalarType::Int> tf;
101+
Tensor self = tf.make({2}, {5, 6});
102+
Tensor out = tf.zeros({2});
103+
104+
ET_EXPECT_KERNEL_FAILURE(
105+
context_, op_fmod_tensor_out(self, tf.make({2}, {2, 0}), out));
106+
ET_EXPECT_KERNEL_FAILURE(context_, op_fmod_scalar_out(self, 0, out));
107+
}
108+
109+
TEST_F(OpFmodTest, IntegerScalarIsCastToTensorDtype) {
110+
TensorFactory<ScalarType::Char> tf;
111+
Tensor self = tf.make({3}, {5, -7, 100});
112+
Tensor out = tf.zeros({3});
113+
114+
// Like ATen, the scalar takes the tensor's dtype: 255 is -1 and 257 is 1 as
115+
// int8.
116+
op_fmod_scalar_out(self, 255, out);
117+
EXPECT_TENSOR_EQ(out, tf.zeros({3}));
118+
op_fmod_scalar_out(self, 257, out);
119+
EXPECT_TENSOR_EQ(out, tf.zeros({3}));
120+
121+
// 256 is 0 as int8, and 2^32 is 0 as int32.
122+
ET_EXPECT_KERNEL_FAILURE(context_, op_fmod_scalar_out(self, 256, out));
123+
TensorFactory<ScalarType::Byte> tf_byte;
124+
Tensor out_byte = tf_byte.zeros({2});
125+
ET_EXPECT_KERNEL_FAILURE(
126+
context_, op_fmod_scalar_out(tf_byte.make({2}, {5, 200}), 256, out_byte));
127+
TensorFactory<ScalarType::Int> tf_int;
128+
Tensor out_int = tf_int.zeros({2});
129+
ET_EXPECT_KERNEL_FAILURE(
130+
context_,
131+
op_fmod_scalar_out(
132+
tf_int.make({2}, {5, 6}), Scalar(int64_t{1} << 32), out_int));
133+
}
134+
49135
TEST_F(OpFmodTest, ScalarSmokeTest) {
50136
TensorFactory<ScalarType::Float> tfFloat;
51137
std::vector<float> a(18);

‎shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -635,6 +635,7 @@ ATEN_OPS = (
635635
"//executorch/kernels/portable/cpu/util:broadcast_util",
636636
"//executorch/kernels/portable/cpu/util:dtype_util",
637637
"//executorch/kernels/portable/cpu/util:elementwise_util",
638+
"//executorch/kernels/portable/cpu/util:math_util",
638639
],
639640
),
640641
op_target(

0 commit comments

Comments
 (0)