Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
110 changes: 104 additions & 6 deletions kernels/portable/cpu/test/vec_ops_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@

#include <executorch/kernels/portable/cpu/vec_ops.h>

#include <algorithm>
#include <cmath>
#include <limits>
#include <vector>

Expand All @@ -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};
Expand Down
103 changes: 97 additions & 6 deletions kernels/portable/cpu/vec_ops.h
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
Expand Down Expand Up @@ -29,16 +29,107 @@
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.
float mins[4] = {x[0], x[0], x[0], x[0]};
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
Expand Down
18 changes: 10 additions & 8 deletions kernels/quantized/cpu/op_choose_qparams.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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(
Expand All @@ -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);
Expand Down
4 changes: 4 additions & 0 deletions kernels/quantized/test/op_choose_qparams_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Loading