AVX-512 Pack32To16 / pstore cleanup and triangular stack allocation limit guard libeigen/eigen!3365 Co-authored-by: Antonio Sanchez <cantonios@google.com>
diff --git a/Eigen/src/Core/arch/AVX512/PacketMath.h b/Eigen/src/Core/arch/AVX512/PacketMath.h index 7d31c46..77a4489 100644 --- a/Eigen/src/Core/arch/AVX512/PacketMath.h +++ b/Eigen/src/Core/arch/AVX512/PacketMath.h
@@ -733,20 +733,7 @@ // Helper function for bit packing snippet of low precision comparison. // It packs the flags from 32x16 to 16x16. -EIGEN_STRONG_INLINE __m256i Pack32To16(const Packet16f& rf) { - // Split data into small pieces and handle with AVX instructions - // to guarantee internal order of vector. - // Operation: - // dst[15:0] := Saturate16(rf[31:0]) - // dst[31:16] := Saturate16(rf[63:32]) - // ... - // dst[255:240] := Saturate16(rf[255:224]) - __m256i lo = _mm256_castps_si256(extract256<0>(rf)); - __m256i hi = _mm256_castps_si256(extract256<1>(rf)); - __m128i result_lo = _mm_packs_epi32(_mm256_extractf128_si256(lo, 0), _mm256_extractf128_si256(lo, 1)); - __m128i result_hi = _mm_packs_epi32(_mm256_extractf128_si256(hi, 0), _mm256_extractf128_si256(hi, 1)); - return _mm256_insertf128_si256(_mm256_castsi128_si256(result_lo), result_hi, 1); -} +EIGEN_STRONG_INLINE __m256i Pack32To16(const Packet16f& rf) { return _mm512_cvtepi32_epi16(_mm512_castps_si512(rf)); } template <> EIGEN_STRONG_INLINE Packet16f pisnan(const Packet16f& a) { @@ -3091,21 +3078,13 @@ template <> EIGEN_STRONG_INLINE void pstore<numext::int16_t, Packet16s>(numext::int16_t* out, const Packet16s& x) { EIGEN_DEBUG_ALIGNED_STORE -#if defined(EIGEN_VECTORIZE_AVX512F) && defined(EIGEN_VECTORIZE_AVX512VL) - _mm256_store_epi32(out, x); -#else _mm256_store_si256(reinterpret_cast<__m256i*>(out), x); -#endif } template <> EIGEN_STRONG_INLINE void pstore<numext::int16_t, Packet8s>(numext::int16_t* out, const Packet8s& x) { EIGEN_DEBUG_ALIGNED_STORE -#if defined(EIGEN_VECTORIZE_AVX512F) && defined(EIGEN_VECTORIZE_AVX512VL) - _mm256_store_epi32(out, x); -#else _mm_store_si128(reinterpret_cast<__m128i*>(out), x); -#endif } template <>
diff --git a/Eigen/src/Core/products/GeneralMatrixMatrixTriangular.h b/Eigen/src/Core/products/GeneralMatrixMatrixTriangular.h index 8edb29d..e19d4fc 100644 --- a/Eigen/src/Core/products/GeneralMatrixMatrixTriangular.h +++ b/Eigen/src/Core/products/GeneralMatrixMatrixTriangular.h
@@ -163,7 +163,11 @@ gebp_kernel<LhsScalar, RhsScalar, Index, ResMapper, mr, nr, ConjLhs, ConjRhs> gebp_kernel1; gebp_kernel<LhsScalar, RhsScalar, Index, BufferMapper, mr, nr, ConjLhs, ConjRhs> gebp_kernel2; - Matrix<ResScalar, BlockSize, BlockSize, ColMajor> buffer; + static constexpr int BufferCapacity = + (EIGEN_STACK_ALLOCATION_LIMIT == 0 || BlockSize * BlockSize * sizeof(ResScalar) <= EIGEN_STACK_ALLOCATION_LIMIT) + ? int(BlockSize) + : Dynamic; + Matrix<ResScalar, BufferCapacity, BufferCapacity, ColMajor> buffer(BlockSize, BlockSize); // let's process the block per panel of actual_mc x BlockSize, // again, each is split into three parts, etc.
diff --git a/Eigen/src/Core/products/TriangularMatrixMatrix.h b/Eigen/src/Core/products/TriangularMatrixMatrix.h index a8a4e49..04a657f 100644 --- a/Eigen/src/Core/products/TriangularMatrixMatrix.h +++ b/Eigen/src/Core/products/TriangularMatrixMatrix.h
@@ -90,7 +90,12 @@ ei_declare_aligned_stack_constructed_variable(Scalar, blockA, sizeA, blocking.blockA()); ei_declare_aligned_stack_constructed_variable(Scalar, blockB, sizeB, blocking.blockB()); - Matrix<Scalar, SmallPanelWidth, SmallPanelWidth, LhsStorageOrder> triangularBuffer; + static constexpr int PanelCapacity = + (EIGEN_STACK_ALLOCATION_LIMIT == 0 || + SmallPanelWidth * SmallPanelWidth * sizeof(Scalar) <= EIGEN_STACK_ALLOCATION_LIMIT) + ? int(SmallPanelWidth) + : Dynamic; + Matrix<Scalar, PanelCapacity, PanelCapacity, LhsStorageOrder> triangularBuffer(SmallPanelWidth, SmallPanelWidth); triangularBuffer.setZero(); EIGEN_IF_CONSTEXPR ((Mode & ZeroDiag) == ZeroDiag) triangularBuffer.diagonal().setZero(); @@ -218,7 +223,12 @@ ei_declare_aligned_stack_constructed_variable(Scalar, blockA, sizeA, blocking.blockA()); ei_declare_aligned_stack_constructed_variable(Scalar, blockB, sizeB, blocking.blockB()); - Matrix<Scalar, SmallPanelWidth, SmallPanelWidth, RhsStorageOrder> triangularBuffer; + static constexpr int PanelCapacity = + (EIGEN_STACK_ALLOCATION_LIMIT == 0 || + SmallPanelWidth * SmallPanelWidth * sizeof(Scalar) <= EIGEN_STACK_ALLOCATION_LIMIT) + ? int(SmallPanelWidth) + : Dynamic; + Matrix<Scalar, PanelCapacity, PanelCapacity, RhsStorageOrder> triangularBuffer(SmallPanelWidth, SmallPanelWidth); triangularBuffer.setZero(); EIGEN_IF_CONSTEXPR ((Mode & ZeroDiag) == ZeroDiag) triangularBuffer.diagonal().setZero();
diff --git a/Eigen/src/Core/products/TriangularSolverMatrix.h b/Eigen/src/Core/products/TriangularSolverMatrix.h index 59454b6..0ec89dd 100644 --- a/Eigen/src/Core/products/TriangularSolverMatrix.h +++ b/Eigen/src/Core/products/TriangularSolverMatrix.h
@@ -250,7 +250,13 @@ Index otherStride, Index cols) { Map<Matrix<Scalar, Dynamic, Dynamic, ColMajor>, Unaligned, OuterStride<>> rest(other, size, cols, OuterStride<>(otherStride)); - Matrix<Scalar, Dynamic, Dynamic, ColMajor, Traits::WorkspaceRows, PacketSize> padded(size, Index(PacketSize)); + static constexpr int PaddedBytes = Traits::WorkspaceRows * PacketSize * int(sizeof(Scalar)); + static constexpr int PaddedMaxRows = + (EIGEN_STACK_ALLOCATION_LIMIT == 0 || PaddedBytes <= EIGEN_STACK_ALLOCATION_LIMIT) ? int(Traits::WorkspaceRows) + : Dynamic; + static constexpr int PaddedMaxCols = + (EIGEN_STACK_ALLOCATION_LIMIT == 0 || PaddedBytes <= EIGEN_STACK_ALLOCATION_LIMIT) ? int(PacketSize) : Dynamic; + Matrix<Scalar, Dynamic, Dynamic, ColMajor, PaddedMaxRows, PaddedMaxCols> padded(size, Index(PacketSize)); padded.leftCols(cols) = rest; padded.rightCols(PacketSize - cols).setZero(); solve<1>(size, a, inverse, padded.data(), size);
diff --git a/benchmarks/Core/bench_packet_ops.cpp b/benchmarks/Core/bench_packet_ops.cpp index 19f6ad3..a84387d 100644 --- a/benchmarks/Core/bench_packet_ops.cpp +++ b/benchmarks/Core/bench_packet_ops.cpp
@@ -148,6 +148,39 @@ BENCHMARK(BM_Ploadquad<float>)->Name("Ploadquad_float"); BENCHMARK(BM_Ploadquad<double>)->Name("Ploadquad_double"); +// ---- pcmp_eq ---- + +template <typename Scalar> +void BM_PcmpEq(benchmark::State& state) { + using Packet = typename packet_traits<Scalar>::type; + constexpr int N = packet_traits<Scalar>::size; + Scalar a_in[N], b_in[N], out[N]; + for (int i = 0; i < N; ++i) { + a_in[i] = static_cast<Scalar>(i); + b_in[i] = static_cast<Scalar>(i % 2 == 0 ? i : i + 1); + } + Packet a = ploadu<Packet>(a_in); + Packet b = ploadu<Packet>(b_in); + for (auto _ : state) { + benchmark::DoNotOptimize(a); + benchmark::DoNotOptimize(b); + Packet r = internal::pcmp_eq(a, b); + benchmark::DoNotOptimize(r); + } + pstoreu(out, internal::pcmp_eq(a, b)); + for (int i = 0; i < N; ++i) { + const numext::uint16_t bits = numext::bit_cast<numext::uint16_t>(out[i]); + const numext::uint16_t expected = (i % 2 == 0) ? 0xFFFFu : 0x0000u; + if (bits != expected) { + state.SkipWithError("pcmp_eq mismatch"); + return; + } + } + state.SetItemsProcessed(state.iterations() * N); +} +BENCHMARK(BM_PcmpEq<Eigen::half>)->Name("PcmpEq_half"); +BENCHMARK(BM_PcmpEq<Eigen::bfloat16>)->Name("PcmpEq_bfloat16"); + // ---- 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:
diff --git a/test/stack_allocation_limit.cpp b/test/stack_allocation_limit.cpp index fcae862..13cd4ba 100644 --- a/test/stack_allocation_limit.cpp +++ b/test/stack_allocation_limit.cpp
@@ -37,7 +37,32 @@ VERIFY(!Vector::Constant(1e200).isApprox(Vector::Constant(2e200))); } +// Triangular product and solver kernels declare internal panel buffers (TriangularMatrixMatrix, +// GeneralMatrixMatrixTriangular, TriangularSolverMatrix); they must fall back to dynamic allocation when +// EIGEN_STACK_ALLOCATION_LIMIT is smaller than the panel size. +template <typename Scalar> +void check_triangular_product_scratch() { + using Mat = Matrix<Scalar, Dynamic, Dynamic>; + Mat a = Mat::Random(16, 16); + Mat b = Mat::Random(16, 16); + Mat c = Mat::Zero(16, 16); + c.noalias() += a.template triangularView<Lower>() * b; + c.noalias() += b * a.template triangularView<Upper>(); + c.template triangularView<Lower>() += a * a.adjoint(); + // Use a non-multiple-of-PacketSize RHS column count to also exercise TriangularSolverMatrix::solve_padded. + Mat rhs = Mat::Random(16, 3); + a.diagonal().array() += Scalar(4); + Mat x = a.template triangularView<Upper>().solve(rhs); + VERIFY_IS_APPROX(a.template triangularView<Upper>() * x, rhs); + Matrix<Scalar, Dynamic, Dynamic, RowMajor> a_row = a; + Mat x_row = a_row.template triangularView<Upper>().solve(rhs); + VERIFY_IS_APPROX(a_row.template triangularView<Upper>() * x_row, rhs); +} + EIGEN_DECLARE_TEST(stack_allocation_limit) { CALL_SUBTEST(check_stack_allocation_is_disabled()); CALL_SUBTEST(check_approx_comparison_temporaries()); + CALL_SUBTEST(check_triangular_product_scratch<float>()); + CALL_SUBTEST(check_triangular_product_scratch<double>()); + CALL_SUBTEST(check_triangular_product_scratch<std::complex<double>>()); }