Use Accelerate for quantized SDPA prefill on Apple (#22760)

extension/llm/custom_ops/op_sdpa_impl.h

@@ -62,6 +62,22 @@ struct MaybeQuantizedMatrixData {
         dtype(dtype_) {}
 };
 
+void dequantize_per_channel_optimized(
+    const int8_t* in_data,
+    const float* scales_data,
+    const int8_t* zero_points_data,
+    float* out_data,
+    int64_t quant_min,
+    int64_t quant_max,
+    size_t outer_size,
+    size_t in_outer_stride,
+    size_t out_outer_stride,
+    size_t num_channels,
+    size_t in_channel_stride,
+    size_t out_channel_stride,
+    size_t channel_size,
+    size_t qparams_stride);
+
 template <typename accum_t>
 void _q_at_k_gemm(
     const int64_t q_m,
@@ -81,6 +97,55 @@ void _q_at_k_gemm(
       "q and k must be int8, float, half, or bfloat16");
   if (q_data.dtype == ScalarType::Char) {
     if constexpr (std::is_same<accum_t, float>::value) {
+      if (widen_scratch != nullptr) {
+        accum_t* k_f32 = widen_scratch;
+        accum_t* q_f32 = widen_scratch + k_n * qk_k;
+        dequantize_per_channel_optimized(
+            static_cast<const int8_t*>(k_data.data),
+            k_data.scales,
+            k_data.zero_points,
+            k_f32,
+            -128,
+            127,
+            1,
+            0,
+            0,
+            k_n,
+            k_stride_n,
+            qk_k,
+            qk_k,
+            k_data.scales_stride);
+        dequantize_per_channel_optimized(
+            static_cast<const int8_t*>(q_data.data),
+            q_data.scales,
+            q_data.zero_points,
+            q_f32,
+            -128,
+            127,
+            1,
+            0,
+            0,
+            q_m,
+            q_stride_m,
+            qk_k,
+            qk_k,
+            q_data.scales_stride);
+        ::executorch::cpublas::gemm(
+            ::executorch::cpublas::TransposeType::Transpose,
+            ::executorch::cpublas::TransposeType::NoTranspose,
+            k_n,
+            q_m,
+            qk_k,
+            static_cast<accum_t>(1),
+            k_f32,
+            qk_k,
+            q_f32,
+            qk_k,
+            static_cast<accum_t>(0),
+            qk_data,
+            k_n);
+        return;
+      }
       int a_stride_m_tmp, b_stride_n_tmp;
       auto kernel = torchao::kernels::cpu::quantized_matmul::
           get_int8_a_int8_b_channelwise_qmatmul(
@@ -940,16 +1005,23 @@ void cpu_flash_attention(
   // Scratch for widening q@K.T to fp32 (see _q_at_k_gemm): one K block plus one
   // q block. qBlockSize cannot exceed qSplitSize, so include the runtime bounds
   // that determine whether any block can use the widened path.
-  const bool widen_qk =
+  const bool widen_reduced_qk =
       std::is_same<scalar_t, ::executorch::aten::BFloat16>::value &&
       ::executorch::cpublas::gemm_uses_blas() &&
       headSize <= kMaxHeadSizeForWidenedQK &&
       qSplitSize >= kMinQBlockForWidenedQK;
+#if defined(__APPLE__)
+  const bool dequantize_qk = is_quantized_sdpa &&
+      ::executorch::cpublas::gemm_uses_blas() && qSplitSize > 4;
+#else
+  const bool dequantize_qk = false;
+#endif
+  const bool use_qk_conversion_scratch = widen_reduced_qk || dequantize_qk;
   int64_t size_per_thread_widen =
-      widen_qk ? (kvSplitSize + qSplitSize) * headSize : 0;
+      use_qk_conversion_scratch ? (kvSplitSize + qSplitSize) * headSize : 0;
   std::unique_ptr<char[]> allocated_buf_widen;
   accum_t* widen_buf = nullptr;
-  if (widen_qk) {
+  if (use_qk_conversion_scratch) {
     int64_t size_widen_bytes =
         size_per_thread_widen * num_thread * sizeof(accum_t);
     Result<void*> scratch_widen = ctx.allocate_temp(size_widen_bytes, 64);
@@ -1109,8 +1181,10 @@ void cpu_flash_attention(
             k_sub_matrix_data,
             kStrideN,
             qk_data,
-            (widen_qk && qBlockSize >= kMinQBlockForWidenedQK) ? widen_ptr
-                                                               : nullptr);
+            ((widen_reduced_qk && qBlockSize >= kMinQBlockForWidenedQK) ||
+             (dequantize_qk && qBlockSize > 4))
+                ? widen_ptr
+                : nullptr);
 
         // There are 4 cases that is_causal has to cover to fill
         // not-attendable-position with -inf

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