From 28e2b1790fbf85f53c1b4386e7bbc00128ac1a40 Mon Sep 17 00:00:00 2001 From: Chen Fu <1316708+chenfucn@users.noreply.github.com> Date: Mon, 19 Dec 2022 09:39:19 -0800 Subject: [PATCH] Moving MLAS threaded QGEMM packing buffer from stack to heap (#14002) ### Description MLAS QGEMM kernel need memory buffer for packing of source tensors. This change moves these buffers from stack to heap ### Motivation and Context MLAS QGEMM kernels have packing buffers on the stack since the beginning of time. Emerging hardware demands larger and larger buffers, causing potential stack overflow problems down the road. This change moves these buffers from stack to the heap. This change also introduces a thread initializer per kernel. For instance, in the new AMX instruction set (support coming), we need to initialize the tile registers per thread. This requirement can be easily satisfied by tapping into this change. Co-authored-by: Chen Fu --- onnxruntime/core/mlas/lib/mlasi.h | 48 ++++++++++++++ onnxruntime/core/mlas/lib/platform.cpp | 9 ++- onnxruntime/core/mlas/lib/qgemm.h | 66 ++++++++++++++++--- .../core/mlas/lib/qgemm_kernel_sse.cpp | 1 + .../core/mlas/lib/qgemm_kernel_wasmsimd.cpp | 1 + 5 files changed, 115 insertions(+), 10 deletions(-) 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;