From 0e091153494ca2defd285a323a89bc1b8c847a89 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 15 Sep 2026 09:46:53 -0700 Subject: [PATCH] Define quantization for non-finite values (#22811) Summary: Host replay sanitizers exposed undefined behavior when non-finite quantization results were converted to integers before clamping. Define deterministic behavior: NaN maps to the clamped zero point, positive infinity to quant_max, and negative infinity to quant_min. Clamp before integer conversion and keep scalar and ARM exceptional-value handling consistent, with regression coverage for non-finite inputs, finite overflow, zero-point clamping, and per-channel quantization. Adapted from the repair proposed in D119101166. Edits were made in fbcode; the repository commit hook mirrored them into xplat, and the mirrored changes are retained. Authored with assistance from Devmate. Reviewed By: rascani Differential Revision: D120017028 --- kernels/quantized/cpu/op_quantize.cpp | 73 +++++++++----- kernels/quantized/test/op_quantize_test.cpp | 101 ++++++++++++++++++++ 2 files changed, 150 insertions(+), 24 deletions(-) diff --git a/kernels/quantized/cpu/op_quantize.cpp b/kernels/quantized/cpu/op_quantize.cpp index e52b9a371e6..72f08435665 100644 --- a/kernels/quantized/cpu/op_quantize.cpp +++ b/kernels/quantized/cpu/op_quantize.cpp @@ -92,6 +92,25 @@ void check_quantize_per_tensor_args( } // namespace +template +T quantize_val_with_inv_scale( + float inv_scale, + int64_t zero_point, + K value, + int64_t quant_min, + int64_t quant_max) { + const float scaled_value = static_cast(inv_scale * value); + // Map NaN to the zero point and clamp infinities before integer conversion. + // Double precision keeps int32 bounds exactly representable while clamping. + const double qvalue = std::isnan(scaled_value) + ? static_cast(zero_point) + : static_cast(zero_point) + std::nearbyint(scaled_value); + const double clamped_value = std::min( + static_cast(quant_max), + std::max(static_cast(quant_min), qvalue)); + return static_cast(clamped_value); +} + template T quantize_val( double scale, @@ -99,15 +118,12 @@ T quantize_val( K value, int64_t quant_min, int64_t quant_max) { - int64_t qvalue; - float inv_scale = 1.0f / static_cast(scale); - qvalue = static_cast( - static_cast(zero_point) + - std::nearbyint(static_cast(inv_scale * value))); - - qvalue = std::max(qvalue, quant_min); - qvalue = std::min(qvalue, quant_max); - return static_cast(qvalue); + return quantize_val_with_inv_scale( + 1.0f / static_cast(scale), + zero_point, + value, + quant_min, + quant_max); } #if defined(__aarch64__) || defined(__ARM_NEON__) @@ -180,7 +196,8 @@ void quantize_arm( const float32x4_t vin0123 = vld1q_f32(in + i); const float32x4_t vin4567 = vld1q_f32(in + i + 4); - // Multiply by inv_scale and round + // FCVTNS maps NaN to zero and saturates infinities before adding + // zero_point. const int32x4_t v0123_rounded = vcvtnq_s32_f32(vmulq_f32(vin0123, vinv_scale)); const int32x4_t v4567_rounded = @@ -201,16 +218,15 @@ void quantize_arm( // Handle remaining elements with proper quant_min/quant_max clamping for (; i < N; ++i) { - float val = in[i] * inv_scale; - int32_t qval = static_cast(std::nearbyint(val)) + zero_point; - qval = std::max(quant_min, std::min(quant_max, qval)); - out[i] = static_cast(qval); + out[i] = quantize_val_with_inv_scale( + inv_scale, zero_point, in[i], quant_min, quant_max); } #else // ARMv7: Use magic float rounding const int32x4_t voffset = vdupq_n_s32(zero_point - 0x4B400000); const float32x4_t vmagic_float = vdupq_n_f32(12582912.0f); + const float32x4_t vmagic_float_maximum = vdupq_n_f32(4194304.0f); int64_t i = 0; // Process 8 elements at a time @@ -218,14 +234,25 @@ void quantize_arm( const float32x4_t vin0123 = vld1q_f32(in + i); const float32x4_t vin4567 = vld1q_f32(in + i + 4); + const float32x4_t vscaled0123 = vmulq_f32(vin0123, vinv_scale); + const float32x4_t vscaled4567 = vmulq_f32(vin4567, vinv_scale); + const uint32x4_t valid = vandq_u32( + vcleq_f32(vabsq_f32(vscaled0123), vmagic_float_maximum), + vcleq_f32(vabsq_f32(vscaled4567), vmagic_float_maximum)); + const uint32x2_t valid_pairs = + vand_u32(vget_low_u32(valid), vget_high_u32(valid)); + if (vget_lane_u32(vpmin_u32(valid_pairs, valid_pairs), 0) != UINT32_MAX) { + for (int64_t j = 0; j < 8; ++j) { + out[i + j] = quantize_val_with_inv_scale( + inv_scale, zero_point, in[i + j], quant_min, quant_max); + } + continue; + } + const int32x4_t vraw0123 = vaddq_s32( - voffset, - vreinterpretq_s32_f32( - vaddq_f32(vmagic_float, vmulq_f32(vin0123, vinv_scale)))); + voffset, vreinterpretq_s32_f32(vaddq_f32(vmagic_float, vscaled0123))); const int32x4_t vraw4567 = vaddq_s32( - voffset, - vreinterpretq_s32_f32( - vaddq_f32(vmagic_float, vmulq_f32(vin4567, vinv_scale)))); + voffset, vreinterpretq_s32_f32(vaddq_f32(vmagic_float, vscaled4567))); const int16x8_t vraw01234567 = vcombine_s16(vqmovn_s32(vraw0123), vqmovn_s32(vraw4567)); @@ -237,10 +264,8 @@ void quantize_arm( // Handle remaining elements with proper quant_min/quant_max clamping for (; i < N; ++i) { - float val = in[i] * inv_scale; - int32_t qval = static_cast(std::nearbyint(val)) + zero_point; - qval = std::max(quant_min, std::min(quant_max, qval)); - out[i] = static_cast(qval); + out[i] = quantize_val_with_inv_scale( + inv_scale, zero_point, in[i], quant_min, quant_max); } #endif } diff --git a/kernels/quantized/test/op_quantize_test.cpp b/kernels/quantized/test/op_quantize_test.cpp index 79629b461e4..97d3bb43285 100644 --- a/kernels/quantized/test/op_quantize_test.cpp +++ b/kernels/quantized/test/op_quantize_test.cpp @@ -14,6 +14,7 @@ #include #include +#include using namespace ::testing; using executorch::aten::ScalarType; @@ -182,6 +183,106 @@ TEST(OpQuantizeOutTest, TestOutOfBounds) { EXPECT_TENSOR_EQ(out, expected); } +template +void test_non_finite_values() { + TensorFactory tf; + TensorFactory tfo; + const float nan = std::numeric_limits::quiet_NaN(); + const float inf = std::numeric_limits::infinity(); + // One SIMD block and a remainder containing each non-finite value. + Tensor input = + tf.make({11}, {nan, inf, -inf, 0, 1, -1, 0.25, 0.75, nan, inf, -inf}); + Tensor out = tfo.zeros({11}); + Tensor expected = tfo.make({11}, {5, 10, 0, 5, 7, 3, 5, 7, 5, 10, 0}); + + quantize_per_tensor_out(input, 0.5, 5, 0, 10, OUTPUT_DTYPE, out); + + EXPECT_TENSOR_EQ(out, expected); +} + +TEST(OpQuantizeOutTest, NonFiniteValuesHaveDeterministicOutputs) { + test_non_finite_values(); + test_non_finite_values(); + test_non_finite_values(); + test_non_finite_values(); + test_non_finite_values(); + test_non_finite_values(); + test_non_finite_values(); + test_non_finite_values(); +} + +TEST(OpQuantizeOutTest, NaNMapsToClampedZeroPoint) { + TensorFactory tf; + TensorFactory tfo; + Tensor input = tf.full({9}, std::numeric_limits::quiet_NaN()); + Tensor out = tfo.zeros({9}); + + for (const int64_t zero_point : {-20, 20}) { + SCOPED_TRACE(zero_point); + Tensor expected = tfo.full({9}, zero_point < 0 ? -10 : 10); + quantize_per_tensor_out( + input, 0.5, zero_point, -10, 10, ScalarType::Char, out); + EXPECT_TENSOR_EQ(out, expected); + } +} + +TEST(OpQuantizeOutTest, LargeFiniteValuesSaturateBeforeIntegerConversion) { + TensorFactory tf; + const float largest = std::numeric_limits::max(); + Tensor input = tf.make( + {10}, + {largest, + -largest, + 1e20f, + -1e20f, + 1e10f, + -1e10f, + 0, + 1, + largest, + -largest}); + TensorFactory tf_char; + Tensor out_char = tf_char.zeros({10}); + Tensor expected_char = + tf_char.make({10}, {10, -10, 10, -10, 10, -10, 5, 7, 10, -10}); + + quantize_per_tensor_out(input, 0.5, 5, -10, 10, ScalarType::Char, out_char); + EXPECT_TENSOR_EQ(out_char, expected_char); + + TensorFactory tf_int; + constexpr int32_t min = std::numeric_limits::min(); + constexpr int32_t max = std::numeric_limits::max(); + Tensor out_int = tf_int.zeros({10}); + Tensor expected_int = + tf_int.make({10}, {max, min, max, min, max, min, 5, 7, max, min}); + + quantize_per_tensor_out(input, 0.5, 5, min, max, ScalarType::Int, out_int); + EXPECT_TENSOR_EQ(out_int, expected_int); +} + +TEST(OpQuantizeOutTest, PerChannelNonFiniteValuesHaveDeterministicOutputs) { + TensorFactory tf; + TensorFactory tf_double; + TensorFactory tf_long; + TensorFactory tfo; + const float nan = std::numeric_limits::quiet_NaN(); + const float inf = std::numeric_limits::infinity(); + Tensor input = + tf.make({2, 11}, {nan, inf, -inf, 0, 1, -1, 0.25, 0.75, nan, inf, -inf, + nan, inf, -inf, 0, 1, -1, 0.25, 0.75, nan, inf, -inf}); + Tensor scale = tf_double.make({2}, {0.5, 1.0}); + Tensor zero_point = tf_long.make({2}, {5, -5}); + Tensor out = tfo.zeros({2, 11}); + Tensor expected = + tfo.make({2, 11}, {5, 10, -10, 5, 7, 3, 5, 7, 5, 10, -10, + -5, 10, -10, -5, -4, -6, -5, -4, -5, 10, -10}); + + quantize_per_channel_out( + input, scale, zero_point, 0, -10, 10, ScalarType::Char, out); + + EXPECT_TENSOR_EQ(out, expected); +} + TEST(OpQuantizeOutTest, QuantizePerChannel) { TensorFactory tf_float; TensorFactory tf_double;