Define quantization for non-finite values (pytorch#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

Pull Request resolved: pytorch#22811

Author: metascroy

kernels/quantized/cpu/op_quantize.cpp

@@ -92,22 +92,38 @@ void check_quantize_per_tensor_args(
 
 } // namespace
 
+template <typename T, typename K>
+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<float>(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<double>(zero_point)
+      : static_cast<double>(zero_point) + std::nearbyint(scaled_value);
+  const double clamped_value = std::min(
+      static_cast<double>(quant_max),
+      std::max(static_cast<double>(quant_min), qvalue));
+  return static_cast<T>(clamped_value);
+}
+
 template <typename T, typename K>
 T quantize_val(
     double scale,
     int64_t zero_point,
     K value,
     int64_t quant_min,
     int64_t quant_max) {
-  int64_t qvalue;
-  float inv_scale = 1.0f / static_cast<float>(scale);
-  qvalue = static_cast<int64_t>(
-      static_cast<int32_t>(zero_point) +
-      std::nearbyint(static_cast<float>(inv_scale * value)));
-
-  qvalue = std::max<int64_t>(qvalue, quant_min);
-  qvalue = std::min<int64_t>(qvalue, quant_max);
-  return static_cast<T>(qvalue);
+  return quantize_val_with_inv_scale<T>(
+      1.0f / static_cast<float>(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,31 +218,41 @@ 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<int32_t>(std::nearbyint(val)) + zero_point;
-    qval = std::max(quant_min, std::min(quant_max, qval));
-    out[i] = static_cast<T>(qval);
+    out[i] = quantize_val_with_inv_scale<T>(
+        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
   for (; i + 8 <= N; i += 8) {
     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<T>(
+            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<int32_t>(std::nearbyint(val)) + zero_point;
-    qval = std::max(quant_min, std::min(quant_max, qval));
-    out[i] = static_cast<T>(qval);
+    out[i] = quantize_val_with_inv_scale<T>(
+        inv_scale, zero_point, in[i], quant_min, quant_max);
   }
 #endif
 }

kernels/quantized/test/op_quantize_test.cpp

@@ -14,6 +14,7 @@
 #include <executorch/test/utils/DeathTest.h>
 
 #include <gtest/gtest.h>
+#include <limits>
 
 using namespace ::testing;
 using executorch::aten::ScalarType;
@@ -182,6 +183,106 @@ TEST(OpQuantizeOutTest, TestOutOfBounds) {
   EXPECT_TENSOR_EQ(out, expected);
 }
 
+template <ScalarType INPUT_DTYPE, ScalarType OUTPUT_DTYPE>
+void test_non_finite_values() {
+  TensorFactory<INPUT_DTYPE> tf;
+  TensorFactory<OUTPUT_DTYPE> tfo;
+  const float nan = std::numeric_limits<float>::quiet_NaN();
+  const float inf = std::numeric_limits<float>::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<ScalarType::Float, ScalarType::Char>();
+  test_non_finite_values<ScalarType::Float, ScalarType::Byte>();
+  test_non_finite_values<ScalarType::Float, ScalarType::Short>();
+  test_non_finite_values<ScalarType::Float, ScalarType::UInt16>();
+  test_non_finite_values<ScalarType::Float, ScalarType::Bits16>();
+  test_non_finite_values<ScalarType::Float, ScalarType::Int>();
+  test_non_finite_values<ScalarType::Half, ScalarType::Char>();
+  test_non_finite_values<ScalarType::Double, ScalarType::Char>();
+}
+
+TEST(OpQuantizeOutTest, NaNMapsToClampedZeroPoint) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Char> tfo;
+  Tensor input = tf.full({9}, std::numeric_limits<float>::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<ScalarType::Float> tf;
+  const float largest = std::numeric_limits<float>::max();
+  Tensor input = tf.make(
+      {10},
+      {largest,
+       -largest,
+       1e20f,
+       -1e20f,
+       1e10f,
+       -1e10f,
+       0,
+       1,
+       largest,
+       -largest});
+  TensorFactory<ScalarType::Char> 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<ScalarType::Int> tf_int;
+  constexpr int32_t min = std::numeric_limits<int32_t>::min();
+  constexpr int32_t max = std::numeric_limits<int32_t>::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<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Double> tf_double;
+  TensorFactory<ScalarType::Long> tf_long;
+  TensorFactory<ScalarType::Char> tfo;
+  const float nan = std::numeric_limits<float>::quiet_NaN();
+  const float inf = std::numeric_limits<float>::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<ScalarType::Float> tf_float;
   TensorFactory<ScalarType::Double> tf_double;