diff --git a/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h b/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h index 722cdc8e3..e3cfd9ce4 100644 --- a/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h +++ b/Eigen/src/Core/arch/SME/GeneralBlockPanelKernel.h @@ -21,31 +21,191 @@ namespace internal { // --------------------------------------------------------------------------- // Streaming vector length and tile geometry. // -// The micro-kernel is organised around a logical kSmeMr x kSmeNr output -// block, packed depth-major (kSmeMr contiguous floats per depth step). Those -// dimensions are compile-time constants: they feed gebp_traits (cache -// blocking) and the packers, none of which can depend on a runtime value. +// The micro-kernel is organised around a logical mr x nr output block, packed +// depth-major (mr contiguous scalars per depth step). Those dimensions are +// compile-time constants: they feed gebp_traits (cache blocking) and the +// packers, none of which can depend on a runtime value. // -// The *physical* tiling of that block onto ZA.S tiles, on the other hand, is -// driven by the runtime streaming vector length. An fp32 ZA.S tile is -// svlw x svlw, where svlw = svcntsw() is the number of 32-bit elements in a -// streaming vector. There are only 4 ZA.S tiles, so the block is covered by -// up to a 2x2 grid of svlw x svlw tiles, iterated in sub-block passes when the -// grid is smaller than the block (and predicated down to it when larger). +// The *physical* tiling of that block onto ZA tiles, on the other hand, is +// driven by the runtime streaming vector length. A ZA tile of Scalar is +// svl x svl, where svl is the number of Scalars in a streaming vector +// (svcntsw() for fp32, svcntsd() for fp64). The block is covered by up to a +// 2x2 grid of svl x svl tiles, iterated in sub-block passes when the grid is +// smaller than the block (and predicated down to it when larger). +// +// fp32 uses the 4 ZA.S tiles, so the 2x2 grid is all of ZA. fp64 uses ZA.D, of +// which there are 8, and deliberately leaves tiles 4-7 idle: a 2x2 grid loads 2 +// packed vectors per side per depth step to feed 4 FMOPAs, i.e. 64 bytes of +// packed panel per FMOPA at either element width, and FMOPA issues at the same +// rate for both. A 2x4 grid over all eight needs a quarter less panel traffic +// per FMOPA and still measures 0.92-1.00x of the 2x2 on Apple M4, so the wider +// block is not worth its L1 footprint. // // This translation unit must be built without -msve-vector-bits (scalable/VLA // mode); see the guard in ConfigureVectorization.h for the rationale. -// Everything below derives lane counts/predicates from the runtime svlw; when a +// Everything below derives lane counts/predicates from the runtime svl; when a // block matches the tile grid exactly, the micro-kernel additionally switches // to a hand-scheduled multi-vector-load loop (see sme_process). // --------------------------------------------------------------------------- -// Logical micro-kernel block (LHS/RHS panel widths): a full 2x2 ZA-tile grid -// at SVL=512; other SVLs tile the block at runtime. If a future SVL ever -// justifies a larger block, these two constants are the only knobs -- but -// don't grow them speculatively, a doubled block measures slower at SVL=512. -static constexpr int kSmeMr = 32; -static constexpr int kSmeNr = 32; +// The per-element-width half of the ACLE surface: everything the kernel needs +// that is selected by the element type alone rather than by an argument. +// Operations that can be overloaded on their arguments are free functions +// below, and the ones taking a ZA tile number take it as a template parameter +// because the underlying instructions encode it as an immediate. +template +struct sme_traits; + +template <> +struct sme_traits { + using Vec = svfloat32_t; + using Vec2 = svfloat32x2_t; + using Vec4 = svfloat32x4_t; + // ZA.S tiles. + static constexpr int kNumTiles = 4; + static EIGEN_ALWAYS_INLINE int svl() __arm_streaming_compatible { return static_cast(svcntsw()); } + static EIGEN_ALWAYS_INLINE svbool_t whilelt(int begin, int end) __arm_streaming { + return svwhilelt_b32(static_cast(begin), static_cast(end)); + } + static EIGEN_ALWAYS_INLINE svbool_t ptrue() __arm_streaming { return svptrue_b32(); } + static EIGEN_ALWAYS_INLINE svcount_t ptrue_c() __arm_streaming { return svptrue_c32(); } + static EIGEN_ALWAYS_INLINE Vec dup(float x) __arm_streaming { return svdup_f32(x); } +}; + +#ifdef EIGEN_VECTORIZE_SME_F64F64 +template <> +struct sme_traits { + using Vec = svfloat64_t; + using Vec2 = svfloat64x2_t; + using Vec4 = svfloat64x4_t; + // ZA.D tiles. + static constexpr int kNumTiles = 8; + static EIGEN_ALWAYS_INLINE int svl() __arm_streaming_compatible { return static_cast(svcntsd()); } + static EIGEN_ALWAYS_INLINE svbool_t whilelt(int begin, int end) __arm_streaming { + return svwhilelt_b64(static_cast(begin), static_cast(end)); + } + static EIGEN_ALWAYS_INLINE svbool_t ptrue() __arm_streaming { return svptrue_b64(); } + static EIGEN_ALWAYS_INLINE svcount_t ptrue_c() __arm_streaming { return svptrue_c64(); } + static EIGEN_ALWAYS_INLINE Vec dup(double x) __arm_streaming { return svdup_f64(x); } +}; +#endif + +// Contiguous predicated load/store, fused multiply-add and multiply. +static EIGEN_ALWAYS_INLINE svfloat32_t sme_ld1(svbool_t pg, const float* p) __arm_streaming { return svld1_f32(pg, p); } +static EIGEN_ALWAYS_INLINE void sme_st1(svbool_t pg, float* p, svfloat32_t v) __arm_streaming { svst1_f32(pg, p, v); } +static EIGEN_ALWAYS_INLINE svfloat32x2_t sme_ld1_x2(svcount_t pn, const float* p) __arm_streaming { + return svld1_f32_x2(pn, p); +} +static EIGEN_ALWAYS_INLINE svfloat32x4_t sme_ld1_x4(svcount_t pn, const float* p) __arm_streaming { + return svld1_f32_x4(pn, p); +} +static EIGEN_ALWAYS_INLINE svfloat32_t sme_mla(svbool_t pg, svfloat32_t acc, svfloat32_t a, + svfloat32_t b) __arm_streaming { + return svmla_f32_x(pg, acc, a, b); +} +static EIGEN_ALWAYS_INLINE svfloat32_t sme_mul(svbool_t pg, svfloat32_t a, svfloat32_t b) __arm_streaming { + return svmul_f32_x(pg, a, b); +} +template +static EIGEN_ALWAYS_INLINE svfloat32_t sme_get(svfloat32x2_t v) __arm_streaming { + return svget2_f32(v, Lane); +} +template +static EIGEN_ALWAYS_INLINE svfloat32_t sme_get(svfloat32x4_t v) __arm_streaming { + return svget4_f32(v, Lane); +} + +// ZA tile access. The tile number is an instruction immediate, hence a template +// parameter; the slice number is a register operand and stays a value. +template +static EIGEN_ALWAYS_INLINE void sme_ld1_hor_za(uint32_t slice, svbool_t pg, + const float* p) __arm_streaming __arm_inout("za") { + svld1_hor_za32(Tile, slice, pg, p); +} +template +static EIGEN_ALWAYS_INLINE svfloat32_t sme_read_hor_za(svfloat32_t zero, svbool_t pg, + uint32_t slice) __arm_streaming __arm_inout("za") { + return svread_hor_za32_f32_m(zero, pg, Tile, slice); +} +template +static EIGEN_ALWAYS_INLINE svfloat32_t sme_read_ver_za(svfloat32_t zero, svbool_t pg, + uint32_t slice) __arm_streaming __arm_inout("za") { + return svread_ver_za32_f32_m(zero, pg, Tile, slice); +} +template +static EIGEN_ALWAYS_INLINE void sme_mopa(svbool_t pm, svbool_t pn, svfloat32_t a, + svfloat32_t b) __arm_streaming __arm_inout("za") { + svmopa_za32_f32_m(Tile, pm, pn, a, b); +} + +#ifdef EIGEN_VECTORIZE_SME_F64F64 +static EIGEN_ALWAYS_INLINE svfloat64_t sme_ld1(svbool_t pg, const double* p) __arm_streaming { + return svld1_f64(pg, p); +} +static EIGEN_ALWAYS_INLINE void sme_st1(svbool_t pg, double* p, svfloat64_t v) __arm_streaming { svst1_f64(pg, p, v); } +static EIGEN_ALWAYS_INLINE svfloat64x2_t sme_ld1_x2(svcount_t pn, const double* p) __arm_streaming { + return svld1_f64_x2(pn, p); +} +static EIGEN_ALWAYS_INLINE svfloat64x4_t sme_ld1_x4(svcount_t pn, const double* p) __arm_streaming { + return svld1_f64_x4(pn, p); +} +static EIGEN_ALWAYS_INLINE svfloat64_t sme_mla(svbool_t pg, svfloat64_t acc, svfloat64_t a, + svfloat64_t b) __arm_streaming { + return svmla_f64_x(pg, acc, a, b); +} +static EIGEN_ALWAYS_INLINE svfloat64_t sme_mul(svbool_t pg, svfloat64_t a, svfloat64_t b) __arm_streaming { + return svmul_f64_x(pg, a, b); +} +template +static EIGEN_ALWAYS_INLINE svfloat64_t sme_get(svfloat64x2_t v) __arm_streaming { + return svget2_f64(v, Lane); +} +template +static EIGEN_ALWAYS_INLINE svfloat64_t sme_get(svfloat64x4_t v) __arm_streaming { + return svget4_f64(v, Lane); +} + +template +static EIGEN_ALWAYS_INLINE void sme_ld1_hor_za(uint32_t slice, svbool_t pg, + const double* p) __arm_streaming __arm_inout("za") { + svld1_hor_za64(Tile, slice, pg, p); +} +template +static EIGEN_ALWAYS_INLINE svfloat64_t sme_read_hor_za(svfloat64_t zero, svbool_t pg, + uint32_t slice) __arm_streaming __arm_inout("za") { + return svread_hor_za64_f64_m(zero, pg, Tile, slice); +} +template +static EIGEN_ALWAYS_INLINE svfloat64_t sme_read_ver_za(svfloat64_t zero, svbool_t pg, + uint32_t slice) __arm_streaming __arm_inout("za") { + return svread_ver_za64_f64_m(zero, pg, Tile, slice); +} +template +static EIGEN_ALWAYS_INLINE void sme_mopa(svbool_t pm, svbool_t pn, svfloat64_t a, + svfloat64_t b) __arm_streaming __arm_inout("za") { + svmopa_za64_f64_m(Tile, pm, pn, a, b); +} +#endif // EIGEN_VECTORIZE_SME_F64F64 + +// Logical micro-kernel block (LHS/RHS panel widths): a full 2x2 ZA-tile grid at +// the 512-bit design point, where a streaming vector holds 64 / sizeof(Scalar) +// scalars. Other SVLs tile the block at runtime. If a future SVL ever +// justifies a larger block, this is the only knob -- but don't grow it +// speculatively, a doubled block measures slower at SVL=512. +static constexpr int kSmeDesignVectorBytes = 64; + +template +struct sme_block { + static constexpr int mr = 2 * kSmeDesignVectorBytes / int(sizeof(Scalar)); + static constexpr int nr = mr; +}; + +static constexpr int kSmeMr = sme_block::mr; +static constexpr int kSmeNr = sme_block::nr; +#ifdef EIGEN_VECTORIZE_SME_F64F64 +static constexpr int kSmeMrD = sme_block::mr; +static constexpr int kSmeNrD = sme_block::nr; +#endif // min() usable from streaming functions (numext::mini lacks the // __arm_streaming_compatible attribute). @@ -58,101 +218,129 @@ static EIGEN_ALWAYS_INLINE T sme_min(T a, T b) __arm_streaming_compatible { // packed panel of width `width`, for the depth sub-range [k0, k1). Both dst and // src are indexed by the absolute depth index k (dst[k*width+off], // src[k*src_stride+off]); the caller offsets `src` to the region's column base -// and `dst` to the panel base. Generalised over the runtime svlw: the panel is -// covered in svlw-wide column chunks, each streamed over the depth sub-range. +// and `dst` to the panel base. Generalised over the runtime svl: the panel is +// covered in svl-wide column chunks, each streamed over the depth sub-range. // The chunk loop is outermost so each chunk's predicate is computed once instead // of per depth step (the runtime chunk count keeps the compiler from hoisting it // on its own). The symm packers reuse this for the diagonal-split direct/ // transposed regions (a contiguous depth sub-range at a depth offset). -template -static EIGEN_ALWAYS_INLINE void sve_copy_panel_range(float* EIGEN_RESTRICT dst, const float* EIGEN_RESTRICT src, +template +static EIGEN_ALWAYS_INLINE void sve_copy_panel_range(Scalar* EIGEN_RESTRICT dst, const Scalar* EIGEN_RESTRICT src, Index src_stride, Index k0, Index k1, int width) __arm_streaming { - const int svlw = static_cast(svcntsw()); - for (int off = 0; off < width; off += svlw) { - const int w = sme_min(width - off, svlw); - const svbool_t pred = svwhilelt_b32(uint32_t(0), uint32_t(w)); + const int svl = sme_traits::svl(); + for (int off = 0; off < width; off += svl) { + const int w = sme_min(width - off, svl); + const svbool_t pred = sme_traits::whilelt(0, w); for (Index k = k0; k < k1; ++k) { - svst1_f32(pred, &dst[k * width + off], svld1_f32(pred, &src[k * src_stride + off])); + sme_st1(pred, &dst[k * width + off], sme_ld1(pred, &src[k * src_stride + off])); } } } // Copy the full depth [0, depth): thin wrapper used by the (non-symm) gemm // packers, which always pack a whole panel. -template -static EIGEN_ALWAYS_INLINE void sve_copy_panel(float* EIGEN_RESTRICT dst, const float* EIGEN_RESTRICT src, +template +static EIGEN_ALWAYS_INLINE void sve_copy_panel(Scalar* EIGEN_RESTRICT dst, const Scalar* EIGEN_RESTRICT src, Index src_stride, Index depth, int width) __arm_streaming { sve_copy_panel_range(dst, src, src_stride, Index(0), depth, width); } // Transpose-pack `width` source rows into depth-major packed output using ZA's -// 2D store as a free transpose, for the depth sub-range [k0, k1): a svlw x svlw -// block of source (svlw rows x svlw depth) is loaded as horizontal ZA slices, +// 2D store as a free transpose, for the depth sub-range [k0, k1): a svl x svl +// block of source (svl rows x svl depth) is loaded as horizontal ZA slices, // then read back as vertical slices, which emits it depth-major. Row-groups of -// svlw rows are processed two at a time through ZA tiles 0 and 1: ZA is not +// svl rows are processed two at a time through ZA tiles 0 and 1: ZA is not // renamed, so a single tile would stall every load pass on the previous read // pass (write-after-read); two tiles in flight keep the phases independent. -// Trailing row-groups (when width is not a multiple of 2*svlw) use tile 0 with +// Trailing row-groups (when width is not a multiple of 2*svl) use tile 0 with // predicated rows. Both dst and src are indexed by the absolute depth index k: // dst[k*width + r] = src[r*src_stride + k], k in [k0,k1), r in [0,width). // The symm packers reuse this for the diagonal-split transposed/direct regions -// (a depth sub-range at a depth offset, with a tail-panel width < kSmeMr). -template -static EIGEN_ALWAYS_INLINE void sme_transpose_pack_range(float* EIGEN_RESTRICT dst, const float* EIGEN_RESTRICT src, +// (a depth sub-range at a depth offset, with a tail-panel width < mr). +template +static EIGEN_ALWAYS_INLINE void sme_transpose_pack_range(Scalar* EIGEN_RESTRICT dst, const Scalar* EIGEN_RESTRICT src, Index src_stride, Index k0, Index k1, int width) __arm_streaming __arm_inout("za") { - const svfloat32_t zero = svdup_f32(0.f); - const svbool_t pg_all = svptrue_b32(); - const int svlw = static_cast(svcntsw()); + using Traits = sme_traits; + const typename Traits::Vec zero = Traits::dup(Scalar(0)); + const svbool_t pg_all = Traits::ptrue(); + const int svl = Traits::svl(); - for (Index k = k0; k < k1; k += svlw) { - const int dk = static_cast(sme_min(k1 - k, Index(svlw))); - const svbool_t pg_d = svwhilelt_b32(uint32_t(0), uint32_t(dk)); + for (Index k = k0; k < k1; k += svl) { + const int dk = static_cast(sme_min(k1 - k, Index(svl))); + const svbool_t pg_d = Traits::whilelt(0, dk); int r0 = 0; // Pairs of full row-groups: tiles 0 and 1 in flight. - for (; r0 + 2 * svlw <= width; r0 += 2 * svlw) { - for (int r = 0; r < svlw; ++r) { - svld1_hor_za32(0, uint32_t(r), pg_d, &src[(r0 + r) * src_stride + k]); - svld1_hor_za32(1, uint32_t(r), pg_d, &src[(r0 + svlw + r) * src_stride + k]); + for (; r0 + 2 * svl <= width; r0 += 2 * svl) { + for (int r = 0; r < svl; ++r) { + sme_ld1_hor_za<0>(uint32_t(r), pg_d, &src[(r0 + r) * src_stride + k]); + sme_ld1_hor_za<1>(uint32_t(r), pg_d, &src[(r0 + svl + r) * src_stride + k]); } for (int c = 0; c < dk; ++c) { - svst1_f32(pg_all, &dst[(k + c) * width + r0], svread_ver_za32_f32_m(zero, pg_all, 0, uint32_t(c))); - svst1_f32(pg_all, &dst[(k + c) * width + r0 + svlw], svread_ver_za32_f32_m(zero, pg_all, 1, uint32_t(c))); + sme_st1(pg_all, &dst[(k + c) * width + r0], sme_read_ver_za<0>(zero, pg_all, uint32_t(c))); + sme_st1(pg_all, &dst[(k + c) * width + r0 + svl], sme_read_ver_za<1>(zero, pg_all, uint32_t(c))); } } - // Trailing row-groups (at most two svlw-wide passes remain, since the pair - // loop consumed all multiples of 2*svlw): predicate down to the remaining + // Trailing row-groups (at most two svl-wide passes remain, since the pair + // loop consumed all multiples of 2*svl): predicate down to the remaining // rows. A single `if` would drop rows when a tail width lands in - // (svlw, 2*svlw); a loop handles any leftover. - for (; r0 < width; r0 += svlw) { - const int rg = sme_min(width - r0, svlw); - const svbool_t pg_r = svwhilelt_b32(uint32_t(0), uint32_t(rg)); + // (svl, 2*svl); a loop handles any leftover. + for (; r0 < width; r0 += svl) { + const int rg = sme_min(width - r0, svl); + const svbool_t pg_r = Traits::whilelt(0, rg); for (int r = 0; r < rg; ++r) { - svld1_hor_za32(0, uint32_t(r), pg_d, &src[(r0 + r) * src_stride + k]); + sme_ld1_hor_za<0>(uint32_t(r), pg_d, &src[(r0 + r) * src_stride + k]); } for (int c = 0; c < dk; ++c) { - svst1_f32(pg_r, &dst[(k + c) * width + r0], svread_ver_za32_f32_m(zero, pg_r, 0, uint32_t(c))); + sme_st1(pg_r, &dst[(k + c) * width + r0], sme_read_ver_za<0>(zero, pg_r, uint32_t(c))); } } } } -// Transpose-pack a whole kSmeMr-wide panel over the full depth [0, depth): thin -// wrapper used by the (non-symm) gemm packers. -template -static EIGEN_ALWAYS_INLINE void sme_transpose_pack(float* EIGEN_RESTRICT dst, const float* EIGEN_RESTRICT src, - Index src_stride, Index depth) __arm_streaming __arm_inout("za") { - static_assert(kSmeMr == kSmeNr, "SME transpose pack assumes square panels"); - sme_transpose_pack_range(dst, src, src_stride, Index(0), depth, kSmeMr); +// Transpose-pack a whole `width`-wide panel over the full depth [0, depth): +// thin wrapper used by the (non-symm) gemm packers. +template +static EIGEN_ALWAYS_INLINE void sme_transpose_pack(Scalar* EIGEN_RESTRICT dst, const Scalar* EIGEN_RESTRICT src, + Index src_stride, Index depth, + int width) __arm_streaming __arm_inout("za") { + sme_transpose_pack_range(dst, src, src_stride, Index(0), depth, width); } -template -static EIGEN_ALWAYS_INLINE void scalar_tail_pack(float* EIGEN_RESTRICT dst_panel, const float* EIGEN_RESTRICT src, - Index src_stride, Index depth, Index tail) __arm_streaming { - for (Index k = 0; k < depth; ++k) { - for (Index i = 0; i < tail; ++i) { - dst_panel[k * tail + i] = src[i * src_stride + k]; +// Transposing copy for a panel narrower than the pack width: +// dst_panel[k*tail + i] = src[i*src_stride + k]. +// Kept outside the caller's __arm_locally_streaming region: it needs neither SVE +// nor ZA, and streaming mode runs scalar floating-point ~40x slower on Apple M4. +// Outside it the source rows are contiguous in k, so PacketSize of them +// transpose in register as in sme_pack_rhs_fallback; a product with cols < nr is +// packed entirely here. +template +static void tail_transpose_pack(Scalar* EIGEN_RESTRICT dst_panel, const Scalar* EIGEN_RESTRICT src, Index src_stride, + Index depth, Index tail) { + using Packet = typename packet_traits::type; + constexpr int PacketSize = int(packet_traits::size); + const Index peeled_tail = (tail / Index(PacketSize)) * Index(PacketSize); + const Index peeled_depth = (depth / Index(PacketSize)) * Index(PacketSize); + + Index i = 0; + for (; i < peeled_tail; i += Index(PacketSize)) { + Index k = 0; + for (; k < peeled_depth; k += Index(PacketSize)) { + PacketBlock block; + for (int p = 0; p < PacketSize; ++p) { + block.packet[p] = ploadu(src + (i + Index(p)) * src_stride + k); + } + ptranspose(block); + for (int p = 0; p < PacketSize; ++p) { + pstoreu(dst_panel + (k + Index(p)) * tail + i, block.packet[p]); + } } + for (; k < depth; ++k) { + for (Index p = 0; p < Index(PacketSize); ++p) dst_panel[k * tail + i + p] = src[(i + p) * src_stride + k]; + } + } + for (; i < tail; ++i) { + for (Index k = 0; k < depth; ++k) dst_panel[k * tail + i] = src[i * src_stride + k]; } } @@ -207,19 +395,18 @@ struct sme_mapper_has_direct_access { // for ColMajor tensor and blas mappers, but not for RowMajor mappers, whose // packets run along the storage-inner second index. RowMajor dispatches pass // vectorise = false and take the scalar element loop. -template -void sme_pack_lhs_fallback(float* dst_base, const DataMapper& lhs, Index depth, Index rows, Index dst_stride, +template +void sme_pack_lhs_fallback(Scalar* dst_base, const DataMapper& lhs, Index depth, Index rows, Index dst_stride, Index dst_offset, bool vectorise) { - typedef typename packet_traits::type Packet; - constexpr Index PacketSize = Index(packet_traits::size); - constexpr int MR = kSmeMr; + using Packet = typename packet_traits::type; + constexpr Index PacketSize = Index(packet_traits::size); for (Index i = 0; i < rows; i += MR) { const Index w = numext::mini(rows - i, Index(MR)); - float* dst_panel = PanelMode ? dst_base + i * dst_stride + dst_offset * w : dst_base + i * depth; + Scalar* dst_panel = PanelMode ? dst_base + i * dst_stride + dst_offset * w : dst_base + i * depth; const Index peeled_w = vectorise ? (w / PacketSize) * PacketSize : Index(0); for (Index k = 0; k < depth; ++k) { - float* dst_row = dst_panel + k * w; + Scalar* dst_row = dst_panel + k * w; Index r = 0; for (; r < peeled_w; r += PacketSize) { pstoreu(dst_row + r, lhs.template loadPacket(i + r, k)); @@ -231,47 +418,51 @@ void sme_pack_lhs_fallback(float* dst_base, const DataMapper& lhs, Index depth, } } +// The PacketSize column sub-mappers one packed column group loads from. +template +EIGEN_ALWAYS_INLINE std::array sme_column_mappers( + const DataMapper& rhs, Index col, std::index_sequence) { + return {{rhs.getLinearMapper(0, col + Index(Is))...}}; +} + // RHS fallback, mirroring sme_pack_lhs_fallback (including the vectorise // contract: LinearMapper packets must advance the first (depth) index). The // packed layout wants consecutive columns contiguous while the mapper's // packets run along the depth k, so PacketSize columns are loaded as packets // along k and transposed in-register (the same LinearMapper + ptranspose // scheme as the generic gemm_pack_rhs). -template -void sme_pack_rhs_fallback(float* dst_base, const DataMapper& rhs, Index depth, Index cols, Index dst_stride, +template +void sme_pack_rhs_fallback(Scalar* dst_base, const DataMapper& rhs, Index depth, Index cols, Index dst_stride, Index dst_offset, bool vectorise) { - typedef typename packet_traits::type Packet; - typedef typename DataMapper::LinearMapper LinearMapper; - constexpr Index PacketSize = Index(packet_traits::size); - static_assert(PacketSize == 4, "the transpose micro-block below is written for 4-wide NEON packets"); - constexpr int NR = kSmeNr; - const Index peeled_depth = (depth / PacketSize) * PacketSize; + using Packet = typename packet_traits::type; + using LinearMapper = typename DataMapper::LinearMapper; + constexpr int PacketSize = int(packet_traits::size); + const Index peeled_depth = (depth / Index(PacketSize)) * Index(PacketSize); for (Index j = 0; j < cols; j += NR) { const Index w = numext::mini(cols - j, Index(NR)); - float* dst_panel = PanelMode ? dst_base + j * dst_stride + dst_offset * w : dst_base + j * depth; - const Index peeled_w = vectorise ? (w / PacketSize) * PacketSize : Index(0); + Scalar* dst_panel = PanelMode ? dst_base + j * dst_stride + dst_offset * w : dst_base + j * depth; + const Index peeled_w = vectorise ? (w / Index(PacketSize)) * Index(PacketSize) : Index(0); Index c = 0; - for (; c < peeled_w; c += PacketSize) { - const LinearMapper dm0 = rhs.getLinearMapper(0, j + c + 0); - const LinearMapper dm1 = rhs.getLinearMapper(0, j + c + 1); - const LinearMapper dm2 = rhs.getLinearMapper(0, j + c + 2); - const LinearMapper dm3 = rhs.getLinearMapper(0, j + c + 3); + for (; c < peeled_w; c += Index(PacketSize)) { + // Loop-invariant in k, but not hoisted out of the k loop by the compiler + // for a mapper that returns its sub-mappers by value -- which is the hot + // path here: tensor contractions pack through TensorContractionSubMapper. + const std::array dm = + sme_column_mappers(rhs, j + c, std::make_index_sequence{}); Index k = 0; - for (; k < peeled_depth; k += PacketSize) { - PacketBlock block; - block.packet[0] = dm0.template loadPacket(k); - block.packet[1] = dm1.template loadPacket(k); - block.packet[2] = dm2.template loadPacket(k); - block.packet[3] = dm3.template loadPacket(k); + for (; k < peeled_depth; k += Index(PacketSize)) { + PacketBlock block; + for (int p = 0; p < PacketSize; ++p) { + block.packet[p] = dm[p].template loadPacket(k); + } ptranspose(block); - pstoreu(dst_panel + (k + 0) * w + c, block.packet[0]); - pstoreu(dst_panel + (k + 1) * w + c, block.packet[1]); - pstoreu(dst_panel + (k + 2) * w + c, block.packet[2]); - pstoreu(dst_panel + (k + 3) * w + c, block.packet[3]); + for (int p = 0; p < PacketSize; ++p) { + pstoreu(dst_panel + (k + Index(p)) * w + c, block.packet[p]); + } } for (; k < depth; ++k) { - for (Index p = 0; p < PacketSize; ++p) { + for (Index p = 0; p < Index(PacketSize); ++p) { dst_panel[k * w + c + p] = rhs(k, j + c + p); } } @@ -289,50 +480,53 @@ void sme_pack_rhs_fallback(float* dst_base, const DataMapper& rhs, Index depth, // packet/element fallback. Tag-dispatched so &m(0,0) is only compiled for // lvalue mappers. UsePacketPath records whether the mapper's packets advance // the index the fallback needs, independently of its direct-access category. -template -EIGEN_ALWAYS_INLINE void sme_dispatch_pack(DirectFn direct, FallbackFn fallback, float* block, const DataMapper& m, +template +EIGEN_ALWAYS_INLINE void sme_dispatch_pack(DirectFn direct, FallbackFn fallback, Scalar* block, const DataMapper& m, Index depth, Index n, Index stride, Index offset, std::true_type /* direct access */) { if (sme_mapper_incr(m) == 1) { - const float* src = (n > 0 && depth > 0) ? &m(0, 0) : nullptr; + const Scalar* src = (n > 0 && depth > 0) ? &m(0, 0) : nullptr; direct(block, src, m.stride(), depth, n, stride, offset); } else { fallback(block, m, depth, n, stride, offset, UsePacketPath); } } -template -EIGEN_ALWAYS_INLINE void sme_dispatch_pack(DirectFn, FallbackFn fallback, float* block, const DataMapper& m, +template +EIGEN_ALWAYS_INLINE void sme_dispatch_pack(DirectFn, FallbackFn fallback, Scalar* block, const DataMapper& m, Index depth, Index n, Index stride, Index offset, std::false_type /* no direct access */) { fallback(block, m, depth, n, stride, offset, UsePacketPath); } /***************************************************************************** - * gebp_traits specialization for SME (float x float) + * gebp_traits specializations for SME (float x float, double x double) * - * Overrides mr and nr so that: + * Override mr and nr so that: * - gemm_pack_lhs receives Pack1 = mr, creating uniform LHS panels * - gemm_pack_rhs receives nr, creating uniform RHS panels * - mc is rounded to a multiple of mr, nc to a multiple of nr * - Cache blocking (kc, mc, nc) is recomputed accordingly * - * We provide custom gemm_pack_lhs/gemm_pack_rhs specializations for float, - * so both ColMajor and RowMajor source matrices produce an identical, + * We provide custom gemm_pack_lhs/gemm_pack_rhs specializations for both + * scalars, so both ColMajor and RowMajor source matrices produce an identical, * simple packed format that the SME kernel consumes. * * Mixed-scalar products (e.g. MatrixXf * MatrixXcf) also instantiate * gemm_pack_lhs, but with Pack1/nr from the generic * gebp_traits> (mr=6, nr=4) and are consumed by the * generic gebp_kernel, not the SME one. So the specializations below pin - * Pack1/nr_ to kSmeMr/kSmeNr: only the instantiation that feeds the SME + * Pack1/nr_ to the SME block sizes: only the instantiation that feeds the SME * gebp_kernel matches; mixed-scalar ones fall through to the generic template. - * This is load-bearing: it relies on no other float consumer instantiating the - * packer with mr == kSmeMr (holds today -- generic float traits give mr <= 12). - * The kernel side is self-checking (the SME gebp_kernel static_asserts - * mr/nr == kSmeMr/kSmeNr, so a float traits change breaks the build instead of - * silently mispairing packer and kernel); the packer side is enforced by the - * static_asserts below for the in-tree mixed-scalar traits (downstream code - * instantiating the packers with hand-picked mr/nr remains uncovered). + * This is load-bearing: it relies on no other consumer of the same scalar + * instantiating the packer with mr == the SME block size (holds today -- + * generic float traits give mr <= 12). The kernel side is self-checking (the + * SME gebp_kernel static_asserts mr/nr against the block sizes, so a traits + * change breaks the build instead of silently mispairing packer and kernel); + * the packer side is enforced by the static_asserts below for the in-tree + * mixed-scalar traits (downstream code instantiating the packers with + * hand-picked mr/nr remains uncovered). *****************************************************************************/ template <> @@ -347,7 +541,7 @@ class gebp_traits>::mr) != kSmeMr, @@ -357,31 +551,44 @@ static_assert(int(gebp_traits, float>::nr) != kSmeNr, "gebp_traits, float>::nr collides with kSmeNr: the SME gemm_pack_rhs would silently " "emit SME panel layout for the generic gebp_kernel"); +#ifdef EIGEN_VECTORIZE_SME_F64F64 +template <> +class gebp_traits + : public gebp_traits { + public: + // As above, only the register-block sizes are overridden. + static constexpr int mr = kSmeMrD; + static constexpr int nr = kSmeNrD; +}; + +static_assert(int(gebp_traits>::mr) != kSmeMrD, + "gebp_traits>::mr collides with kSmeMrD: the SME gemm_pack_lhs would silently " + "emit SME panel layout for the generic gebp_kernel"); +static_assert(int(gebp_traits, double>::nr) != kSmeNrD, + "gebp_traits, double>::nr collides with kSmeNrD: the SME gemm_pack_rhs would silently " + "emit SME panel layout for the generic gebp_kernel"); +#endif + /***************************************************************************** - * gemm_pack_lhs specialization for SME (float, ColMajor) + * gemm_pack_lhs for SME (ColMajor source) * - * Packs the LHS matrix into uniform panels of width mr = kSmeMr. - * Each depth step k writes exactly MR contiguous floats. + * Packs the LHS matrix into uniform panels of width mr. + * Each depth step k writes exactly MR contiguous scalars. *****************************************************************************/ -// Pack1 is pinned to kSmeMr (rather than left open) so this specialization -// only matches consumers that actually feed the SME gebp_kernel -- see -// "Mixed-scalar products" in the gebp_traits doc comment above. -template -struct gemm_pack_lhs { - typedef float Scalar; - // Conjugate is deliberately ignored: conj is the identity for real scalars, - // and Conjugate=true instantiations do occur (e.g. the SYMM above-diagonal - // transposed pack). A complex port of these packers must actually conjugate. +template +struct sme_pack_lhs_colmajor { + // Conjugate is deliberately ignored by the specializations below: conj is the + // identity for real scalars, and Conjugate=true instantiations do occur (e.g. + // the SYMM above-diagonal transposed pack). A complex port of these packers + // must actually conjugate. static_assert(!NumTraits::IsComplex, "the SME packers only support real scalars"); - __arm_locally_streaming static void pack_lhs_colmajor(Scalar* dst_base, const Scalar* EIGEN_RESTRICT src, - Index src_stride, Index depth, Index rows, Index dst_stride, - Index dst_offset) { - constexpr int MR = kSmeMr; + __arm_locally_streaming static void pack_direct(Scalar* dst_base, const Scalar* EIGEN_RESTRICT src, Index src_stride, + Index depth, Index rows, Index dst_stride, Index dst_offset) { const Index peeled_rows = (rows / MR) * MR; - // Full panels of width MR, streamed in svlw-wide predicated chunks. + // Full panels of width MR, streamed in svl-wide predicated chunks. for (Index i = 0; i < peeled_rows; i += MR) { Scalar* dst_panel = PanelMode ? dst_base + i * dst_stride + dst_offset * MR : dst_base + i * depth; sve_copy_panel(dst_panel, src + i, src_stride, depth, MR); @@ -403,7 +610,7 @@ struct gemm_pack_lhs(&pack_lhs_colmajor, &sme_pack_lhs_fallback, blockA, lhs, + sme_dispatch_pack(&pack_direct, &sme_pack_lhs_fallback, blockA, lhs, depth, rows, stride, offset, bool_constant::value>{}); } @@ -414,38 +621,42 @@ struct gemm_pack_lhs -struct gemm_pack_lhs { - typedef float Scalar; - // See the ColMajor specialization above: Conjugate is ignored, sound only - // for real scalars. +template +struct sme_pack_lhs_rowmajor { + // See sme_pack_lhs_colmajor: Conjugate is ignored, sound only for real scalars. static_assert(!NumTraits::IsComplex, "the SME packers only support real scalars"); - __arm_locally_streaming __arm_new("za") static void pack_lhs_rowmajor(Scalar* dst_base, - const Scalar* EIGEN_RESTRICT src, - Index src_stride, Index depth, Index rows, - Index dst_stride, Index dst_offset) { - constexpr int MR = kSmeMr; - const Index peeled_rows = (rows / MR) * MR; - + __arm_locally_streaming __arm_new("za") static void pack_full_panels(Scalar* dst_base, + const Scalar* EIGEN_RESTRICT src, + Index src_stride, Index depth, Index peeled_rows, + Index dst_stride, Index dst_offset) { for (Index i = 0; i < peeled_rows; i += MR) { Scalar* dst_panel = PanelMode ? dst_base + i * dst_stride + dst_offset * MR : dst_base + i * depth; - sme_transpose_pack(dst_panel, src + i * src_stride, src_stride, depth); + sme_transpose_pack(dst_panel, src + i * src_stride, src_stride, depth, MR); + } + } + + static void pack_direct(Scalar* dst_base, const Scalar* EIGEN_RESTRICT src, Index src_stride, Index depth, Index rows, + Index dst_stride, Index dst_offset) { + const Index peeled_rows = (rows / MR) * MR; + + if (peeled_rows > 0) { + pack_full_panels(dst_base, src, src_stride, depth, peeled_rows, dst_stride, dst_offset); } // Row tail (rows - peeled_rows in [1, MR-1]). This branch runs at most // once per pack_lhs call with < MR rows and would need a partial-ZA-tile // dance to vectorise; total copies are < MR * depth per call, which is - // noise vs the main packer's workload, so scalar is the simple choice. + // noise vs the main packer's workload, so scalar is the simple choice -- + // taken outside the streaming region above (see tail_transpose_pack). if (peeled_rows < rows) { const Index tail = rows - peeled_rows; Scalar* dst_panel = PanelMode ? dst_base + peeled_rows * dst_stride + dst_offset * tail : dst_base + peeled_rows * depth; - scalar_tail_pack(dst_panel, src + peeled_rows * src_stride, src_stride, depth, tail); + tail_transpose_pack(dst_panel, src + peeled_rows * src_stride, src_stride, depth, tail); } } @@ -457,50 +668,53 @@ struct gemm_pack_lhs(&pack_lhs_rowmajor, &sme_pack_lhs_fallback, blockA, lhs, - depth, rows, stride, offset, + sme_dispatch_pack(&pack_direct, &sme_pack_lhs_fallback, blockA, + lhs, depth, rows, stride, offset, bool_constant::value>{}); } }; /***************************************************************************** - * gemm_pack_rhs specialization for SME (float, ColMajor) -- SME in-ZA - * transpose, mirroring the RowMajor LHS packer. + * gemm_pack_rhs for SME (ColMajor source) -- SME in-ZA transpose, mirroring + * the RowMajor LHS packer. * - * Packs the RHS matrix into panels of width nr = kSmeNr. ColMajor source has + * Packs the RHS matrix into panels of width nr. ColMajor source has * columns contiguous; we load NR columns as horizontal ZA slices and then * read verticals to produce depth-major packed output. *****************************************************************************/ -// nr_ is pinned to kSmeNr (rather than left open) so this specialization -// only matches consumers that actually feed the SME gebp_kernel -- see -// "Mixed-scalar products" in the gebp_traits doc comment above. -template -struct gemm_pack_rhs { - typedef float Scalar; - // See gemm_pack_lhs above: Conjugate is ignored, sound only for real scalars. +template +struct sme_pack_rhs_colmajor { + // See sme_pack_lhs_colmajor: Conjugate is ignored, sound only for real scalars. static_assert(!NumTraits::IsComplex, "the SME packers only support real scalars"); - __arm_locally_streaming __arm_new("za") static void pack_rhs_colmajor(Scalar* dst_base, - const Scalar* EIGEN_RESTRICT src, - Index src_stride, Index depth, Index cols, - Index dst_stride, Index dst_offset) { - constexpr int NR = kSmeNr; - const Index peeled_cols = (cols / NR) * NR; - + __arm_locally_streaming __arm_new("za") static void pack_full_panels(Scalar* dst_base, + const Scalar* EIGEN_RESTRICT src, + Index src_stride, Index depth, Index peeled_cols, + Index dst_stride, Index dst_offset) { for (Index j = 0; j < peeled_cols; j += NR) { Scalar* dst_panel = PanelMode ? dst_base + j * dst_stride + dst_offset * NR : dst_base + j * depth; - sme_transpose_pack(dst_panel, src + j * src_stride, src_stride, depth); + sme_transpose_pack(dst_panel, src + j * src_stride, src_stride, depth, NR); + } + } + + static void pack_direct(Scalar* dst_base, const Scalar* EIGEN_RESTRICT src, Index src_stride, Index depth, Index cols, + Index dst_stride, Index dst_offset) { + const Index peeled_cols = (cols / NR) * NR; + + if (peeled_cols > 0) { + pack_full_panels(dst_base, src, src_stride, depth, peeled_cols, dst_stride, dst_offset); } // Col tail (cols - peeled_cols in [1, NR-1]). Same reasoning as the LHS // RowMajor packer's row tail: runs at most once per call, < NR cols, not - // worth the partial-ZA-tile handling. + // worth the partial-ZA-tile handling, and taken outside the streaming + // region above (see tail_transpose_pack). if (peeled_cols < cols) { const Index tail = cols - peeled_cols; Scalar* dst_panel = PanelMode ? dst_base + peeled_cols * dst_stride + dst_offset * tail : dst_base + peeled_cols * depth; - scalar_tail_pack(dst_panel, src + peeled_cols * src_stride, src_stride, depth, tail); + tail_transpose_pack(dst_panel, src + peeled_cols * src_stride, src_stride, depth, tail); } } @@ -511,25 +725,21 @@ struct gemm_pack_rhs(&pack_rhs_colmajor, &sme_pack_rhs_fallback, blockB, rhs, + sme_dispatch_pack(&pack_direct, &sme_pack_rhs_fallback, blockB, rhs, depth, cols, stride, offset, bool_constant::value>{}); } }; // RowMajor RHS packer -- streaming SVE copy (mirrors the ColMajor LHS packer). -// Rows are contiguous in the source, so each depth-step is NR contiguous fp32. -// nr_ pinned to kSmeNr -- see the ColMajor specialization above. -template -struct gemm_pack_rhs { - typedef float Scalar; - // See gemm_pack_lhs above: Conjugate is ignored, sound only for real scalars. +// Rows are contiguous in the source, so each depth-step is NR contiguous scalars. +template +struct sme_pack_rhs_rowmajor { + // See sme_pack_lhs_colmajor: Conjugate is ignored, sound only for real scalars. static_assert(!NumTraits::IsComplex, "the SME packers only support real scalars"); - __arm_locally_streaming static void pack_rhs_rowmajor(Scalar* dst_base, const Scalar* EIGEN_RESTRICT src, - Index src_stride, Index depth, Index cols, Index dst_stride, - Index dst_offset) { - constexpr int NR = kSmeNr; + __arm_locally_streaming static void pack_direct(Scalar* dst_base, const Scalar* EIGEN_RESTRICT src, Index src_stride, + Index depth, Index cols, Index dst_stride, Index dst_offset) { const Index peeled_cols = (cols / NR) * NR; for (Index j = 0; j < peeled_cols; j += NR) { @@ -553,60 +763,120 @@ struct gemm_pack_rhs(&pack_rhs_rowmajor, &sme_pack_rhs_fallback, blockB, rhs, - depth, cols, stride, offset, + sme_dispatch_pack(&pack_direct, &sme_pack_rhs_fallback, blockB, + rhs, depth, cols, stride, offset, bool_constant::value>{}); } }; +// Pack1/nr_ are pinned to the SME block sizes (rather than left open) so these +// specializations only match consumers that actually feed the SME gebp_kernel +// -- see "Mixed-scalar products" in the gebp_traits doc comment above. +template +struct gemm_pack_lhs + : sme_pack_lhs_colmajor {}; + +template +struct gemm_pack_lhs + : sme_pack_lhs_rowmajor {}; + +template +struct gemm_pack_rhs + : sme_pack_rhs_colmajor {}; + +template +struct gemm_pack_rhs + : sme_pack_rhs_rowmajor {}; + +#ifdef EIGEN_VECTORIZE_SME_F64F64 +template +struct gemm_pack_lhs + : sme_pack_lhs_colmajor {}; + +template +struct gemm_pack_lhs + : sme_pack_lhs_rowmajor {}; + +template +struct gemm_pack_rhs + : sme_pack_rhs_colmajor {}; + +template +struct gemm_pack_rhs + : sme_pack_rhs_rowmajor {}; +#endif + /***************************************************************************** - * sme_store_za_tile -- Store one ZA.S tile back to C with alpha scaling. + * sme_store_za_tile -- Store one ZA tile back to C with alpha scaling. * * `pw` is the row-predicate width for this tile, `cw` the col-predicate width - * (both <= the runtime svlw). + * (both <= the runtime svl). *****************************************************************************/ -template -EIGEN_ALWAYS_INLINE void sme_store_za_tile(float* EIGEN_RESTRICT C, Index C_stride_row, Index C_stride_col, float alpha, - Index row_start, int pw, Index col_start, +template +EIGEN_ALWAYS_INLINE void sme_store_za_tile(Scalar* EIGEN_RESTRICT C, Index C_stride_row, Index C_stride_col, + Scalar alpha, Index row_start, int pw, Index col_start, int cw) __arm_streaming __arm_inout("za") { - const svbool_t pg_m = svwhilelt_b32((uint32_t)0, (uint32_t)pw); - const svbool_t pg_n = svwhilelt_b32((uint32_t)0, (uint32_t)cw); + using Traits = sme_traits; + using Vec = typename Traits::Vec; + const svbool_t pg_m = Traits::whilelt(0, pw); + const svbool_t pg_n = Traits::whilelt(0, cw); // FMLA and FADD have equal latency/throughput on ARMv9 cores, and // multiplying by alpha=1.0 is exact in IEEE-754 so the FMLA form is // bit-identical to FADD in that case. A single unconditional FMLA // keeps the store compact and measures no worse (and a few percent // better on small matrices, where the branch would otherwise disrupt // instruction scheduling). - const svfloat32_t vzero = svdup_f32(0.f); - const svfloat32_t valpha = svdup_f32(alpha); + const Vec vzero = Traits::dup(Scalar(0)); + const Vec valpha = Traits::dup(alpha); + // Two C slices are loaded before either is stored: a C line the caller wrote + // from non-streaming code just before the kernel does not forward across the + // mode switch on Apple M4, and a serial load/store pays that latency per slice. + // C = A*B meets the condition on every call, since evalTo zeroes the + // destination first. SVE vectors are sizeless, hence the spelled-out pair. if (C_stride_row == 1) { // Column-major C: extract vertical slices (columns of the ZA tile) - for (int ci = 0; ci < cw; ++ci) { - svfloat32_t vres = svread_ver_za32_f32_m(vzero, pg_m, TileId, (uint32_t)ci); - float* pC = C + row_start + (col_start + ci) * C_stride_col; - svfloat32_t vc = svld1_f32(pg_m, pC); - svst1_f32(pg_m, pC, svmla_f32_x(pg_m, vc, vres, valpha)); + int ci = 0; + for (; ci + 2 <= cw; ci += 2) { + Scalar* p0 = C + row_start + (col_start + ci) * C_stride_col; + Scalar* p1 = p0 + C_stride_col; + Vec c0 = sme_ld1(pg_m, p0); + Vec c1 = sme_ld1(pg_m, p1); + sme_st1(pg_m, p0, sme_mla(pg_m, c0, sme_read_ver_za(vzero, pg_m, (uint32_t)ci), valpha)); + sme_st1(pg_m, p1, sme_mla(pg_m, c1, sme_read_ver_za(vzero, pg_m, (uint32_t)(ci + 1)), valpha)); + } + if (ci < cw) { + Scalar* pC = C + row_start + (col_start + ci) * C_stride_col; + Vec vc = sme_ld1(pg_m, pC); + sme_st1(pg_m, pC, sme_mla(pg_m, vc, sme_read_ver_za(vzero, pg_m, (uint32_t)ci), valpha)); } } else if (C_stride_col == 1) { // Row-major C: extract horizontal slices (rows of the ZA tile) - for (int ri = 0; ri < pw; ++ri) { - svfloat32_t vres = svread_hor_za32_f32_m(vzero, pg_n, TileId, (uint32_t)ri); - float* pC = C + (row_start + ri) * C_stride_row + col_start; - svfloat32_t vc = svld1_f32(pg_n, pC); - svst1_f32(pg_n, pC, svmla_f32_x(pg_n, vc, vres, valpha)); + int ri = 0; + for (; ri + 2 <= pw; ri += 2) { + Scalar* p0 = C + (row_start + ri) * C_stride_row + col_start; + Scalar* p1 = p0 + C_stride_row; + Vec c0 = sme_ld1(pg_n, p0); + Vec c1 = sme_ld1(pg_n, p1); + sme_st1(pg_n, p0, sme_mla(pg_n, c0, sme_read_hor_za(vzero, pg_n, (uint32_t)ri), valpha)); + sme_st1(pg_n, p1, sme_mla(pg_n, c1, sme_read_hor_za(vzero, pg_n, (uint32_t)(ri + 1)), valpha)); + } + if (ri < pw) { + Scalar* pC = C + (row_start + ri) * C_stride_row + col_start; + Vec vc = sme_ld1(pg_n, pC); + sme_st1(pg_n, pC, sme_mla(pg_n, vc, sme_read_hor_za(vzero, pg_n, (uint32_t)ri), valpha)); } } else { // General stride: extract rows to temp buffer, scatter to C. scratch - // holds one ZA row; every caller passes cw <= min(svlw, kSmeNr) (a tile - // never spans more than the logical block), so kSmeNr is a static - // bound independent of the runtime svlw. - float scratch[kSmeNr]; + // holds one ZA row; every caller passes cw <= min(svl, nr) (a tile + // never spans more than the logical block), so nr is a static + // bound independent of the runtime svl. + Scalar scratch[sme_block::nr]; for (int ri = 0; ri < pw; ++ri) { - svfloat32_t vres = svread_hor_za32_f32_m(vzero, pg_n, TileId, (uint32_t)ri); - vres = svmul_f32_x(pg_n, vres, valpha); - svst1_f32(pg_n, scratch, vres); + Vec vres = sme_read_hor_za(vzero, pg_n, (uint32_t)ri); + vres = sme_mul(pg_n, vres, valpha); + sme_st1(pg_n, scratch, vres); for (int ci = 0; ci < cw; ++ci) { C[(row_start + ri) * C_stride_row + (col_start + ci) * C_stride_col] += scratch[ci]; } @@ -615,7 +885,7 @@ EIGEN_ALWAYS_INLINE void sme_store_za_tile(float* EIGEN_RESTRICT C, Index C_stri } /***************************************************************************** - * sme_store_2x2_grid -- store the (up to) 2x2 grid of svlw x svlw ZA tiles. + * sme_store_2x2_grid -- store the (up to) 2x2 grid of svl x svl ZA tiles. * * Tile layout: 0 = (row-lo, col-lo) 1 = (row-lo, col-hi) * 2 = (row-hi, col-lo) 3 = (row-hi, col-hi) @@ -626,122 +896,126 @@ EIGEN_ALWAYS_INLINE void sme_store_za_tile(float* EIGEN_RESTRICT C, Index C_stri * pattern repeats across blocks). *****************************************************************************/ -template -EIGEN_ALWAYS_INLINE void sme_store_2x2_grid(float* EIGEN_RESTRICT C, Index C_stride_row, Index C_stride_col, - float alpha, Index row_start, int rlo, int rhi, Index col_start, int clo, +template +EIGEN_ALWAYS_INLINE void sme_store_2x2_grid(Scalar* EIGEN_RESTRICT C, Index C_stride_row, Index C_stride_col, + Scalar alpha, Index row_start, int rlo, int rhi, Index col_start, int clo, int chi) __arm_streaming __arm_inout("za") { - const int svlw = static_cast(svcntsw()); - sme_store_za_tile<0>(C, C_stride_row, C_stride_col, alpha, row_start, rlo, col_start, clo); + const int svl = sme_traits::svl(); + sme_store_za_tile(C, C_stride_row, C_stride_col, alpha, row_start, rlo, col_start, clo); if (chi > 0) { - sme_store_za_tile<1>(C, C_stride_row, C_stride_col, alpha, row_start, rlo, col_start + svlw, chi); + sme_store_za_tile(C, C_stride_row, C_stride_col, alpha, row_start, rlo, col_start + svl, chi); } if (rhi > 0) { - sme_store_za_tile<2>(C, C_stride_row, C_stride_col, alpha, row_start + svlw, rhi, col_start, clo); + sme_store_za_tile(C, C_stride_row, C_stride_col, alpha, row_start + svl, rhi, col_start, clo); if (chi > 0) { - sme_store_za_tile<3>(C, C_stride_row, C_stride_col, alpha, row_start + svlw, rhi, col_start + svlw, chi); + sme_store_za_tile(C, C_stride_row, C_stride_col, alpha, row_start + svl, rhi, col_start + svl, chi); } } } // One depth step's worth of the exact-match grid: the four FMOPAs that take the // lo/hi halves of a packed A column and a packed B column and accumulate the -// 2x2 ZA-tile outer product. `all` is svptrue_b32() because this is only used -// on the exact-match path, where the block fills the grid (pg == svptrue_b32), -// so factoring it out is identical to the inline form. -static EIGEN_ALWAYS_INLINE void outer_product_2x2(svfloat32_t a_lo, svfloat32_t a_hi, svfloat32_t b_lo, - svfloat32_t b_hi) __arm_streaming __arm_inout("za") { - const svbool_t all = svptrue_b32(); - svmopa_za32_f32_m(0, all, all, a_lo, b_lo); - svmopa_za32_f32_m(1, all, all, a_lo, b_hi); - svmopa_za32_f32_m(2, all, all, a_hi, b_lo); - svmopa_za32_f32_m(3, all, all, a_hi, b_hi); +// 2x2 ZA-tile outer product. `all` is the all-true predicate because this is +// only used on the exact-match path, where the block fills the grid, so +// factoring it out is identical to the inline form. +template +static EIGEN_ALWAYS_INLINE void outer_product_2x2( + typename sme_traits::Vec a_lo, typename sme_traits::Vec a_hi, typename sme_traits::Vec b_lo, + typename sme_traits::Vec b_hi) __arm_streaming __arm_inout("za") { + const svbool_t all = sme_traits::ptrue(); + sme_mopa<0>(all, all, a_lo, b_lo); + sme_mopa<1>(all, all, a_lo, b_hi); + sme_mopa<2>(all, all, a_hi, b_lo); + sme_mopa<3>(all, all, a_hi, b_hi); } /***************************************************************************** * sme_process -- micro-kernel for one pw x cw output block. * - * Tiles the block into svlw x svlw ZA tiles, processed in passes of up to a 2x2 - * tile grid (the 4-ZA-tile budget): several (2*svlw) x (2*svlw) sub-block passes - * when the grid is smaller than the block, tiles predicated down to the block - * width when it is larger. blA/blB are packed depth-major with depth-strides - * pw and cw respectively. + * Tiles the block into svl x svl ZA tiles, processed in passes of up to a 2x2 + * tile grid: several (2*svl) x (2*svl) sub-block passes when the grid is + * smaller than the block, tiles predicated down to the block width when it is + * larger. blA/blB are packed depth-major with depth-strides pw and cw + * respectively. * - * When the block matches the tile grid exactly (pw == cw == 2 * svlw), the + * When the block matches the tile grid exactly (pw == cw == 2 * svl), the * packed rows are also contiguous across depth steps, enabling the - * hand-scheduled loop below: per 4 unrolled depth steps, 2 svld1_f32_x4 per + * hand-scheduled loop below: per 4 unrolled depth steps, 2 x4 loads per * side (each spanning 2 depth steps) feed 16 FMOPAs -- a 1:1 compute:load * ratio at the vector level. All other geometries use predicated * per-depth-step loads. *****************************************************************************/ -template -EIGEN_ALWAYS_INLINE void sme_process(float* EIGEN_RESTRICT C, Index C_stride_row, Index C_stride_col, - const float* EIGEN_RESTRICT blA, const float* EIGEN_RESTRICT blB, Index depth, - float alpha, Index row_start, int pw, Index col_start, +template +EIGEN_ALWAYS_INLINE void sme_process(Scalar* EIGEN_RESTRICT C, Index C_stride_row, Index C_stride_col, + const Scalar* EIGEN_RESTRICT blA, const Scalar* EIGEN_RESTRICT blB, Index depth, + Scalar alpha, Index row_start, int pw, Index col_start, int cw) __arm_streaming __arm_inout("za") { - const int svlw = static_cast(svcntsw()); + using Traits = sme_traits; + using Vec = typename Traits::Vec; + const int svl = Traits::svl(); - for (int rt = 0; rt < pw; rt += 2 * svlw) { - const int rpw = sme_min(pw - rt, 2 * svlw); - const int rlo = sme_min(rpw, svlw); - const int rhi = rpw - rlo; // >= 0; > 0 only when rpw > svlw, in which case rlo == svlw - const svbool_t pg_rlo = svwhilelt_b32((uint32_t)0, (uint32_t)rlo); - const svbool_t pg_rhi = svwhilelt_b32((uint32_t)0, (uint32_t)rhi); + for (int rt = 0; rt < pw; rt += 2 * svl) { + const int rpw = sme_min(pw - rt, 2 * svl); + const int rlo = sme_min(rpw, svl); + const int rhi = rpw - rlo; // >= 0; > 0 only when rpw > svl, in which case rlo == svl + const svbool_t pg_rlo = Traits::whilelt(0, rlo); + const svbool_t pg_rhi = Traits::whilelt(0, rhi); - for (int ct = 0; ct < cw; ct += 2 * svlw) { - const int cpw = sme_min(cw - ct, 2 * svlw); - const int clo = sme_min(cpw, svlw); + for (int ct = 0; ct < cw; ct += 2 * svl) { + const int cpw = sme_min(cw - ct, 2 * svl); + const int clo = sme_min(cpw, svl); const int chi = cpw - clo; - const svbool_t pg_clo = svwhilelt_b32((uint32_t)0, (uint32_t)clo); - const svbool_t pg_chi = svwhilelt_b32((uint32_t)0, (uint32_t)chi); + const svbool_t pg_clo = Traits::whilelt(0, clo); + const svbool_t pg_chi = Traits::whilelt(0, chi); svzero_za(); - if (pw == 2 * svlw && cw == 2 * svlw) { + if (pw == 2 * svl && cw == 2 * svl) { // The block is exactly one full-grid patch (single pass, rt == ct == - // 0, rlo == rhi == clo == chi == svlw), so a packed row is the + // 0, rlo == rhi == clo == chi == svl), so a packed row is the // patch's slice and rows are contiguous across depth steps: x4 loads // each span 2 of them, e.g. va_01 = [d0 lo, d0 hi, d1 lo, d1 hi]. - const svcount_t pn = svptrue_c32(); + const svcount_t pn = Traits::ptrue_c(); const Index depth_4 = (depth / 4) * 4; Index k = 0; for (; k < depth_4; k += 4) { - svfloat32x4_t va_01 = svld1_f32_x4(pn, &blA[k * pw]); - svfloat32x4_t vb_01 = svld1_f32_x4(pn, &blB[k * cw]); + typename Traits::Vec4 va_01 = sme_ld1_x4(pn, &blA[k * pw]); + typename Traits::Vec4 vb_01 = sme_ld1_x4(pn, &blB[k * cw]); // d0 - outer_product_2x2(svget4_f32(va_01, 0), svget4_f32(va_01, 1), svget4_f32(vb_01, 0), svget4_f32(vb_01, 1)); + outer_product_2x2(sme_get<0>(va_01), sme_get<1>(va_01), sme_get<0>(vb_01), sme_get<1>(vb_01)); // d1 - outer_product_2x2(svget4_f32(va_01, 2), svget4_f32(va_01, 3), svget4_f32(vb_01, 2), svget4_f32(vb_01, 3)); + outer_product_2x2(sme_get<2>(va_01), sme_get<3>(va_01), sme_get<2>(vb_01), sme_get<3>(vb_01)); - svfloat32x4_t va_23 = svld1_f32_x4(pn, &blA[(k + 2) * pw]); - svfloat32x4_t vb_23 = svld1_f32_x4(pn, &blB[(k + 2) * cw]); + typename Traits::Vec4 va_23 = sme_ld1_x4(pn, &blA[(k + 2) * pw]); + typename Traits::Vec4 vb_23 = sme_ld1_x4(pn, &blB[(k + 2) * cw]); // d2 - outer_product_2x2(svget4_f32(va_23, 0), svget4_f32(va_23, 1), svget4_f32(vb_23, 0), svget4_f32(vb_23, 1)); + outer_product_2x2(sme_get<0>(va_23), sme_get<1>(va_23), sme_get<0>(vb_23), sme_get<1>(vb_23)); // d3 - outer_product_2x2(svget4_f32(va_23, 2), svget4_f32(va_23, 3), svget4_f32(vb_23, 2), svget4_f32(vb_23, 3)); + outer_product_2x2(sme_get<2>(va_23), sme_get<3>(va_23), sme_get<2>(vb_23), sme_get<3>(vb_23)); } // Depth tail: one x2 load per side per step. for (; k < depth; ++k) { - svfloat32x2_t va = svld1_f32_x2(pn, &blA[k * pw]); - svfloat32x2_t vb = svld1_f32_x2(pn, &blB[k * cw]); - outer_product_2x2(svget2_f32(va, 0), svget2_f32(va, 1), svget2_f32(vb, 0), svget2_f32(vb, 1)); + typename Traits::Vec2 va = sme_ld1_x2(pn, &blA[k * pw]); + typename Traits::Vec2 vb = sme_ld1_x2(pn, &blB[k * cw]); + outer_product_2x2(sme_get<0>(va), sme_get<1>(va), sme_get<0>(vb), sme_get<1>(vb)); } } else { for (Index k = 0; k < depth; ++k) { - svfloat32_t a_lo = svld1_f32(pg_rlo, &blA[k * pw + rt]); - svfloat32_t b_lo = svld1_f32(pg_clo, &blB[k * cw + ct]); - svmopa_za32_f32_m(0, pg_rlo, pg_clo, a_lo, b_lo); - svfloat32_t b_hi = svdup_f32(0.f); + Vec a_lo = sme_ld1(pg_rlo, &blA[k * pw + rt]); + Vec b_lo = sme_ld1(pg_clo, &blB[k * cw + ct]); + sme_mopa<0>(pg_rlo, pg_clo, a_lo, b_lo); + Vec b_hi = Traits::dup(Scalar(0)); if (chi > 0) { - b_hi = svld1_f32(pg_chi, &blB[k * cw + ct + svlw]); - svmopa_za32_f32_m(1, pg_rlo, pg_chi, a_lo, b_hi); + b_hi = sme_ld1(pg_chi, &blB[k * cw + ct + svl]); + sme_mopa<1>(pg_rlo, pg_chi, a_lo, b_hi); } if (rhi > 0) { - svfloat32_t a_hi = svld1_f32(pg_rhi, &blA[k * pw + rt + svlw]); - svmopa_za32_f32_m(2, pg_rhi, pg_clo, a_hi, b_lo); + Vec a_hi = sme_ld1(pg_rhi, &blA[k * pw + rt + svl]); + sme_mopa<2>(pg_rhi, pg_clo, a_hi, b_lo); if (chi > 0) { - svmopa_za32_f32_m(3, pg_rhi, pg_chi, a_hi, b_hi); + sme_mopa<3>(pg_rhi, pg_chi, a_hi, b_hi); } } } @@ -753,12 +1027,12 @@ EIGEN_ALWAYS_INLINE void sme_process(float* EIGEN_RESTRICT C, Index C_stride_row } } -template +template EIGEN_DONT_INLINE __arm_locally_streaming __arm_new("za") void sme_gebp_impl( - float* C, Index C_stride_row, Index C_stride_col, const float* blockA, const float* blockB, Index rows, Index depth, - Index cols, float alpha, Index strideA, Index strideB, Index offsetA, Index offsetB) { - constexpr int MR = kSmeMr; - constexpr int NR = kSmeNr; + Scalar* C, Index C_stride_row, Index C_stride_col, const Scalar* blockA, const Scalar* blockB, Index rows, + Index depth, Index cols, Scalar alpha, Index strideA, Index strideB, Index offsetA, Index offsetB) { + constexpr int MR = sme_block::mr; + constexpr int NR = sme_block::nr; // Column-outer, row-inner: keeps blB (one kc × NR panel) hot in L1 while // smaller blA tiles stream from L2. The outer GOTO loop in @@ -769,26 +1043,26 @@ EIGEN_DONT_INLINE __arm_locally_streaming __arm_new("za") void sme_gebp_impl( // blocks are tiled and predicated inside the generic path. for (Index j = 0; j < cols; j += NR) { const int cw = static_cast(sme_min(cols - j, Index(NR))); - const float* blB = blockB + j * strideB + offsetB * cw; + const Scalar* blB = blockB + j * strideB + offsetB * cw; for (Index i = 0; i < rows; i += MR) { const int pw = static_cast(sme_min(rows - i, Index(MR))); - const float* blA = blockA + i * strideA + offsetA * pw; + const Scalar* blA = blockA + i * strideA + offsetA * pw; sme_process(C, C_stride_row, C_stride_col, blA, blB, depth, alpha, i, pw, j, cw); } } } -template -struct gebp_kernel { - typedef float Scalar; - typedef float ResScalar; +template +struct sme_gebp_kernel { + using ResScalar = Scalar; EIGEN_DONT_INLINE void operator()(const DataMapper& res, const Scalar* blockA, const Scalar* blockB, Index rows, Index depth, Index cols, ResScalar alpha, Index strideA = -1, Index strideB = -1, Index offsetA = 0, Index offsetB = 0) { - static_assert(!ConjugateLhs && !ConjugateRhs, "SME fp32 kernel does not support conjugation"); - static_assert(mr == kSmeMr && nr == kSmeNr, "SME fp32 kernel expects kSmeMr/kSmeNr-sized packed panels"); + static_assert(!ConjugateLhs && !ConjugateRhs, "the SME kernel does not support conjugation"); + static_assert(mr == sme_block::mr && nr == sme_block::nr, + "the SME kernel expects packed panels of the SME block width"); if (strideA == -1) strideA = depth; if (strideB == -1) strideB = depth; @@ -804,6 +1078,28 @@ struct gebp_kernel +struct gebp_kernel + : sme_gebp_kernel {}; + +#ifdef EIGEN_VECTORIZE_SME_F64F64 +template +struct gebp_kernel + : sme_gebp_kernel {}; +#endif + +// sme_has_gebp_kernel (products/GeneralBlockPanelKernel.h) drives the cache +// blocking and the GEMM loop order, and is declared before this header. A pair +// listed there but not specialized here would be packed and blocked for SME and +// then handed to the generic kernel. +static_assert(sme_has_gebp_kernel::value, "the SME float kernel is not advertised to the GEMM driver"); +#ifdef EIGEN_VECTORIZE_SME_F64F64 +static_assert(sme_has_gebp_kernel::value, "the SME double kernel is not advertised to the GEMM driver"); +#else +static_assert(!sme_has_gebp_kernel::value, + "double is advertised to the GEMM driver without FEAT_SME_F64F64 to implement it"); +#endif + // --------------------------------------------------------------------------- // Selfadjoint (SYMM) packers. // @@ -811,12 +1107,12 @@ struct gebp_kernel= col is the // stored triangle. It reads that half directly and mirrors the other half -// (conjugation is an identity for float): +// (conjugation is an identity for real scalars): // full(row,col) = (row >= col) ? m(row,col) : m(col,row) // // Regions wholly below or above the diagonal use the normal dense copy or @@ -837,25 +1133,30 @@ struct gebp_kernel -EIGEN_DONT_INLINE __arm_locally_streaming __arm_new("za") void sme_symm_pack_panels(float* block, - const float* EIGEN_RESTRICT base, - Index stride, Index depth, - Index outer, Index k2) { - static_assert(kSmeMr == kSmeNr, "the shared SYMM packer assumes square panels"); - constexpr int PACK = kSmeMr; +// Depth-region boundaries for the panel at outer offset `j`, all clamped to +// [0, depth]: the diagonal splits it into a transposed head [0, t_end), a +// straddle band [t_end, s_end) and a direct tail [s_end, depth). +template +static EIGEN_ALWAYS_INLINE void sme_symm_panel_regions(Index j, int w, Index depth, Index k2, Index& t_end, + Index& s_end) __arm_streaming_compatible { + const Index raw_t = j - k2, raw_s = j + Index(w) - k2; + t_end = raw_t <= 0 ? Index(0) : sme_min(raw_t, depth); + s_end = raw_s <= 0 ? Index(0) : sme_min(raw_s, depth); +} + +// The two dense regions of every panel, which are ordinary copies or ZA +// transposes of the stored triangle. ColM selects the ColMajor operand. +template +EIGEN_DONT_INLINE __arm_locally_streaming __arm_new("za") void sme_symm_pack_dense_regions( + Scalar* block, const Scalar* EIGEN_RESTRICT base, Index stride, Index depth, Index outer, Index k2) { + constexpr int PACK = sme_block::mr; constexpr bool ColM = (StorageOrder == ColMajor); - const int svlw = static_cast(svcntsw()); for (Index j = 0; j < outer; j += PACK) { const int w = static_cast(sme_min(outer - j, Index(PACK))); - float* dst = block + j * depth; // depth-major panel of width w - - // Depth-region boundaries (all clamped to [0, depth]). - const Index raw_t = j - k2, raw_s = j + Index(w) - k2; - const Index t_end = raw_t <= 0 ? Index(0) : sme_min(raw_t, depth); // transposed [0, t_end) - const Index s_end = raw_s <= 0 ? Index(0) : sme_min(raw_s, depth); // straddle [t_end, s_end) - // direct [s_end, depth) + Scalar* dst = block + j * depth; // depth-major panel of width w + Index t_end, s_end; + sme_symm_panel_regions(j, w, depth, k2, t_end, s_end); // Transposed region: full(k2+k, j+c) = m(j+c, k2+k). if (t_end > 0) { @@ -873,59 +1174,93 @@ EIGEN_DONT_INLINE __arm_locally_streaming __arm_new("za") void sme_symm_pack_pan sve_copy_panel_range(dst, base + k2 * stride + j, stride, s_end, depth, w); } } - // Straddle band: the diagonal crosses the panel at c* = (k2+k) - j (in - // [0, w) for every k in the band), splitting each depth step into a - // direct head (c < c*: m(k2+k, j+c)) and a mirrored tail (c >= c*: - // m(j+c, k2+k); at c == c* both name the diagonal element). One side is - // contiguous in c -- the head for RowMajor, the tail for ColMajor -- and - // is copied with predicated vectors; the other walks the stride scalar. + } +} + +// The diagonal band of every panel: the diagonal crosses at c* = (k2+k) - j +// (in [0, w) throughout the band), so each depth step splits into a direct head +// (c < c*: m(k2+k, j+c)) and a mirrored tail (c >= c*: m(j+c, k2+k); at c == c* +// both name the diagonal element). +// +// Kept out of the streaming region above for the reason tail_transpose_pack gives, +// at the cost of a second pass over the panels: it is scalar floating-point, +// and fusing it made the float SYMM packers 2-11x slower. +template +EIGEN_DONT_INLINE void sme_symm_pack_straddle(Scalar* block, const Scalar* EIGEN_RESTRICT base, Index stride, + Index depth, Index outer, Index k2) { + constexpr int PACK = sme_block::mr; + constexpr bool ColM = (StorageOrder == ColMajor); + + for (Index j = 0; j < outer; j += PACK) { + const int w = static_cast(numext::mini(outer - j, Index(PACK))); + Scalar* dst = block + j * depth; + Index t_end, s_end; + sme_symm_panel_regions(j, w, depth, k2, t_end, s_end); + for (Index k = t_end; k < s_end; ++k) { const Index row = k2 + k; const int cs = static_cast(row - j); - float* dst_row = dst + k * w; + Scalar* dst_row = dst + k * w; EIGEN_IF_CONSTEXPR (ColM) { - const float* head = base + row + j * stride; // m(row, j+c): stride-strided + const Scalar* head = base + row + j * stride; // m(row, j+c): stride-strided for (int c = 0; c < cs; ++c, head += stride) dst_row[c] = *head; - const float* tail = base + j + row * stride; // m(j+c, row): contiguous - for (int c = cs; c < w; c += svlw) { - const svbool_t pred = svwhilelt_b32(uint32_t(c), uint32_t(w)); - svst1_f32(pred, dst_row + c, svld1_f32(pred, tail + c)); - } + const Scalar* tail = base + j + row * stride; // m(j+c, row): contiguous + for (int c = cs; c < w; ++c) dst_row[c] = tail[c]; } else { - const float* head = base + row * stride + j; // m(row, j+c): contiguous - for (int c = 0; c < cs; c += svlw) { - const svbool_t pred = svwhilelt_b32(uint32_t(c), uint32_t(cs)); - svst1_f32(pred, dst_row + c, svld1_f32(pred, head + c)); - } - const float* tail = base + (j + Index(cs)) * stride + row; // m(j+c, row): stride-strided + const Scalar* head = base + row * stride + j; // m(row, j+c): contiguous + for (int c = 0; c < cs; ++c) dst_row[c] = head[c]; + const Scalar* tail = base + (j + Index(cs)) * stride + row; // m(j+c, row): stride-strided for (int c = cs; c < w; ++c, tail += stride) dst_row[c] = *tail; } } } } -// symm_pack_lhs/rhs SME specializations: emit the uniform kSmeMr/kSmeNr panels +// Packer shared by the LHS (k2 == 0) and RHS symm specializations. +template +EIGEN_DONT_INLINE void sme_symm_pack_panels(Scalar* block, const Scalar* EIGEN_RESTRICT base, Index stride, Index depth, + Index outer, Index k2) { + static_assert(sme_block::mr == sme_block::nr, "the shared SYMM packer assumes square panels"); + sme_symm_pack_dense_regions(block, base, stride, depth, outer, k2); + sme_symm_pack_straddle(block, base, stride, depth, outer, k2); +} + +// symm_pack_lhs/rhs SME specializations: emit the uniform mr/nr panels // sme_gebp_impl reads. Pack1/nr pinned exactly as gemm_pack_lhs/rhs above. -template -struct symm_pack_lhs { - typedef float Scalar; +template +struct sme_symm_pack_lhs { // Note: generic symm_pack_lhs's "cols" is the depth extent, and the LHS // block is diagonal-anchored (base = &lhs(k2,k2)), so its depth offset is 0. EIGEN_DONT_INLINE void operator()(Scalar* blockA, const Scalar* lhs_, Index lhsStride, Index cols, Index rows) const { - sme_symm_pack_panels(blockA, lhs_, lhsStride, cols, rows, Index(0)); + sme_symm_pack_panels(blockA, lhs_, lhsStride, cols, rows, Index(0)); } }; -template -struct symm_pack_rhs { - typedef float Scalar; +template +struct sme_symm_pack_rhs { // Note: generic symm_pack_rhs's "rows" is the depth extent (end_k = k2 + rows), not a row count. EIGEN_DONT_INLINE void operator()(Scalar* blockB, const Scalar* rhs_, Index rhsStride, Index rows, Index cols, Index k2) const { - sme_symm_pack_panels(blockB, rhs_, rhsStride, rows, cols, k2); + sme_symm_pack_panels(blockB, rhs_, rhsStride, rows, cols, k2); } }; +template +struct symm_pack_lhs : sme_symm_pack_lhs { +}; + +template +struct symm_pack_rhs : sme_symm_pack_rhs {}; + +#ifdef EIGEN_VECTORIZE_SME_F64F64 +template +struct symm_pack_lhs + : sme_symm_pack_lhs {}; + +template +struct symm_pack_rhs : sme_symm_pack_rhs {}; +#endif + } // namespace internal } // namespace Eigen diff --git a/Eigen/src/Core/products/GeneralBlockPanelKernel.h b/Eigen/src/Core/products/GeneralBlockPanelKernel.h index 1d3f8c583..30de6108a 100644 --- a/Eigen/src/Core/products/GeneralBlockPanelKernel.h +++ b/Eigen/src/Core/products/GeneralBlockPanelKernel.h @@ -127,6 +127,19 @@ inline void manage_caching_sizes(Action action, std::ptrdiff_t* l1, std::ptrdiff * \sa setCpuCacheSizes */ #ifdef EIGEN_VECTORIZE_SME +// True for the scalar pairs the SME gebp_kernel specializes (see +// arch/SME/GeneralBlockPanelKernel.h, which static_asserts that it agrees with +// this list); every other pair keeps Eigen's generic kernel, packers, cache +// blocking and GEMM loop order. +template +struct sme_has_gebp_kernel : std::false_type {}; +template <> +struct sme_has_gebp_kernel : std::true_type {}; +#ifdef EIGEN_VECTORIZE_SME_F64F64 +template <> +struct sme_has_gebp_kernel : std::true_type {}; +#endif + // Overridable SME packed-panel budgets. The defaults are empirically tuned // fp32 working-set limits for Apple M4 — heuristic budgets, not generic ARM64 // cache defaults; redefine them to retune for other SME implementations. @@ -162,8 +175,11 @@ void evaluateProductBlockingSizesHeuristicForSme(Index& k, Index& m, Index& n) { #endif // Keep kc large enough to amortize SME setup and accumulation, but cap very - // deep products to avoid too many result store passes. - k = (numext::mini)(k, sme_max_kc); + // deep products to avoid too many result store passes. The cap is a scalar + // count tuned for fp32; scale it by the scalar width so every element type + // gets the same packed-panel byte budget. + const Index max_kc = (numext::maxi)(Index(1), sme_max_kc * Index(sizeof(float)) / Index(sizeof(LhsScalar))); + k = (numext::mini)(k, max_kc); // Bound the packed RHS strip so very wide matrices do not allocate an // unbounded blockB panel. @@ -267,9 +283,9 @@ void evaluateProductBlockingSizesHeuristic(Index& k, Index& m, Index& n, Index n if ((numext::maxi)(k, (numext::maxi)(m, n)) < 48) return; #ifdef EIGEN_VECTORIZE_SME - // Only float×float uses the SME kernel; other scalar pairs run the generic - // kernel below and would thrash L1/L2 with the SME-sized budgets. - if (std::is_same::value && std::is_same::value) { + // Only the scalar pairs the SME kernel specializes use the SME budgets; + // the others run the generic kernel below and would thrash L1/L2 with them. + EIGEN_IF_CONSTEXPR ((sme_has_gebp_kernel::value)) { evaluateProductBlockingSizesHeuristicForSme(k, m, n); return; } diff --git a/Eigen/src/Core/products/GeneralMatrixMatrix.h b/Eigen/src/Core/products/GeneralMatrixMatrix.h index 3b66e173f..8bb2b28d7 100644 --- a/Eigen/src/Core/products/GeneralMatrixMatrix.h +++ b/Eigen/src/Core/products/GeneralMatrixMatrix.h @@ -228,11 +228,13 @@ struct general_matrix_matrix_product::value, + gemm_pack_rhs_first_loop_policy, gemm_pack_lhs_first_loop_policy>; #else using SequentialGemmLoop = gemm_pack_lhs_first_loop_policy; #endif diff --git a/Eigen/src/Core/util/ConfigureVectorization.h b/Eigen/src/Core/util/ConfigureVectorization.h index 0bbc0ab69..020a32c33 100644 --- a/Eigen/src/Core/util/ConfigureVectorization.h +++ b/Eigen/src/Core/util/ConfigureVectorization.h @@ -462,6 +462,16 @@ extern "C" { "EIGEN_ARM64_USE_SME must be built without -msve-vector-bits (scalable/VLA mode): a fixed SVE vector length pins the kernel to one runtime streaming SVL and silently miscomputes at any other." #endif +// Double-precision outer products (FMOPA into a ZA.D tile) need the optional +// FEAT_SME_F64F64, which each compiler reports differently: GCC defines the ACLE +// macro, clang defines no macro but gates the builtin on the target feature. +// Both halves are needed -- clang otherwise accepts svmopa_za64_f64_m without the +// feature, so a missed gate faults at run time rather than at build time. +#if !defined(EIGEN_ARM64_NO_SME_F64F64) && \ + (defined(__ARM_FEATURE_SME_F64F64) || EIGEN_HAS_BUILTIN(__builtin_sme_svmopa_za64_f64_m)) +#define EIGEN_VECTORIZE_SME_F64F64 +#endif + #elif EIGEN_ARCH_RISCV #if defined(__riscv_zfh) diff --git a/ci/build.linux.gitlab-ci.yml b/ci/build.linux.gitlab-ci.yml index ae520f340..2e81100dd 100644 --- a/ci/build.linux.gitlab-ci.yml +++ b/ci/build.linux.gitlab-ci.yml @@ -429,9 +429,12 @@ build:linux:cross:arm:clang-14:default: # line (ci/scripts/build.linux.script.sh), so each -D value must be a single # shell word: multi-token flags use CMake's ';' list separator, with no # spaces or embedded quotes (cf. the arm / generic-vector jobs). + # +sme-f64f64 (FEAT_SME_F64F64) enables the double-precision GEMM kernel: + # GCC refuses to compile its fp64 FMOPA without it, and qemu's `max` CPU + # implements the extension, so the test jobs below exercise that kernel. EIGEN_CI_ADDITIONAL_ARGS: >- -DEIGEN_TEST_SME=ON - -DEIGEN_TEST_CUSTOM_CXX_FLAGS=-march=armv9.2-a+sme2 + -DEIGEN_TEST_CUSTOM_CXX_FLAGS=-march=armv9.2-a+sme2+sme-f64f64 -DCMAKE_CROSSCOMPILING_EMULATOR=qemu-aarch64;-L;/usr/aarch64-linux-gnu # SME ACLE (arm_sme.h, __arm_streaming, +sme2) requires GCC >= 14. # The build stage only cross-compiles and never invokes the emulator, so @@ -440,8 +443,9 @@ build:linux:cross:arm:clang-14:default: # product_symm/product_trsolve/product_trmm/product_syrk/mixingtypes exercise # the SME packers through their non-dense-GEMM consumers (selfadjoint, # triangular-solve, triangular-product, rank-update, mixed-scalar) that - # product_sme can't reach. Explicit list, not `buildtests`: cross-building - # every test under the SME flags overruns the runner timeout. + # product_sme can't reach; each has double subtests, so they cover both SME + # scalar types. Explicit list, not `buildtests`: cross-building every test + # under the SME flags overruns the runner timeout. EIGEN_CI_BUILD_TARGET: product_sme product_symm product_trsolve product_trmm product_syrk mixingtypes # Inherit the 2xlarge runner from .build:linux:cross: GitLab's open-source # program bills small and large runners the same, so don't downsize. @@ -475,6 +479,28 @@ build:linux:cross:sme:gcc-14:full: EIGEN_CI_BUILD_TARGET: buildtests timeout: 2h +# FEAT_SME_F64F64 is optional, and every job above builds with it, so nothing +# else compiles the configuration where it is absent: the #else static_assert +# that keeps double off the GEMM driver, and double on the generic kernel inside +# an SME build. GCC hard-errors on the fp64 FMOPA without the feature, so this +# is also what catches an fp64 SME intrinsic escaping the +# EIGEN_VECTORIZE_SME_F64F64 gate. The EIGEN_ARM64_NO_SME_F64F64 opt-out reaches +# the same branch, so one job covers both ways in. product_sme alone: its double +# subtest is written to run in either configuration, and one target keeps the +# job cheap. +build:linux:cross:sme:gcc-14:no-f64f64: + extends: .build:linux:cross:sme + variables: + EIGEN_CI_C_COMPILER: gcc-14 + EIGEN_CI_CXX_COMPILER: g++-14 + EIGEN_CI_CROSS_C_COMPILER: aarch64-linux-gnu-gcc-14 + EIGEN_CI_CROSS_CXX_COMPILER: aarch64-linux-gnu-g++-14 + # The anchor's arguments, less +sme-f64f64. + EIGEN_CI_ADDITIONAL_ARGS: >- + -DEIGEN_TEST_SME=ON + -DEIGEN_TEST_CUSTOM_CXX_FLAGS=-march=armv9.2-a+sme2 + -DCMAKE_CROSSCOMPILING_EMULATOR=qemu-aarch64;-L;/usr/aarch64-linux-gnu + EIGEN_CI_BUILD_TARGET: product_sme build:linux:cross:aarch64:gcc-10:default: @@ -728,3 +754,9 @@ build:linux:riscv64:gcc-15:default:affected: build:linux:cross:sme:gcc-14:full:affected: extends: build:linux:cross:sme:gcc-14:full rules: !reference [.rules:libeigen:affected-tests:sve-sme, rules] + +# One product_sme build, so a change to the gate is caught by the affected tier +# rather than only by the nightly run. +build:linux:cross:sme:gcc-14:no-f64f64:affected: + extends: build:linux:cross:sme:gcc-14:no-f64f64 + rules: !reference [.rules:libeigen:affected-tests:sve-sme, rules] diff --git a/ci/test.linux.gitlab-ci.yml b/ci/test.linux.gitlab-ci.yml index c5d881248..0b644a1ac 100644 --- a/ci/test.linux.gitlab-ci.yml +++ b/ci/test.linux.gitlab-ci.yml @@ -470,6 +470,18 @@ test:linux:sme2048:gcc-14:default:official: variables: QEMU_CPU: max,sme=on,sme2048=on,sme128=off,sme256=off,sme512=off,sme1024=off +# The no-FEAT_SME_F64F64 build, at one SVL: double has to reach the generic +# kernel and still be correct inside an SME build, where packet traits, +# alignment and cache blocking all differ from a plain NEON build. Its subtest +# skips the SME-only symm and mapper cases there, so the run is short. One SVL +# is enough -- the generic kernel it exercises is SVL-independent. +test:linux:sme512:gcc-14:no-f64f64:official: + extends: .test:linux:sme:gcc-14:default + needs: [ build:linux:cross:sme:gcc-14:no-f64f64 ] + variables: + QEMU_CPU: max,sme=on,sme512=on,sme128=off,sme256=off,sme1024=off,sme2048=off + EIGEN_CI_CTEST_REGEX: product_sme + # Full nightly SME test run: the complete Official + Unsupported ctest suite # against the full build above, at a single representative SVL (512: the logical # block matches the ZA tile grid exactly, exercising the multi-vector-load fast diff --git a/doc/PreprocessorDirectives.dox b/doc/PreprocessorDirectives.dox index 75db150bd..0de389e0e 100644 --- a/doc/PreprocessorDirectives.dox +++ b/doc/PreprocessorDirectives.dox @@ -124,8 +124,12 @@ run time. However, these assertions do cost time and can thus be turned off. - \b \c EIGEN_DEFAULT_L1_CACHE_SIZE - Sets the default L1 cache size that is used in Eigen's GEBP kernel when the correct cache size cannot be determined at runtime. - \b \c EIGEN_DEFAULT_L2_CACHE_SIZE - Sets the default L2 cache size that is used in Eigen's GEBP kernel when the correct cache size cannot be determined at runtime. - \b \c EIGEN_DEFAULT_L3_CACHE_SIZE - Sets the default L3 cache size that is used in Eigen's GEBP kernel when the correct cache size cannot be determined at runtime. - - \b \c EIGEN_SME_MAX_KC - Maximum depth (k) blocking size used by the ARM SME GEMM kernel. The default (2048) is - empirically tuned for Apple M4; override to retune for other SME implementations. + - \b \c EIGEN_SME_MAX_KC - Maximum depth (k) blocking size used by the ARM SME GEMM kernel, expressed in \c float + elements and scaled by the scalar width for wider types. The default (2048) is empirically tuned for Apple M4; + override to retune for other SME implementations. + - \b \c EIGEN_ARM64_NO_SME_F64F64 - Disables the double-precision ARM SME GEMM kernel, which needs the optional + FEAT_SME_F64F64 extension. %Eigen enables it whenever the compiler reports that extension; define this macro when + the build target is wider than the run target, so that \c double keeps the generic kernel. - \b \c EIGEN_SME_PACKED_RHS_BUDGET_BYTES - Byte budget for the packed RHS panel in the ARM SME GEMM kernel. The default (32 MB) is empirically tuned for Apple M4. - \b \c EIGEN_SME_LHS_WORKING_SET_BUDGET_BYTES - Byte budget for the LHS working set in the ARM SME GEMM kernel. diff --git a/doc/TopicVectorization.dox b/doc/TopicVectorization.dox index 6a5a7ced5..129eb6125 100644 --- a/doc/TopicVectorization.dox +++ b/doc/TopicVectorization.dox @@ -20,8 +20,11 @@ On CPUs, %Eigen provides vectorized kernels for the following instruction sets: ARM / AArch64NEON; SVE and SME as opt-in backends SVE requires \c EIGEN_ARM64_USE_SVE and a fixed vector length (-msve-vector-bits=N); the SME backend, enabled with \c EIGEN_ARM64_USE_SME, - accelerates matrix products and must be built \em without -msve-vector-bits, since - a fixed length would pin the kernels to one runtime streaming vector length + accelerates \c float and \c double matrix products and must be built \em without + -msve-vector-bits, since a fixed length would pin the kernels to one runtime + streaming vector length. Double precision additionally needs the optional FEAT_SME_F64F64 + extension (+sme-f64f64, or a -mcpu that implies it); without it \c double + keeps the generic kernel PowerPCAltiVec, VSX, MMA IBM Z (s390x)ZVector MIPSMSA diff --git a/test/product_sme.cpp b/test/product_sme.cpp index 82f761246..2465ea30b 100644 --- a/test/product_sme.cpp +++ b/test/product_sme.cpp @@ -9,6 +9,10 @@ // SME GEMM kernel tests. // Requires compiler flags: -march=armv9.2-a+sme2 and -DEIGEN_ARM64_USE_SME. +// Double precision additionally needs FEAT_SME_F64F64 (+sme-f64f64, or a -mcpu +// that implies it); without it EIGEN_VECTORIZE_SME_F64F64 is undefined and +// double keeps the generic kernel, so the double subtest packs its cases +// through that path instead. #include "product.h" @@ -23,10 +27,28 @@ "the typical CMake invocation)." #endif -using SmeColMajorMatF = Matrix; -using SmeRowMajorMatF = Matrix; -using SmeColMajorStridedMatF = Map>; -using SmeRowMajorStridedMatF = Map>; +template +using SmeColMajorMat = Matrix; +template +using SmeRowMajorMat = Matrix; +template +using SmeVector = Matrix; +template +using SmeColMajorStridedMat = Map, 0, Stride>; +template +using SmeRowMajorStridedMat = Map, 0, Stride>; + +// The logical micro-kernel block width for Scalar: kSmeMr for float, kSmeMrD +// for double. Sizes below are expressed in terms of it so each scalar sweeps +// its own block and ZA-tile boundaries. +template +static constexpr int sme_mr() { + return internal::sme_block::mr; +} +template +static constexpr int sme_nr() { + return internal::sme_block::nr; +} template static void verify_strided_result(int n, ResultMat& storage, const Stride& stride) { @@ -43,44 +65,48 @@ static void verify_strided_result(int n, ResultMat& storage, const Stride +template static void test_general_strided_result(int n) { // General-stride C path: InputMat selects the source packers, while both C // strides are non-unit so sme_store_za_tile uses scalar scatter. - SmeColMajorMatF storage = SmeColMajorMatF::Zero(2 * n, n); - verify_strided_result( + SmeColMajorMat storage = SmeColMajorMat::Zero(2 * n, n); + verify_strided_result, SmeColMajorStridedMat>( n, storage, Stride(/*outer=*/2 * n, /*inner=*/2)); // Padding rows skipped by the strided Map should not be touched. for (int i = 0; i < n; ++i) { for (int j = 0; j < n; ++j) { - VERIFY(storage(2 * i + 1, j) == float(0)); + VERIFY(storage(2 * i + 1, j) == Scalar(0)); } } } +template static void test_rowmajor_strided_result(int n) { // RowMajor C path: inner stride is one, with padded columns after the Map. - SmeRowMajorMatF storage = SmeRowMajorMatF::Zero(n, 2 * n); - verify_strided_result( + SmeRowMajorMat storage = SmeRowMajorMat::Zero(n, 2 * n); + verify_strided_result, SmeRowMajorMat, SmeRowMajorStridedMat>( n, storage, Stride(/*outer=*/2 * n, /*inner=*/1)); // Padding columns skipped by the strided Map should not be touched. for (int i = 0; i < n; ++i) { for (int j = n; j < 2 * n; ++j) { - VERIFY(storage(i, j) == float(0)); + VERIFY(storage(i, j) == Scalar(0)); } } } +// Exercise the kc split path just above the SME blocking heuristic's depth cap +// (sme_max_kc in GeneralBlockPanelKernel.h, scaled by the scalar width). +template static void test_deep_k_split() { constexpr int rows = 64; - constexpr int depth = 2050; + const int depth = int(2 * (Index(EIGEN_SME_MAX_KC) * Index(sizeof(float)) / Index(sizeof(Scalar)))) + 2; constexpr int cols = 64; - SmeColMajorMatF A = SmeColMajorMatF::Random(rows, depth); - SmeColMajorMatF B = SmeColMajorMatF::Random(depth, cols); - SmeColMajorMatF C = SmeColMajorMatF::Random(rows, cols); - SmeColMajorMatF c_before = C; + SmeColMajorMat A = SmeColMajorMat::Random(rows, depth); + SmeColMajorMat B = SmeColMajorMat::Random(depth, cols); + SmeColMajorMat C = SmeColMajorMat::Random(rows, cols); + SmeColMajorMat c_before = C; C.noalias() += A * B; @@ -102,81 +128,86 @@ static void test_deep_k_split() { // A distinctive marker for buffer cells the packer must leave untouched, and // for the unused triangle of a lower-triangular operand. Random values live in // [-1, 1], so it never collides with a real packed value. -static constexpr float kPackSentinel = 98765.0f; +template +static Scalar pack_sentinel() { + return Scalar(98765); +} // Lower-triangular n x n operand plus the dense symmetric reference the packer -// must emit. The unused triangle is filled with kPackSentinel so a packer that +// must emit. The unused triangle is filled with the sentinel so a packer that // copies the dense matrix and never mirrors fails VERIFY_IS_EQUAL. // product_selfadjoint_matrix stores the valid triangle where row >= col // (after the Upper/RowMajor xor), so the packer must read stored(row,col) // below the diagonal and stored(col,row) above it. -template -static void make_lower_stored_symmetric(Index n, Matrix& stored, - Matrix& full) { - full = Matrix::Random(n, n); - full = ((full + full.transpose()) * 0.5f).eval(); - stored = Matrix::Constant(n, n, kPackSentinel); +template +static void make_lower_stored_symmetric(Index n, Matrix& stored, + Matrix& full) { + full = Matrix::Random(n, n); + full = ((full + full.transpose()) * Scalar(0.5)).eval(); + stored = Matrix::Constant(n, n, pack_sentinel()); for (Index i = 0; i < n; ++i) for (Index j = 0; j <= i; ++j) stored(i, j) = full(i, j); } // LHS SYMM packer: a square selfadjoint diagonal block of size n, packed into -// uniform kSmeMr-wide depth-major panels. Reference: full(i+r, k). -template +// uniform mr-wide depth-major panels. Reference: full(i+r, k). +template static void verify_symm_pack_lhs(Index n) { - const Index MR = internal::kSmeMr; - Matrix stored, full; - make_lower_stored_symmetric(n, stored, full); + const Index MR = sme_mr(); + Matrix stored, full; + make_lower_stored_symmetric(n, stored, full); - VectorXf packed = VectorXf::Constant(n * n, kPackSentinel); - VectorXf ref = VectorXf::Constant(n * n, kPackSentinel); + SmeVector packed = SmeVector::Constant(n * n, pack_sentinel()); + SmeVector ref = SmeVector::Constant(n * n, pack_sentinel()); for (Index i = 0; i < n; i += MR) { const Index w = numext::mini(MR, n - i); for (Index k = 0; k < n; ++k) for (Index r = 0; r < w; ++r) ref[i * n + k * w + r] = full(i + r, k); } - internal::symm_pack_lhs pack; + internal::symm_pack_lhs(), 1, StorageOrder> pack; pack(packed.data(), stored.data(), stored.outerStride(), /*cols(depth)=*/n, /*rows=*/n); VERIFY_IS_EQUAL(packed, ref); } // RHS SYMM packer: a depth block [k2, k2 + depth) x cols columns of an N x N -// selfadjoint matrix, packed into kSmeNr-wide depth-major panels. Reference: +// selfadjoint matrix, packed into nr-wide depth-major panels. Reference: // full(k2 + k, j + c). A k2 > 0 offset makes the transposed region non-empty, // so partial-width panels reach the two-pass transpose. -template +template static void verify_symm_pack_rhs(Index N, Index depth, Index cols, Index k2) { eigen_assert(k2 + depth <= N && cols <= N); - const Index NR = internal::kSmeNr; - Matrix stored, full; - make_lower_stored_symmetric(N, stored, full); + const Index NR = sme_nr(); + Matrix stored, full; + make_lower_stored_symmetric(N, stored, full); - VectorXf packed = VectorXf::Constant(cols * depth, kPackSentinel); - VectorXf ref = VectorXf::Constant(cols * depth, kPackSentinel); + SmeVector packed = SmeVector::Constant(cols * depth, pack_sentinel()); + SmeVector ref = SmeVector::Constant(cols * depth, pack_sentinel()); for (Index j = 0; j < cols; j += NR) { const Index w = numext::mini(NR, cols - j); for (Index k = 0; k < depth; ++k) for (Index c = 0; c < w; ++c) ref[j * depth + k * w + c] = full(k2 + k, j + c); } - internal::symm_pack_rhs pack; + internal::symm_pack_rhs(), StorageOrder> pack; pack(packed.data(), stored.data(), stored.outerStride(), /*rows(depth)=*/depth, /*cols=*/cols, k2); VERIFY_IS_EQUAL(packed, ref); } +template static void test_symm_pack() { // The last panel width sweeps a range of partial widths; at each SVL the // two-pass trailing transpose (the if->loop fix) fires when a partial width - // leaves a trailing row-group remainder in (svlw, 2*svlw). The spread below - // hits that for svlw in {4, 8, 16, 32, 64} (SVL 128..2048). + // leaves a trailing row-group remainder in (svl, 2*svl). The spread below + // hits that for svl in {2, 4, 8, 16, 32, 64} -- fp32 SVL 128..2048 and the + // fp64 lane counts, which are half of those. const int sizes[] = {1, 5, 7, 17, 31, 32, 33, 37, 39, 45, 48, 49, 55, 57, 63, 64, 65, 79, 96, 97}; for (int n : sizes) { - verify_symm_pack_lhs(n); - verify_symm_pack_lhs(n); + verify_symm_pack_lhs(n); + verify_symm_pack_lhs(n); // RHS, single depth block anchored at the diagonal (k2 == 0). - verify_symm_pack_rhs(n, n, n, 0); - verify_symm_pack_rhs(n, n, n, 0); + verify_symm_pack_rhs(n, n, n, 0); + verify_symm_pack_rhs(n, n, n, 0); } // RHS depth blocks offset from the diagonal (k2 > 0): the transposed region is @@ -189,8 +220,8 @@ static void test_symm_pack() { {100, 32, 39, 16}, {100, 24, 39, 32}, {100, 40, 64, 8}, {100, 39, 39, 33}, {128, 57, 57, 40}, {128, 33, 45, 60}, }; for (const RhsCase& c : rhs_cases) { - verify_symm_pack_rhs(c.N, c.depth, c.cols, c.k2); - verify_symm_pack_rhs(c.N, c.depth, c.cols, c.k2); + verify_symm_pack_rhs(c.N, c.depth, c.cols, c.k2); + verify_symm_pack_rhs(c.N, c.depth, c.cols, c.k2); } } @@ -207,53 +238,57 @@ static void test_symm_pack() { // Minimal stand-ins for by-value sub-mappers. ColMajor packets advance the // first index, while RowMajor packets follow the normal storage-inner second // index. operator() returns by value so both take the no-direct-access dispatch. +template struct ByValueColMajorLhsMapper { - const float* data; + const Scalar* data; Index stride; // element(i, k) = data[i + k * stride], contiguous in i - float operator()(Index i, Index k) const { return data[i + k * stride]; } + Scalar operator()(Index i, Index k) const { return data[i + k * stride]; } template EIGEN_ALWAYS_INLINE Packet loadPacket(Index i, Index k) const { return internal::ploadu(data + i + k * stride); } }; +template struct ByValueRowMajorLhsMapper { - const float* data; + const Scalar* data; Index stride; // element(i, k) = data[i * stride + k], contiguous in k - float operator()(Index i, Index k) const { return data[i * stride + k]; } + Scalar operator()(Index i, Index k) const { return data[i * stride + k]; } template EIGEN_ALWAYS_INLINE Packet loadPacket(Index i, Index k) const { return internal::ploadu(data + i * stride + k); } }; +template struct ByValueColMajorRhsMapper { struct LinearMapper { - const float* p; // &element(0, col); contiguous in depth - float operator()(Index k) const { return p[k]; } + const Scalar* p; // &element(0, col); contiguous in depth + Scalar operator()(Index k) const { return p[k]; } template EIGEN_ALWAYS_INLINE Packet loadPacket(Index k) const { return internal::ploadu(p + k); } }; - const float* data; + const Scalar* data; Index stride; // element(k, col) = data[k + col * stride], contiguous in k - float operator()(Index k, Index col) const { return data[k + col * stride]; } + Scalar operator()(Index k, Index col) const { return data[k + col * stride]; } LinearMapper getLinearMapper(Index k, Index col) const { return LinearMapper{data + k + col * stride}; } }; +template struct ByValueRowMajorRhsMapper { struct LinearMapper { - const float* p; // &element(row, col); packet offsets advance columns - float operator()(Index offset) const { return p[offset]; } + const Scalar* p; // &element(row, col); packet offsets advance columns + Scalar operator()(Index offset) const { return p[offset]; } template EIGEN_ALWAYS_INLINE Packet loadPacket(Index offset) const { return internal::ploadu(p + offset); } }; - const float* data; + const Scalar* data; Index stride; // element(k, col) = data[k * stride + col], contiguous in col - float operator()(Index k, Index col) const { return data[k * stride + col]; } + Scalar operator()(Index k, Index col) const { return data[k * stride + col]; } LinearMapper getLinearMapper(Index k, Index col) const { return LinearMapper{data + k * stride + col}; } }; @@ -261,7 +296,7 @@ struct ByValueRowMajorRhsMapper { // dst_panel formula in sme_pack_lhs_fallback. template static Index packed_len(Index outer, Index depth, Index unit, Index dst_stride, Index dst_offset) { - // `outer` is rows (LHS) or cols (RHS); `unit` is kSmeMr or kSmeNr. + // `outer` is rows (LHS) or cols (RHS); `unit` is the panel width mr or nr. if (!PanelMode) return outer * depth; Index end = 0; for (Index i = 0; i < outer; i += unit) { @@ -271,11 +306,11 @@ static Index packed_len(Index outer, Index depth, Index unit, Index dst_stride, return end; } -template -static void fill_lhs_ref(VectorXf& ref, const MatrixType& V, Index rows, Index depth, Index dst_stride, +template +static void fill_lhs_ref(SmeVector& ref, const MatrixType& V, Index rows, Index depth, Index dst_stride, Index dst_offset) { - const Index MR = internal::kSmeMr; - ref.setConstant(kPackSentinel); + const Index MR = sme_mr(); + ref.setConstant(pack_sentinel()); for (Index i = 0; i < rows; i += MR) { const Index w = numext::mini(MR, rows - i); const Index base = PanelMode ? i * dst_stride + dst_offset * w : i * depth; @@ -284,11 +319,11 @@ static void fill_lhs_ref(VectorXf& ref, const MatrixType& V, Index rows, Index d } } -template -static void fill_rhs_ref(VectorXf& ref, const MatrixType& V, Index cols, Index depth, Index dst_stride, +template +static void fill_rhs_ref(SmeVector& ref, const MatrixType& V, Index cols, Index depth, Index dst_stride, Index dst_offset) { - const Index NR = internal::kSmeNr; - ref.setConstant(kPackSentinel); + const Index NR = sme_nr(); + ref.setConstant(pack_sentinel()); for (Index j = 0; j < cols; j += NR) { const Index w = numext::mini(NR, cols - j); const Index base = PanelMode ? j * dst_stride + dst_offset * w : j * depth; @@ -300,12 +335,12 @@ static void fill_rhs_ref(VectorXf& ref, const MatrixType& V, Index cols, Index d // Inner-strided blas mapper LHS: element(i, k) laid out with inner stride // `incr`. ColMajor takes the vectorised gather path; RowMajor takes the scalar // path (its packets would run along depth, not rows). -template +template static void verify_fallback_lhs_strided(Index rows, Index depth, Index incr) { - typedef internal::blas_data_mapper Mapper; - MatrixXf V = MatrixXf::Random(rows, depth); + using Mapper = internal::blas_data_mapper; + Matrix V = Matrix::Random(rows, depth); const Index mstride = (StorageOrder == ColMajor ? rows : depth) * incr; - VectorXf buf = VectorXf::Zero((StorageOrder == ColMajor ? depth : rows) * mstride + incr); + SmeVector buf = SmeVector::Zero((StorageOrder == ColMajor ? depth : rows) * mstride + incr); for (Index k = 0; k < depth; ++k) for (Index i = 0; i < rows; ++i) buf[StorageOrder == ColMajor ? i * incr + k * mstride : k * incr + i * mstride] = V(i, k); @@ -313,12 +348,12 @@ static void verify_fallback_lhs_strided(Index rows, Index depth, Index incr) { const Index dst_stride = PanelMode ? depth + 5 : 0; const Index dst_offset = PanelMode ? 3 : 0; - const Index len = packed_len(rows, depth, internal::kSmeMr, dst_stride, dst_offset); - VectorXf packed = VectorXf::Constant(len, kPackSentinel); - VectorXf ref(len); - fill_lhs_ref(ref, V, rows, depth, dst_stride, dst_offset); + const Index len = packed_len(rows, depth, sme_mr(), dst_stride, dst_offset); + SmeVector packed = SmeVector::Constant(len, pack_sentinel()); + SmeVector ref(len); + fill_lhs_ref(ref, V, rows, depth, dst_stride, dst_offset); - internal::gemm_pack_lhs::type, + internal::gemm_pack_lhs(), 1, typename internal::packet_traits::type, StorageOrder, false, PanelMode> pack; pack(packed.data(), mapper, depth, rows, dst_stride, dst_offset); @@ -327,22 +362,22 @@ static void verify_fallback_lhs_strided(Index rows, Index depth, Index incr) { // By-value LHS mappers exercise both packet directions. RowMajor must stay // scalar because its packets advance depth rather than rows. -template +template static void verify_fallback_lhs_byvalue(Index rows, Index depth) { - typedef Matrix MatrixType; - typedef typename std::conditional::type - Mapper; + using MatrixType = Matrix; + using Mapper = typename std::conditional, + ByValueRowMajorLhsMapper>::type; MatrixType V = MatrixType::Random(rows, depth); Mapper mapper{V.data(), V.outerStride()}; const Index dst_stride = PanelMode ? depth + 5 : 0; const Index dst_offset = PanelMode ? 3 : 0; - const Index len = packed_len(rows, depth, internal::kSmeMr, dst_stride, dst_offset); - VectorXf packed = VectorXf::Constant(len, kPackSentinel); - VectorXf ref(len); - fill_lhs_ref(ref, V, rows, depth, dst_stride, dst_offset); + const Index len = packed_len(rows, depth, sme_mr(), dst_stride, dst_offset); + SmeVector packed = SmeVector::Constant(len, pack_sentinel()); + SmeVector ref(len); + fill_lhs_ref(ref, V, rows, depth, dst_stride, dst_offset); - internal::gemm_pack_lhs::type, + internal::gemm_pack_lhs(), 1, typename internal::packet_traits::type, StorageOrder, false, PanelMode> pack; pack(packed.data(), mapper, depth, rows, dst_stride, dst_offset); @@ -351,12 +386,12 @@ static void verify_fallback_lhs_byvalue(Index rows, Index depth) { // Inner-strided blas mapper RHS: element(k, col) with inner stride `incr`. // ColMajor takes the vectorised transpose path; RowMajor takes the scalar path. -template +template static void verify_fallback_rhs_strided(Index depth, Index cols, Index incr) { - typedef internal::blas_data_mapper Mapper; - MatrixXf V = MatrixXf::Random(depth, cols); + using Mapper = internal::blas_data_mapper; + Matrix V = Matrix::Random(depth, cols); const Index mstride = (StorageOrder == ColMajor ? depth : cols) * incr; - VectorXf buf = VectorXf::Zero((StorageOrder == ColMajor ? cols : depth) * mstride + incr); + SmeVector buf = SmeVector::Zero((StorageOrder == ColMajor ? cols : depth) * mstride + incr); for (Index col = 0; col < cols; ++col) for (Index k = 0; k < depth; ++k) buf[StorageOrder == ColMajor ? k * incr + col * mstride : col * incr + k * mstride] = V(k, col); @@ -364,109 +399,106 @@ static void verify_fallback_rhs_strided(Index depth, Index cols, Index incr) { const Index dst_stride = PanelMode ? depth + 5 : 0; const Index dst_offset = PanelMode ? 3 : 0; - const Index len = packed_len(cols, depth, internal::kSmeNr, dst_stride, dst_offset); - VectorXf packed = VectorXf::Constant(len, kPackSentinel); - VectorXf ref(len); - fill_rhs_ref(ref, V, cols, depth, dst_stride, dst_offset); + const Index len = packed_len(cols, depth, sme_nr(), dst_stride, dst_offset); + SmeVector packed = SmeVector::Constant(len, pack_sentinel()); + SmeVector ref(len); + fill_rhs_ref(ref, V, cols, depth, dst_stride, dst_offset); - internal::gemm_pack_rhs pack; + internal::gemm_pack_rhs(), StorageOrder, false, PanelMode> pack; pack(packed.data(), mapper, depth, cols, dst_stride, dst_offset); VERIFY_IS_EQUAL(packed, ref); } // By-value RHS mappers likewise cover both packet directions. RowMajor packets // advance columns, so the depth-oriented transpose fallback must stay scalar. -template +template static void verify_fallback_rhs_byvalue(Index depth, Index cols) { - typedef Matrix MatrixType; - typedef typename std::conditional::type - Mapper; + using MatrixType = Matrix; + using Mapper = typename std::conditional, + ByValueRowMajorRhsMapper>::type; MatrixType V = MatrixType::Random(depth, cols); Mapper mapper{V.data(), V.outerStride()}; const Index dst_stride = PanelMode ? depth + 5 : 0; const Index dst_offset = PanelMode ? 3 : 0; - const Index len = packed_len(cols, depth, internal::kSmeNr, dst_stride, dst_offset); - VectorXf packed = VectorXf::Constant(len, kPackSentinel); - VectorXf ref(len); - fill_rhs_ref(ref, V, cols, depth, dst_stride, dst_offset); + const Index len = packed_len(cols, depth, sme_nr(), dst_stride, dst_offset); + SmeVector packed = SmeVector::Constant(len, pack_sentinel()); + SmeVector ref(len); + fill_rhs_ref(ref, V, cols, depth, dst_stride, dst_offset); - internal::gemm_pack_rhs pack; + internal::gemm_pack_rhs(), StorageOrder, false, PanelMode> pack; pack(packed.data(), mapper, depth, cols, dst_stride, dst_offset); VERIFY_IS_EQUAL(packed, ref); } +template static void test_mapper_fallback() { - const int widths[] = {4, 5, 32, 33, 65}; // rows/cols around 4 and 32 - const int depths[] = {1, 3, 8, 35}; // depth remainders 1..3 and larger + const int MR = sme_mr(); + const int widths[] = {4, 5, MR, MR + 1, 2 * MR + 1}; // rows/cols around 4 and the panel width + const int depths[] = {1, 3, 8, 35}; // depth remainders 1..3 and larger for (int n : widths) { for (int d : depths) { for (int incr : {2, 3}) { - verify_fallback_lhs_strided(n, d, incr); - verify_fallback_lhs_strided(n, d, incr); - verify_fallback_lhs_strided(n, d, incr); // scalar path - verify_fallback_lhs_strided(n, d, incr); - verify_fallback_rhs_strided(d, n, incr); - verify_fallback_rhs_strided(d, n, incr); - verify_fallback_rhs_strided(d, n, incr); // scalar path - verify_fallback_rhs_strided(d, n, incr); + verify_fallback_lhs_strided(n, d, incr); + verify_fallback_lhs_strided(n, d, incr); + verify_fallback_lhs_strided(n, d, incr); // scalar path + verify_fallback_lhs_strided(n, d, incr); + verify_fallback_rhs_strided(d, n, incr); + verify_fallback_rhs_strided(d, n, incr); + verify_fallback_rhs_strided(d, n, incr); // scalar path + verify_fallback_rhs_strided(d, n, incr); } - verify_fallback_lhs_byvalue(n, d); - verify_fallback_lhs_byvalue(n, d); - verify_fallback_lhs_byvalue(n, d); - verify_fallback_lhs_byvalue(n, d); - verify_fallback_rhs_byvalue(d, n); - verify_fallback_rhs_byvalue(d, n); - verify_fallback_rhs_byvalue(d, n); - verify_fallback_rhs_byvalue(d, n); + verify_fallback_lhs_byvalue(n, d); + verify_fallback_lhs_byvalue(n, d); + verify_fallback_lhs_byvalue(n, d); + verify_fallback_lhs_byvalue(n, d); + verify_fallback_rhs_byvalue(d, n); + verify_fallback_rhs_byvalue(d, n); + verify_fallback_rhs_byvalue(d, n); + verify_fallback_rhs_byvalue(d, n); } } } -EIGEN_DECLARE_TEST(product_sme) { - // Square edge cases around the block and tile boundaries (the block is - // kSmeMr x kSmeNr and a ZA tile is svlw x svlw, so the sizes below land - // just on/off the intra-tile splits and the block tails at SVL=512). - CALL_SUBTEST_1(product(Matrix(1, 1))); - CALL_SUBTEST_1(product(Matrix(15, 15))); - CALL_SUBTEST_1(product(Matrix(16, 16))); - CALL_SUBTEST_1(product(Matrix(17, 17))); - CALL_SUBTEST_1(product(Matrix(31, 31))); - CALL_SUBTEST_1(product(Matrix(33, 33))); - CALL_SUBTEST_1(product(Matrix(63, 63))); - CALL_SUBTEST_1(product(Matrix(64, 64))); - CALL_SUBTEST_1(product(Matrix(65, 65))); +// --------------------------------------------------------------------------- +// Product-level coverage, swept relative to the scalar's own block width. +// --------------------------------------------------------------------------- - // Thin / wide rectangular cases (M x 1, 1 x N) - CALL_SUBTEST_1(product(Matrix(32, 1))); - CALL_SUBTEST_1(product(Matrix(1, 32))); - CALL_SUBTEST_1(product(Matrix(1, 64))); - CALL_SUBTEST_1(product(Matrix(64, 1))); +// Sizes that land just on and off the block tails and the intra-block ZA-tile +// splits. MR/2 is the tile side at the SVL=512 design point. +template +static std::vector sme_edge_sizes() { + const int MR = sme_mr(); + return {1, MR / 2 - 1, MR / 2, MR / 2 + 1, MR - 1, MR, MR + 1, 2 * MR - 1, 2 * MR, 2 * MR + 1}; +} - // Non-float scalar smoke tests: SME only specializes fp32, so these prove - // unsupported scalar types still route through the generic product path. - CALL_SUBTEST_2(product(Matrix(33, 17))); - CALL_SUBTEST_3(product(Matrix, Dynamic, Dynamic>(33, 17))); +template +static void test_products() { + const int MR = sme_mr(); - // Non-square cases that exercise tail paths for both M and N - CALL_SUBTEST_1(product(Matrix(17, 65))); - CALL_SUBTEST_1(product(Matrix(65, 17))); - CALL_SUBTEST_1(product(Matrix(15, 63))); - CALL_SUBTEST_1(product(Matrix(33, 7))); - CALL_SUBTEST_1(product(Matrix(7, 33))); - CALL_SUBTEST_1(product(Matrix(128, 3))); - CALL_SUBTEST_1(product(Matrix(3, 128))); + // Square edge cases around the block and tile boundaries. + for (int n : sme_edge_sizes()) product(SmeColMajorMat(n, n)); - // Exercise the kc split path just above the SME blocking heuristic's depth - // cap (sme_max_kc in GeneralBlockPanelKernel.h). - test_deep_k_split(); - test_symm_pack(); - test_mapper_fallback(); + // Thin / wide rectangular cases (M x 1, 1 x N) and non-square cases that + // exercise tail paths for both M and N. + product(SmeColMajorMat(MR, 1)); + product(SmeColMajorMat(1, MR)); + product(SmeColMajorMat(1, 2 * MR)); + product(SmeColMajorMat(2 * MR, 1)); + product(SmeColMajorMat(MR + 1, 2 * MR + 1)); + product(SmeColMajorMat(2 * MR + 1, MR + 1)); + product(SmeColMajorMat(MR - 1, 2 * MR - 1)); + product(SmeColMajorMat(MR + 1, 7)); + product(SmeColMajorMat(7, MR + 1)); + product(SmeColMajorMat(4 * MR, 3)); + product(SmeColMajorMat(3, 4 * MR)); + + test_deep_k_split(); // Random sizes for (int i = 0; i < g_repeat; i++) { - CALL_SUBTEST_1(product(Matrix(internal::random(1, EIGEN_TEST_MAX_SIZE), - internal::random(1, EIGEN_TEST_MAX_SIZE)))); + product(SmeColMajorMat(internal::random(1, EIGEN_TEST_MAX_SIZE), + internal::random(1, EIGEN_TEST_MAX_SIZE))); } // Exercise the RowMajor packers and RowMajor result path. When the input @@ -476,40 +508,57 @@ EIGEN_DECLARE_TEST(product_sme) { // - the RowMajor RHS packer (gemm_pack_rhs<..., RowMajor>) // - the RowMajor-C dispatch in GeneralMatrixMatrix.h (which transposes // the computation: C^T = B^T * A^T). - CALL_SUBTEST_1(product(SmeRowMajorMatF(15, 15))); - CALL_SUBTEST_1(product(SmeRowMajorMatF(16, 16))); - CALL_SUBTEST_1(product(SmeRowMajorMatF(17, 17))); - CALL_SUBTEST_1(product(SmeRowMajorMatF(31, 31))); - CALL_SUBTEST_1(product(SmeRowMajorMatF(32, 32))); - CALL_SUBTEST_1(product(SmeRowMajorMatF(33, 33))); - CALL_SUBTEST_1(product(SmeRowMajorMatF(64, 64))); - CALL_SUBTEST_1(product(SmeRowMajorMatF(65, 65))); - CALL_SUBTEST_1(product(SmeRowMajorMatF(17, 65))); - CALL_SUBTEST_1(product(SmeRowMajorMatF(65, 17))); + for (int n : sme_edge_sizes()) { + if (n > 1) product(SmeRowMajorMat(n, n)); + } + product(SmeRowMajorMat(MR + 1, 2 * MR + 1)); + product(SmeRowMajorMat(2 * MR + 1, MR + 1)); for (int i = 0; i < g_repeat; i++) { - CALL_SUBTEST_1(product( - SmeRowMajorMatF(internal::random(1, EIGEN_TEST_MAX_SIZE), internal::random(1, EIGEN_TEST_MAX_SIZE)))); + product(SmeRowMajorMat(internal::random(1, EIGEN_TEST_MAX_SIZE), + internal::random(1, EIGEN_TEST_MAX_SIZE))); } // Exercise the general-stride branch of sme_store_za_tile: fires when both // C_stride_row != 1 and C_stride_col != 1, e.g. a Map with an // explicit non-unit inner stride. product.h never builds such a result, so // without this subtest the scalar-scatter path is effectively untested. - for (int n : {15, 16, 17, 31, 32, 33, 63, 64, 65}) { - test_general_strided_result(n); - test_general_strided_result(n); - test_rowmajor_strided_result(n); + for (int n : sme_edge_sizes()) { + if (n < 2) continue; + test_general_strided_result>(n); + test_general_strided_result>(n); + test_rowmajor_strided_result(n); } // Row-LHS x Row-RHS -> Col-C: the one LHS/RHS/C storage combination that // product.h's transpose-style expressions never build directly (it always // flips one side of the multiplication). The code paths are the same as // other combinations via Eigen's dispatch, but exercise them explicitly. - for (int n : {15, 16, 17, 31, 32, 33, 63, 64, 65}) { - Matrix A = Matrix::Random(n, n); - Matrix B = Matrix::Random(n, n); - SmeColMajorMatF C = SmeColMajorMatF::Zero(n, n); + for (int n : sme_edge_sizes()) { + if (n < 2) continue; + SmeRowMajorMat A = SmeRowMajorMat::Random(n, n); + SmeRowMajorMat B = SmeRowMajorMat::Random(n, n); + SmeColMajorMat C = SmeColMajorMat::Zero(n, n); C.noalias() += A * B; VERIFY_IS_APPROX(C, (A.lazyProduct(B)).eval()); } } + +EIGEN_DECLARE_TEST(product_sme) { + CALL_SUBTEST_1(test_products()); + CALL_SUBTEST_1(test_symm_pack()); + CALL_SUBTEST_1(test_mapper_fallback()); + + // double only reaches the SME kernel and packers with FEAT_SME_F64F64; the + // product sweep is meaningful either way, but the packed-layout tests name + // specializations that only exist when it is available. + CALL_SUBTEST_2(test_products()); +#ifdef EIGEN_VECTORIZE_SME_F64F64 + CALL_SUBTEST_2(test_symm_pack()); + CALL_SUBTEST_2(test_mapper_fallback()); +#endif + + // Scalar types SME does not specialize: these prove they still route through + // the generic product path. + CALL_SUBTEST_3(product(Matrix, Dynamic, Dynamic>(33, 17))); + CALL_SUBTEST_3(product(Matrix, Dynamic, Dynamic>(33, 17))); +}