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

MLAS: enable threading for quantized GEMMs (#2844)

Browse source
This commit is contained in:
Tracy Sharpe 2020-01-15 19:25:40 -08:00 committed by GitHub
parent 5db8543018
commit 928b6bb210
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 707 additions and 271 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
onnxruntime/contrib_ops/cpu/nchwc_ops.cc
View file

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

68
onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm
View file

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

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

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

734
onnxruntime/core/mlas/lib/qgemm.cpp
View file

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

30
onnxruntime/core/mlas/lib/snchwc.cpp
View file

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

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

@ -46,7 +46,6 @@ MlasExecuteThreaded(
}
#endif
//
// Fallback to OpenMP or a serialized implementation.
//

68
onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S
View file

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

5
onnxruntime/core/providers/cpu/math/gemm.h
View file

@ -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<Tensor>(0);
const auto* W = context->Input<Tensor>(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<T>(activation_, y_data, M * N, leaky_relu_alpha_);

3
onnxruntime/core/providers/cpu/math/matmul.cc
View file

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

1
onnxruntime/core/providers/cpu/math/matmul_integer.cc
View file

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

12
onnxruntime/core/providers/cpu/nn/conv.cc
View file

@ -78,7 +78,7 @@ Status Conv<T>::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<T>::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<float>::Compute(OpKernelContext* context) const {
auto* Ydata = Y->template MutableData<float>();
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<float>::Compute(OpKernelContext* context) const {
static_cast<size_t>(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<float>::Compute(OpKernelContext* context) const {
Bdata,
static_cast<float*>(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<float>::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);

3
onnxruntime/core/providers/cpu/nn/conv_integer.cc
View file

@ -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<int>(output_image_size),
nullptr);
thread_pool);
}
Xdata += X_offset * conv_attrs_.group;

9
onnxruntime/core/providers/cpu/nn/conv_transpose.cc
View file

@ -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<T>::Compute(OpKernelContext* context) const {
template <typename T>
Status ConvTranspose<T>::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<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_
Xdata + group_id * X_offset,
0,
col_buffer_data,
tp);
thread_pool);
// Col2im
math::Col2im<T, CPUMathUtil, StorageOrder::NCHW>(
@ -136,7 +135,7 @@ Status ConvTranspose<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_
Xdata + group_id * X_offset,
0,
col_buffer_data,
tp);
thread_pool);
// Col2im
math::Col2imNd<T, CPUMathUtil, StorageOrder::NCHW>(
@ -155,7 +154,7 @@ Status ConvTranspose<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_
}
if (p.B != nullptr) {
auto Ymatrix = EigenMatrixMap<T>(Ydata, output_size, p.num_output_channels);
auto Bvec = ConstEigenVectorMap<T>(p.B->template Data<T>(), p.num_output_channels);
Ymatrix.rowwise() += Bvec.transpose();

1
onnxruntime/core/providers/cpu/nn/pool.cc
View file

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