diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index 0f365159ac..56d417b7e6 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -2062,3 +2062,51 @@ MlasReadTimeStampCounter(void) #endif #endif } + +// +// Aligned buffer for GEMM packing, etc. +// + + +constexpr size_t ThreadedBufAlignment = 64; +extern thread_local size_t ThreadedBufSize; +#ifdef _MSC_VER +extern thread_local std::unique_ptr ThreadedBufHolder; +#else +extern thread_local std::unique_ptr ThreadedBufHolder; +#endif + +MLAS_FORCEINLINE +constexpr size_t +UpAlignSize(size_t size) +{ + size = (size + ThreadedBufAlignment - 1) / ThreadedBufAlignment; + return size * ThreadedBufAlignment; +} + + +MLAS_FORCEINLINE +void +MlasThreadedBufAlloc(size_t size) +{ + if (size > ThreadedBufSize) { +#ifdef _MSC_VER + ThreadedBufHolder.reset( + reinterpret_cast(_aligned_malloc(size, ThreadedBufAlignment))); +#elif (__STDC_VERSION__ >= 201112L) && !defined(__APPLE__) + ThreadedBufHolder.reset( + reinterpret_cast(aligned_alloc(ThreadedBufAlignment, size))); +#else + // aligned_alloc unavailable macos 10.14 or earlier + void* ptr; + int err = posix_memalign(&ptr, ThreadedBufAlignment, size); + if (err != 0) { + ptr = nullptr; + } + ThreadedBufHolder.reset(reinterpret_cast(ptr)); +#endif + + ThreadedBufSize = size; + } +} + diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp index 29e3b9ceb4..33e1489a10 100644 --- a/onnxruntime/core/mlas/lib/platform.cpp +++ b/onnxruntime/core/mlas/lib/platform.cpp @@ -478,4 +478,11 @@ MlasPlatformU8S8Overflow( return p.GemmU8U8Dispatch != p.GemmU8S8Dispatch; } -#endif \ No newline at end of file +#endif + +thread_local size_t ThreadedBufSize = 0; +#ifdef _MSC_VER +thread_local std::unique_ptr ThreadedBufHolder(nullptr, &_aligned_free); +#else +thread_local std::unique_ptr ThreadedBufHolder(nullptr, &free); +#endif diff --git a/onnxruntime/core/mlas/lib/qgemm.h b/onnxruntime/core/mlas/lib/qgemm.h index 8a1856eb71..c34a916941 100644 --- a/onnxruntime/core/mlas/lib/qgemm.h +++ b/onnxruntime/core/mlas/lib/qgemm.h @@ -34,6 +34,7 @@ Abstract: #include #include +#include // // Define the default striding parameters used for the quantized integer @@ -222,6 +223,28 @@ MlasGemmQuantScaleSumBuffer( return MlasGemmQuantScaleSumBuffer(SumBuffer, SumBuffer, N, Scale); } +template +MLAS_FORCEINLINE +void +MlasGemmQuantThreadInit() +{ + constexpr MLAS_GEMM_QUANT_STRIDES Strides = KernelType::Strides; + constexpr size_t packASize = + UpAlignSize(Strides.M * Strides.K * sizeof(typename KernelType::PackedAType)); + constexpr size_t packBSize = + UpAlignSize(Strides.N * Strides.K * sizeof(typename KernelType::PackedBType)); + constexpr size_t rowSumSize = UpAlignSize(Strides.M * sizeof(int32_t)); + constexpr size_t colSumSize = UpAlignSize(Strides.N * sizeof(int32_t)); + constexpr size_t zpbSize = UpAlignSize(Strides.N * sizeof(int32_t)); + + constexpr MLAS_GEMM_QUANT_STRIDES PackedStrides = KernelType::PackedStrides; + constexpr size_t packedASize = + UpAlignSize(PackedStrides.M * PackedStrides.K * sizeof(typename KernelType::PackedAType)); + + constexpr size_t bufsize = std::max(packASize + packBSize, packedASize) + rowSumSize + colSumSize + zpbSize; + + MlasThreadedBufAlloc(bufsize); +} template void @@ -261,13 +284,28 @@ Return Value: --*/ { constexpr MLAS_GEMM_QUANT_STRIDES Strides = KernelType::Strides; + constexpr size_t packASize = + UpAlignSize(Strides.M * Strides.K * sizeof(typename KernelType::PackedAType)); + constexpr size_t packBSize = + UpAlignSize(Strides.N * Strides.K * sizeof(typename KernelType::PackedBType)); + constexpr size_t rowSumSize = UpAlignSize(Strides.M * sizeof(int32_t)); + constexpr size_t colSumSize = UpAlignSize(Strides.N * sizeof(int32_t)); - MLAS_DECLSPEC_ALIGN(typename KernelType::PackedAType PanelA[Strides.M * Strides.K], 64); - MLAS_DECLSPEC_ALIGN(typename KernelType::PackedBType PanelB[Strides.N * Strides.K], 64); + MlasGemmQuantThreadInit(); + + uint8_t* p = ThreadedBufHolder.get(); + typename KernelType::PackedAType* PanelA = + reinterpret_cast(p); + p += packASize; + typename KernelType::PackedBType* PanelB = + reinterpret_cast(p); + p += packBSize; + int32_t* RowSumBuffer = reinterpret_cast(p); + p += rowSumSize; + int32_t* ColumnSumBuffer = reinterpret_cast(p); + p += colSumSize; + int32_t* ZeroPointBBuffer = reinterpret_cast(p); - MLAS_DECLSPEC_ALIGN(int32_t RowSumBuffer[Strides.M], 64); - MLAS_DECLSPEC_ALIGN(int32_t ColumnSumBuffer[Strides.N], 64); - MLAS_DECLSPEC_ALIGN(int32_t ZeroPointBBuffer[Strides.N], 64); const size_t K = Shape->K; @@ -497,12 +535,22 @@ Return Value: --*/ { constexpr MLAS_GEMM_QUANT_STRIDES Strides = KernelType::PackedStrides; + constexpr size_t packASize = + UpAlignSize(Strides.M * Strides.K * sizeof(typename KernelType::PackedAType)); + constexpr size_t rowSumSize = UpAlignSize(Strides.M * sizeof(int32_t)); + constexpr size_t colSumSize = UpAlignSize(Strides.N * sizeof(int32_t)); - MLAS_DECLSPEC_ALIGN(typename KernelType::PackedAType PanelA[Strides.M * Strides.K], 64); + MlasGemmQuantThreadInit(); - MLAS_DECLSPEC_ALIGN(int32_t RowSumBuffer[Strides.M], 64); - MLAS_DECLSPEC_ALIGN(int32_t ColumnSumBuffer[Strides.N], 64); - MLAS_DECLSPEC_ALIGN(int32_t ZeroPointBBuffer[Strides.N], 64); + uint8_t* p = ThreadedBufHolder.get(); + typename KernelType::PackedAType* PanelA = + reinterpret_cast(p); + p += packASize; + int32_t* RowSumBuffer = reinterpret_cast(p); + p += rowSumSize; + int32_t* ColumnSumBuffer = reinterpret_cast(p); + p += colSumSize; + int32_t* ZeroPointBBuffer = reinterpret_cast(p); const size_t K = Shape->K; diff --git a/onnxruntime/core/mlas/lib/qgemm_kernel_sse.cpp b/onnxruntime/core/mlas/lib/qgemm_kernel_sse.cpp index 06f936f50c..65c9e2b5ae 100644 --- a/onnxruntime/core/mlas/lib/qgemm_kernel_sse.cpp +++ b/onnxruntime/core/mlas/lib/qgemm_kernel_sse.cpp @@ -26,6 +26,7 @@ struct MLAS_GEMM_U8X8_KERNEL_SSE static constexpr size_t PackedK = 2; static constexpr MLAS_GEMM_QUANT_STRIDES Strides{ 12, 128, 128 }; + static constexpr MLAS_GEMM_QUANT_STRIDES PackedStrides{0, 0, 0}; }; constexpr size_t MLAS_GEMM_U8X8_KERNEL_SSE::PackedK; diff --git a/onnxruntime/core/mlas/lib/qgemm_kernel_wasmsimd.cpp b/onnxruntime/core/mlas/lib/qgemm_kernel_wasmsimd.cpp index f85fc929f7..1f33d77adf 100644 --- a/onnxruntime/core/mlas/lib/qgemm_kernel_wasmsimd.cpp +++ b/onnxruntime/core/mlas/lib/qgemm_kernel_wasmsimd.cpp @@ -36,6 +36,7 @@ struct MLAS_GEMM_U8X8_KERNEL_WASMSIMD static constexpr size_t PackedK = 2; static constexpr MLAS_GEMM_QUANT_STRIDES Strides{ 12, 128, 128 }; + static constexpr MLAS_GEMM_QUANT_STRIDES PackedStrides{0, 0, 0}; }; constexpr size_t MLAS_GEMM_U8X8_KERNEL_WASMSIMD::PackedK;