From eafce2f0ce73a62deb29bbcb9b841e7c395e4ab3 Mon Sep 17 00:00:00 2001 From: Anu Oguntayo Date: Tue, 7 Jul 2026 14:02:27 +0100 Subject: [PATCH] mlas: add missing s390x ZVECTOR kernel headers to fix build The previous commit bdf348c13ac230923d203eec7325d5c49eaa5fb2 introduced MLAS support including the s390x backend, but accidentally omitted the required internal header files. This causes a compilation failure on s390x: fatal error: SgemmKernelZVECTOR.h: No such file or directory This PR restores the build by adding the following missing headers from upstream (microsoft/onnxruntime): - FgemmKernelZVECTOR.h - SgemmKernelZVECTOR.h This is the same class of bug as #29422 (loongarch64), which was fixed by PR #29423 for that architecture only. --- 3rdparty/mlas/lib/s390x/FgemmKernelZVECTOR.h | 333 +++++++++++++++++++ 3rdparty/mlas/lib/s390x/SgemmKernelZVECTOR.h | 139 ++++++++ 2 files changed, 472 insertions(+) create mode 100644 3rdparty/mlas/lib/s390x/FgemmKernelZVECTOR.h create mode 100644 3rdparty/mlas/lib/s390x/SgemmKernelZVECTOR.h diff --git a/3rdparty/mlas/lib/s390x/FgemmKernelZVECTOR.h b/3rdparty/mlas/lib/s390x/FgemmKernelZVECTOR.h new file mode 100644 index 0000000000..af779747a2 --- /dev/null +++ b/3rdparty/mlas/lib/s390x/FgemmKernelZVECTOR.h @@ -0,0 +1,333 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + FgemmKernelZVECTOR.h + +Abstract: + + This module implements the kernels for the single/double precision matrix/matrix + multiply operation (DGEMM/SGEMM). + +--*/ + +#include "mlasi.h" +#if defined(SINGLE) +#define MLAS_FLOATTYPE MLAS_FLOAT32X4 +#define MLAS_GEMMTYPE float +#define MLAS_LOAD_FLOAT MlasLoadFloat32x4 +#define MLAS_ZERO_FLOAT MlasZeroFloat32x4 +#define MLAS_STORE_FLOAT MlasStoreFloat32x4 +#define MLAS_EXTRACT_FLOAT MlasExtractLaneFloat32x4 +#define MLAS_MUL_FLOAT MlasMultiplyFloat32x4 +#define MLAS_MULADD_FLOAT MlasMultiplyAddFloat32x4 +#define MLAS_BROADCAST_FLOAT MlasBroadcastFloat32x4 +#else +#define MLAS_FLOATTYPE MLAS_FLOAT64X2 +#define MLAS_GEMMTYPE double +#define MLAS_LOAD_FLOAT MlasLoadFloat64x2 +#define MLAS_ZERO_FLOAT MlasZeroFloat64x2 +#define MLAS_STORE_FLOAT MlasStoreFloat64x2 +#define MLAS_EXTRACT_FLOAT MlasExtractLaneFloat64x2 +#define MLAS_MUL_FLOAT MlasMultiplyFloat64x2 +#define MLAS_MULADD_FLOAT MlasMultiplyAddFloat64x2 +#define MLAS_BROADCAST_FLOAT MlasBroadcastFloat64x2 +#endif +// +// Templates to ensure that a loop is unrolled. +// + +template +struct MlasLoopUnrollStep +{ + template + MLAS_FORCEINLINE + static + void + Step( + IterationArgs&&... Arguments + ) + { + IterationType::template Iteration(Arguments...); + MlasLoopUnrollStep::template Step(Arguments...); + } +}; + +template +struct MlasLoopUnrollStep +{ + template + MLAS_FORCEINLINE + static + void + Step( + IterationArgs&&... + ) + { + // Terminate the loop. + } +}; + +template +struct MlasLoopUnroll +{ + template + MLAS_FORCEINLINE + void + operator()( + IterationArgs&&... Arguments + ) + { + MlasLoopUnrollStep::template Step(Arguments...); + } +}; + +// +// Templates used with loop unrolling to perform an action on one row of the +// output. +// + +struct MlasFgemmZeroAccumulators +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOATTYPE Accumulators[RowCount][4] + ) + { + Accumulators[Row][0] = MLAS_ZERO_FLOAT(); + Accumulators[Row][1] = MLAS_ZERO_FLOAT(); + Accumulators[Row][2] = MLAS_ZERO_FLOAT(); + Accumulators[Row][3] = MLAS_ZERO_FLOAT(); + } +}; + +struct MlasFgemmLoadAElements +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOATTYPE AElements[RowCount], + const MLAS_GEMMTYPE* A, + size_t lda + ) + { + AElements[Row] = MLAS_LOAD_FLOAT(A + Row * lda); + } +}; + +struct MlasFgemmBroadcastAElements +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOATTYPE ABroadcast[RowCount], + const MLAS_GEMMTYPE* A, + size_t lda + ) + { + ABroadcast[Row] = MLAS_BROADCAST_FLOAT(A + Row * lda); + } +}; + +template +struct MlasFgemmSplatAElements +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOATTYPE AElements[RowCount], + MLAS_FLOATTYPE ABroadcast[RowCount] + ) + { + ABroadcast[Row] = vec_splat(AElements[Row], Lane); + } +}; + +struct MlasFgemmMultiplyAddRow +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOATTYPE Accumulators[RowCount][4], + MLAS_FLOATTYPE ABroadcast[RowCount], + MLAS_FLOATTYPE BElements[4] + ) + { + Accumulators[Row][0] = MLAS_MULADD_FLOAT(ABroadcast[Row], BElements[0], Accumulators[Row][0]); + Accumulators[Row][1] = MLAS_MULADD_FLOAT(ABroadcast[Row], BElements[1], Accumulators[Row][1]); + Accumulators[Row][2] = MLAS_MULADD_FLOAT(ABroadcast[Row], BElements[2], Accumulators[Row][2]); + Accumulators[Row][3] = MLAS_MULADD_FLOAT(ABroadcast[Row], BElements[3], Accumulators[Row][3]); + } +}; + +template +MLAS_FORCEINLINE +void +MlasFgemmComputeBlock( + MLAS_FLOATTYPE Accumulators[RowCount][4], + MLAS_FLOATTYPE ABroadcast[RowCount], + const MLAS_GEMMTYPE* B + ) +{ + MLAS_FLOATTYPE BElements[4]; +#if defined(SINGLE) + BElements[0] = MLAS_LOAD_FLOAT(B); + BElements[1] = MLAS_LOAD_FLOAT(B + 4); + BElements[2] = MLAS_LOAD_FLOAT(B + 8); + BElements[3] = MLAS_LOAD_FLOAT(B + 12); +#else + BElements[0] = MLAS_LOAD_FLOAT(B); + BElements[1] = MLAS_LOAD_FLOAT(B + 2); + BElements[2] = MLAS_LOAD_FLOAT(B + 4); + BElements[3] = MLAS_LOAD_FLOAT(B + 6); +#endif + + MlasLoopUnroll()(Accumulators, ABroadcast, BElements); +} + +struct MlasFgemmMultiplyAlphaRow +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOATTYPE Accumulators[4], + MLAS_FLOATTYPE AlphaBroadcast + ) + { + Accumulators[Index] = MLAS_MUL_FLOAT(Accumulators[Index], AlphaBroadcast); + } +}; + +struct MlasFgemmMultiplyAlphaAddRow +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOATTYPE Accumulators[4], + MLAS_FLOATTYPE AlphaBroadcast, + const MLAS_GEMMTYPE* C + ) + { +#if defined(SINGLE) + Accumulators[Index] = MLAS_MULADD_FLOAT(Accumulators[Index], + AlphaBroadcast, MLAS_LOAD_FLOAT(C + Index * 4)); +#else + Accumulators[Index] = MLAS_MULADD_FLOAT(Accumulators[Index], + AlphaBroadcast, MLAS_LOAD_FLOAT(C + Index * 2)); +#endif + } +}; + +struct MlasFgemmStoreRow +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOATTYPE Accumulators[4], + MLAS_GEMMTYPE* C + ) + { +#if defined(SINGLE) + MLAS_STORE_FLOAT(C + Index * 4, Accumulators[Index]); +#else + MLAS_STORE_FLOAT(C + Index * 2, Accumulators[Index]); +#endif + } +}; + +template +struct MlasFgemmStoreVector +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOATTYPE Accumulators[RowCount][4], + MLAS_GEMMTYPE* C, + size_t ldc, + MLAS_FLOATTYPE AlphaBroadcast, + bool ZeroMode + ) + { + MLAS_GEMMTYPE* c = C + Row * ldc; + + if (ZeroMode) { + MlasLoopUnroll()(Accumulators[Row], AlphaBroadcast); + } else { + MlasLoopUnroll()(Accumulators[Row], AlphaBroadcast, c); + } + + MlasLoopUnroll()(Accumulators[Row], c); + + // + // Shift down any unaligned elements to the bottom for further processing. + // + + if (VectorCount < 4) { + Accumulators[Row][0] = Accumulators[Row][VectorCount]; + } + } +}; + +struct MlasFgemmMultiplyAlphaTrailing +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOATTYPE Accumulators[RowCount][4], + MLAS_FLOATTYPE AlphaBroadcast + ) + { + Accumulators[Row][0] = MLAS_MUL_FLOAT(Accumulators[Row][0], AlphaBroadcast); + } +}; + +template +struct MlasFgemmStoreScalar +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOATTYPE Accumulators[RowCount][4], + MLAS_GEMMTYPE* C, + size_t ldc, + bool ZeroMode + ) + { + MLAS_GEMMTYPE* c = C + Row * ldc + Lane; + MLAS_GEMMTYPE Value = MLAS_EXTRACT_FLOAT(Accumulators[Row][0]); + + if (!ZeroMode) { + Value += *c; + } + + *c = Value; + } +}; + diff --git a/3rdparty/mlas/lib/s390x/SgemmKernelZVECTOR.h b/3rdparty/mlas/lib/s390x/SgemmKernelZVECTOR.h new file mode 100644 index 0000000000..83c66eb626 --- /dev/null +++ b/3rdparty/mlas/lib/s390x/SgemmKernelZVECTOR.h @@ -0,0 +1,139 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + SgemmKernelZVECTOR.h + +Abstract: + + This module implements the kernels for the single precision matrix/matrix + multiply operation (SGEMM). + +--*/ + +#include "FgemmKernelZVECTOR.h" + +template +MLAS_FORCEINLINE +size_t +MlasSgemmProcessCount( + const float* A, + const float* B, + float* C, + size_t CountK, + size_t CountN, + size_t lda, + size_t ldc, + MLAS_FLOAT32X4 AlphaBroadcast, + bool ZeroMode + ) +{ + do { + + const float* a = A; + size_t k = CountK; + + MLAS_FLOAT32X4 Accumulators[RowCount][4]; + MLAS_FLOAT32X4 AElements[RowCount]; + MLAS_FLOAT32X4 ABroadcast[RowCount]; + + // + // Clear the block accumulators. + // + + MlasLoopUnroll()(Accumulators); + + // + // Compute the output block. + // + + while (k >= 4) { + + MlasLoopUnroll()(AElements, a, lda); + + MlasLoopUnroll>()(AElements, ABroadcast); + MlasFgemmComputeBlock(Accumulators, ABroadcast, B); + + MlasLoopUnroll>()(AElements, ABroadcast); + MlasFgemmComputeBlock(Accumulators, ABroadcast, B + 16); + + MlasLoopUnroll>()(AElements, ABroadcast); + MlasFgemmComputeBlock(Accumulators, ABroadcast, B + 32); + + MlasLoopUnroll>()(AElements, ABroadcast); + MlasFgemmComputeBlock(Accumulators, ABroadcast, B + 48); + + a += 4; + B += 16 * 4; + k -= 4; + } + + while (k > 0) { + + MlasLoopUnroll()(ABroadcast, a, lda); + MlasFgemmComputeBlock(Accumulators, ABroadcast, B); + + a += 1; + B += 16; + k -= 1; + } + + if (CountN >= 16) { + + // + // Store the entire output block. + // + + MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); + + } else { + + // + // Store the partial output block. + // + + if (CountN >= 12) { + MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); + } else if (CountN >= 8) { + MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); + } else if (CountN >= 4) { + MlasLoopUnroll>()(Accumulators, C, ldc, AlphaBroadcast, ZeroMode); + } + + // + // Store the remaining unaligned columns. + // + + C += (CountN & ~3); + CountN &= 3; + + if (CountN > 0) { + + MlasLoopUnroll()(Accumulators, AlphaBroadcast); + + MlasLoopUnroll>()(Accumulators, C, ldc, ZeroMode); + + if (CountN >= 2) { + MlasLoopUnroll>()(Accumulators, C, ldc, ZeroMode); + } + + if (CountN >= 3) { + MlasLoopUnroll>()(Accumulators, C, ldc, ZeroMode); + } + } + + break; + } + + C += 16; + CountN -= 16; + + } while (CountN > 0); + + return RowCount; +} +