saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-07 17:15:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Moving MLAS threaded QGEMM packing buffer from stack to heap (#14002)

Browse source 
### 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 <fuchen@microsoft.com>
This commit is contained in:
Chen Fu 2022-12-19 09:39:19 -08:00 committed by GitHub
parent fba09faf5b
commit 28e2b1790f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 115 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
onnxruntime/core/mlas/lib/mlasi.h
View file

@ -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<uint8_t, decltype(&_aligned_free)> ThreadedBufHolder;
#else
extern thread_local std::unique_ptr<uint8_t, decltype(&free)> 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<uint8_t*>(_aligned_malloc(size, ThreadedBufAlignment)));
#elif (__STDC_VERSION__ >= 201112L) && !defined(__APPLE__)
ThreadedBufHolder.reset(
reinterpret_cast<uint8_t*>(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<uint8_t*>(ptr));
#endif
ThreadedBufSize = size;
}
}

9
onnxruntime/core/mlas/lib/platform.cpp
View file

@ -478,4 +478,11 @@ MlasPlatformU8S8Overflow(
return p.GemmU8U8Dispatch != p.GemmU8S8Dispatch;
}
#endif
#endif
thread_local size_t ThreadedBufSize = 0;
#ifdef _MSC_VER
thread_local std::unique_ptr<uint8_t, decltype(&_aligned_free)> ThreadedBufHolder(nullptr, &_aligned_free);
#else
thread_local std::unique_ptr<uint8_t, decltype(&free)> ThreadedBufHolder(nullptr, &free);
#endif

66
onnxruntime/core/mlas/lib/qgemm.h
View file

@ -34,6 +34,7 @@ Abstract:
#include <sstream>
#include <string>
#include <cstdlib>
//
// Define the default striding parameters used for the quantized integer
@ -222,6 +223,28 @@ MlasGemmQuantScaleSumBuffer(
return MlasGemmQuantScaleSumBuffer(SumBuffer, SumBuffer, N, Scale);
}
template<typename KernelType>
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<typename KernelType>
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<KernelType>();
uint8_t* p = ThreadedBufHolder.get();
typename KernelType::PackedAType* PanelA =
reinterpret_cast<typename KernelType::PackedAType*>(p);
p += packASize;
typename KernelType::PackedBType* PanelB =
reinterpret_cast<typename KernelType::PackedBType*>(p);
p += packBSize;
int32_t* RowSumBuffer = reinterpret_cast<int32_t*>(p);
p += rowSumSize;
int32_t* ColumnSumBuffer = reinterpret_cast<int32_t*>(p);
p += colSumSize;
int32_t* ZeroPointBBuffer = reinterpret_cast<int32_t*>(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<KernelType>();
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<typename KernelType::PackedAType*>(p);
p += packASize;
int32_t* RowSumBuffer = reinterpret_cast<int32_t*>(p);
p += rowSumSize;
int32_t* ColumnSumBuffer = reinterpret_cast<int32_t*>(p);
p += colSumSize;
int32_t* ZeroPointBBuffer = reinterpret_cast<int32_t*>(p);
const size_t K = Shape->K;

1
onnxruntime/core/mlas/lib/qgemm_kernel_sse.cpp
View file

@ -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;

1
onnxruntime/core/mlas/lib/qgemm_kernel_wasmsimd.cpp
View file

@ -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;