| // SPDX-FileCopyrightText: The Eigen Authors |
| // SPDX-License-Identifier: MPL-2.0 |
| #ifndef EIGEN_SME_VECTOR_KERNELS_H |
| #define EIGEN_SME_VECTOR_KERNELS_H |
| |
| // IWYU pragma: private |
| #include "../../InternalHeaderCheck.h" |
| |
| namespace Eigen { |
| namespace internal { |
| |
| #if EIGEN_COMP_CLANG |
| #define EIGEN_SME_VECTOR_UNROLL4 _Pragma("unroll") |
| #elif EIGEN_COMP_GNUC |
| #define EIGEN_SME_VECTOR_UNROLL4 _Pragma("GCC unroll 4") |
| #else |
| #define EIGEN_SME_VECTOR_UNROLL4 |
| #endif |
| |
| // FPCR.RMode == 0b10 (roundTowardNegative), the one mode where (+0) + (-0) is -0. |
| static EIGEN_ALWAYS_INLINE bool sme_rounds_toward_negative() { |
| std::uint64_t fpcr; |
| asm volatile("mrs %0, fpcr" : "=r"(fpcr)); |
| return ((fpcr >> 22) & 3) == 2; |
| } |
| |
| template <typename Scalar, typename Index> |
| __arm_new("za") __arm_locally_streaming |
| EIGEN_DONT_INLINE void sme_axpy(Index n, const Scalar* x, Scalar* y, Scalar alpha, Index prefix) { |
| using Traits = sme_packet_traits<Scalar>; |
| using Vec = typename Traits::type; |
| const Index lanes = Traits::size(); |
| const auto pn = Traits::ptrue_c(); |
| const auto a = pset1<Vec>(alpha); |
| Index i = 0; |
| for (; i < prefix; i += prefix - i < lanes ? prefix - i : lanes) { |
| auto active = Traits::whilelt(i, prefix); |
| pstoreu(active, y + i, pmadd(active, ploadu(active, x + i), a, ploadu(active, y + i))); |
| } |
| for (; i <= n - 16 * lanes; i += 16 * lanes) { |
| EIGEN_SME_VECTOR_UNROLL4 |
| for (int k = 0; k < 4; ++k) { |
| auto xv = ploadu_x4(pn, x + i + k * 4 * lanes); |
| sme_write_za_vg1x4(k, ploadu_x4(pn, y + i + k * 4 * lanes)); |
| sme_madd_za_vg1x4(k, xv, a); |
| } |
| EIGEN_SME_VECTOR_UNROLL4 |
| for (int k = 0; k < 4; ++k) pstoreu_x4(pn, y + i + k * 4 * lanes, sme_read_za_vg1x4<Scalar>(k)); |
| } |
| for (; i <= n - 4 * lanes; i += 4 * lanes) { |
| auto xv = ploadu_x4(pn, x + i); |
| sme_write_za_vg1x4(0, ploadu_x4(pn, y + i)); |
| sme_madd_za_vg1x4(0, xv, a); |
| pstoreu_x4(pn, y + i, sme_read_za_vg1x4<Scalar>(0)); |
| } |
| for (; i < n; i += n - i < lanes ? n - i : lanes) { |
| auto tail = Traits::whilelt(i, n); |
| pstoreu(tail, y + i, pmadd(tail, ploadu(tail, x + i), a, ploadu(tail, y + i))); |
| } |
| } |
| |
| template <typename Scalar, typename Index> |
| __arm_new("za") __arm_locally_streaming EIGEN_DONT_INLINE Scalar sme_dot(Index n, const Scalar* x, const Scalar* y) { |
| using Traits = sme_packet_traits<Scalar>; |
| using Vec = typename Traits::type; |
| const Index lanes = Traits::size(); |
| const auto pg = Traits::ptrue(); |
| const auto pn = Traits::ptrue_c(); |
| // -0 is the additive identity except under roundTowardNegative, so a zero sum gets the IEEE sign. |
| const auto negative_zero = pset1<Vec>(Scalar(-0.0)); |
| EIGEN_SME_VECTOR_UNROLL4 |
| for (int k = 0; k < 4; ++k) |
| sme_write_za_vg1x4(k, pcreate(negative_zero, negative_zero, negative_zero, negative_zero)); |
| Index i = 0; |
| for (; i <= n - 16 * lanes; i += 16 * lanes) { |
| EIGEN_SME_VECTOR_UNROLL4 |
| for (int k = 0; k < 4; ++k) |
| sme_madd_za_vg1x4(k, ploadu_x4(pn, x + i + k * 4 * lanes), ploadu_x4(pn, y + i + k * 4 * lanes)); |
| } |
| for (; i <= n - 4 * lanes; i += 4 * lanes) sme_madd_za_vg1x4(0, ploadu_x4(pn, x + i), ploadu_x4(pn, y + i)); |
| auto tail_accumulator = negative_zero; |
| for (; i < n; i += n - i < lanes ? n - i : lanes) { |
| auto tail = Traits::whilelt(i, n); |
| tail_accumulator = pmadd_m(tail, ploadu(tail, x + i), ploadu(tail, y + i), tail_accumulator); |
| } |
| auto sum = tail_accumulator; |
| EIGEN_SME_VECTOR_UNROLL4 |
| for (int k = 0; k < 4; ++k) { |
| auto v = sme_read_za_vg1x4<Scalar>(k); |
| sum = padd(pg, sum, padd(pg, padd(pg, pget<0>(v), pget<1>(v)), padd(pg, pget<2>(v), pget<3>(v)))); |
| } |
| return predux(pg, sum); |
| } |
| |
| // y += alpha * ZA group k, updated in ZA like the accumulation. |
| template <typename Scalar, typename Vec> |
| static EIGEN_ALWAYS_INLINE void sme_gemv_update(unsigned int k, svcount_t active, Scalar* y, |
| Vec alpha) __arm_streaming __arm_inout("za") { |
| const auto sum = sme_read_za_vg1x4<Scalar>(k); |
| sme_write_za_vg1x4(k, ploadu_x4(active, y)); |
| sme_madd_za_vg1x4(k, sum, alpha); |
| pstoreu_x4(active, y, sme_read_za_vg1x4<Scalar>(k)); |
| } |
| |
| template <typename Scalar, typename Index> |
| __arm_new("za") __arm_locally_streaming |
| EIGEN_DONT_INLINE void sme_gemv(Index rows, Index cols, const Scalar* a, Index stride, const Scalar* x, Scalar* y, |
| Scalar alpha, Index block_cols) { |
| using Traits = sme_packet_traits<Scalar>; |
| using Vec = typename Traits::type; |
| const Index lanes = Traits::size(); |
| const auto scale = pset1<Vec>(alpha); |
| // Match the generic GEMV's scaled column batches to bound the unscaled sums. |
| for (Index first = 0; first < cols;) { |
| const Index end = first + (cols - first < block_cols ? cols - first : block_cols); |
| // Four predicated chunks per row block keep four independent FMLA chains, also in the final partial block. |
| // An inactive chunk addresses the block start and accesses no memory. |
| for (Index i = 0; i < rows;) { |
| const Index remaining = rows - i; |
| const auto p0 = Traits::whilelt_c4(i, rows), p1 = Traits::whilelt_c4(i + std::int64_t(4) * lanes, rows), |
| p2 = Traits::whilelt_c4(i + std::int64_t(8) * lanes, rows), |
| p3 = Traits::whilelt_c4(i + std::int64_t(12) * lanes, rows); |
| const Index o1 = remaining > 4 * lanes ? 4 * lanes : 0, o2 = remaining > 8 * lanes ? 8 * lanes : 0, |
| o3 = remaining > 12 * lanes ? 12 * lanes : 0; |
| svzero_za(); |
| for (Index j = first; j < end; ++j) { |
| const auto b = pset1<Vec>(x[j]); |
| const Scalar* column = a + i + j * stride; |
| sme_madd_za_vg1x4(0, ploadu_x4(p0, column), b); |
| sme_madd_za_vg1x4(1, ploadu_x4(p1, column + o1), b); |
| sme_madd_za_vg1x4(2, ploadu_x4(p2, column + o2), b); |
| sme_madd_za_vg1x4(3, ploadu_x4(p3, column + o3), b); |
| } |
| sme_gemv_update(0, p0, y + i, scale); |
| sme_gemv_update(1, p1, y + i + o1, scale); |
| sme_gemv_update(2, p2, y + i + o2, scale); |
| sme_gemv_update(3, p3, y + i + o3, scale); |
| // Bound the final increment to avoid signed overflow with a 32-bit Index near its maximum. |
| i += remaining < 16 * lanes ? remaining : 16 * lanes; |
| } |
| first = end; |
| } |
| } |
| |
| #undef EIGEN_SME_VECTOR_UNROLL4 |
| |
| template <typename Scalar> |
| struct sme_vector_scalar : false_type {}; |
| template <> |
| struct sme_vector_scalar<float> : true_type {}; |
| #ifdef EIGEN_VECTORIZE_SME_F64F64 |
| template <> |
| struct sme_vector_scalar<double> : true_type {}; |
| #endif |
| |
| template <typename Xpr> |
| struct sme_vector_access : bool_constant<(sme_vector_scalar<typename traits<Xpr>::Scalar>::value && |
| bool(Xpr::IsVectorAtCompileTime) && has_direct_access<Xpr>::value)> {}; |
| |
| template <typename Lhs, typename Rhs> |
| struct sme_dot_supported : bool_constant<sme_vector_access<Lhs>::value && sme_vector_access<Rhs>::value && |
| is_same<typename traits<Lhs>::Scalar, typename traits<Rhs>::Scalar>::value> {}; |
| |
| // Fixed crossovers, fitted on Apple M4 (SVL 512): below them the NEON/SME memory handoff outweighs the |
| // streaming kernels. Core's cache estimates, including setCpuCacheSizes overrides, supply the others. |
| static constexpr std::size_t kSmeDotMinBytes = 16384; // per operand |
| static constexpr std::size_t kSmeAxpyMinBytes = 4096; // per operand |
| static constexpr Index kSmeGemvMinRows = 128; |
| static constexpr Index kSmeGemvMinCols = 4; |
| // A GEMV whose RHS reaches kSmeGemvWideRhsBytes needs only kSmeGemvWideMinBytes of matrix, not Core's L1 estimate. |
| static constexpr std::size_t kSmeGemvWideRhsBytes = 128; |
| static constexpr std::size_t kSmeGemvWideMinBytes = 32768; |
| |
| // Keep cache queries out of callers so the small vector fallback can inline. |
| // Per operand, the L1 estimate (divided by L1Divisor) is the lower crossover and the L2 estimate the upper one. |
| template <typename Scalar, int L1Divisor = 1> |
| EIGEN_DONT_INLINE bool sme_vector_size_suitable(Index size) { |
| std::ptrdiff_t l1, l2, l3; |
| manage_caching_sizes(GetAction, &l1, &l2, &l3); |
| return l1 > 0 && l2 > 0 && size >= 0 && |
| static_cast<std::size_t>(size) >= (static_cast<std::size_t>(l1) - 1) / (L1Divisor * sizeof(Scalar)) + 1 && |
| static_cast<std::size_t>(size) <= static_cast<std::size_t>(l2) / sizeof(Scalar); |
| } |
| |
| template <typename Lhs, typename Rhs> |
| struct default_inner_product_impl<Lhs, Rhs, true, std::enable_if_t<sme_dot_supported<Lhs, Rhs>::value>> |
| : default_inner_product_impl<Lhs, Rhs, true, false_type> { |
| using Base = default_inner_product_impl<Lhs, Rhs, true, false_type>; |
| using Scalar = typename traits<Lhs>::Scalar; |
| static EIGEN_STRONG_INLINE Scalar run(const MatrixBase<Lhs>& lhs, const MatrixBase<Rhs>& rhs) { |
| inner_product_assert<Lhs, Rhs>::run(lhs.derived(), rhs.derived()); |
| // DOT has no streaming stores for a following NEON consumer: half L1 per operand suffices. |
| if (lhs.size() >= Index(kSmeDotMinBytes / sizeof(Scalar)) && sme_vector_size_suitable<Scalar, 2>(lhs.size()) && |
| lhs.innerStride() == 1 && rhs.innerStride() == 1) { |
| sme_fpsr_guard status; |
| const Scalar result = sme_dot(lhs.size(), lhs.derived().data(), rhs.derived().data()); |
| // Under roundTowardNegative the -0 seeds turn +0 sums into -0. |
| if (result != Scalar(0) || !sme_rounds_toward_negative()) return result; |
| } |
| return Base::run(lhs, rhs); |
| } |
| }; |
| |
| template <typename Dst, typename Src> |
| EIGEN_STRONG_INLINE bool sme_try_axpy(Dst& dst, const Src& src, typename traits<Src>::Scalar alpha) { |
| using Scalar = typename traits<Src>::Scalar; |
| eigen_assert(dst.rows() == src.rows() && dst.cols() == src.cols()); |
| if (dst.size() < Index(kSmeAxpyMinBytes / sizeof(Scalar)) || !sme_vector_size_suitable<Scalar>(dst.size()) || |
| dst.innerStride() != 1 || src.innerStride() != 1) |
| return false; |
| const std::uintptr_t dst_address = reinterpret_cast<std::uintptr_t>(dst.data()); |
| const std::uintptr_t src_address = reinterpret_cast<std::uintptr_t>(src.data()); |
| const std::uintptr_t distance = dst_address > src_address ? dst_address - src_address : src_address - dst_address; |
| // Exact aliasing is coefficient-wise; partial overlap must retain the default traversal. |
| if (distance != 0 && distance / sizeof(Scalar) < static_cast<std::uintptr_t>(dst.size())) return false; |
| // Align streaming stores to 64 bytes without changing Eigen's allocation alignment. |
| sme_fpsr_guard status; |
| sme_axpy(dst.size(), src.data(), dst.data(), alpha, first_aligned<64>(dst.data(), dst.size())); |
| return true; |
| } |
| |
| template <typename Dst, typename Scalar, typename Lhs, typename Rhs> |
| struct Assignment< |
| Dst, CwiseBinaryOp<scalar_product_op<Scalar, Scalar>, Lhs, Rhs>, add_assign_op<Scalar, Scalar>, Dense2Dense, |
| std::enable_if_t<sme_vector_access<Dst>::value && |
| (sme_vector_access<Lhs>::value || sme_vector_access<Rhs>::value) && |
| blas_traits<CwiseBinaryOp<scalar_product_op<Scalar, Scalar>, Lhs, Rhs>>::HasScalarFactor && |
| is_same<Scalar, typename traits<Dst>::Scalar>::value>> { |
| using Source = CwiseBinaryOp<scalar_product_op<Scalar, Scalar>, Lhs, Rhs>; |
| using BlasTraits = blas_traits<Source>; |
| static EIGEN_STRONG_INLINE void run(Dst& dst, const Source& src, const add_assign_op<Scalar, Scalar>& func) { |
| // A direct operand restricts extraction to one scale, preserving nested products' evaluation order. |
| if (!sme_try_axpy(dst, BlasTraits::extract(src), BlasTraits::extractScalarFactor(src))) |
| Assignment<Dst, Source, add_assign_op<Scalar, Scalar>, Dense2Dense, false_type>::run(dst, src, func); |
| } |
| }; |
| |
| #ifndef EIGEN_USE_BLAS |
| template <typename Scalar, typename Index> |
| EIGEN_STRONG_INLINE bool sme_gemv_size_suitable(Index rows, Index cols) { |
| if (rows < Index(kSmeGemvMinRows) || cols < Index(kSmeGemvMinCols)) return false; |
| if (cols >= Index(kSmeGemvWideRhsBytes / sizeof(Scalar)) && |
| cols > Index(kSmeGemvWideMinBytes / sizeof(Scalar) - 1) / rows) |
| return true; |
| // Thin products amortize the NEON/SME handoff once the matrix reaches Core's L1 estimate. |
| const std::ptrdiff_t l1 = l1CacheSize(); |
| return l1 > 0 && static_cast<std::size_t>(cols) > |
| (static_cast<std::size_t>(l1) - 1) / sizeof(Scalar) / static_cast<std::size_t>(rows); |
| } |
| |
| #define EIGEN_SME_GEMV_SPECIALIZATION(Scalar) \ |
| template <typename Index, bool ConjugateLhs, bool ConjugateRhs> \ |
| struct general_matrix_vector_product<Index, Scalar, const_blas_data_mapper<Scalar, Index, ColMajor>, ColMajor, \ |
| ConjugateLhs, Scalar, const_blas_data_mapper<Scalar, Index, RowMajor>, \ |
| ConjugateRhs, Specialized> { \ |
| static EIGEN_STRONG_INLINE void run(Index rows, Index cols, \ |
| const const_blas_data_mapper<Scalar, Index, ColMajor>& lhs, \ |
| const const_blas_data_mapper<Scalar, Index, RowMajor>& rhs, Scalar* res, \ |
| Index resIncr, Scalar alpha) { \ |
| if (sme_gemv_size_suitable<Scalar>(rows, cols) && rhs.stride() == 1 && resIncr == 1) { \ |
| const std::ptrdiff_t l1 = l1CacheSize(); \ |
| const Index block_cols = cols < 128 ? cols \ |
| : (l1 > 0 && static_cast<std::size_t>(lhs.stride()) <= \ |
| (static_cast<std::size_t>(l1) - 1) / sizeof(Scalar) \ |
| ? Index(16) \ |
| : Index(4)); \ |
| /* BLAS contract: alpha == 0 leaves the result unchanged. */ \ |
| if (alpha == Scalar(0)) return; \ |
| sme_fpsr_guard status; \ |
| sme_gemv(rows, cols, lhs.data(), lhs.stride(), rhs.data(), res, alpha, block_cols); \ |
| } else { \ |
| general_matrix_vector_product<Index, Scalar, const_blas_data_mapper<Scalar, Index, ColMajor>, ColMajor, \ |
| ConjugateLhs, Scalar, const_blas_data_mapper<Scalar, Index, RowMajor>, \ |
| ConjugateRhs, BuiltIn>::run(rows, cols, lhs, rhs, res, resIncr, alpha); \ |
| } \ |
| } \ |
| }; |
| EIGEN_SME_GEMV_SPECIALIZATION(float) |
| #ifdef EIGEN_VECTORIZE_SME_F64F64 |
| EIGEN_SME_GEMV_SPECIALIZATION(double) |
| #endif |
| #undef EIGEN_SME_GEMV_SPECIALIZATION |
| #endif |
| |
| } // namespace internal |
| } // namespace Eigen |
| |
| #endif |
