diff --git a/kernels/portable/cpu/test/vec_ops_test.cpp b/kernels/portable/cpu/test/vec_ops_test.cpp index 9611acd9d2c..0f5b96aaf38 100644 --- a/kernels/portable/cpu/test/vec_ops_test.cpp +++ b/kernels/portable/cpu/test/vec_ops_test.cpp @@ -8,6 +8,8 @@ #include +#include +#include #include #include @@ -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 {}; + +TEST_P(VecMinMaxfTest, EveryExtremumPositionAndUnalignedInput) { + const size_t size = GetParam(); + std::vector 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::quiet_NaN(); + std::vector 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::infinity(); + std::vector 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 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}; diff --git a/kernels/portable/cpu/vec_ops.h b/kernels/portable/cpu/vec_ops.h index 16a9e939e59..3b1f117a42a 100644 --- a/kernels/portable/cpu/vec_ops.h +++ b/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 +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(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(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(x, size, min_out, max_out); } /// Add each element of `x` and `y` into the corresponding element of `z`. All diff --git a/kernels/quantized/cpu/op_choose_qparams.cpp b/kernels/quantized/cpu/op_choose_qparams.cpp index acb8e100af6..837a0b0268a 100644 --- a/kernels/quantized/cpu/op_choose_qparams.cpp +++ b/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(); // 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); diff --git a/kernels/quantized/test/op_choose_qparams_test.cpp b/kernels/quantized/test/op_choose_qparams_test.cpp index d03d4cba9f6..3fe6b59571d 100644 --- a/kernels/quantized/test/op_choose_qparams_test.cpp +++ b/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(); +} + TEST(OpChooseQparamsPerTokenAsymmetricTensorOutTest, Float) { et_pal_init(); TensorFactory tf_float;