SVE: Assorted optimizations libeigen/eigen!2928
diff --git a/Eigen/src/Core/arch/SVE/PacketMath.h b/Eigen/src/Core/arch/SVE/PacketMath.h index f1ad4a7..d01cbed 100644 --- a/Eigen/src/Core/arch/SVE/PacketMath.h +++ b/Eigen/src/Core/arch/SVE/PacketMath.h
@@ -94,9 +94,7 @@ template <> EIGEN_STRONG_INLINE PacketXi plset<PacketXi>(const numext::int32_t& a) { - numext::int32_t c[packet_traits<numext::int32_t>::size]; - for (int i = 0; i < packet_traits<numext::int32_t>::size; i++) c[i] = i; - return svadd_s32_x(svptrue_b32(), pset1<PacketXi>(a), svld1_s32(svptrue_b32(), c)); + return svindex_s32(a, 1); } template <> @@ -218,17 +216,24 @@ template <> EIGEN_STRONG_INLINE PacketXi ploaddup<PacketXi>(const numext::int32_t* from) { - svuint32_t indices = svindex_u32(0, 1); // index {base=0, base+step=1, base+step*2, ...} - indices = svzip1_u32(indices, indices); // index in the format {a0, a0, a1, a1, a2, a2, ...} - return svld1_gather_u32index_s32(svptrue_b32(), from, indices); + // Load the size/2 values this reads into the low half and interleave them + // with themselves: svzip1 only consumes the low halves of its operands. + // The predicate is exact rather than svptrue -- ploaddup may only touch + // size/2 elements, and a wider one would read past the end of the input. + constexpr uint64_t kHalf = uint64_t(packet_traits<numext::int32_t>::size) / 2; + svint32_t lo = svld1_s32(svwhilelt_b32(uint64_t(0), kHalf), from); + return svzip1_s32(lo, lo); } template <> EIGEN_STRONG_INLINE PacketXi ploadquad<PacketXi>(const numext::int32_t* from) { - svuint32_t indices = svindex_u32(0, 1); // index {base=0, base+step=1, base+step*2, ...} - indices = svzip1_u32(indices, indices); // index in the format {a0, a0, a1, a1, a2, a2, ...} - indices = svzip1_u32(indices, indices); // index in the format {a0, a0, a0, a0, a1, a1, a1, a1, ...} - return svld1_gather_u32index_s32(svptrue_b32(), from, indices); + // As ploaddup, one zip further: size/4 values, each repeated four times. + // At the smallest vector length size/4 rounds to zero, where one element + // still has to be read. + constexpr uint64_t kQuarter = numext::maxi(uint64_t(packet_traits<numext::int32_t>::size) / 4, uint64_t(1)); + svint32_t lo = svld1_s32(svwhilelt_b32(uint64_t(0), kQuarter), from); + lo = svzip1_s32(lo, lo); + return svzip1_s32(lo, lo); } template <> @@ -279,32 +284,13 @@ template <> EIGEN_STRONG_INLINE numext::int32_t predux_mul<PacketXi>(const PacketXi& a) { - EIGEN_STATIC_ASSERT((EIGEN_ARM64_SVE_VL % 128 == 0), EIGEN_INTERNAL_ERROR_PLEASE_FILE_A_BUG_REPORT); - - // Multiply the vector by its reverse + // Multiply the vector by its reverse. svint32_t prod = svmul_s32_x(svptrue_b32(), a, svrev_s32(a)); - svint32_t half_prod; - // Extract the high half of the vector. Depending on the VL more reductions need to be done - EIGEN_IF_CONSTEXPR (EIGEN_ARM64_SVE_VL >= 2048) { - half_prod = svtbl_s32(prod, svindex_u32(32, 1)); - prod = svmul_s32_x(svptrue_b32(), prod, half_prod); - } - EIGEN_IF_CONSTEXPR (EIGEN_ARM64_SVE_VL >= 1024) { - half_prod = svtbl_s32(prod, svindex_u32(16, 1)); - prod = svmul_s32_x(svptrue_b32(), prod, half_prod); - } - EIGEN_IF_CONSTEXPR (EIGEN_ARM64_SVE_VL >= 512) { - half_prod = svtbl_s32(prod, svindex_u32(8, 1)); - prod = svmul_s32_x(svptrue_b32(), prod, half_prod); - } - EIGEN_IF_CONSTEXPR (EIGEN_ARM64_SVE_VL >= 256) { - half_prod = svtbl_s32(prod, svindex_u32(4, 1)); - prod = svmul_s32_x(svptrue_b32(), prod, half_prod); - } - // Last reduction - half_prod = svtbl_s32(prod, svindex_u32(2, 1)); - prod = svmul_s32_x(svptrue_b32(), prod, half_prod); + // Reduce with interleave-and-multiply. + // NOTE: Skip the final reduction since it is already handled by `rev` above. + for (int n = unpacket_traits<PacketXi>::size; n > 2; n >>= 1) + prod = svmul_s32_x(svptrue_b32(), svzip1_s32(prod, prod), svzip2_s32(prod, prod)); // The reduction is done to the first element. return pfirst<PacketXi>(prod); @@ -322,16 +308,16 @@ template <int N> EIGEN_DEVICE_FUNC inline void ptranspose(PacketBlock<PacketXi, N>& kernel) { - int buffer[packet_traits<numext::int32_t>::size * N] = {0}; - int i = 0; - - PacketXi stride_index = svindex_s32(0, N); - - for (i = 0; i < N; i++) { - svst1_scatter_s32index_s32(svptrue_b32(), buffer + i, stride_index, kernel.packet[i]); - } - for (i = 0; i < N; i++) { - kernel.packet[i] = svld1_s32(svptrue_b32(), buffer + i * packet_traits<numext::int32_t>::size); + EIGEN_STATIC_ASSERT((N & (N - 1)) == 0, EIGEN_INTERNAL_ERROR_PLEASE_FILE_A_BUG_REPORT); + for (int stride = N / 2; stride > 0; stride >>= 1) { + for (int block = 0; block < N; block += 2 * stride) { + for (int k = 0; k < stride; ++k) { + PacketXi lo = svzip1_s32(kernel.packet[block + k], kernel.packet[block + k + stride]); + PacketXi hi = svzip2_s32(kernel.packet[block + k], kernel.packet[block + k + stride]); + kernel.packet[block + k] = lo; + kernel.packet[block + k + stride] = hi; + } + } } } @@ -568,17 +554,24 @@ template <> EIGEN_STRONG_INLINE PacketXf ploaddup<PacketXf>(const float* from) { - svuint32_t indices = svindex_u32(0, 1); // index {base=0, base+step=1, base+step*2, ...} - indices = svzip1_u32(indices, indices); // index in the format {a0, a0, a1, a1, a2, a2, ...} - return svld1_gather_u32index_f32(svptrue_b32(), from, indices); + // Load the size/2 values this reads into the low half and interleave them + // with themselves: svzip1 only consumes the low halves of its operands. + // The predicate is exact rather than svptrue -- ploaddup may only touch + // size/2 elements, and a wider one would read past the end of the input. + constexpr uint64_t kHalf = uint64_t(packet_traits<float>::size) / 2; + svfloat32_t lo = svld1_f32(svwhilelt_b32(uint64_t(0), kHalf), from); + return svzip1_f32(lo, lo); } template <> EIGEN_STRONG_INLINE PacketXf ploadquad<PacketXf>(const float* from) { - svuint32_t indices = svindex_u32(0, 1); // index {base=0, base+step=1, base+step*2, ...} - indices = svzip1_u32(indices, indices); // index in the format {a0, a0, a1, a1, a2, a2, ...} - indices = svzip1_u32(indices, indices); // index in the format {a0, a0, a0, a0, a1, a1, a1, a1, ...} - return svld1_gather_u32index_f32(svptrue_b32(), from, indices); + // As ploaddup, one zip further: size/4 values, each repeated four times. + // At the smallest vector length size/4 rounds to zero, where one element + // still has to be read. + constexpr uint64_t kQuarter = numext::maxi(uint64_t(packet_traits<float>::size) / 4, uint64_t(1)); + svfloat32_t lo = svld1_f32(svwhilelt_b32(uint64_t(0), kQuarter), from); + lo = svzip1_f32(lo, lo); + return svzip1_f32(lo, lo); } template <> @@ -635,34 +628,15 @@ // Other reduction functions: // mul -// Only works for SVE Vls multiple of 128 template <> EIGEN_STRONG_INLINE float predux_mul<PacketXf>(const PacketXf& a) { - EIGEN_STATIC_ASSERT((EIGEN_ARM64_SVE_VL % 128 == 0), EIGEN_INTERNAL_ERROR_PLEASE_FILE_A_BUG_REPORT); - // Multiply the vector by its reverse + // Multiply the vector by its reverse. svfloat32_t prod = svmul_f32_x(svptrue_b32(), a, svrev_f32(a)); - svfloat32_t half_prod; - // Extract the high half of the vector. Depending on the VL more reductions need to be done - EIGEN_IF_CONSTEXPR (EIGEN_ARM64_SVE_VL >= 2048) { - half_prod = svtbl_f32(prod, svindex_u32(32, 1)); - prod = svmul_f32_x(svptrue_b32(), prod, half_prod); - } - EIGEN_IF_CONSTEXPR (EIGEN_ARM64_SVE_VL >= 1024) { - half_prod = svtbl_f32(prod, svindex_u32(16, 1)); - prod = svmul_f32_x(svptrue_b32(), prod, half_prod); - } - EIGEN_IF_CONSTEXPR (EIGEN_ARM64_SVE_VL >= 512) { - half_prod = svtbl_f32(prod, svindex_u32(8, 1)); - prod = svmul_f32_x(svptrue_b32(), prod, half_prod); - } - EIGEN_IF_CONSTEXPR (EIGEN_ARM64_SVE_VL >= 256) { - half_prod = svtbl_f32(prod, svindex_u32(4, 1)); - prod = svmul_f32_x(svptrue_b32(), prod, half_prod); - } - // Last reduction - half_prod = svtbl_f32(prod, svindex_u32(2, 1)); - prod = svmul_f32_x(svptrue_b32(), prod, half_prod); + // Reduce with interleave-and-multiply. + // NOTE: Skip the final reduction since it is already handled by `rev` above. + for (int n = unpacket_traits<PacketXf>::size; n > 2; n >>= 1) + prod = svmul_f32_x(svptrue_b32(), svzip1_f32(prod, prod), svzip2_f32(prod, prod)); // The reduction is done to the first element. return pfirst<PacketXf>(prod); @@ -680,17 +654,16 @@ template <int N> EIGEN_DEVICE_FUNC inline void ptranspose(PacketBlock<PacketXf, N>& kernel) { - EIGEN_ALIGN_MAX float buffer[packet_traits<float>::size * N] = {}; - int i = 0; - - PacketXi stride_index = svindex_s32(0, N); - - for (i = 0; i < N; i++) { - svst1_scatter_s32index_f32(svptrue_b32(), buffer + i, stride_index, kernel.packet[i]); - } - - for (i = 0; i < N; i++) { - kernel.packet[i] = svld1_f32(svptrue_b32(), buffer + i * packet_traits<float>::size); + EIGEN_STATIC_ASSERT((N & (N - 1)) == 0, EIGEN_INTERNAL_ERROR_PLEASE_FILE_A_BUG_REPORT); + for (int stride = N / 2; stride > 0; stride >>= 1) { + for (int block = 0; block < N; block += 2 * stride) { + for (int k = 0; k < stride; ++k) { + PacketXf lo = svzip1_f32(kernel.packet[block + k], kernel.packet[block + k + stride]); + PacketXf hi = svzip2_f32(kernel.packet[block + k], kernel.packet[block + k + stride]); + kernel.packet[block + k] = lo; + kernel.packet[block + k + stride] = hi; + } + } } } @@ -995,18 +968,17 @@ return svaddv_f64(svptrue_b64(), a); } -// Only works for SVE VLs that are a multiple of 128. template <> EIGEN_STRONG_INLINE double predux_mul<PacketXd>(const PacketXd& a) { - EIGEN_STATIC_ASSERT((EIGEN_ARM64_SVE_VL % 128 == 0), EIGEN_INTERNAL_ERROR_PLEASE_FILE_A_BUG_REPORT); - // Multiplying by the reverse pairs lane i with lane n-1-i, leaving every - // product of a pair in both halves; halving the span each round then folds - // the halves together. At VL = 128 there are two lanes and the first multiply - // has already combined them. + // Multiply the vector by its reverse. svfloat64_t prod = svmul_f64_x(svptrue_b64(), a, svrev_f64(a)); - for (int span = unpacket_traits<PacketXd>::size / 2; span >= 2; span >>= 1) { - prod = svmul_f64_x(svptrue_b64(), prod, svtbl_f64(prod, svindex_u64(span, 1))); - } + + // Reduce with interleave-and-multiply. + // NOTE: Skip the final reduction since it is already handled by `rev` above. + for (int n = unpacket_traits<PacketXd>::size; n > 2; n >>= 1) + prod = svmul_f64_x(svptrue_b64(), svzip1_f64(prod, prod), svzip2_f64(prod, prod)); + + // The reduction is done to the first element. return pfirst<PacketXd>(prod); } @@ -1022,17 +994,16 @@ template <int N> EIGEN_DEVICE_FUNC inline void ptranspose(PacketBlock<PacketXd, N>& kernel) { - EIGEN_ALIGN_MAX double buffer[packet_traits<double>::size * N] = {}; - int i = 0; - - svint64_t stride_index = svindex_s64(0, N); - - for (i = 0; i < N; i++) { - svst1_scatter_s64index_f64(svptrue_b64(), buffer + i, stride_index, kernel.packet[i]); - } - - for (i = 0; i < N; i++) { - kernel.packet[i] = svld1_f64(svptrue_b64(), buffer + i * packet_traits<double>::size); + EIGEN_STATIC_ASSERT((N & (N - 1)) == 0, EIGEN_INTERNAL_ERROR_PLEASE_FILE_A_BUG_REPORT); + for (int stride = N / 2; stride > 0; stride >>= 1) { + for (int block = 0; block < N; block += 2 * stride) { + for (int k = 0; k < stride; ++k) { + PacketXd lo = svzip1_f64(kernel.packet[block + k], kernel.packet[block + k + stride]); + PacketXd hi = svzip2_f64(kernel.packet[block + k], kernel.packet[block + k + stride]); + kernel.packet[block + k] = lo; + kernel.packet[block + k + stride] = hi; + } + } } }
diff --git a/Eigen/src/Core/util/ConfigureVectorization.h b/Eigen/src/Core/util/ConfigureVectorization.h index ea65c3f..f3b1fab 100644 --- a/Eigen/src/Core/util/ConfigureVectorization.h +++ b/Eigen/src/Core/util/ConfigureVectorization.h
@@ -482,6 +482,11 @@ // to ensure a fixed length is set #if defined __ARM_FEATURE_SVE_BITS #define EIGEN_ARM64_SVE_VL __ARM_FEATURE_SVE_BITS + +// Architecture-mandated length constraints. +static_assert((EIGEN_ARM64_SVE_VL >= 128) && (EIGEN_ARM64_SVE_VL <= 2048) && + ((EIGEN_ARM64_SVE_VL & (EIGEN_ARM64_SVE_VL - 1)) == 0), + "SVE vector length must be 2^n for some n in [7, 11]"); #else #error "Eigen requires a fixed SVE vector length but EIGEN_ARM64_SVE_VL is not set." #endif
diff --git a/benchmarks/Core/CMakeLists.txt b/benchmarks/Core/CMakeLists.txt index f777291..5549b21 100644 --- a/benchmarks/Core/CMakeLists.txt +++ b/benchmarks/Core/CMakeLists.txt
@@ -17,6 +17,7 @@ eigen_add_benchmark(bench_stable_norm bench_stable_norm.cpp) eigen_add_benchmark(bench_cwise_math bench_cwise_math.cpp) eigen_add_benchmark(bench_int64_ops bench_int64_ops.cpp) +eigen_add_benchmark(bench_packet_ops bench_packet_ops.cpp) eigen_add_benchmark(bench_broadcasting bench_broadcasting.cpp) eigen_add_benchmark(bench_block_ops bench_block_ops.cpp) eigen_add_benchmark(bench_map bench_map.cpp)
diff --git a/benchmarks/Core/bench_packet_ops.cpp b/benchmarks/Core/bench_packet_ops.cpp new file mode 100644 index 0000000..9d78923 --- /dev/null +++ b/benchmarks/Core/bench_packet_ops.cpp
@@ -0,0 +1,205 @@ +// Benchmarks for the PacketMath implementations of `plset`, `ploaddup`, +// `ploadquad`, `predux_mul`, and `ptranspose`, at the packet-op level and +// shared across whichever architecture backend the build targets. To +// compare against a prior implementation, build and run this same file +// against the Eigen checkout in question -- it only calls the public +// `Eigen::internal` packet API, so it is source-compatible with whatever +// PacketMath.h happens to provide. +// SPDX-FileCopyrightText: The Eigen Authors +// SPDX-License-Identifier: MPL-2.0 + +#include <benchmark/benchmark.h> +#include <Eigen/Core> + +#include <cmath> +#include <cstdint> +#include <type_traits> + +namespace Eigen { +namespace { + +using internal::packet_traits; +using internal::PacketBlock; +using internal::pfirst; +using internal::ploadu; +using internal::pstoreu; + +template <typename Packet, int N> +EIGEN_DONT_INLINE void call_ptranspose(PacketBlock<Packet, N>& kernel) { + internal::ptranspose(kernel); +} + +// ---- plset ---- + +template <typename Scalar> +void BM_Plset(benchmark::State& state) { + using Packet = typename packet_traits<Scalar>::type; + constexpr int N = packet_traits<Scalar>::size; + Scalar a = Scalar(7); + Scalar out[N]; + + pstoreu(out, internal::plset<Packet>(a)); + for (int i = 0; i < N; ++i) { + if (out[i] != static_cast<Scalar>(a + i)) { + state.SkipWithError("Plset: materialized result does not match scalar reference"); + return; + } + } + + benchmark::DoNotOptimize(a); + for (auto _ : state) { + pstoreu(out, internal::plset<Packet>(a)); + benchmark::DoNotOptimize(out); + } +} +BENCHMARK(BM_Plset<numext::int32_t>)->Name("Plset_int32"); +BENCHMARK(BM_Plset<float>)->Name("Plset_float"); +BENCHMARK(BM_Plset<double>)->Name("Plset_double"); + +// ---- ploaddup ---- + +template <typename Scalar> +void BM_Ploaddup(benchmark::State& state) { + using Packet = typename packet_traits<Scalar>::type; + constexpr int N = packet_traits<Scalar>::size; + Scalar in[N]; + for (int i = 0; i < N; ++i) in[i] = static_cast<Scalar>(i); + Scalar out[N]; + + pstoreu(out, internal::ploaddup<Packet>(in)); + for (int i = 0; i < N; ++i) { + if (out[i] != in[i / 2]) { + state.SkipWithError("Ploaddup: materialized result does not match scalar reference"); + return; + } + } + + benchmark::DoNotOptimize(in); + for (auto _ : state) { + pstoreu(out, internal::ploaddup<Packet>(in)); + benchmark::DoNotOptimize(out); + } +} +BENCHMARK(BM_Ploaddup<numext::int32_t>)->Name("Ploaddup_int32"); +BENCHMARK(BM_Ploaddup<float>)->Name("Ploaddup_float"); +BENCHMARK(BM_Ploaddup<double>)->Name("Ploaddup_double"); + +// ---- ploadquad ---- + +template <typename Scalar> +void BM_Ploadquad(benchmark::State& state) { + using Packet = typename packet_traits<Scalar>::type; + constexpr int N = packet_traits<Scalar>::size; + Scalar in[N]; + for (int i = 0; i < N; ++i) in[i] = static_cast<Scalar>(i); + Scalar out[N]; + + pstoreu(out, internal::ploadquad<Packet>(in)); + for (int i = 0; i < N; ++i) { + if (out[i] != in[i / 4]) { + state.SkipWithError("Ploadquad: materialized result does not match scalar reference"); + return; + } + } + + benchmark::DoNotOptimize(in); + for (auto _ : state) { + pstoreu(out, internal::ploadquad<Packet>(in)); + benchmark::DoNotOptimize(out); + } +} +BENCHMARK(BM_Ploadquad<numext::int32_t>)->Name("Ploadquad_int32"); +BENCHMARK(BM_Ploadquad<float>)->Name("Ploadquad_float"); +BENCHMARK(BM_Ploadquad<double>)->Name("Ploadquad_double"); + +// ---- predux_mul ---- +// Inputs are chosen so the true product is exactly representable regardless +// of the order the reduction folds lanes together, at every vector length: +// powers of two multiply without rounding, and centering the exponents around +// zero keeps both the inputs and the product in range for float/double even +// at N == 64; multiplying by 1 is exact and overflow-free for int32. + +template <typename Scalar> +void fill_redux_mul_input(Scalar (&in)[packet_traits<Scalar>::size], Scalar& expected) { + constexpr int N = packet_traits<Scalar>::size; + if constexpr (std::is_integral<Scalar>::value) { + for (int i = 0; i < N; ++i) in[i] = Scalar(1); + in[0] = Scalar(-3); + expected = Scalar(-3); + if (N > 1) { + in[N - 1] = Scalar(2); + expected = Scalar(-6); + } + } else { + for (int i = 0; i < N; ++i) in[i] = std::ldexp(Scalar(1), i - N / 2); + expected = std::ldexp(Scalar(1), -N / 2); + } +} + +template <typename Scalar> +void BM_ReduxMul(benchmark::State& state) { + using Packet = typename packet_traits<Scalar>::type; + constexpr int N = packet_traits<Scalar>::size; + Scalar in[N]; + Scalar expected; + fill_redux_mul_input<Scalar>(in, expected); + Packet a = ploadu<Packet>(in); + + if (internal::predux_mul<Packet>(a) != expected) { + state.SkipWithError("ReduxMul: materialized result does not match scalar reference"); + return; + } + + benchmark::DoNotOptimize(a); + for (auto _ : state) benchmark::DoNotOptimize(internal::predux_mul<Packet>(a)); +} +BENCHMARK(BM_ReduxMul<numext::int32_t>)->Name("ReduxMul_int32"); +BENCHMARK(BM_ReduxMul<float>)->Name("ReduxMul_float"); +BENCHMARK(BM_ReduxMul<double>)->Name("ReduxMul_double"); + +// ---- ptranspose ---- +// Benchmarked at N == the type's packet width, i.e. a full square transpose, +// so the expected result is simply new_packet[i][k] == old_packet[k][i]. +// Correctness is checked once on a scratch kernel -- ptranspose applied +// repeatedly toggles between the original and transposed state, so checking +// after the timed loop would depend on the (unpredictable) iteration count. +// +// call_ptranspose is deliberately EIGEN_DONT_INLINE: real callers inline +// ptranspose on register-resident packets, but this wrapper forces the block +// through memory, so this benchmark understates the win; transposeInPlace is +// the macro benchmark for the inlined case. + +template <typename Scalar> +void BM_Ptranspose(benchmark::State& state) { + using Packet = typename packet_traits<Scalar>::type; + constexpr int N = packet_traits<Scalar>::size; + Scalar in[N * N]; + for (int i = 0; i < N * N; ++i) in[i] = static_cast<Scalar>(i); + + PacketBlock<Packet, N> check; + for (int i = 0; i < N; ++i) check.packet[i] = ploadu<Packet>(in + i * N); + call_ptranspose<Packet, N>(check); + for (int i = 0; i < N; ++i) { + Scalar row[N]; + pstoreu(row, check.packet[i]); + for (int k = 0; k < N; ++k) { + if (row[k] != in[k * N + i]) { + state.SkipWithError("Ptranspose: materialized result does not match scalar reference"); + return; + } + } + } + + PacketBlock<Packet, N> kernel; + for (int i = 0; i < N; ++i) kernel.packet[i] = ploadu<Packet>(in + i * N); + for (auto _ : state) { + call_ptranspose<Packet, N>(kernel); + benchmark::DoNotOptimize(pfirst<Packet>(kernel.packet[0])); + } +} +BENCHMARK(BM_Ptranspose<numext::int32_t>)->Name("Ptranspose_int32"); +BENCHMARK(BM_Ptranspose<float>)->Name("Ptranspose_float"); +BENCHMARK(BM_Ptranspose<double>)->Name("Ptranspose_double"); + +} // namespace +} // namespace Eigen