Vectorize min + max and add fused minmax (pytorch#22759) (pytorch#22759)

Summary:

Replace the iterator-returning `std::min_element` / `std::max_element`
implementations of `torch::executor::vec_minf` and `vec_maxf` with four
independent value-reduction lanes. A shared compile-time implementation
computes only the requested extrema, with a narrowly scoped Clang
vectorization hint for NEON/SSE2 targets and no fast-math requirement.

Add `torch::executor::vec_minmaxf(const float* x, size_t size, float*
min_out, float* max_out)` to compute both extrema together. Wire the
per-tensor and both serial/parallel per-token `choose_qparams` paths to
the fused helper, eliminating their separate minimum and maximum scans
without changing scale/zero-point calculations.

Differential Revision: D119391738

Author: JakeStevens

kernels/portable/cpu/test/vec_ops_test.cpp

@@ -8,6 +8,8 @@
 
 #include <executorch/kernels/portable/cpu/vec_ops.h>
 
+#include <algorithm>
+#include <cmath>
 #include <limits>
 #include <vector>
 
@@ -16,23 +18,119 @@
 using namespace ::testing;
 
 TEST(VecMinfTest, Smoke) {
-  // No need to be super thorough since we know this is implemented with
-  // std::min_element(). Just show that it's hooked up correctly.
-
   constexpr size_t kNumVals = 5;
   float x[kNumVals] = {1.1, -2.2, 0, -1234.5, 10.0};
   EXPECT_EQ(torch::executor::vec_minf(x, kNumVals), -1234.5);
 }
 
 TEST(VecMaxfTest, Smoke) {
-  // No need to be super thorough since we know this is implemented with
-  // std::max_element(). Just show that it's hooked up correctly.
-
   constexpr size_t kNumVals = 5;
   float x[kNumVals] = {1.1, -2.2, 0, -1234.5, 10.0};
   EXPECT_EQ(torch::executor::vec_maxf(x, kNumVals), 10.0);
 }
 
+namespace {
+
+void expect_extrema(const float* x, size_t size, float min, float max) {
+  float fused_min = 123.0f;
+  float fused_max = -456.0f;
+  torch::executor::vec_minmaxf(x, size, &fused_min, &fused_max);
+  for (float actual : {torch::executor::vec_minf(x, size), fused_min}) {
+    if (std::isnan(min)) {
+      EXPECT_TRUE(std::isnan(actual));
+    } else {
+      EXPECT_EQ(actual, min);
+      EXPECT_EQ(std::signbit(actual), std::signbit(min));
+    }
+  }
+  for (float actual : {torch::executor::vec_maxf(x, size), fused_max}) {
+    if (std::isnan(max)) {
+      EXPECT_TRUE(std::isnan(actual));
+    } else {
+      EXPECT_EQ(actual, max);
+      EXPECT_EQ(std::signbit(actual), std::signbit(max));
+    }
+  }
+}
+
+class VecMinMaxfTest : public TestWithParam<size_t> {};
+
+TEST_P(VecMinMaxfTest, EveryExtremumPositionAndUnalignedInput) {
+  const size_t size = GetParam();
+  std::vector<float> storage(size + 1, 7.0f);
+  float* x = storage.data() + 1;
+  for (size_t i = 0; i < size; ++i) {
+    SCOPED_TRACE(i);
+    x[i] = -13.0f;
+    expect_extrema(x, size, -13.0f, size == 1 ? -13.0f : 7.0f);
+    x[i] = 19.0f;
+    expect_extrema(x, size, size == 1 ? 19.0f : 7.0f, 19.0f);
+    x[i] = 7.0f;
+  }
+}
+
+TEST_P(VecMinMaxfTest, LeadingNaNPersistsAndLaterNaNsAreIgnored) {
+  const size_t size = GetParam();
+  const float nan = std::numeric_limits<float>::quiet_NaN();
+  std::vector<float> x(size, 3.0f);
+  for (size_t i = 0; i < size; ++i) {
+    SCOPED_TRACE(i);
+    x[i] = nan;
+    expect_extrema(x.data(), size, i == 0 ? nan : 3.0f, i == 0 ? nan : 3.0f);
+    x[i] = 3.0f;
+  }
+  std::fill(x.begin(), x.end(), nan);
+  expect_extrema(x.data(), size, nan, nan);
+}
+
+TEST_P(VecMinMaxfTest, Infinities) {
+  const size_t size = GetParam();
+  const float inf = std::numeric_limits<float>::infinity();
+  std::vector<float> x(size, inf);
+  expect_extrema(x.data(), size, inf, inf);
+  std::fill(x.begin(), x.end(), -inf);
+  expect_extrema(x.data(), size, -inf, -inf);
+  if (size > 1) {
+    x.back() = inf;
+    expect_extrema(x.data(), size, -inf, inf);
+  }
+}
+
+TEST_P(VecMinMaxfTest, FirstSignedZeroWins) {
+  const size_t size = GetParam();
+  for (float first_zero : {0.0f, -0.0f}) {
+    for (size_t first = 0; first < size; ++first) {
+      SCOPED_TRACE(first);
+      for (float other : {-1.0f, 1.0f}) {
+        std::vector<float> x(size, other);
+        x[first] = first_zero;
+        for (size_t j = first + 1; j < size; ++j) {
+          x[j] = -first_zero;
+        }
+        expect_extrema(
+            x.data(),
+            size,
+            first == 0 || other > 0 ? first_zero : other,
+            first == 0 || other < 0 ? first_zero : other);
+      }
+    }
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    VectorBoundaries,
+    VecMinMaxfTest,
+    Values(1, 2, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 32, 33, 127, 128, 129));
+
+TEST(VecMinMaxfTest, OutputsMayAliasInput) {
+  float x[] = {4.0f, -9.0f, 2.0f, 11.0f, -3.0f};
+  torch::executor::vec_minmaxf(x, 5, &x[0], &x[1]);
+  EXPECT_EQ(x[0], -9.0f);
+  EXPECT_EQ(x[1], 11.0f);
+}
+
+} // namespace
+
 TEST(VecAddfTest, Smoke) {
   constexpr size_t kNumVals = 5;
   float in1[kNumVals] = {1, 2, 3, 4, 5};

kernels/portable/cpu/vec_ops.h

@@ -29,16 +29,104 @@
 namespace torch {
 namespace executor {
 
-/// Returns the minimum element of the array at `x`, which must have `size`
-/// elements.
+namespace internal {
+
+template <bool ComputeMin, bool ComputeMax>
+inline void
+vec_minmaxf_impl(const float* x, size_t size, float* min_out, float* max_out) {
+  float min_value = x[0];
+  float max_value = x[0];
+  size_t i = 1;
+  if (size >= 4) {
+    // Seeding every lane with x[0] preserves a leading NaN and ignores later
+    // NaNs.
+    ET_UNUSED float mins[4] = {x[0], x[0], x[0], x[0]};
+    ET_UNUSED float maxs[4] = {x[0], x[0], x[0], x[0]};
+    const size_t vector_end = size - size % 4;
+    for (i = 0; i < vector_end; i += 4) {
+      // Keep the independent lanes together even under size optimization.
+#if defined(__clang__) && (defined(__ARM_NEON) || defined(__SSE2__))
+#pragma clang loop vectorize_width(4) interleave_count(1) unroll(disable)
+#endif
+      for (size_t j = 0; j < 4; ++j) {
+        if constexpr (ComputeMin) {
+          mins[j] = std::min(mins[j], x[i + j]);
+        }
+        if constexpr (ComputeMax) {
+          maxs[j] = std::max(maxs[j], x[i + j]);
+        }
+      }
+    }
+    if constexpr (ComputeMin) {
+      min_value =
+          std::min(std::min(mins[0], mins[1]), std::min(mins[2], mins[3]));
+    }
+    if constexpr (ComputeMax) {
+      max_value =
+          std::max(std::max(maxs[0], maxs[1]), std::max(maxs[2], maxs[3]));
+    }
+  }
+  for (; i < size; ++i) {
+    if constexpr (ComputeMin) {
+      min_value = std::min(min_value, x[i]);
+    }
+    if constexpr (ComputeMax) {
+      max_value = std::max(max_value, x[i]);
+    }
+  }
+
+  // Lane reduction can reorder equal signed zeros. Preserve the first zero.
+  if (size >= 4 &&
+      ((ComputeMin && min_value == 0.0f) ||
+       (ComputeMax && max_value == 0.0f))) {
+    for (size_t j = 0; j < size; ++j) {
+      if (x[j] == 0.0f) {
+        if constexpr (ComputeMin) {
+          if (min_value == 0.0f) {
+            min_value = x[j];
+          }
+        }
+        if constexpr (ComputeMax) {
+          if (max_value == 0.0f) {
+            max_value = x[j];
+          }
+        }
+        break;
+      }
+    }
+  }
+  if constexpr (ComputeMin) {
+    *min_out = min_value;
+  }
+  if constexpr (ComputeMax) {
+    *max_out = max_value;
+  }
+}
+
+} // namespace internal
+
+/// Returns the minimum element of the nonempty array at `x`, which must have
+/// `size` elements.
 inline float vec_minf(const float* x, size_t size) {
-  return *std::min_element(x, x + size);
+  float minimum;
+  internal::vec_minmaxf_impl<true, false>(x, size, &minimum, nullptr);
+  return minimum;
 }
 
-/// Returns the maximum element of the array at `x`, which must have `size`
-/// elements.
+/// Returns the maximum element of the nonempty array at `x`, which must have
+/// `size` elements.
 inline float vec_maxf(const float* x, size_t size) {
-  return *std::max_element(x, x + size);
+  float maximum;
+  internal::vec_minmaxf_impl<false, true>(x, size, nullptr, &maximum);
+  return maximum;
+}
+
+/// Writes the minimum and maximum of the nonempty array at `x`, which must have
+/// `size` elements. `min_out` and `max_out` must point to distinct valid
+/// floats.
+inline void
+vec_minmaxf(const float* x, size_t size, float* min_out, float* max_out) {
+  internal::vec_minmaxf_impl<true, true>(x, size, min_out, max_out);
 }
 
 /// Add each element of `x` and `y` into the corresponding element of `z`. All

kernels/quantized/cpu/op_choose_qparams.cpp

@@ -179,8 +179,9 @@ void choose_qparams(
     Tensor& zero_point_out) {
   const float* x_fp32 = input.const_data_ptr<float>();
   // Compute x_min, x_max and q_params (scale, zero_point)
-  float min = torch::executor::vec_minf(x_fp32, input.numel());
-  float max = torch::executor::vec_maxf(x_fp32, input.numel());
+  float min;
+  float max;
+  torch::executor::vec_minmaxf(x_fp32, input.numel(), &min, &max);
 
   double scale;
   int32_t zero_point;
@@ -216,8 +217,10 @@ void choose_qparams_per_token(
         0, num_tokens, 1, [&](const int64_t begin, const int64_t end) {
           for (int64_t i = begin; i < end; i++) {
             const float* token_data = x_fp32 + i * token_dim_size;
-            float min = torch::executor::vec_minf(token_data, token_dim_size);
-            float max = torch::executor::vec_maxf(token_data, token_dim_size);
+            float min;
+            float max;
+            torch::executor::vec_minmaxf(
+                token_data, token_dim_size, &min, &max);
             double scale;
             int32_t zero_point;
             calculate_scale_and_zero_point(
@@ -228,10 +231,9 @@ void choose_qparams_per_token(
         });
   } else {
     for (auto i = 0; i < num_tokens; i++) {
-      // vec_minf uses std::min_element. Check if it actually
-      // gets vectorized.
-      float min = torch::executor::vec_minf(x_fp32, token_dim_size);
-      float max = torch::executor::vec_maxf(x_fp32, token_dim_size);
+      float min;
+      float max;
+      torch::executor::vec_minmaxf(x_fp32, token_dim_size, &min, &max);
       double scale;
       int32_t zero_point;
       calculate_scale_and_zero_point(min, max, qmin, qmax, scale, zero_point);

kernels/quantized/test/op_choose_qparams_test.cpp

@@ -50,6 +50,10 @@ void test_dtype() {
   EXPECT_TENSOR_EQ(zero_point_out, expected_zero_point);
 }
 
+TEST(OpChooseQparamsTensorOutTest, Byte) {
+  test_dtype<ScalarType::Byte>();
+}
+
 TEST(OpChooseQparamsPerTokenAsymmetricTensorOutTest, Float) {
   et_pal_init();
   TensorFactory<ScalarType::Float> tf_float;