From 90642e7eacdf018cf26b633ec0c7ad1a1072e41a Mon Sep 17 00:00:00 2001 From: Tracy Sharpe <42477615+tracysh@users.noreply.github.com> Date: Wed, 17 Mar 2021 09:22:55 -0700 Subject: [PATCH] MLAS: more code cleanup (#7036) Change int32_t->ptrdiff_t when interacting with the threadpool. Migrate more code from MlasMaskMoveAvx->MlasMaskMoveTableAvx. Update more code to use FUNCTION_ENTRY macro. --- onnxruntime/core/mlas/inc/mlas.h | 3 +- .../core/mlas/lib/amd64/DgemmKernelCommon.inc | 2 -- .../mlas/lib/amd64/FgemmKernelAvxCommon.inc | 14 +++----- .../mlas/lib/amd64/FgemmKernelFma3Common.inc | 8 ++--- .../core/mlas/lib/amd64/SgemmKernelCommon.inc | 2 -- onnxruntime/core/mlas/lib/compute.cpp | 10 +++--- onnxruntime/core/mlas/lib/convolve.cpp | 34 +++++++------------ onnxruntime/core/mlas/lib/dgemm.cpp | 24 ++++++------- onnxruntime/core/mlas/lib/mlasi.h | 16 ++++----- onnxruntime/core/mlas/lib/platform.cpp | 6 ---- onnxruntime/core/mlas/lib/qgemm.cpp | 20 +++++------ onnxruntime/core/mlas/lib/sgemm.cpp | 24 ++++++------- onnxruntime/core/mlas/lib/snchwc.cpp | 16 ++++----- onnxruntime/core/mlas/lib/threading.cpp | 6 ++-- .../core/mlas/lib/x86_64/DgemmKernelCommon.h | 2 -- .../mlas/lib/x86_64/FgemmKernelAvxCommon.h | 12 ++----- .../core/mlas/lib/x86_64/FgemmKernelCommon.h | 9 +++-- .../mlas/lib/x86_64/FgemmKernelFma3Common.h | 6 ++-- .../core/mlas/lib/x86_64/SgemmKernelCommon.h | 2 -- .../core/mlas/lib/x86_64/SgemmKernelM1Avx.S | 3 +- .../lib/x86_64/SgemmKernelM1TransposeBAvx.S | 3 +- .../lib/x86_64/SgemmTransposePackB16x4Avx.S | 3 +- .../lib/x86_64/SgemmTransposePackB16x4Sse2.S | 3 +- 23 files changed, 95 insertions(+), 133 deletions(-) diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index aae34514bb..9a23bfc561 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -17,6 +17,7 @@ Abstract: #pragma once +#include #include #include @@ -362,7 +363,7 @@ struct MLAS_CONV_PARAMETERS { size_t OutputSize; size_t K; MLAS_CONV_ALGORITHM Algorithm; - int32_t ThreadCount; + ptrdiff_t ThreadCount; union { struct { CBLAS_TRANSPOSE TransB; diff --git a/onnxruntime/core/mlas/lib/amd64/DgemmKernelCommon.inc b/onnxruntime/core/mlas/lib/amd64/DgemmKernelCommon.inc index d1c650db87..52ee715643 100644 --- a/onnxruntime/core/mlas/lib/amd64/DgemmKernelCommon.inc +++ b/onnxruntime/core/mlas/lib/amd64/DgemmKernelCommon.inc @@ -23,7 +23,6 @@ FgemmElementShift EQU 3 FgemmElementSize EQU (1 SHL FgemmElementShift) FgemmElementPtr EQU QWORD PTR FgemmElementBcst EQU QWORD BCST -FgemmMaskMoveVector EQU MlasMaskMoveAvx64 ; ; Define the typed instructions for double precision. @@ -41,7 +40,6 @@ vmovapf EQU vmovapd vmovsf EQU vmovsd vmovupf EQU vmovupd vmulpf EQU vmulpd -vpcmpgtf EQU vpcmpgtq vxorpf EQU vxorpd INCLUDE FgemmKernelCommon.inc diff --git a/onnxruntime/core/mlas/lib/amd64/FgemmKernelAvxCommon.inc b/onnxruntime/core/mlas/lib/amd64/FgemmKernelAvxCommon.inc index dac48c23db..5b65c37bb9 100644 --- a/onnxruntime/core/mlas/lib/amd64/FgemmKernelAvxCommon.inc +++ b/onnxruntime/core/mlas/lib/amd64/FgemmKernelAvxCommon.inc @@ -17,7 +17,7 @@ ; ;-- - EXTERN FgemmMaskMoveVector:NEAR + EXTERN MlasMaskMoveTableAvx:NEAR ; ; Macro Description: @@ -319,15 +319,9 @@ ENDIF add rbp,FgemmYmmElementCount ; correct for over-subtract above OutputMasked1xNBlock: - mov FgemmKernelFrame.CountN[rsp],rbp -IF FgemmElementSize EQ 4 - vbroadcastsf xmm0,FgemmElementPtr FgemmKernelFrame.CountN[rsp] -ELSE - vmovddup xmm0,FgemmElementPtr FgemmKernelFrame.CountN[rsp] -ENDIF - vpcmpgtf xmm1,xmm0,XMMWORD PTR [FgemmMaskMoveVector+16] - vpcmpgtf xmm0,xmm0,XMMWORD PTR [FgemmMaskMoveVector] - vinsertf128 ymm0,ymm0,xmm1,1 + neg rbp + lea rcx,MlasMaskMoveTableAvx+8*4 + vmovdqu ymm0,YMMWORD PTR [rcx+rbp*FgemmElementSize] test r15b,r15b ; ZeroMode? jnz StoreMasked1xNBlock EmitIfCountGE RowCount, 1, diff --git a/onnxruntime/core/mlas/lib/amd64/FgemmKernelFma3Common.inc b/onnxruntime/core/mlas/lib/amd64/FgemmKernelFma3Common.inc index a7f3ad9fc2..298b9dca81 100644 --- a/onnxruntime/core/mlas/lib/amd64/FgemmKernelFma3Common.inc +++ b/onnxruntime/core/mlas/lib/amd64/FgemmKernelFma3Common.inc @@ -17,7 +17,7 @@ ; ;-- - EXTERN FgemmMaskMoveVector:NEAR + EXTERN MlasMaskMoveTableAvx:NEAR ; ; Macro Description: @@ -359,9 +359,9 @@ ENDIF add rbp,FgemmYmmElementCount ; correct for over-subtract above OutputMasked1xNBlock: - mov FgemmKernelFrame.CountN[rsp],rbp - vbroadcastsf ymm0,FgemmElementPtr FgemmKernelFrame.CountN[rsp] - vpcmpgtf ymm0,ymm0,YMMWORD PTR [FgemmMaskMoveVector] + neg rbp + lea rcx,MlasMaskMoveTableAvx+8*4 + vmovdqu ymm0,YMMWORD PTR [rcx+rbp*FgemmElementSize] test r15b,r15b ; ZeroMode? jnz MultiplyAlphaMasked1xNBlock EmitIfCountGE RowCount, 1, diff --git a/onnxruntime/core/mlas/lib/amd64/SgemmKernelCommon.inc b/onnxruntime/core/mlas/lib/amd64/SgemmKernelCommon.inc index be474e8040..f7bc7870bb 100644 --- a/onnxruntime/core/mlas/lib/amd64/SgemmKernelCommon.inc +++ b/onnxruntime/core/mlas/lib/amd64/SgemmKernelCommon.inc @@ -23,7 +23,6 @@ FgemmElementShift EQU 2 FgemmElementSize EQU (1 SHL FgemmElementShift) FgemmElementPtr EQU DWORD PTR FgemmElementBcst EQU DWORD BCST -FgemmMaskMoveVector EQU MlasMaskMoveAvx ; ; Define the typed instructions for single precision. @@ -41,7 +40,6 @@ vmovapf EQU vmovaps vmovsf EQU vmovss vmovupf EQU vmovups vmulpf EQU vmulps -vpcmpgtf EQU vpcmpgtd vxorpf EQU vxorps INCLUDE FgemmKernelCommon.inc diff --git a/onnxruntime/core/mlas/lib/compute.cpp b/onnxruntime/core/mlas/lib/compute.cpp index 04b9cc3a1b..af5ee95ac8 100644 --- a/onnxruntime/core/mlas/lib/compute.cpp +++ b/onnxruntime/core/mlas/lib/compute.cpp @@ -69,7 +69,7 @@ MLAS_INTERNAL_DATA const float MlasMinimumF32Value = std::numeric_limits: // struct MLAS_SOFTMAX_WORK_BLOCK { - int32_t ThreadCountN; + ptrdiff_t ThreadCountN; bool LogSoftmax; const float* Input; float* Output; @@ -801,7 +801,7 @@ Return Value: void MlasComputeSoftmaxThreaded( void* Context, - int32_t Index + ptrdiff_t Index ) /*++ @@ -971,10 +971,10 @@ Return Value: // another thread. // - int32_t ThreadCountN = MlasGetMaximumThreadCount(ThreadPool); + ptrdiff_t ThreadCountN = MlasGetMaximumThreadCount(ThreadPool); if (size_t(ThreadCountN) > N) { - ThreadCountN = int32_t(N); + ThreadCountN = ptrdiff_t(N); } constexpr size_t MinimumElementsPerThread = 16384; @@ -982,7 +982,7 @@ Return Value: size_t BlockCount = ((N * D) / MinimumElementsPerThread) + 1; if (size_t(ThreadCountN) > BlockCount) { - ThreadCountN = int32_t(BlockCount); + ThreadCountN = ptrdiff_t(BlockCount); } WorkBlock.ThreadCountN = ThreadCountN; diff --git a/onnxruntime/core/mlas/lib/convolve.cpp b/onnxruntime/core/mlas/lib/convolve.cpp index 72141607cf..86f174ac1c 100644 --- a/onnxruntime/core/mlas/lib/convolve.cpp +++ b/onnxruntime/core/mlas/lib/convolve.cpp @@ -39,7 +39,7 @@ struct MLAS_CONV_WORK_BLOCK { size_t StartN; size_t CountN; } Segments[MLAS_MAXIMUM_THREAD_COUNT]; - int32_t TargetThreadCount; + ptrdiff_t TargetThreadCount; }; void @@ -609,7 +609,7 @@ Return Value: void MlasConvOperationThreaded( void* Context, - int32_t Index + ptrdiff_t Index ) /*++ @@ -645,7 +645,7 @@ Return Value: void MlasConvGemmDirectThreaded( void* Context, - int32_t Index + ptrdiff_t Index ) /*++ @@ -677,21 +677,13 @@ Return Value: const size_t GroupCount = Parameters->GroupCount; const size_t BatchGroupCount = Parameters->BatchCount * GroupCount; - const size_t TargetThreadCount = WorkBlock->TargetThreadCount; - - const size_t BatchGroupCountPerThread = BatchGroupCount / TargetThreadCount; - const size_t BatchGroupCountExtra = BatchGroupCount % TargetThreadCount; - size_t BatchGroupStart; - size_t BatchGroupEnd; + size_t BatchGroupRemaining; - if (uint32_t(Index) < BatchGroupCountExtra) { - BatchGroupStart = (BatchGroupCountPerThread + 1) * Index; - BatchGroupEnd = BatchGroupStart + BatchGroupCountPerThread + 1; - } else { - BatchGroupStart = BatchGroupCountPerThread * Index + BatchGroupCountExtra; - BatchGroupEnd = BatchGroupStart + BatchGroupCountPerThread; - } + MlasPartitionWork(Index, WorkBlock->TargetThreadCount, BatchGroupCount, + &BatchGroupStart, &BatchGroupRemaining); + + size_t BatchGroupEnd = BatchGroupStart + BatchGroupRemaining; // // Iterate over the batch and groups allocated to this thread. @@ -889,10 +881,10 @@ Return Value: const size_t BatchGroupCount = BatchCount * GroupCount; - int32_t TargetThreadCount = MlasGetMaximumThreadCount(ThreadPool); + ptrdiff_t TargetThreadCount = MlasGetMaximumThreadCount(ThreadPool); if (size_t(TargetThreadCount) >= BatchGroupCount) { - TargetThreadCount = int32_t(BatchGroupCount); + TargetThreadCount = ptrdiff_t(BatchGroupCount); } MLAS_CONV_WORK_BLOCK WorkBlock; @@ -1213,16 +1205,16 @@ Return Value: // threaded path. // - int32_t TargetThreadCount; + ptrdiff_t TargetThreadCount; double Complexity = double(FilterCount) * double(OutputSize) * double(K); if (Complexity < double(MLAS_SGEMM_THREAD_COMPLEXITY * MLAS_MAXIMUM_THREAD_COUNT)) { - TargetThreadCount = int32_t(Complexity / double(MLAS_SGEMM_THREAD_COMPLEXITY)) + 1; + TargetThreadCount = ptrdiff_t(Complexity / double(MLAS_SGEMM_THREAD_COMPLEXITY)) + 1; } else { TargetThreadCount = MLAS_MAXIMUM_THREAD_COUNT; } - int32_t MaximumThreadCount = MlasGetMaximumThreadCount(ThreadPool); + ptrdiff_t MaximumThreadCount = MlasGetMaximumThreadCount(ThreadPool); if (TargetThreadCount >= MaximumThreadCount) { TargetThreadCount = MaximumThreadCount; diff --git a/onnxruntime/core/mlas/lib/dgemm.cpp b/onnxruntime/core/mlas/lib/dgemm.cpp index 7d327aa483..666856a6b6 100644 --- a/onnxruntime/core/mlas/lib/dgemm.cpp +++ b/onnxruntime/core/mlas/lib/dgemm.cpp @@ -32,8 +32,8 @@ Abstract: // struct MLAS_DGEMM_WORK_BLOCK { - int32_t ThreadCountM; - int32_t ThreadCountN; + ptrdiff_t ThreadCountM; + ptrdiff_t ThreadCountN; CBLAS_TRANSPOSE TransA; CBLAS_TRANSPOSE TransB; size_t M; @@ -751,7 +751,7 @@ Return Value: void MlasDgemmThreaded( void* Context, - int32_t ThreadId + ptrdiff_t ThreadId ) /*++ @@ -774,11 +774,11 @@ Return Value: { const auto* WorkBlock = (MLAS_DGEMM_WORK_BLOCK*)Context; - const int32_t ThreadCountM = WorkBlock->ThreadCountM; - const int32_t ThreadCountN = WorkBlock->ThreadCountN; + const ptrdiff_t ThreadCountM = WorkBlock->ThreadCountM; + const ptrdiff_t ThreadCountN = WorkBlock->ThreadCountN; - const int32_t ThreadIdM = ThreadId / ThreadCountN; - const int32_t ThreadIdN = ThreadId % ThreadCountN; + const ptrdiff_t ThreadIdM = ThreadId / ThreadCountN; + const ptrdiff_t ThreadIdN = ThreadId % ThreadCountN; // // Partition the operation along the M dimension. @@ -864,15 +864,15 @@ Return Value: const double Complexity = double(M) * double(N) * double(K); - int32_t TargetThreadCount; + ptrdiff_t TargetThreadCount; if (Complexity < double(MLAS_DGEMM_THREAD_COMPLEXITY * MlasPlatform.MaximumThreadCount)) { - TargetThreadCount = int32_t(Complexity / double(MLAS_DGEMM_THREAD_COMPLEXITY)) + 1; + TargetThreadCount = ptrdiff_t(Complexity / double(MLAS_DGEMM_THREAD_COMPLEXITY)) + 1; } else { TargetThreadCount = MlasPlatform.MaximumThreadCount; } - int32_t MaximumThreadCount = MlasGetMaximumThreadCount(ThreadPool); + ptrdiff_t MaximumThreadCount = MlasGetMaximumThreadCount(ThreadPool); if (TargetThreadCount >= MaximumThreadCount) { TargetThreadCount = MaximumThreadCount; @@ -891,7 +891,7 @@ Return Value: MLAS_DGEMM_STRIDEN_THREAD_ALIGN; if (size_t(TargetThreadCount) > BlockedN) { - TargetThreadCount = int32_t(BlockedN); + TargetThreadCount = ptrdiff_t(BlockedN); } WorkBlock->ThreadCountM = 1; @@ -900,7 +900,7 @@ Return Value: } else { if (size_t(TargetThreadCount) > M) { - TargetThreadCount = int32_t(M); + TargetThreadCount = ptrdiff_t(M); } WorkBlock->ThreadCountM = TargetThreadCount; diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index ad56b2ec22..305a0ddded 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -733,9 +733,9 @@ struct MLAS_PLATFORM { MLAS_QUANTIZE_LINEAR_U8_KERNEL* QuantizeLinearU8Kernel; uint32_t NchwcBlockSize; uint32_t PreferredBufferAlignment; - uint32_t MaximumThreadCount; + int32_t MaximumThreadCount; #else - static constexpr uint32_t MaximumThreadCount = MLAS_MAXIMUM_THREAD_COUNT; + static constexpr int32_t MaximumThreadCount = MLAS_MAXIMUM_THREAD_COUNT; #endif #if defined(MLAS_TARGET_ARM64) @@ -753,19 +753,19 @@ typedef void (MLAS_THREADED_ROUTINE)( void* Context, - int32_t Index + ptrdiff_t Index ); void MlasExecuteThreaded( MLAS_THREADED_ROUTINE* ThreadedRoutine, void* Context, - int32_t Iterations, + ptrdiff_t Iterations, MLAS_THREADPOOL* ThreadPool ); inline -int32_t +ptrdiff_t MlasGetMaximumThreadCount( MLAS_THREADPOOL* ThreadPool ) @@ -786,8 +786,8 @@ MlasGetMaximumThreadCount( inline void MlasPartitionWork( - int32_t ThreadId, - int32_t ThreadCount, + ptrdiff_t ThreadId, + ptrdiff_t ThreadCount, size_t TotalWork, size_t* WorkIndex, size_t* WorkRemaining @@ -796,7 +796,7 @@ MlasPartitionWork( const size_t WorkPerThread = TotalWork / ThreadCount; const size_t WorkPerThreadExtra = TotalWork % ThreadCount; - if (uint32_t(ThreadId) < WorkPerThreadExtra) { + if (size_t(ThreadId) < WorkPerThreadExtra) { *WorkIndex = (WorkPerThread + 1) * ThreadId; *WorkRemaining = WorkPerThread + 1; } else { diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp index 43dab620a2..7e66e0b2f7 100644 --- a/onnxruntime/core/mlas/lib/platform.cpp +++ b/onnxruntime/core/mlas/lib/platform.cpp @@ -41,12 +41,6 @@ MLAS_PLATFORM MlasPlatform; MLAS_INTERNAL_DATA MLAS_DECLSPEC_ALIGN(const uint32_t MlasMaskMoveAvx[8], 32) = { 0, 1, 2, 3, 4, 5, 6, 7 }; -// -// Stores a vector to build a conditional load/store mask for vmaskmovpd. -// - -MLAS_INTERNAL_DATA MLAS_DECLSPEC_ALIGN(const uint64_t MlasMaskMoveAvx64[4], 32) = { 0, 1, 2, 3 }; - // // Stores a table of AVX vmaskmovps/vmaskmovpd load/store masks. // diff --git a/onnxruntime/core/mlas/lib/qgemm.cpp b/onnxruntime/core/mlas/lib/qgemm.cpp index f4a1868db9..a6c83e5951 100644 --- a/onnxruntime/core/mlas/lib/qgemm.cpp +++ b/onnxruntime/core/mlas/lib/qgemm.cpp @@ -85,8 +85,8 @@ MlasGemmU8X8GetDispatch( // struct MLAS_GEMM_U8X8_WORK_BLOCK { - int32_t ThreadCountM; - int32_t ThreadCountN; + ptrdiff_t ThreadCountM; + ptrdiff_t ThreadCountN; const MLAS_GEMM_U8X8_PARAMETERS* Parameters; }; @@ -2660,7 +2660,7 @@ const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8X8DispatchDefault = { void MlasGemmU8X8Threaded( void* Context, - int32_t ThreadId + ptrdiff_t ThreadId ) /*++ @@ -2684,8 +2684,8 @@ Return Value: const auto* WorkBlock = (MLAS_GEMM_U8X8_WORK_BLOCK*)Context; const auto* Parameters = WorkBlock->Parameters; - const int32_t ThreadIdM = ThreadId / WorkBlock->ThreadCountN; - const int32_t ThreadIdN = ThreadId % WorkBlock->ThreadCountN; + const ptrdiff_t ThreadIdM = ThreadId / WorkBlock->ThreadCountN; + const ptrdiff_t ThreadIdN = ThreadId % WorkBlock->ThreadCountN; // // Partition the operation along the M dimension. @@ -2771,15 +2771,15 @@ Return Value: const double Complexity = double(M) * double(N) * double(K); - int32_t TargetThreadCount; + ptrdiff_t TargetThreadCount; if (Complexity < double(MLAS_QGEMM_THREAD_COMPLEXITY * MlasPlatform.MaximumThreadCount)) { - TargetThreadCount = int32_t(Complexity / double(MLAS_QGEMM_THREAD_COMPLEXITY)) + 1; + TargetThreadCount = ptrdiff_t(Complexity / double(MLAS_QGEMM_THREAD_COMPLEXITY)) + 1; } else { TargetThreadCount = MlasPlatform.MaximumThreadCount; } - int32_t MaximumThreadCount = MlasGetMaximumThreadCount(ThreadPool); + ptrdiff_t MaximumThreadCount = MlasGetMaximumThreadCount(ThreadPool); if (TargetThreadCount >= MaximumThreadCount) { TargetThreadCount = MaximumThreadCount; @@ -2802,7 +2802,7 @@ Return Value: MLAS_QGEMM_STRIDEN_THREAD_ALIGN; if (size_t(TargetThreadCount) > BlockedN) { - TargetThreadCount = int32_t(BlockedN); + TargetThreadCount = ptrdiff_t(BlockedN); } WorkBlock.ThreadCountM = 1; @@ -2811,7 +2811,7 @@ Return Value: } else { if (size_t(TargetThreadCount) > M) { - TargetThreadCount = int32_t(M); + TargetThreadCount = ptrdiff_t(M); } WorkBlock.ThreadCountM = TargetThreadCount; diff --git a/onnxruntime/core/mlas/lib/sgemm.cpp b/onnxruntime/core/mlas/lib/sgemm.cpp index 996a18ff9b..51abd18c27 100644 --- a/onnxruntime/core/mlas/lib/sgemm.cpp +++ b/onnxruntime/core/mlas/lib/sgemm.cpp @@ -32,8 +32,8 @@ Abstract: // struct MLAS_SGEMM_WORK_BLOCK { - int32_t ThreadCountM; - int32_t ThreadCountN; + ptrdiff_t ThreadCountM; + ptrdiff_t ThreadCountN; CBLAS_TRANSPOSE TransA; CBLAS_TRANSPOSE TransB; size_t M; @@ -1206,7 +1206,7 @@ Return Value: void MlasSgemmThreaded( void* Context, - int32_t ThreadId + ptrdiff_t ThreadId ) /*++ @@ -1229,11 +1229,11 @@ Return Value: { const auto* WorkBlock = (MLAS_SGEMM_WORK_BLOCK*)Context; - const int32_t ThreadCountM = WorkBlock->ThreadCountM; - const int32_t ThreadCountN = WorkBlock->ThreadCountN; + const ptrdiff_t ThreadCountM = WorkBlock->ThreadCountM; + const ptrdiff_t ThreadCountN = WorkBlock->ThreadCountN; - const int32_t ThreadIdM = ThreadId / ThreadCountN; - const int32_t ThreadIdN = ThreadId % ThreadCountN; + const ptrdiff_t ThreadIdM = ThreadId / ThreadCountN; + const ptrdiff_t ThreadIdN = ThreadId % ThreadCountN; // // Partition the operation along the M dimension. @@ -1331,15 +1331,15 @@ Return Value: const double Complexity = double(M) * double(N) * double(K); - int32_t TargetThreadCount; + ptrdiff_t TargetThreadCount; if (Complexity < double(MLAS_SGEMM_THREAD_COMPLEXITY * MlasPlatform.MaximumThreadCount)) { - TargetThreadCount = int32_t(Complexity / double(MLAS_SGEMM_THREAD_COMPLEXITY)) + 1; + TargetThreadCount = ptrdiff_t(Complexity / double(MLAS_SGEMM_THREAD_COMPLEXITY)) + 1; } else { TargetThreadCount = MlasPlatform.MaximumThreadCount; } - int32_t MaximumThreadCount = MlasGetMaximumThreadCount(ThreadPool); + ptrdiff_t MaximumThreadCount = MlasGetMaximumThreadCount(ThreadPool); if (TargetThreadCount >= MaximumThreadCount) { TargetThreadCount = MaximumThreadCount; @@ -1358,7 +1358,7 @@ Return Value: MLAS_SGEMM_STRIDEN_THREAD_ALIGN; if (size_t(TargetThreadCount) > BlockedN) { - TargetThreadCount = int32_t(BlockedN); + TargetThreadCount = ptrdiff_t(BlockedN); } WorkBlock->ThreadCountM = 1; @@ -1367,7 +1367,7 @@ Return Value: } else { if (size_t(TargetThreadCount) > M) { - TargetThreadCount = int32_t(M); + TargetThreadCount = ptrdiff_t(M); } WorkBlock->ThreadCountM = TargetThreadCount; diff --git a/onnxruntime/core/mlas/lib/snchwc.cpp b/onnxruntime/core/mlas/lib/snchwc.cpp index 23240d1909..a30ee219af 100644 --- a/onnxruntime/core/mlas/lib/snchwc.cpp +++ b/onnxruntime/core/mlas/lib/snchwc.cpp @@ -23,7 +23,7 @@ Abstract: struct MLAS_NCHWC_WORK_BLOCK { - int32_t tids; + ptrdiff_t tids; size_t BatchCount; size_t InputChannels; size_t InputShape[2]; @@ -350,7 +350,7 @@ template void MlasNchwcThreaded( void* Context, - int32_t Index + ptrdiff_t Index ) { AlgorithmType((decltype(AlgorithmType::WorkBlock))Context).Execute(Index); @@ -542,7 +542,7 @@ struct MLAS_NCHWC_GROUPED_CONV_ALGORITHM : MLAS_NCHWC_CONV_ALGORITHM FilterCount = std::min(FilterSetSize, (OutputChannels / BlockSize) - FilterSet * FilterSetSize); } - void PrepareWork(int32_t Index) + void PrepareWork(ptrdiff_t Index) { const size_t TotalWork = BatchCount * GroupCount * FilterSetCount * OutputHeight; @@ -652,7 +652,7 @@ struct MLAS_NCHWC_CONV_NCHWC_ALGORITHM : MLAS_NCHWC_GROUPED_CONV_ALGORITHM { } - void Execute(int32_t Index) + void Execute(ptrdiff_t Index) { // // Setup the convolution state based on the thread index. @@ -762,7 +762,7 @@ struct MLAS_NCHWC_CONV_NCHW_ALGORITHM : MLAS_NCHWC_GROUPED_CONV_ALGORITHM { } - void Execute(int32_t Index) + void Execute(ptrdiff_t Index) { // // Setup the convolution state based on the thread index. @@ -862,7 +862,7 @@ struct MLAS_NCHWC_CONV_POINTWISE_ALGORITHM : MLAS_NCHWC_GROUPED_CONV_ALGORITHM { } - void Execute(int32_t Index) + void Execute(ptrdiff_t Index) { // // Setup the convolution state based on the thread index. @@ -970,7 +970,7 @@ struct MLAS_NCHWC_CONV_DEPTHWISE_ALGORITHM : MLAS_NCHWC_CONV_ALGORITHM { } - void Execute(int32_t Index) + void Execute(ptrdiff_t Index) { const size_t GroupBlockCount = ((GroupCount + BlockSize - 1) / BlockSize); @@ -1105,7 +1105,7 @@ struct MLAS_NCHWC_POOL_ALGORITHM : MLAS_NCHWC_NN_ALGORITHM { } - void Execute(int32_t Index) + void Execute(ptrdiff_t Index) { const size_t TotalWork = ((BatchCount * InputChannels + BlockSize - 1) / BlockSize) * OutputHeight; diff --git a/onnxruntime/core/mlas/lib/threading.cpp b/onnxruntime/core/mlas/lib/threading.cpp index e8cb071982..a0a4867f8b 100644 --- a/onnxruntime/core/mlas/lib/threading.cpp +++ b/onnxruntime/core/mlas/lib/threading.cpp @@ -20,7 +20,7 @@ void MlasExecuteThreaded( MLAS_THREADED_ROUTINE* ThreadedRoutine, void* Context, - int32_t Iterations, + ptrdiff_t Iterations, MLAS_THREADPOOL* ThreadPool ) { @@ -46,7 +46,7 @@ MlasExecuteThreaded( #ifdef _OPENMP #pragma omp parallel for #endif - for (int32_t tid = 0; tid < Iterations; tid++) { + for (ptrdiff_t tid = 0; tid < Iterations; tid++) { ThreadedRoutine(Context, tid); } #else @@ -55,7 +55,7 @@ MlasExecuteThreaded( // MLAS_THREADPOOL::TrySimpleParallelFor(ThreadPool, Iterations, [&](ptrdiff_t tid) { - ThreadedRoutine(Context, static_cast(tid)); + ThreadedRoutine(Context, tid); }); #endif } diff --git a/onnxruntime/core/mlas/lib/x86_64/DgemmKernelCommon.h b/onnxruntime/core/mlas/lib/x86_64/DgemmKernelCommon.h index 0728abfb62..ba53a91b68 100644 --- a/onnxruntime/core/mlas/lib/x86_64/DgemmKernelCommon.h +++ b/onnxruntime/core/mlas/lib/x86_64/DgemmKernelCommon.h @@ -21,7 +21,6 @@ Abstract: .equ .LFgemmElementShift, 3 .equ .LFgemmElementSize, 1 << .LFgemmElementShift - .equ .LFgemmMaskMoveVector, C_UNDERSCORE(MlasMaskMoveAvx64) #include "FgemmKernelCommon.h" @@ -42,7 +41,6 @@ FGEMM_TYPED_INSTRUCTION(vmovapf, vmovapd) FGEMM_TYPED_INSTRUCTION(vmovsf, vmovsd) FGEMM_TYPED_INSTRUCTION(vmovupf, vmovupd) FGEMM_TYPED_INSTRUCTION(vmulpf, vmulpd) -FGEMM_TYPED_INSTRUCTION(vpcmpgtf, vpcmpgtq) FGEMM_TYPED_INSTRUCTION(vxorpf, vxorpd) .macro vfmadd231pf_bcst DestReg, SrcReg, Address diff --git a/onnxruntime/core/mlas/lib/x86_64/FgemmKernelAvxCommon.h b/onnxruntime/core/mlas/lib/x86_64/FgemmKernelAvxCommon.h index d84cc54953..69c8d17e27 100644 --- a/onnxruntime/core/mlas/lib/x86_64/FgemmKernelAvxCommon.h +++ b/onnxruntime/core/mlas/lib/x86_64/FgemmKernelAvxCommon.h @@ -316,15 +316,9 @@ Implicit Arguments: add r9,.LFgemmYmmElementCount # correct for over-subtract above .LOutputMasked1xNBlock\@: - vmovd xmm0,r9d -.if .LFgemmElementSize == 4 - vshufps xmm0,xmm0,xmm0,0 -.else - vmovddup xmm0,xmm0 -.endif - vpcmpgtf xmm1,xmm0,XMMWORD PTR .LFgemmMaskMoveVector[rip+16] - vpcmpgtf xmm0,xmm0,XMMWORD PTR .LFgemmMaskMoveVector[rip] - vinsertf128 ymm0,ymm0,xmm1,1 + neg r9 + lea rdi,C_UNDERSCORE(MlasMaskMoveTableAvx)[rip+8*4] + vmovdqu ymm0,YMMWORD PTR [rdi+r9*.LFgemmElementSize] test r15b,r15b # ZeroMode? jnz .LStoreMasked1xNBlock\@ EmitIfCountGE \RowCount\(), 1, "vmaskmovpf ymm8,ymm0,YMMWORD PTR [rdx]" diff --git a/onnxruntime/core/mlas/lib/x86_64/FgemmKernelCommon.h b/onnxruntime/core/mlas/lib/x86_64/FgemmKernelCommon.h index 941cb8d1de..f3e8890fd3 100644 --- a/onnxruntime/core/mlas/lib/x86_64/FgemmKernelCommon.h +++ b/onnxruntime/core/mlas/lib/x86_64/FgemmKernelCommon.h @@ -19,11 +19,10 @@ Abstract: // Stack frame layout for the floating point kernels. // - .equ .LFgemmKernelFrame_SavedR12, -40 - .equ .LFgemmKernelFrame_SavedR13, -32 - .equ .LFgemmKernelFrame_SavedR14, -24 - .equ .LFgemmKernelFrame_alpha, -16 - .equ .LFgemmKernelFrame_mask, -8 + .equ .LFgemmKernelFrame_SavedR12, -32 + .equ .LFgemmKernelFrame_SavedR13, -24 + .equ .LFgemmKernelFrame_SavedR14, -16 + .equ .LFgemmKernelFrame_alpha, -8 .equ .LFgemmKernelFrame_SavedR15, 0 .equ .LFgemmKernelFrame_SavedRbx, 8 .equ .LFgemmKernelFrame_SavedRbp, 16 diff --git a/onnxruntime/core/mlas/lib/x86_64/FgemmKernelFma3Common.h b/onnxruntime/core/mlas/lib/x86_64/FgemmKernelFma3Common.h index 3c26e053b7..77fc1f1d9f 100644 --- a/onnxruntime/core/mlas/lib/x86_64/FgemmKernelFma3Common.h +++ b/onnxruntime/core/mlas/lib/x86_64/FgemmKernelFma3Common.h @@ -354,9 +354,9 @@ Implicit Arguments: add r9,.LFgemmYmmElementCount # correct for over-subtract above .LOutputMasked1xNBlock\@: - mov [rsp+.LFgemmKernelFrame_mask],r9 - vbroadcastsf ymm0,[rsp+.LFgemmKernelFrame_mask] - vpcmpgtf ymm0,ymm0,YMMWORD PTR .LFgemmMaskMoveVector[rip] + neg r9 + lea rdi,C_UNDERSCORE(MlasMaskMoveTableAvx)[rip+8*4] + vmovdqu ymm0,YMMWORD PTR [rdi+r9*.LFgemmElementSize] test r15b,r15b # ZeroMode? jnz .LMultiplyAlphaMasked1xNBlock\@ EmitIfCountGE \RowCount\(), 1, "vmaskmovpf ymm4,ymm0,YMMWORD PTR [rdx]" diff --git a/onnxruntime/core/mlas/lib/x86_64/SgemmKernelCommon.h b/onnxruntime/core/mlas/lib/x86_64/SgemmKernelCommon.h index 20eb07710a..5802028788 100644 --- a/onnxruntime/core/mlas/lib/x86_64/SgemmKernelCommon.h +++ b/onnxruntime/core/mlas/lib/x86_64/SgemmKernelCommon.h @@ -21,7 +21,6 @@ Abstract: .equ .LFgemmElementShift, 2 .equ .LFgemmElementSize, 1 << .LFgemmElementShift - .equ .LFgemmMaskMoveVector, C_UNDERSCORE(MlasMaskMoveAvx) #include "FgemmKernelCommon.h" @@ -42,7 +41,6 @@ FGEMM_TYPED_INSTRUCTION(vmovapf, vmovaps) FGEMM_TYPED_INSTRUCTION(vmovsf, vmovss) FGEMM_TYPED_INSTRUCTION(vmovupf, vmovups) FGEMM_TYPED_INSTRUCTION(vmulpf, vmulps) -FGEMM_TYPED_INSTRUCTION(vpcmpgtf, vpcmpgtd) FGEMM_TYPED_INSTRUCTION(vxorpf, vxorps) .macro vfmadd231pf_bcst DestReg, SrcReg, Address diff --git a/onnxruntime/core/mlas/lib/x86_64/SgemmKernelM1Avx.S b/onnxruntime/core/mlas/lib/x86_64/SgemmKernelM1Avx.S index 3a02845a5f..5c759847e2 100644 --- a/onnxruntime/core/mlas/lib/x86_64/SgemmKernelM1Avx.S +++ b/onnxruntime/core/mlas/lib/x86_64/SgemmKernelM1Avx.S @@ -56,8 +56,7 @@ Return Value: --*/ - .globl C_UNDERSCORE(MlasSgemmKernelM1Avx) -C_UNDERSCORE(MlasSgemmKernelM1Avx): + FUNCTION_ENTRY MlasSgemmKernelM1Avx push rbx shl r9,2 # convert ldb to bytes diff --git a/onnxruntime/core/mlas/lib/x86_64/SgemmKernelM1TransposeBAvx.S b/onnxruntime/core/mlas/lib/x86_64/SgemmKernelM1TransposeBAvx.S index b3db39989d..b205c3d6d7 100644 --- a/onnxruntime/core/mlas/lib/x86_64/SgemmKernelM1TransposeBAvx.S +++ b/onnxruntime/core/mlas/lib/x86_64/SgemmKernelM1TransposeBAvx.S @@ -56,8 +56,7 @@ Return Value: --*/ - .globl C_UNDERSCORE(MlasSgemmKernelM1TransposeBAvx) -C_UNDERSCORE(MlasSgemmKernelM1TransposeBAvx): + FUNCTION_ENTRY MlasSgemmKernelM1TransposeBAvx push rbx shl r9,2 # convert ldb to bytes diff --git a/onnxruntime/core/mlas/lib/x86_64/SgemmTransposePackB16x4Avx.S b/onnxruntime/core/mlas/lib/x86_64/SgemmTransposePackB16x4Avx.S index 0f92de07d8..644077838a 100644 --- a/onnxruntime/core/mlas/lib/x86_64/SgemmTransposePackB16x4Avx.S +++ b/onnxruntime/core/mlas/lib/x86_64/SgemmTransposePackB16x4Avx.S @@ -108,8 +108,7 @@ Return Value: --*/ - .globl C_UNDERSCORE(MlasSgemmTransposePackB16x4Avx) -C_UNDERSCORE(MlasSgemmTransposePackB16x4Avx): + FUNCTION_ENTRY MlasSgemmTransposePackB16x4Avx shl rdx,2 # convert ldb to bytes TransposePackB8x4BlockAvx 0*4 diff --git a/onnxruntime/core/mlas/lib/x86_64/SgemmTransposePackB16x4Sse2.S b/onnxruntime/core/mlas/lib/x86_64/SgemmTransposePackB16x4Sse2.S index e39133c794..d3ef6f3237 100644 --- a/onnxruntime/core/mlas/lib/x86_64/SgemmTransposePackB16x4Sse2.S +++ b/onnxruntime/core/mlas/lib/x86_64/SgemmTransposePackB16x4Sse2.S @@ -47,8 +47,7 @@ Return Value: --*/ - .globl C_UNDERSCORE(MlasSgemmTransposePackB16x4Sse) -C_UNDERSCORE(MlasSgemmTransposePackB16x4Sse): + FUNCTION_ENTRY MlasSgemmTransposePackB16x4Sse shl rdx,2 # convert ldb to bytes mov ecx,4 # transpose four 4x4 blocks