From 928b6bb21010b06c446248fa89ee95f11e61dded Mon Sep 17 00:00:00 2001 From: Tracy Sharpe <42477615+tracysh@users.noreply.github.com> Date: Wed, 15 Jan 2020 19:25:40 -0800 Subject: [PATCH] MLAS: enable threading for quantized GEMMs (#2844) --- onnxruntime/contrib_ops/cpu/nchwc_ops.cc | 1 - .../mlas/lib/amd64/QgemmU8S8KernelAvx2.asm | 68 +- onnxruntime/core/mlas/lib/mlasi.h | 42 + onnxruntime/core/mlas/lib/qgemm.cpp | 734 ++++++++++++++---- onnxruntime/core/mlas/lib/snchwc.cpp | 30 +- onnxruntime/core/mlas/lib/threading.cpp | 1 - .../mlas/lib/x86_64/QgemmU8S8KernelAvx2.S | 68 +- onnxruntime/core/providers/cpu/math/gemm.h | 5 +- onnxruntime/core/providers/cpu/math/matmul.cc | 3 +- .../core/providers/cpu/math/matmul_integer.cc | 1 - onnxruntime/core/providers/cpu/nn/conv.cc | 12 +- .../core/providers/cpu/nn/conv_integer.cc | 3 +- .../core/providers/cpu/nn/conv_transpose.cc | 9 +- onnxruntime/core/providers/cpu/nn/pool.cc | 1 - 14 files changed, 707 insertions(+), 271 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/nchwc_ops.cc b/onnxruntime/contrib_ops/cpu/nchwc_ops.cc index 325217a1f8..14e0a83384 100644 --- a/onnxruntime/contrib_ops/cpu/nchwc_ops.cc +++ b/onnxruntime/contrib_ops/cpu/nchwc_ops.cc @@ -1,7 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/framework/op_kernel_context_internal.h" #include "nchwc_ops.h" #include "core/mlas/inc/mlas.h" diff --git a/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm b/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm index f701fb623f..a4d9ec8468 100644 --- a/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm +++ b/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm @@ -170,6 +170,8 @@ ProcessNextRowM4: vpxor xmm1,xmm1,xmm1 vpxor xmm2,xmm2,xmm2 vpxor xmm3,xmm3,xmm3 + lea r13,[r8+r8*2] ; compute ldb * 3 + lea rax,[r11+r11*2] ; compute output stride * 3 mov rdx,rsi mov rcx,rdi lea rsi,[rsi+r8*4] ; advance next matrix A by 4 rows @@ -179,16 +181,14 @@ ProcessNextRowM4: jb ProcessRemainingColumnsM4 ProcessNextColumnLoopM4: - lea rax,[rdx+r8*2] ; compute matrix A plus 2 rows vmovdqu ymm4,YMMWORD PTR [rdx] vmovdqu ymm5,YMMWORD PTR [rdx+r8] - vmovdqu ymm6,YMMWORD PTR [rax] - vmovdqu ymm7,YMMWORD PTR [rax+r8] - lea rax,[rcx+r11*2] ; compute matrix D plus 2 rows + vmovdqu ymm6,YMMWORD PTR [rdx+r8*2] + vmovdqu ymm7,YMMWORD PTR [rdx+r13] vmovdqu YMMWORD PTR [rcx],ymm4 vmovdqu YMMWORD PTR [rcx+r11],ymm5 - vmovdqu YMMWORD PTR [rax],ymm6 - vmovdqu YMMWORD PTR [rax+r11],ymm7 + vmovdqu YMMWORD PTR [rcx+r11*2],ymm6 + vmovdqu YMMWORD PTR [rcx+rax],ymm7 vpmaddubsw ymm4,ymm4,ymm9 ; horizontal byte+byte=word per row vpaddw ymm0,ymm0,ymm4 ; add words to row accumulators vpmaddubsw ymm5,ymm5,ymm9 @@ -207,16 +207,14 @@ ProcessRemainingColumnsM4: jz ReduceRowSumVectorM4 test bl,16 ; (CountK & 16) != 0? jz CopyRemainingCountKLessThan16M4 - lea rax,[rdx+r8*2] ; compute matrix A plus 2 rows vmovdqu xmm4,XMMWORD PTR [rdx] vmovdqu xmm5,XMMWORD PTR [rdx+r8] - vmovdqu xmm6,XMMWORD PTR [rax] - vmovdqu xmm7,XMMWORD PTR [rax+r8] - lea rax,[rcx+r11*2] ; compute matrix D plus 2 rows + vmovdqu xmm6,XMMWORD PTR [rdx+r8*2] + vmovdqu xmm7,XMMWORD PTR [rdx+r13] vmovdqu XMMWORD PTR [rcx],xmm4 vmovdqu XMMWORD PTR [rcx+r11],xmm5 - vmovdqu XMMWORD PTR [rax],xmm6 - vmovdqu XMMWORD PTR [rax+r11],xmm7 + vmovdqu XMMWORD PTR [rcx+r11*2],xmm6 + vmovdqu XMMWORD PTR [rcx+rax],xmm7 vpmaddubsw xmm4,xmm4,xmm9 ; horizontal byte+byte=word per row vpaddw ymm0,ymm0,ymm4 ; add words to row accumulators vpmaddubsw xmm5,xmm5,xmm9 @@ -239,14 +237,13 @@ CopyRemainingCountKLessThan16M4: mov rbp,rsp ; GemmU8S8CopyPackAFrame.PaddedMatrixAData test bl,8 ; (CountK & 8) != 0? jz CopyRemainingCountKLessThan8M4 - lea r13,[rdx+r8*2] ; compute matrix A plus 2 rows mov rax,QWORD PTR [rdx] mov QWORD PTR [rbp],rax mov rax,QWORD PTR [rdx+r8] mov QWORD PTR [rbp+16],rax - mov rax,QWORD PTR [r13] + mov rax,QWORD PTR [rdx+r8*2] mov QWORD PTR [rbp+32],rax - mov rax,QWORD PTR [r13+r8] + mov rax,QWORD PTR [rdx+r13] mov QWORD PTR [rbp+48],rax add rdx,8 add rbp,8 ; advance padded buffer destination @@ -254,14 +251,13 @@ CopyRemainingCountKLessThan16M4: CopyRemainingCountKLessThan8M4: test bl,4 ; (CountK & 4) != 0? jz CopyRemainingCountKLessThan4M4 - lea r13,[rdx+r8*2] ; compute matrix A plus 2 rows mov eax,DWORD PTR [rdx] mov DWORD PTR [rbp],eax mov eax,DWORD PTR [rdx+r8] mov DWORD PTR [rbp+16],eax - mov eax,DWORD PTR [r13] + mov eax,DWORD PTR [rdx+r8*2] mov DWORD PTR [rbp+32],eax - mov eax,DWORD PTR [r13+r8] + mov eax,DWORD PTR [rdx+r13] mov DWORD PTR [rbp+48],eax add rdx,4 add rbp,4 ; advance padded buffer destination @@ -269,14 +265,13 @@ CopyRemainingCountKLessThan8M4: CopyRemainingCountKLessThan4M4: test bl,2 ; (CountK & 2) != 0? jz CopyRemainingCountKLessThan2M4 - lea r13,[rdx+r8*2] ; compute matrix A plus 2 rows movzx eax,WORD PTR [rdx] mov WORD PTR [rbp],ax movzx eax,WORD PTR [rdx+r8] mov WORD PTR [rbp+16],ax - movzx eax,WORD PTR [r13] + movzx eax,WORD PTR [rdx+r8*2] mov WORD PTR [rbp+32],ax - movzx eax,WORD PTR [r13+r8] + movzx eax,WORD PTR [rdx+r13] mov WORD PTR [rbp+48],ax add rdx,2 add rbp,2 ; advance padded buffer destination @@ -284,14 +279,13 @@ CopyRemainingCountKLessThan4M4: CopyRemainingCountKLessThan2M4: test bl,1 ; (CountK & 1) != 0? jz ProcessPaddedMatrixADataM4 - lea r13,[rdx+r8*2] ; compute matrix A plus 2 rows movzx eax,BYTE PTR [rdx] mov BYTE PTR [rbp],al movzx eax,BYTE PTR [rdx+r8] mov BYTE PTR [rbp+16],al - movzx eax,BYTE PTR [r13] + movzx eax,BYTE PTR [rdx+r8*2] mov BYTE PTR [rbp+32],al - movzx eax,BYTE PTR [r13+r8] + movzx eax,BYTE PTR [rdx+r13] mov BYTE PTR [rbp+48],al ; @@ -507,6 +501,7 @@ ExitRoutine: END_PROLOGUE mov rsi,rdx + lea rdi,[r8+r8*2] ; compute ldb * 3 mov r10,GemmU8S8CopyPackBFrame.CountK[rsp] mov r11,GemmU8S8CopyPackBFrame.ColumnSumVector[rsp] vpbroadcastw ymm7,WORD PTR GemmU8S8CopyPackBFrame.offa[rsp] @@ -532,11 +527,10 @@ ProcessNextColumnN16: jb ProcessRemainingRowsN16 ProcessNextRowLoopN16: - lea rax,[rdx+r8*2] ; compute matrix B plus 2 rows vmovdqu xmm2,XMMWORD PTR [rdx] ; load 4 rows vmovdqu xmm3,XMMWORD PTR [rdx+r8] - vmovdqu xmm4,XMMWORD PTR [rax] - vmovdqu xmm5,XMMWORD PTR [rax+r8] + vmovdqu xmm4,XMMWORD PTR [rdx+r8*2] + vmovdqu xmm5,XMMWORD PTR [rdx+rdi] lea rdx,[rdx+r8*4] ; advance matrix B by 4 rows InterleaveRowDataN16: @@ -630,14 +624,13 @@ ProcessNextRowLoopNUnaligned: mov rbp,rsp ; GemmU8S8CopyPackBFrame.PaddedMatrixBData test r9b,8 ; (CountN & 8) != 0? jz CopyRemainingCountNLessThan8K4 - lea rdi,[rdx+r8*2] ; compute matrix B plus 2 rows mov rax,QWORD PTR [rdx] mov QWORD PTR [rbp],rax mov rax,QWORD PTR [rdx+r8] mov QWORD PTR [rbp+16],rax - mov rax,QWORD PTR [rdi] + mov rax,QWORD PTR [rdx+r8*2] mov QWORD PTR [rbp+32],rax - mov rax,QWORD PTR [rdi+r8] + mov rax,QWORD PTR [rdx+rdi] mov QWORD PTR [rbp+48],rax add rdx,8 ; advance matrix B add rbp,8 ; advance padded buffer destination @@ -645,14 +638,13 @@ ProcessNextRowLoopNUnaligned: CopyRemainingCountNLessThan8K4: test r9b,4 ; (CountN & 4) != 0? jz CopyRemainingCountNLessThan4K4 - lea rdi,[rdx+r8*2] ; compute matrix B plus 2 rows mov eax,DWORD PTR [rdx] mov DWORD PTR [rbp],eax mov eax,DWORD PTR [rdx+r8] mov DWORD PTR [rbp+16],eax - mov eax,DWORD PTR [rdi] + mov eax,DWORD PTR [rdx+r8*2] mov DWORD PTR [rbp+32],eax - mov eax,DWORD PTR [rdi+r8] + mov eax,DWORD PTR [rdx+rdi] mov DWORD PTR [rbp+48],eax add rdx,4 ; advance matrix B add rbp,4 ; advance padded buffer destination @@ -660,14 +652,13 @@ CopyRemainingCountNLessThan8K4: CopyRemainingCountNLessThan4K4: test r9b,2 ; (CountN & 2) != 0? jz CopyRemainingCountNLessThan2K4 - lea rdi,[rdx+r8*2] ; compute matrix B plus 2 rows movzx eax,WORD PTR [rdx] mov WORD PTR [rbp],ax movzx eax,WORD PTR [rdx+r8] mov WORD PTR [rbp+16],ax - movzx eax,WORD PTR [rdi] + movzx eax,WORD PTR [rdx+r8*2] mov WORD PTR [rbp+32],ax - movzx eax,WORD PTR [rdi+r8] + movzx eax,WORD PTR [rdx+rdi] mov WORD PTR [rbp+48],ax add rdx,2 ; advance matrix B add rbp,2 ; advance padded buffer destination @@ -675,14 +666,13 @@ CopyRemainingCountNLessThan4K4: CopyRemainingCountNLessThan2K4: test r9b,1 ; (CountN & 1) != 0? jz ProcessPaddedMatrixBData - lea rdi,[rdx+r8*2] ; compute matrix B plus 2 rows movzx eax,BYTE PTR [rdx] mov BYTE PTR [rbp],al movzx eax,BYTE PTR [rdx+r8] mov BYTE PTR [rbp+16],al - movzx eax,BYTE PTR [rdi] + movzx eax,BYTE PTR [rdx+r8*2] mov BYTE PTR [rbp+32],al - movzx eax,BYTE PTR [rdi+r8] + movzx eax,BYTE PTR [rdx+rdi] mov BYTE PTR [rbp+48],al ProcessPaddedMatrixBData: diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index ff0b038cf2..12ed28672e 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -131,6 +131,7 @@ Abstract: #define MLAS_SGEMM_STRIDEN_THREAD_ALIGN 16 #define MLAS_DGEMM_STRIDEN_THREAD_ALIGN 8 +#define MLAS_QGEMM_STRIDEN_THREAD_ALIGN 16 // // Define the prototypes of the platform optimized routines. @@ -333,6 +334,24 @@ size_t typedef MLAS_GEMM_U8U8_KERNEL* PMLAS_GEMM_U8U8_KERNEL; +typedef +void +(MLASCALL MLAS_GEMM_X8X8_OPERATION)( + size_t M, + size_t N, + size_t K, + const uint8_t* A, + size_t lda, + int16_t offa, + const uint8_t* B, + size_t ldb, + int16_t offb, + int32_t* C, + size_t ldc + ); + +typedef MLAS_GEMM_X8X8_OPERATION* PMLAS_GEMM_X8X8_OPERATION; + typedef void (MLASCALL MLAS_CONV_FLOAT_KERNEL)( @@ -557,6 +576,7 @@ extern "C" { #define MLAS_SGEMM_THREAD_COMPLEXITY (64 * 1024) #define MLAS_DGEMM_THREAD_COMPLEXITY (64 * 1024) +#define MLAS_QGEMM_THREAD_COMPLEXITY (64 * 1024) // // Single-threaded single precision matrix/matrix multiply operation. @@ -660,6 +680,28 @@ MlasGetMaximumThreadCount( #endif } +inline +void +MlasPartitionWork( + int32_t ThreadId, + int32_t ThreadCount, + size_t TotalWork, + size_t* WorkIndex, + size_t* WorkRemaining + ) +{ + const size_t WorkPerThread = TotalWork / ThreadCount; + const size_t WorkPerThreadExtra = TotalWork % ThreadCount; + + if (uint32_t(ThreadId) < WorkPerThreadExtra) { + *WorkIndex = (WorkPerThread + 1) * ThreadId; + *WorkRemaining = WorkPerThread + 1; + } else { + *WorkIndex = WorkPerThread * ThreadId + WorkPerThreadExtra; + *WorkRemaining = WorkPerThread; + } +} + // // Define the missing ARM64 NEON intrinsic macros from arm64_neon.h that enable // cross-compiler support. diff --git a/onnxruntime/core/mlas/lib/qgemm.cpp b/onnxruntime/core/mlas/lib/qgemm.cpp index 776107fc5c..377628c356 100644 --- a/onnxruntime/core/mlas/lib/qgemm.cpp +++ b/onnxruntime/core/mlas/lib/qgemm.cpp @@ -21,13 +21,33 @@ Abstract: // Define the default strides to step through slices of the input matrices. // -#define MLAS_GEMM_U8S8_STRIDEM 24 -#define MLAS_GEMM_U8S8_STRIDEN 256 -#define MLAS_GEMM_U8S8_STRIDEK 128 +#define MLAS_GEMM_X8X8_STRIDEM 24 +#define MLAS_GEMM_X8X8_STRIDEN 256 +#define MLAS_GEMM_X8X8_STRIDEK 128 -#define MLAS_GEMM_U8U8_STRIDEM 24 -#define MLAS_GEMM_U8U8_STRIDEN 256 -#define MLAS_GEMM_U8U8_STRIDEK 128 +// +// Define the parameters to execute segments of a QGEMM operation on worker +// threads. +// + +struct MLAS_GEMM_X8X8_WORK_BLOCK { + PMLAS_GEMM_X8X8_OPERATION GemmX8X8Operation; + size_t M; + size_t N; + size_t K; + const uint8_t* A; + size_t lda; + const uint8_t* B; + size_t ldb; + int32_t* C; + size_t ldc; + int32_t ThreadCountM; + int32_t ThreadCountN; + size_t StrideM; + size_t StrideN; + int16_t offa; + int16_t offb; +}; #ifdef MLAS_TARGET_AMD64_IX86 @@ -1100,6 +1120,448 @@ Return Value: return 1; } +void +MLASCALL +MlasGemmU8S8Operation( + size_t M, + size_t N, + size_t K, + const uint8_t* A, + size_t lda, + int16_t offa, + const uint8_t* B, + size_t ldb, + int16_t offb, + int32_t* C, + size_t ldc + ) +/*++ + +Routine Description: + + This module implements the quantized integer matrix/matrix multiply + operation (QGEMM). + +Arguments: + + M - Supplies the number of rows of matrix A and matrix C. + + N - Supplies the number of columns of matrix B and matrix C. + + K - Supplies the number of columns of matrix A and the number of rows of + matrix B. + + A - Supplies the address of matrix A. + + lda - Supplies the first dimension of matrix A. + + offa - Supplies the zero point offset of matrix A. + + B - Supplies the address of matrix B. + + ldb - Supplies the first dimension of matrix B. + + offb - Supplies the zero point offset of matrix B. + + C - Supplies the address of matrix C. + + ldc - Supplies the first dimension of matrix C. + +Return Value: + + None. + +--*/ +{ + MLAS_DECLSPEC_ALIGN(uint8_t PanelA[MLAS_GEMM_X8X8_STRIDEM * MLAS_GEMM_X8X8_STRIDEK], 64); + MLAS_DECLSPEC_ALIGN(int8_t PanelB[MLAS_GEMM_X8X8_STRIDEN * MLAS_GEMM_X8X8_STRIDEK], 64); + + MLAS_DECLSPEC_ALIGN(int32_t RowSumVector[MLAS_GEMM_X8X8_STRIDEM], 16); + MLAS_DECLSPEC_ALIGN(int32_t ColumnSumVector[MLAS_GEMM_X8X8_STRIDEN], 16); + + size_t StrideM = MLAS_GEMM_X8X8_STRIDEM; + size_t StrideN = MLAS_GEMM_X8X8_STRIDEN; + size_t StrideK = MLAS_GEMM_X8X8_STRIDEK; + +#if defined(MLAS_TARGET_AMD64) + + if (M == 1 && offa == 0 && offb == 0) { + + if (MlasPlatform.GemvU8S8Kernel != nullptr) { + MlasPlatform.GemvU8S8Kernel(A, (const int8_t*)B, C, K, N, ldb); + return; + } + } + +#endif + + // + // Step through each slice of matrix B along the K dimension. + // + + size_t CountK; + + for (size_t k = 0; k < K; k += CountK) { + + CountK = StrideK; + + if (CountK > (K - k)) { + CountK = K - k; + } + + // + // Step through each slice of matrix B along the N dimension. + // + + size_t CountN; + + for (size_t n = 0; n < N; n += CountN) { + + CountN = StrideN; + + if (CountN > (N - n)) { + CountN = N - n; + } + + const int8_t* b = (const int8_t*)B + n + k * ldb; + + MlasPlatform.GemmU8S8CopyPackBRoutine(PanelB, b, ldb, CountN, + CountK, ColumnSumVector, -int16_t(offa)); + + size_t CountM; + + for (size_t m = 0; m < M; m += CountM) { + + CountM = StrideM; + + if (CountM > (M - m)) { + CountM = M - m; + } + + MlasPlatform.GemmU8S8CopyPackARoutine(PanelA, A + k + m * lda, + lda, CountM, CountK, RowSumVector, -int16_t(offb)); + + uint8_t* pa = PanelA; + int32_t* c = C + n + m * ldc; + + int32_t* RowSums = RowSumVector; + + size_t RowsRemaining = CountM; + size_t RowsHandled; + + size_t QuadCountK = (CountK + 3) / 4; + + while (RowsRemaining > 0) { + + RowsHandled = MlasPlatform.GemmU8S8Kernel(pa, PanelB, c, + QuadCountK, RowsRemaining, CountN, ldc, RowSums, + ColumnSumVector, int32_t(CountK) * offa * offb, k == 0); + + RowsRemaining -= RowsHandled; + c += ldc * RowsHandled; + pa += 4 * QuadCountK * RowsHandled; + RowSums += RowsHandled; + } + } + } + } +} + +void +MLASCALL +MlasGemmU8U8Operation( + size_t M, + size_t N, + size_t K, + const uint8_t* A, + size_t lda, + int16_t offa, + const uint8_t* B, + size_t ldb, + int16_t offb, + int32_t* C, + size_t ldc + ) +/*++ + +Routine Description: + + This module implements the quantized integer matrix/matrix multiply + operation (QGEMM). + +Arguments: + + M - Supplies the number of rows of matrix A and matrix C. + + N - Supplies the number of columns of matrix B and matrix C. + + K - Supplies the number of columns of matrix A and the number of rows of + matrix B. + + A - Supplies the address of matrix A. + + lda - Supplies the first dimension of matrix A. + + offa - Supplies the zero point offset of matrix A. + + B - Supplies the address of matrix B. + + ldb - Supplies the first dimension of matrix B. + + offb - Supplies the zero point offset of matrix B. + + C - Supplies the address of matrix C. + + ldc - Supplies the first dimension of matrix C. + +Return Value: + + None. + +--*/ +{ + MLAS_DECLSPEC_ALIGN(int16_t PanelA[MLAS_GEMM_X8X8_STRIDEM * MLAS_GEMM_X8X8_STRIDEK], 64); + MLAS_DECLSPEC_ALIGN(uint8_t PanelB[MLAS_GEMM_X8X8_STRIDEN * MLAS_GEMM_X8X8_STRIDEK], 64); + + MLAS_DECLSPEC_ALIGN(int32_t RowSumVector[MLAS_GEMM_X8X8_STRIDEM], 16); + MLAS_DECLSPEC_ALIGN(int32_t ColumnSumVector[MLAS_GEMM_X8X8_STRIDEN], 16); + + size_t StrideM = MLAS_GEMM_X8X8_STRIDEM; + size_t StrideN = MLAS_GEMM_X8X8_STRIDEN; + size_t StrideK = MLAS_GEMM_X8X8_STRIDEK; + + // + // Step through each slice of matrix B along the K dimension. + // + + size_t CountK; + + for (size_t k = 0; k < K; k += CountK) { + + CountK = StrideK; + + if (CountK > (K - k)) { + CountK = K - k; + } + + // + // Step through each slice of matrix B along the N dimension. + // + + size_t CountN; + + for (size_t n = 0; n < N; n += CountN) { + + CountN = StrideN; + + if (CountN > (N - n)) { + CountN = N - n; + } + + const uint8_t* b = (const uint8_t*)B + n + k * ldb; + + MlasPlatform.GemmU8U8CopyPackBRoutine(PanelB, b, ldb, CountN, + CountK, ColumnSumVector, -int16_t(offa)); + + size_t CountM; + + for (size_t m = 0; m < M; m += CountM) { + + CountM = StrideM; + + if (CountM > (M - m)) { + CountM = M - m; + } + + MlasPlatform.GemmU8U8CopyPackARoutine(PanelA, A + k + m * lda, + lda, CountM, CountK, RowSumVector, -int16_t(offb)); + + int16_t* pa = PanelA; + int32_t* c = C + n + m * ldc; + + int32_t* RowSums = RowSumVector; + + size_t RowsRemaining = CountM; + size_t RowsHandled; + + size_t PairCountK = (CountK + 1) / 2; + + while (RowsRemaining > 0) { + + RowsHandled = MlasPlatform.GemmU8U8Kernel(pa, PanelB, c, + PairCountK, RowsRemaining, CountN, ldc, RowSums, + ColumnSumVector, int32_t(CountK) * offa * offb, k == 0); + + RowsRemaining -= RowsHandled; + c += ldc * RowsHandled; + pa += 2 * PairCountK * RowsHandled; + RowSums += RowsHandled; + } + } + } + } +} + +void +MlasGemmX8X8Threaded( + void* Context, + int32_t ThreadId + ) +/*++ + +Routine Description: + + This routine is invoked from a worker thread to execute a segment of a + QGEMM operation. + +Arguments: + + Context - Supplies the pointer to the context for the threaded operation. + + ThreadId - Supplies the current index of the threaded operation. + +Return Value: + + None. + +--*/ +{ + const auto* WorkBlock = (MLAS_GEMM_X8X8_WORK_BLOCK*)Context; + + const int32_t ThreadCountM = WorkBlock->ThreadCountM; + const int32_t ThreadCountN = WorkBlock->ThreadCountN; + + const int32_t ThreadIdM = ThreadId / ThreadCountN; + const int32_t ThreadIdN = ThreadId % ThreadCountN; + + // + // Partition the operation along the M dimension. + // + + size_t M = WorkBlock->M; + size_t m; + size_t CountM; + + MlasPartitionWork(ThreadIdM, ThreadCountM, M, &m, &CountM); + + // + // Partition the operation along the N dimension. + // + + size_t N = WorkBlock->N; + size_t n; + size_t CountN; + + const size_t BlockedN = (N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) / + MLAS_QGEMM_STRIDEN_THREAD_ALIGN; + + MlasPartitionWork(ThreadIdN, ThreadCountN, BlockedN, &n, &CountN); + + n *= MLAS_QGEMM_STRIDEN_THREAD_ALIGN; + CountN *= MLAS_QGEMM_STRIDEN_THREAD_ALIGN; + + if (CountN > N - n) { + CountN = N - n; + } + + // + // Dispatch the partitioned operation. + // + + const size_t lda = WorkBlock->lda; + const size_t ldb = WorkBlock->ldb; + const size_t ldc = WorkBlock->ldc; + + const uint8_t* a = WorkBlock->A + m * lda; + const uint8_t* b = WorkBlock->B + n; + int32_t* c = WorkBlock->C + n + m * ldc; + + WorkBlock->GemmX8X8Operation(CountM, CountN, WorkBlock->K, a, lda, + WorkBlock->offa, b, ldb, WorkBlock->offb, c, ldc); +} + +void +MlasGemmX8X8Schedule( + MLAS_GEMM_X8X8_WORK_BLOCK* WorkBlock, + MLAS_THREADPOOL* ThreadPool + ) +/*++ + +Routine Description: + + This module schedules the quantized integer matrix/matrix multiply + operation (QGEMM) across one or more threads. + +Arguments: + + WorkBlock - Supplies the structure containing the GEMM parameters. + + ThreadPool - Supplies the thread pool object to use, else nullptr if the + base library threading support should be used. + +Return Value: + + None. + +--*/ +{ + const size_t M = WorkBlock->M; + const size_t N = WorkBlock->N; + const size_t K = WorkBlock->K; + + // + // Compute the number of target threads given the complexity of the SGEMM + // operation. Small requests should run using the single threaded path. + // + + double Complexity = double(M) * double(N) * double(K); + + int32_t TargetThreadCount; + + if (Complexity < double(MLAS_QGEMM_THREAD_COMPLEXITY * MLAS_MAXIMUM_THREAD_COUNT)) { + TargetThreadCount = int32_t(Complexity / double(MLAS_QGEMM_THREAD_COMPLEXITY)) + 1; + } else { + TargetThreadCount = MLAS_MAXIMUM_THREAD_COUNT; + } + + int32_t MaximumThreadCount = MlasGetMaximumThreadCount(ThreadPool); + + if (TargetThreadCount >= MaximumThreadCount) { + TargetThreadCount = MaximumThreadCount; + } + + // + // Segment the operation across multiple threads. + // + // N.B. Currently, the operation is segmented as a 1D partition, which + // works okay for operations involving skinny matrices. + // + + if (N > M) { + + const size_t BlockedN = (N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) / + MLAS_QGEMM_STRIDEN_THREAD_ALIGN; + + if (size_t(TargetThreadCount) > BlockedN) { + TargetThreadCount = int32_t(BlockedN); + } + + WorkBlock->ThreadCountM = 1; + WorkBlock->ThreadCountN = TargetThreadCount; + + } else { + + if (size_t(TargetThreadCount) > M) { + TargetThreadCount = int32_t(M); + } + + WorkBlock->ThreadCountM = TargetThreadCount; + WorkBlock->ThreadCountN = 1; + } + + MlasExecuteThreaded(MlasGemmX8X8Threaded, WorkBlock, TargetThreadCount, ThreadPool); +} + void MLASCALL MlasGemm( @@ -1116,87 +1578,71 @@ MlasGemm( size_t ldc, MLAS_THREADPOOL* ThreadPool ) +/*++ + +Routine Description: + + This module implements the quantized integer matrix/matrix multiply + operation (QGEMM). + +Arguments: + + M - Supplies the number of rows of matrix A and matrix C. + + N - Supplies the number of columns of matrix B and matrix C. + + K - Supplies the number of columns of matrix A and the number of rows of + matrix B. + + A - Supplies the address of matrix A. + + lda - Supplies the first dimension of matrix A. + + offa - Supplies the zero point offset of matrix A. + + B - Supplies the address of matrix B. + + ldb - Supplies the first dimension of matrix B. + + offb - Supplies the zero point offset of matrix B. + + C - Supplies the address of matrix C. + + ldc - Supplies the first dimension of matrix C. + + ThreadPool - Supplies the thread pool object to use, else nullptr if the + base library threading support should be used. + +Return Value: + + None. + +--*/ { - MLAS_DECLSPEC_ALIGN(uint8_t PanelA[MLAS_GEMM_U8S8_STRIDEM * MLAS_GEMM_U8S8_STRIDEK], 64); - MLAS_DECLSPEC_ALIGN(int8_t PanelB[MLAS_GEMM_U8S8_STRIDEN * MLAS_GEMM_U8S8_STRIDEK], 64); + MLAS_GEMM_X8X8_WORK_BLOCK WorkBlock; - MLAS_DECLSPEC_ALIGN(int32_t RowSumVector[MLAS_GEMM_U8S8_STRIDEM], 16); - MLAS_DECLSPEC_ALIGN(int32_t ColumnSumVector[MLAS_GEMM_U8S8_STRIDEN], 16); + // + // Capture the GEMM parameters to the work block. + // - size_t StrideM = MLAS_GEMM_U8S8_STRIDEM; - size_t StrideN = MLAS_GEMM_U8S8_STRIDEN; - size_t StrideK = MLAS_GEMM_U8S8_STRIDEK; + WorkBlock.M = M; + WorkBlock.N = N; + WorkBlock.K = K; + WorkBlock.A = A; + WorkBlock.lda = lda; + WorkBlock.B = (const uint8_t*)B; + WorkBlock.ldb = ldb; + WorkBlock.C = C; + WorkBlock.ldc = ldc; + WorkBlock.offa = int16_t(offa); + WorkBlock.offb = int16_t(offb); + WorkBlock.GemmX8X8Operation = MlasGemmU8S8Operation; - MLAS_UNREFERENCED_PARAMETER(ThreadPool); + // + // Schedule the operation across a set of worker threads. + // -#if defined(MLAS_TARGET_AMD64) - - if (M == 1 && offa == 0 && offb == 0) { - - if (MlasPlatform.GemvU8S8Kernel != nullptr) { - MlasPlatform.GemvU8S8Kernel(A, B, C, K, N, ldb); - return; - } - } - -#endif - - size_t CountK; - - for (size_t k = 0; k < K; k += CountK) { - - CountK = StrideK; - - if (CountK > (K - k)) { - CountK = K - k; - } - - size_t CountN; - - for (size_t n = 0; n < N; n += CountN) { - - CountN = StrideN; - - if (CountN > (N - n)) { - CountN = N - n; - } - - MlasPlatform.GemmU8S8CopyPackBRoutine(PanelB, B + n + k * ldb, ldb, CountN, CountK, ColumnSumVector, -int16_t(offa)); - - size_t CountM; - - for (size_t m = 0; m < M; m += CountM) { - - CountM = StrideM; - - if (CountM > (M - m)) { - CountM = M - m; - } - - MlasPlatform.GemmU8S8CopyPackARoutine(PanelA, A + k + m * lda, lda, CountM, CountK, RowSumVector, -int16_t(offb)); - - uint8_t* pa = PanelA; - int32_t* c = C + n + m * ldc; - - int32_t* RowSums = RowSumVector; - - size_t RowsRemaining = CountM; - size_t RowsHandled; - - size_t QuadCountK = (CountK + 3) / 4; - - while (RowsRemaining > 0) { - - RowsHandled = MlasPlatform.GemmU8S8Kernel(pa, PanelB, c, QuadCountK, RowsRemaining, CountN, ldc, RowSums, ColumnSumVector, int32_t(CountK) * offa * offb, k == 0); - - RowsRemaining -= RowsHandled; - c += ldc * RowsHandled; - pa += 4 * QuadCountK * RowsHandled; - RowSums += RowsHandled; - } - } - } - } + MlasGemmX8X8Schedule(&WorkBlock, ThreadPool); } void @@ -1215,75 +1661,71 @@ MlasGemm( size_t ldc, MLAS_THREADPOOL* ThreadPool ) +/*++ + +Routine Description: + + This module implements the quantized integer matrix/matrix multiply + operation (QGEMM). + +Arguments: + + M - Supplies the number of rows of matrix A and matrix C. + + N - Supplies the number of columns of matrix B and matrix C. + + K - Supplies the number of columns of matrix A and the number of rows of + matrix B. + + A - Supplies the address of matrix A. + + lda - Supplies the first dimension of matrix A. + + offa - Supplies the zero point offset of matrix A. + + B - Supplies the address of matrix B. + + ldb - Supplies the first dimension of matrix B. + + offb - Supplies the zero point offset of matrix B. + + C - Supplies the address of matrix C. + + ldc - Supplies the first dimension of matrix C. + + ThreadPool - Supplies the thread pool object to use, else nullptr if the + base library threading support should be used. + +Return Value: + + None. + +--*/ { - MLAS_DECLSPEC_ALIGN(int16_t PanelA[MLAS_GEMM_U8U8_STRIDEM * MLAS_GEMM_U8U8_STRIDEK], 64); - MLAS_DECLSPEC_ALIGN(uint8_t PanelB[MLAS_GEMM_U8U8_STRIDEN * MLAS_GEMM_U8U8_STRIDEK], 64); + MLAS_GEMM_X8X8_WORK_BLOCK WorkBlock; - MLAS_DECLSPEC_ALIGN(int32_t RowSumVector[MLAS_GEMM_U8U8_STRIDEM], 16); - MLAS_DECLSPEC_ALIGN(int32_t ColumnSumVector[MLAS_GEMM_U8U8_STRIDEN], 16); + // + // Capture the GEMM parameters to the work block. + // - size_t StrideM = MLAS_GEMM_U8U8_STRIDEM; - size_t StrideN = MLAS_GEMM_U8U8_STRIDEN; - size_t StrideK = MLAS_GEMM_U8U8_STRIDEK; + WorkBlock.M = M; + WorkBlock.N = N; + WorkBlock.K = K; + WorkBlock.A = A; + WorkBlock.lda = lda; + WorkBlock.B = B; + WorkBlock.ldb = ldb; + WorkBlock.C = C; + WorkBlock.ldc = ldc; + WorkBlock.offa = int16_t(offa); + WorkBlock.offb = int16_t(offb); + WorkBlock.GemmX8X8Operation = MlasGemmU8U8Operation; - MLAS_UNREFERENCED_PARAMETER(ThreadPool); + // + // Schedule the operation across a set of worker threads. + // - size_t CountK; - - for (size_t k = 0; k < K; k += CountK) { - - CountK = StrideK; - - if (CountK > (K - k)) { - CountK = K - k; - } - - size_t CountN; - - for (size_t n = 0; n < N; n += CountN) { - - CountN = StrideN; - - if (CountN > (N - n)) { - CountN = N - n; - } - - MlasPlatform.GemmU8U8CopyPackBRoutine(PanelB, B + n + k * ldb, ldb, CountN, CountK, ColumnSumVector, -int16_t(offa)); - - size_t CountM; - - for (size_t m = 0; m < M; m += CountM) { - - CountM = StrideM; - - if (CountM > (M - m)) { - CountM = M - m; - } - - MlasPlatform.GemmU8U8CopyPackARoutine(PanelA, A + k + m * lda, lda, CountM, CountK, RowSumVector, -int16_t(offb)); - - int16_t* pa = PanelA; - int32_t* c = C + n + m * ldc; - - int32_t* RowSums = RowSumVector; - - size_t RowsRemaining = CountM; - size_t RowsHandled; - - size_t PairCountK = (CountK + 1) / 2; - - while (RowsRemaining > 0) { - - RowsHandled = MlasPlatform.GemmU8U8Kernel(pa, PanelB, c, PairCountK, RowsRemaining, CountN, ldc, RowSums, ColumnSumVector, int32_t(CountK) * offa * offb, k == 0); - - RowsRemaining -= RowsHandled; - c += ldc * RowsHandled; - pa += 2 * PairCountK * RowsHandled; - RowSums += RowsHandled; - } - } - } - } + MlasGemmX8X8Schedule(&WorkBlock, ThreadPool); } #endif diff --git a/onnxruntime/core/mlas/lib/snchwc.cpp b/onnxruntime/core/mlas/lib/snchwc.cpp index 843f919f44..7c3193c290 100644 --- a/onnxruntime/core/mlas/lib/snchwc.cpp +++ b/onnxruntime/core/mlas/lib/snchwc.cpp @@ -128,7 +128,7 @@ Routine Description: Arguments: - WorkBlock - Supplies the structure that contains the commong convolution + WorkBlock - Supplies the structure that contains the common convolution and pooling parameters. Dimensions - Supplies the number of dimensions. @@ -344,28 +344,6 @@ struct MLAS_NCHWC_NN_ALGORITHM OutputCountRightPadX(WorkBlock->OutputCountRightPad[WidthShapeIndex]) { } - - static - void - PartitionWork( - int32_t Index, - const MLAS_NCHWC_WORK_BLOCK* WorkBlock, - size_t TotalWork, - size_t* WorkIndex, - size_t* WorkRemaining - ) - { - const size_t WorkPerThread = TotalWork / WorkBlock->tids; - const size_t WorkPerThreadExtra = TotalWork % WorkBlock->tids; - - if (uint32_t(Index) < WorkPerThreadExtra) { - *WorkIndex = (WorkPerThread + 1) * Index; - *WorkRemaining = WorkPerThread + 1; - } else { - *WorkIndex = WorkPerThread * Index + WorkPerThreadExtra; - *WorkRemaining = WorkPerThread; - } - } }; constexpr size_t MLAS_NCHWC_NN_ALGORITHM::HeightShapeIndex; @@ -573,7 +551,7 @@ struct MLAS_NCHWC_GROUPED_CONV_ALGORITHM : MLAS_NCHWC_CONV_ALGORITHM size_t WorkIndex; - PartitionWork(Index, WorkBlock, TotalWork, &WorkIndex, &WorkRemaining); + MlasPartitionWork(Index, WorkBlock->tids, TotalWork, &WorkIndex, &WorkRemaining); // // Extract the current batch, group, filter cluster, and output line @@ -1004,7 +982,7 @@ struct MLAS_NCHWC_CONV_DEPTHWISE_ALGORITHM : MLAS_NCHWC_CONV_ALGORITHM size_t WorkIndex; size_t WorkRemaining; - PartitionWork(Index, WorkBlock, TotalWork, &WorkIndex, &WorkRemaining); + MlasPartitionWork(Index, WorkBlock->tids, TotalWork, &WorkIndex, &WorkRemaining); // // Extract the current batch, group block, and output line from the @@ -1138,7 +1116,7 @@ struct MLAS_NCHWC_POOL_ALGORITHM : MLAS_NCHWC_NN_ALGORITHM size_t WorkIndex; size_t WorkRemaining; - PartitionWork(Index, WorkBlock, TotalWork, &WorkIndex, &WorkRemaining); + MlasPartitionWork(Index, WorkBlock->tids, TotalWork, &WorkIndex, &WorkRemaining); size_t ph = WorkIndex % OutputHeight; const size_t BatchChannel = WorkIndex / OutputHeight; diff --git a/onnxruntime/core/mlas/lib/threading.cpp b/onnxruntime/core/mlas/lib/threading.cpp index ef30de9499..ab797e4312 100644 --- a/onnxruntime/core/mlas/lib/threading.cpp +++ b/onnxruntime/core/mlas/lib/threading.cpp @@ -46,7 +46,6 @@ MlasExecuteThreaded( } #endif - // // Fallback to OpenMP or a serialized implementation. // diff --git a/onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S b/onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S index 3d637fb97e..6a6edd7e80 100644 --- a/onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S +++ b/onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S @@ -130,6 +130,8 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackAAvx2): vpxor xmm1,xmm1,xmm1 vpxor xmm2,xmm2,xmm2 vpxor xmm3,xmm3,xmm3 + lea r13,[r10+r10*2] # compute ldb * 3 + lea rax,[r12+r12*2] # compute output stride * 3 mov rdx,rsi mov rcx,rdi lea rsi,[rsi+r10*4] # advance next matrix A by 4 rows @@ -139,16 +141,14 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackAAvx2): jb .LCopyPackA.ProcessRemainingColumnsM4 .LCopyPackA.ProcessNextColumnLoopM4: - lea rax,[rdx+r10*2] # compute matrix A plus 2 rows vmovdqu ymm4,YMMWORD PTR [rdx] vmovdqu ymm5,YMMWORD PTR [rdx+r10] - vmovdqu ymm6,YMMWORD PTR [rax] - vmovdqu ymm7,YMMWORD PTR [rax+r10] - lea rax,[rcx+r12*2] # compute matrix D plus 2 rows + vmovdqu ymm6,YMMWORD PTR [rdx+r10*2] + vmovdqu ymm7,YMMWORD PTR [rdx+r13] vmovdqu YMMWORD PTR [rcx],ymm4 vmovdqu YMMWORD PTR [rcx+r12],ymm5 - vmovdqu YMMWORD PTR [rax],ymm6 - vmovdqu YMMWORD PTR [rax+r12],ymm7 + vmovdqu YMMWORD PTR [rcx+r12*2],ymm6 + vmovdqu YMMWORD PTR [rcx+rax],ymm7 vpmaddubsw ymm4,ymm4,ymm9 # horizontal byte+byte=word per row vpaddw ymm0,ymm0,ymm4 # add words to row accumulators vpmaddubsw ymm5,ymm5,ymm9 @@ -167,16 +167,14 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackAAvx2): jz .LCopyPackA.ReduceRowSumVectorM4 test bl,16 # (CountK & 16) != 0? jz .LCopyPackA.CopyRemainingCountKLessThan16M4 - lea rax,[rdx+r10*2] # compute matrix A plus 2 rows vmovdqu xmm4,XMMWORD PTR [rdx] vmovdqu xmm5,XMMWORD PTR [rdx+r10] - vmovdqu xmm6,XMMWORD PTR [rax] - vmovdqu xmm7,XMMWORD PTR [rax+r10] - lea rax,[rcx+r12*2] # compute matrix D plus 2 rows + vmovdqu xmm6,XMMWORD PTR [rdx+r10*2] + vmovdqu xmm7,XMMWORD PTR [rdx+r13] vmovdqu XMMWORD PTR [rcx],xmm4 vmovdqu XMMWORD PTR [rcx+r12],xmm5 - vmovdqu XMMWORD PTR [rax],xmm6 - vmovdqu XMMWORD PTR [rax+r12],xmm7 + vmovdqu XMMWORD PTR [rcx+r12*2],xmm6 + vmovdqu XMMWORD PTR [rcx+rax],xmm7 vpmaddubsw xmm4,xmm4,xmm9 # horizontal byte+byte=word per row vpaddw ymm0,ymm0,ymm4 # add words to row accumulators vpmaddubsw xmm5,xmm5,xmm9 @@ -198,14 +196,13 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackAAvx2): lea rbp,.LGemmU8S8CopyPackAFrame_PaddedMatrixAData[rsp] test bl,8 # (CountK & 8) != 0? jz .LCopyPackA.CopyRemainingCountKLessThan8M4 - lea r13,[rdx+r10*2] # compute matrix A plus 2 rows mov rax,QWORD PTR [rdx] mov QWORD PTR [rbp],rax mov rax,QWORD PTR [rdx+r10] mov QWORD PTR [rbp+16],rax - mov rax,QWORD PTR [r13] + mov rax,QWORD PTR [rdx+r10*2] mov QWORD PTR [rbp+32],rax - mov rax,QWORD PTR [r13+r10] + mov rax,QWORD PTR [rdx+r13] mov QWORD PTR [rbp+48],rax add rdx,8 add rbp,8 # advance padded buffer destination @@ -213,14 +210,13 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackAAvx2): .LCopyPackA.CopyRemainingCountKLessThan8M4: test bl,4 # (CountK & 4) != 0? jz .LCopyPackA.CopyRemainingCountKLessThan4M4 - lea r13,[rdx+r10*2] # compute matrix A plus 2 rows mov eax,DWORD PTR [rdx] mov DWORD PTR [rbp],eax mov eax,DWORD PTR [rdx+r10] mov DWORD PTR [rbp+16],eax - mov eax,DWORD PTR [r13] + mov eax,DWORD PTR [rdx+r10*2] mov DWORD PTR [rbp+32],eax - mov eax,DWORD PTR [r13+r10] + mov eax,DWORD PTR [rdx+r13] mov DWORD PTR [rbp+48],eax add rdx,4 add rbp,4 # advance padded buffer destination @@ -228,14 +224,13 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackAAvx2): .LCopyPackA.CopyRemainingCountKLessThan4M4: test bl,2 # (CountK & 2) != 0? jz .LCopyPackA.CopyRemainingCountKLessThan2M4 - lea r13,[rdx+r10*2] # compute matrix A plus 2 rows movzx eax,WORD PTR [rdx] mov WORD PTR [rbp],ax movzx eax,WORD PTR [rdx+r10] mov WORD PTR [rbp+16],ax - movzx eax,WORD PTR [r13] + movzx eax,WORD PTR [rdx+r10*2] mov WORD PTR [rbp+32],ax - movzx eax,WORD PTR [r13+r10] + movzx eax,WORD PTR [rdx+r13] mov WORD PTR [rbp+48],ax add rdx,2 add rbp,2 # advance padded buffer destination @@ -243,14 +238,13 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackAAvx2): .LCopyPackA.CopyRemainingCountKLessThan2M4: test bl,1 # (CountK & 1) != 0? jz .LCopyPackA.ProcessPaddedMatrixADataM4 - lea r13,[rdx+r10*2] # compute matrix A plus 2 rows movzx eax,BYTE PTR [rdx] mov BYTE PTR [rbp],al movzx eax,BYTE PTR [rdx+r10] mov BYTE PTR [rbp+16],al - movzx eax,BYTE PTR [r13] + movzx eax,BYTE PTR [rdx+r10*2] mov BYTE PTR [rbp+32],al - movzx eax,BYTE PTR [r13+r10] + movzx eax,BYTE PTR [rdx+r13] mov BYTE PTR [rbp+48],al // @@ -446,6 +440,7 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): push rbx mov r10,rdx + lea r11,[r10+r10*2] # compute ldb * 3 vpbroadcastw ymm7,WORD PTR .LGemmU8S8CopyPackBFrame_offa[rsp] vpcmpeqw ymm8,ymm8,ymm8 # generate word vector [0xFFFF] vpsrlw ymm8,ymm8,15 # generate word vector [0x0001] @@ -469,11 +464,10 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): jb .LCopyPackB.ProcessRemainingRowsN16 .LCopyPackB.ProcessNextRowLoopN16: - lea rax,[rdx+r10*2] # compute matrix B plus 2 rows vmovdqu xmm2,XMMWORD PTR [rdx] # load 4 rows vmovdqu xmm3,XMMWORD PTR [rdx+r10] - vmovdqu xmm4,XMMWORD PTR [rax] - vmovdqu xmm5,XMMWORD PTR [rax+r10] + vmovdqu xmm4,XMMWORD PTR [rdx+r10*2] + vmovdqu xmm5,XMMWORD PTR [rdx+r11] lea rdx,[rdx+r10*4] # advance matrix B by 4 rows .LCopyPackB.InterleaveRowDataN16: @@ -558,14 +552,13 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): lea rbp,.LGemmU8S8CopyPackBFrame_PaddedMatrixBData[rsp] test cl,8 # (CountN & 8) != 0? jz .LCopyPackB.CopyRemainingCountNLessThan8K4 - lea r11,[rdx+r10*2] # compute matrix B plus 2 rows mov rax,QWORD PTR [rdx] mov QWORD PTR [rbp],rax mov rax,QWORD PTR [rdx+r10] mov QWORD PTR [rbp+16],rax - mov rax,QWORD PTR [r11] + mov rax,QWORD PTR [rdx+r10*2] mov QWORD PTR [rbp+32],rax - mov rax,QWORD PTR [r11+r10] + mov rax,QWORD PTR [rdx+r11] mov QWORD PTR [rbp+48],rax add rdx,8 # advance matrix B add rbp,8 # advance padded buffer destination @@ -573,14 +566,13 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): .LCopyPackB.CopyRemainingCountNLessThan8K4: test cl,4 # (CountN & 4) != 0? jz .LCopyPackB.CopyRemainingCountNLessThan4K4 - lea r11,[rdx+r10*2] # compute matrix B plus 2 rows mov eax,DWORD PTR [rdx] mov DWORD PTR [rbp],eax mov eax,DWORD PTR [rdx+r10] mov DWORD PTR [rbp+16],eax - mov eax,DWORD PTR [r11] + mov eax,DWORD PTR [rdx+r10*2] mov DWORD PTR [rbp+32],eax - mov eax,DWORD PTR [r11+r10] + mov eax,DWORD PTR [rdx+r11] mov DWORD PTR [rbp+48],eax add rdx,4 # advance matrix B add rbp,4 # advance padded buffer destination @@ -588,14 +580,13 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): .LCopyPackB.CopyRemainingCountNLessThan4K4: test cl,2 # (CountN & 2) != 0? jz .LCopyPackB.CopyRemainingCountNLessThan2K4 - lea r11,[rdx+r10*2] # compute matrix B plus 2 rows movzx eax,WORD PTR [rdx] mov WORD PTR [rbp],ax movzx eax,WORD PTR [rdx+r10] mov WORD PTR [rbp+16],ax - movzx eax,WORD PTR [r11] + movzx eax,WORD PTR [rdx+r10*2] mov WORD PTR [rbp+32],ax - movzx eax,WORD PTR [r11+r10] + movzx eax,WORD PTR [rdx+r11] mov WORD PTR [rbp+48],ax add rdx,2 # advance matrix B add rbp,2 # advance padded buffer destination @@ -603,14 +594,13 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): .LCopyPackB.CopyRemainingCountNLessThan2K4: test cl,1 # (CountN & 1) != 0? jz .LCopyPackB.ProcessPaddedMatrixBData - lea r11,[rdx+r10*2] # compute matrix B plus 2 rows movzx eax,BYTE PTR [rdx] mov BYTE PTR [rbp],al movzx eax,BYTE PTR [rdx+r10] mov BYTE PTR [rbp+16],al - movzx eax,BYTE PTR [r11] + movzx eax,BYTE PTR [rdx+r10*2] mov BYTE PTR [rbp+32],al - movzx eax,BYTE PTR [r11+r10] + movzx eax,BYTE PTR [rdx+r11] mov BYTE PTR [rbp+48],al .LCopyPackB.ProcessPaddedMatrixBData: diff --git a/onnxruntime/core/providers/cpu/math/gemm.h b/onnxruntime/core/providers/cpu/math/gemm.h index a309ad144a..227bf46c02 100644 --- a/onnxruntime/core/providers/cpu/math/gemm.h +++ b/onnxruntime/core/providers/cpu/math/gemm.h @@ -8,7 +8,6 @@ #include "core/util/math.h" #include "core/util/math_cpuonly.h" #include "gemm_helper.h" -#include "core/framework/op_kernel_context_internal.h" namespace onnxruntime { @@ -28,7 +27,7 @@ class Gemm : public OpKernel { } Status Compute(OpKernelContext* context) const override { - concurrency::ThreadPool* tp = context->GetOperatorThreadPool(); + concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool(); const auto* X = context->Input(0); const auto* W = context->Input(1); @@ -82,7 +81,7 @@ class Gemm : public OpKernel { // but passing 0 for beta is cheaper and it will ignore any junk in the output buffer B != nullptr ? beta_ : 0, y_data, - tp); + thread_pool); FuseActivation(activation_, y_data, M * N, leaky_relu_alpha_); diff --git a/onnxruntime/core/providers/cpu/math/matmul.cc b/onnxruntime/core/providers/cpu/math/matmul.cc index cff1b1066a..5e4c5fb2a5 100644 --- a/onnxruntime/core/providers/cpu/math/matmul.cc +++ b/onnxruntime/core/providers/cpu/math/matmul.cc @@ -1,8 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/framework/op_kernel_context_internal.h" -#include "core/providers/cpu/math/matmul.h" +#include "core/providers/cpu/math/matmul.h" #include "core/util/math.h" #include "core/util/math_cpuonly.h" #include "matmul_helper.h" diff --git a/onnxruntime/core/providers/cpu/math/matmul_integer.cc b/onnxruntime/core/providers/cpu/math/matmul_integer.cc index b76470e098..26bfaaa013 100644 --- a/onnxruntime/core/providers/cpu/math/matmul_integer.cc +++ b/onnxruntime/core/providers/cpu/math/matmul_integer.cc @@ -2,7 +2,6 @@ // Licensed under the MIT License. #include "core/framework/data_types_internal.h" -#include "core/framework/op_kernel_context_internal.h" #include "core/providers/cpu/math/matmul_integer.h" #include "core/providers/cpu/math/matmul_helper.h" #include "core/util/qmath.h" diff --git a/onnxruntime/core/providers/cpu/nn/conv.cc b/onnxruntime/core/providers/cpu/nn/conv.cc index f7023a98d3..e213db54fc 100644 --- a/onnxruntime/core/providers/cpu/nn/conv.cc +++ b/onnxruntime/core/providers/cpu/nn/conv.cc @@ -78,7 +78,7 @@ Status Conv::Compute(OpKernelContext* context) const { output_shape.GetDims().end()); const size_t kernel_rank = kernel_shape.size(); - concurrency::ThreadPool* tp = context->GetOperatorThreadPool(); + concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool(); for (int image_id = 0; image_id < N; ++image_id) { for (int group_id = 0; group_id < conv_attrs_.group; ++group_id) { @@ -125,7 +125,7 @@ Status Conv::Compute(OpKernelContext* context) const { col_buffer_data, 0, Ydata + group_id * Y_offset, - tp); + thread_pool); } if (B != nullptr) { @@ -186,7 +186,7 @@ Status Conv::Compute(OpKernelContext* context) const { auto* Ydata = Y->template MutableData(); const size_t kernel_rank = kernel_shape.size(); - concurrency::ThreadPool* tp = context->GetOperatorThreadPool(); + concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool(); if (kernel_rank == 2 || kernel_rank == 3) { MLAS_CONV_PARAMETERS Parameters; @@ -205,7 +205,7 @@ Status Conv::Compute(OpKernelContext* context) const { static_cast(M / conv_attrs_.group), &activation_, &WorkingBufferSize, - tp); + thread_pool); auto working_data = WorkingBufferSize > 0 ? alloc->Alloc(sizeof(float) * WorkingBufferSize) : nullptr; BufferUniquePtr working_buffer(working_data, BufferDeleter(alloc)); @@ -216,7 +216,7 @@ Status Conv::Compute(OpKernelContext* context) const { Bdata, static_cast(working_buffer.get()), Ydata, - tp); + thread_pool); } else { const int64_t input_image_size = input_shape.Size(); const int64_t output_image_size = output_shape.Size(); @@ -262,7 +262,7 @@ Status Conv::Compute(OpKernelContext* context) const { col_buffer_data, 0, Ydata + group_id * Y_offset, - tp); + thread_pool); } MlasActivation(&activation_, Ydata, Bdata, M, output_image_size, output_image_size); diff --git a/onnxruntime/core/providers/cpu/nn/conv_integer.cc b/onnxruntime/core/providers/cpu/nn/conv_integer.cc index 7e5951d4ec..53ce6e1070 100644 --- a/onnxruntime/core/providers/cpu/nn/conv_integer.cc +++ b/onnxruntime/core/providers/cpu/nn/conv_integer.cc @@ -90,6 +90,7 @@ Status ConvInteger::Compute(OpKernelContext* context) const { output_shape.GetDims().end()); const size_t kernel_rank = kernel_shape.size(); + concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool(); for (int image_id = 0; image_id < N; ++image_id) { for (int group_id = 0; group_id < conv_attrs_.group; ++group_id) { @@ -139,7 +140,7 @@ Status ConvInteger::Compute(OpKernelContext* context) const { input_offset, Ydata + group_id * Y_offset, static_cast(output_image_size), - nullptr); + thread_pool); } Xdata += X_offset * conv_attrs_.group; diff --git a/onnxruntime/core/providers/cpu/nn/conv_transpose.cc b/onnxruntime/core/providers/cpu/nn/conv_transpose.cc index e025c12979..ab763b8d9f 100644 --- a/onnxruntime/core/providers/cpu/nn/conv_transpose.cc +++ b/onnxruntime/core/providers/cpu/nn/conv_transpose.cc @@ -16,7 +16,6 @@ /* Modifications Copyright (c) Microsoft. */ #include "core/providers/cpu/nn/conv_transpose.h" -#include "core/framework/op_kernel_context_internal.h" #include "core/util/math.h" #include "core/util/math_cpuonly.h" @@ -42,7 +41,7 @@ Status ConvTranspose::Compute(OpKernelContext* context) const { template Status ConvTranspose::DoConvTranspose(OpKernelContext* context, bool dynamic_padding) const { - concurrency::ThreadPool* tp = context->GetOperatorThreadPool(); + concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool(); size_t num_inputs = OpKernel::Node().InputDefs().size(); ConvTransposeAttributes::Prepare p; @@ -87,7 +86,7 @@ Status ConvTranspose::DoConvTranspose(OpKernelContext* context, bool dynamic_ Xdata + group_id * X_offset, 0, col_buffer_data, - tp); + thread_pool); // Col2im math::Col2im( @@ -136,7 +135,7 @@ Status ConvTranspose::DoConvTranspose(OpKernelContext* context, bool dynamic_ Xdata + group_id * X_offset, 0, col_buffer_data, - tp); + thread_pool); // Col2im math::Col2imNd( @@ -155,7 +154,7 @@ Status ConvTranspose::DoConvTranspose(OpKernelContext* context, bool dynamic_ } if (p.B != nullptr) { - + auto Ymatrix = EigenMatrixMap(Ydata, output_size, p.num_output_channels); auto Bvec = ConstEigenVectorMap(p.B->template Data(), p.num_output_channels); Ymatrix.rowwise() += Bvec.transpose(); diff --git a/onnxruntime/core/providers/cpu/nn/pool.cc b/onnxruntime/core/providers/cpu/nn/pool.cc index ec7a5904c7..8bf630eb05 100644 --- a/onnxruntime/core/providers/cpu/nn/pool.cc +++ b/onnxruntime/core/providers/cpu/nn/pool.cc @@ -1,7 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/framework/op_kernel_context_internal.h" #include "core/providers/cpu/nn/pool.h" using namespace ::onnxruntime::common;