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

MLAS: add sgemm weight prepacking (#5183)

Browse source 
Add support to MLAS to prepack weights for the float GEMM. Support for prepacking has been added to MatMul and Attention for this release.
This commit is contained in:
Tracy Sharpe 2020-09-16 08:36:27 -07:00 committed by edgchen1
parent ecf04d23c4
commit b2994492af
11 changed files with 845 additions and 237 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

151
onnxruntime/contrib_ops/cpu/bert/attention.cc
View file

@ -1,7 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "attention.h"
#include "attention_cpu_base.h"
#include "attention_helper.h"
#include "core/framework/tensorprotoutils.h"
#include "core/graph/onnx_protobuf.h"
@ -14,13 +14,33 @@ using onnxruntime::concurrency::ThreadPool;
namespace onnxruntime {
namespace contrib {
// These ops are internal-only, so register outside of onnx
#define REGISTER_KERNEL_TYPED(T) \
ONNX_OPERATOR_TYPED_KERNEL_EX(Attention, kMSDomain, 1, T, kCpuExecutionProvider, \
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
Attention<T>);
REGISTER_KERNEL_TYPED(float)
template <typename T>
class Attention : public OpKernel, public AttentionCPUBase {
public:
explicit Attention(const OpKernelInfo& info);
Status Compute(OpKernelContext* context) const override;
#if !defined(USE_MKLML_FOR_BLAS)
Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override;
#endif
private:
BufferUniquePtr packed_weights_;
size_t packed_weights_size_;
TensorShape weight_shape_;
};
// These ops are internal-only, so register outside of onnx
ONNX_OPERATOR_TYPED_KERNEL_EX(
Attention,
kMSDomain,
1,
float,
kCpuExecutionProvider,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
Attention<float>);
AttentionBase::AttentionBase(const OpKernelInfo& info) {
int64_t num_heads = 0;
@ -116,8 +136,7 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape,
if ((static_cast<int>(mask_dims[0]) == batch_size || static_cast<int>(mask_dims[0]) == 1) && static_cast<int>(mask_dims[1]) == 1) {
// Mask will have same value after propogation, which has same effect as no mask.
mask_index = nullptr;
}
else {
} else {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'mask_index' with raw attention mask shall have shape batch_size x (past_sequence_length + sequence_length)");
}
}
@ -159,15 +178,68 @@ template <typename T>
Attention<T>::Attention(const OpKernelInfo& info) : OpKernel(info), AttentionCPUBase(info) {
}
#if !defined(USE_MKLML_FOR_BLAS)
template <typename T>
Status Attention<T>::PrePack(const Tensor& weights, int input_idx, bool& is_packed) {
is_packed = false;
if (1 != input_idx) {
return Status::OK();
}
weight_shape_ = weights.Shape();
const auto& weights_dims = weight_shape_.GetDims();
if (weights_dims.size() != 2) {
return Status::OK();
}
const size_t hidden_size = static_cast<size_t>(weights_dims[0]);
const size_t hidden_size_x3 = static_cast<size_t>(weights_dims[1]);
const size_t head_size = hidden_size / num_heads_;
// Bail out if the weights shape has an expected shape.
if ((hidden_size == 0) || ((hidden_size % num_heads_) != 0) || (hidden_size_x3 != 3 * hidden_size)) {
return Status::OK();
}
const auto* weights_data = weights.Data<T>();
packed_weights_size_ = MlasGemmPackBSize(head_size, hidden_size);
if (packed_weights_size_ == 0) {
return Status::OK();
}
const size_t loop_len = 3 * num_heads_;
auto alloc = Info().GetAllocator(0, OrtMemTypeDefault);
auto* packed_weights_data = static_cast<uint8_t*>(alloc->Alloc(packed_weights_size_ * loop_len));
packed_weights_ = BufferUniquePtr(packed_weights_data, BufferDeleter(alloc));
for (size_t i = 0; i < loop_len; i++) {
MlasGemmPackB(CblasNoTrans, head_size, hidden_size, weights_data, hidden_size_x3, packed_weights_data);
packed_weights_data += packed_weights_size_;
weights_data += head_size;
}
is_packed = true;
return Status::OK();
}
#endif
template <typename T>
Status Attention<T>::Compute(OpKernelContext* context) const {
const Tensor* input = context->Input<Tensor>(0);
const Tensor* weights = context->Input<Tensor>(1);
const Tensor* weights = packed_weights_ ? nullptr : context->Input<Tensor>(1);
const Tensor* bias = context->Input<Tensor>(2);
const Tensor* mask_index = context->Input<Tensor>(3);
const Tensor* past = context->Input<Tensor>(4);
ORT_RETURN_IF_ERROR(CheckInputs(input->Shape(), weights->Shape(), bias->Shape(), mask_index, past));
ORT_RETURN_IF_ERROR(CheckInputs(input->Shape(),
packed_weights_ ? weight_shape_ : weights->Shape(),
bias->Shape(),
mask_index,
past));
const auto& shape = input->Shape().GetDims();
const int batch_size = static_cast<int>(shape[0]);
@ -187,17 +259,17 @@ Status Attention<T>::Compute(OpKernelContext* context) const {
// gemm_data(BS, 3NH) = input(BS, NH) x weights(NH, 3NH) + bias(3NH)
auto gemm_data = allocator->Alloc(SafeInt<size_t>(batch_size) * sequence_length * 3 * hidden_size * element_size);
BufferUniquePtr gemm_buffer(gemm_data, BufferDeleter(allocator));
auto Q = reinterpret_cast<T*>(gemm_data);
auto K = Q + batch_size * sequence_length * hidden_size;
auto V = K + batch_size * sequence_length * hidden_size;
T* QKV[3] = {Q, K, V};
{
const int loop_len = 3 * batch_size * num_heads_;
const auto input_data = input->template Data<T>();
const auto weights_data = weights->template Data<T>();
const auto bias_data = bias->template Data<T>();
const auto* input_data = input->template Data<T>();
const auto* weights_data = packed_weights_ ? nullptr : weights->template Data<T>();
const auto* bias_data = bias->template Data<T>();
const double cost =
static_cast<double>(sequence_length) * static_cast<double>(head_size) * static_cast<double>(hidden_size);
@ -224,22 +296,39 @@ Status Attention<T>::Compute(OpKernelContext* context) const {
// A: input (BxSxNxH) (B.)S x NH S x NH
// B: weights (NxHx3xNxH) NH x (3.N.)H NH x H
// C: QKV[qkv_index] (3xBxNxSxH) (3.B.N.)S x H S x H
math::GemmEx<float, ThreadPool>(CblasNoTrans, // TransA = no
CblasNoTrans, // TransB = no
sequence_length, // M = S
head_size, // N = H
hidden_size, // K = NH
1.0f, // alpha
input_data + input_offset, // A
hidden_size, // lda = NH
weights_data + weights_offset, // B
3 * hidden_size, // ldb = 3NH
1.0f, // beta
qkv_dest + qkv_offset, // C
head_size, // ldc
nullptr // use single-thread
);
if (packed_weights_) {
const auto* packed_weight =
static_cast<const uint8_t*>(packed_weights_.get()) + packed_weights_size_ * (weights_offset / head_size);
MlasGemm(
CblasNoTrans, // TransA = no
sequence_length, // M = S
head_size, // N = H
hidden_size, // K = NH
1.0f, // alpha
input_data + input_offset, // A
hidden_size, // lda = NH
packed_weight, // B
1.0f, // beta
qkv_dest + qkv_offset, // C
head_size, // ldc
nullptr); // use single-thread
} else {
math::GemmEx<float, ThreadPool>(CblasNoTrans, // TransA = no
CblasNoTrans, // TransB = no
sequence_length, // M = S
head_size, // N = H
hidden_size, // K = NH
1.0f, // alpha
input_data + input_offset, // A
hidden_size, // lda = NH
weights_data + weights_offset, // B
3 * hidden_size, // ldb = 3NH
1.0f, // beta
qkv_dest + qkv_offset, // C
head_size, // ldc
nullptr // use single-thread
);
}
}
});
}

21
onnxruntime/contrib_ops/cpu/bert/attention.h
View file

@ -1,21 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "attention_cpu_base.h"
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
namespace onnxruntime {
namespace contrib {
template <typename T>
class Attention : public OpKernel, public AttentionCPUBase {
public:
explicit Attention(const OpKernelInfo& info);
Status Compute(OpKernelContext* context) const override;
};
} // namespace contrib
} // namespace onnxruntime

2
onnxruntime/contrib_ops/cuda/quantization/attention_quantization.h
View file

@ -6,7 +6,7 @@
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
#include "core/providers/cuda/cuda_common.h"
#include "contrib_ops/cpu/bert/attention.h"
#include "contrib_ops/cpu/bert/attention_base.h"
namespace onnxruntime {
namespace contrib {

35
onnxruntime/core/mlas/inc/mlas.h
View file

@ -128,6 +128,23 @@ MlasGemm(
MLAS_THREADPOOL* ThreadPool
);
void
MLASCALL
MlasGemm(
CBLAS_TRANSPOSE TransA,
size_t M,
size_t N,
size_t K,
float alpha,
const float* A,
size_t lda,
const void* PackedB,
float beta,
float* C,
size_t ldc,
MLAS_THREADPOOL* ThreadPool
);
void
MLASCALL
MlasGemm(
@ -225,6 +242,24 @@ MlasGemm(
// Buffer packing routines.
//
size_t
MLASCALL
MlasGemmPackBSize(
size_t N,
size_t K
);
void
MLASCALL
MlasGemmPackB(
CBLAS_TRANSPOSE TransB,
size_t N,
size_t K,
const float* B,
size_t ldb,
void* PackedB
);
size_t
MLASCALL
MlasGemmPackBSize(

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

@ -139,6 +139,8 @@ Abstract:
#define MLAS_SGEMM_STRIDEN 128
#define MLAS_SGEMM_STRIDEK 128
#define MLAS_SGEMM_PACKED_STRIDEN 128
#define MLAS_SGEMM_PACKED_STRIDEK 256
#define MLAS_DGEMM_STRIDEN 64
#define MLAS_DGEMM_STRIDEK 128

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

@ -157,7 +157,7 @@ Return Value:
for (size_t k = 0; k < K; k += CountK) {
CountK = (std::min)(K - k, Strides.K);
CountK = std::min(K - k, Strides.K);
//
// Step through each slice of matrix B along the N dimension.
@ -167,7 +167,7 @@ Return Value:
for (size_t n = 0; n < N; n += CountN) {
CountN = (std::min)(N - n, Strides.N);
CountN = std::min(N - n, Strides.N);
//
// Copy a panel of matrix B to a local packed buffer.
@ -191,7 +191,7 @@ Return Value:
for (size_t m = 0; m < M; m += CountM) {
CountM = (std::min)(M - m, Strides.M);
CountM = std::min(M - m, Strides.M);
//
// Copy a panel of matrix A to a local packed buffer.
@ -311,7 +311,7 @@ Return Value:
for (size_t k = 0; k < K; k += CountK) {
CountK = (std::min)(K - k, Strides.K);
CountK = std::min(K - k, Strides.K);
const size_t PackedCountK = (CountK + KernelType::PackedK - 1) /
KernelType::PackedK;
@ -328,7 +328,7 @@ Return Value:
for (size_t n = 0; n < N; n += CountN) {
CountN = (std::min)(N - n, Strides.N);
CountN = std::min(N - n, Strides.N);
if (k == 0) {
MlasGemmU8X8ScaleSumBuffer(ColumnSumBuffer, PackedColumnSumBuffer + n,
@ -347,7 +347,7 @@ Return Value:
for (size_t m = 0; m < M; m += CountM) {
CountM = (std::min)(M - m, Strides.M);
CountM = std::min(M - m, Strides.M);
//
// Copy a panel of matrix A to a local packed buffer.
@ -2126,7 +2126,7 @@ Return Value:
for (size_t k = 0; k < K; k += CountK) {
CountK = (std::min)(K - k, StrideK);
CountK = std::min(K - k, StrideK);
//
// Step through each slice of matrix B along the N dimension.
@ -2141,7 +2141,7 @@ Return Value:
constexpr size_t BatchedN = 128;
MLAS_DECLSPEC_ALIGN(int32_t ColumnSumBuffer[BatchedN], 64);
CountN = (std::min)(N - n, BatchedN);
CountN = std::min(N - n, BatchedN);
if (GemmU8X8Operation == &MlasGemmU8X8PackedOperation<MLAS_GEMM_U8S8_KERNEL_AVX2>) {
MLAS_GEMM_U8S8_KERNEL_AVX2::CopyPackB(pb, B + n, ldb, CountN, CountK, ColumnSumBuffer, BIsSigned);

607
onnxruntime/core/mlas/lib/sgemm.cpp
View file

@ -32,21 +32,22 @@ Abstract:
//
struct MLAS_SGEMM_WORK_BLOCK {
int32_t ThreadCountM;
int32_t ThreadCountN;
CBLAS_TRANSPOSE TransA;
CBLAS_TRANSPOSE TransB;
size_t M;
size_t N;
size_t K;
const float* A;
size_t lda;
const float* B;
size_t ldb;
const void* PackedB;
float* C;
size_t ldc;
float alpha;
float beta;
struct SEGMENT {
size_t M;
size_t N;
const float* A;
const float* B;
float* C;
} Segments[MLAS_MAXIMUM_THREAD_COUNT];
};
void
@ -962,8 +963,8 @@ Return Value:
// the A panel needs to be used for transposing.
//
uint32_t StrideN = MLAS_SGEMM_STRIDEN;
uint32_t StrideK = MLAS_SGEMM_STRIDEK;
size_t StrideN = MLAS_SGEMM_STRIDEN;
size_t StrideK = MLAS_SGEMM_STRIDEK;
if (N >= K) {
@ -985,15 +986,10 @@ Return Value:
//
size_t CountN;
size_t CountK;
for (size_t n = 0; n < N; n += CountN) {
CountN = StrideN;
if (CountN > (N - n)) {
CountN = N - n;
}
CountN = std::min(N - n, StrideN);
//
// Multiply the output matrix by beta as needed.
@ -1007,15 +1003,12 @@ Return Value:
// Step through each slice of matrix B along the K dimension.
//
size_t CountK;
bool ZeroMode = (beta == 0.0f);
for (size_t k = 0; k < K; k += CountK) {
CountK = StrideK;
if (CountK > (K - k)) {
CountK = K - k;
}
CountK = std::min(K - k, StrideK);
//
// Copy or transpose a panel of matrix B to a local packed buffer.
@ -1048,11 +1041,7 @@ Return Value:
// Transpose elements from matrix A into a local buffer.
//
size_t RowsTransposed = RowsRemaining;
if (RowsTransposed > MLAS_SGEMM_TRANSA_ROWS) {
RowsTransposed = MLAS_SGEMM_TRANSA_ROWS;
}
size_t RowsTransposed = std::min(RowsRemaining, size_t(MLAS_SGEMM_TRANSA_ROWS));
MlasSgemmTransposeA(PanelA, a, lda, RowsTransposed, CountK);
@ -1074,9 +1063,145 @@ Return Value:
}
void
MlasSgemmOperationThreaded(
MlasSgemmPackedOperation(
CBLAS_TRANSPOSE TransA,
size_t M,
size_t RangeStartN,
size_t RangeCountN,
size_t K,
float alpha,
const float* A,
size_t lda,
const void* PackedB,
size_t AlignedN,
float beta,
float* C,
size_t ldc
)
/*++
Routine Description:
This routine implements the single precision matrix/matrix multiply
operation (SGEMM).
Arguments:
TransA - Supplies the transpose operation for matrix A.
M - Supplies the number of rows of matrix A and matrix C.
RangeStartN - Supplies the starting column from packed matrix B.
RangeCountN - 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.
alpha - Supplies the scalar alpha multiplier (see SGEMM definition).
A - Supplies the address of matrix A.
lda - Supplies the first dimension of matrix A.
PackedB - Supplies the address of packed matrix B.
AlignedN - Supplies the total number of aligned columns for packed matrix B.
ldb - Supplies the first dimension of matrix B.
beta - Supplies the scalar beta multiplier (see SGEMM definition).
C - Supplies the address of matrix C.
ldc - Supplies the first dimension of matrix C.
Return Value:
None.
--*/
{
float PanelA[MLAS_SGEMM_TRANSA_ROWS * MLAS_SGEMM_PACKED_STRIDEK];
//
// Step through each slice of matrix B along the N dimension.
//
size_t CountN;
for (size_t n = 0; n < RangeCountN; n += CountN) {
const size_t SliceStartN = RangeStartN + n;
CountN = std::min(RangeCountN - n, size_t(MLAS_SGEMM_PACKED_STRIDEN));
//
// Multiply the output matrix by beta as needed.
//
if (beta != 0.0f && beta != 1.0f) {
MlasSgemmMultiplyBeta(C + n, M, CountN, ldc, beta);
}
//
// Step through each slice of matrix B along the K dimension.
//
size_t CountK;
bool ZeroMode = (beta == 0.0f);
for (size_t k = 0; k < K; k += CountK) {
CountK = std::min(K - k, size_t(MLAS_SGEMM_PACKED_STRIDEK));
//
// Step through each slice of matrix A along the M dimension.
//
const float* pb = (const float*)PackedB + AlignedN * k + CountK * SliceStartN;
float* c = C + n;
if (TransA == CblasNoTrans) {
MlasSgemmKernelLoop(A + k, pb, c, CountK, M, CountN, lda, ldc, alpha, ZeroMode);
} else {
const float* a = A + k * lda;
size_t RowsRemaining = M;
do {
//
// Transpose elements from matrix A into a local buffer.
//
size_t RowsTransposed = std::min(RowsRemaining, size_t(MLAS_SGEMM_TRANSA_ROWS));
MlasSgemmTransposeA(PanelA, a, lda, RowsTransposed, CountK);
RowsRemaining -= RowsTransposed;
a += RowsTransposed;
//
// Step through the rows of the local buffer.
//
c = MlasSgemmKernelLoop(PanelA, pb, c, CountK, RowsTransposed, CountN, CountK, ldc, alpha, ZeroMode);
} while (RowsRemaining > 0);
}
ZeroMode = false;
}
}
}
void
MlasSgemmThreaded(
void* Context,
int32_t Index
int32_t ThreadId
)
/*++
@ -1089,7 +1214,7 @@ Arguments:
Context - Supplies the pointer to the context for the threaded operation.
Index - Supplies the current index of the threaded operation.
ThreadId - Supplies the current index of the threaded operation.
Return Value:
@ -1097,89 +1222,111 @@ Return Value:
--*/
{
MLAS_SGEMM_WORK_BLOCK* WorkBlock = (MLAS_SGEMM_WORK_BLOCK*)Context;
const auto* WorkBlock = (MLAS_SGEMM_WORK_BLOCK*)Context;
MLAS_SGEMM_WORK_BLOCK::SEGMENT* Segment = &WorkBlock->Segments[Index];
const int32_t ThreadCountM = WorkBlock->ThreadCountM;
const int32_t ThreadCountN = WorkBlock->ThreadCountN;
MlasSgemmOperation(WorkBlock->TransA, WorkBlock->TransB, Segment->M,
Segment->N, WorkBlock->K, WorkBlock->alpha, Segment->A, WorkBlock->lda,
Segment->B, WorkBlock->ldb, WorkBlock->beta, Segment->C,
WorkBlock->ldc);
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 RangeStartM;
size_t RangeCountM;
MlasPartitionWork(ThreadIdM, ThreadCountM, M, &RangeStartM, &RangeCountM);
//
// Partition the operation along the N dimension.
//
size_t N = WorkBlock->N;
size_t RangeStartN;
size_t RangeCountN;
const size_t BlockedN = (N + MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1) /
MLAS_SGEMM_STRIDEN_THREAD_ALIGN;
MlasPartitionWork(ThreadIdN, ThreadCountN, BlockedN, &RangeStartN,
&RangeCountN);
RangeStartN *= MLAS_SGEMM_STRIDEN_THREAD_ALIGN;
RangeCountN *= MLAS_SGEMM_STRIDEN_THREAD_ALIGN;
RangeCountN = std::min(N - RangeStartN, RangeCountN);
//
// Dispatch the partitioned operation.
//
CBLAS_TRANSPOSE TransA = WorkBlock->TransA;
const size_t lda = WorkBlock->lda;
const size_t ldc = WorkBlock->ldc;
const float* A = WorkBlock->A + RangeStartM * ((TransA == CblasNoTrans) ? lda : 1);
float* C = WorkBlock->C + RangeStartM * ldc + RangeStartN;
if (WorkBlock->B != nullptr) {
CBLAS_TRANSPOSE TransB = WorkBlock->TransB;
const size_t ldb = WorkBlock->ldb;
const float* B = WorkBlock->B + RangeStartN * ((TransB == CblasNoTrans) ? 1 : ldb);
MlasSgemmOperation(TransA, TransB, RangeCountM, RangeCountN, WorkBlock->K,
WorkBlock->alpha, A, lda, B, ldb, WorkBlock->beta, C, ldc);
} else {
MlasSgemmPackedOperation(TransA, RangeCountM, RangeStartN, RangeCountN,
WorkBlock->K, WorkBlock->alpha, A, lda, WorkBlock->PackedB,
BlockedN * MLAS_SGEMM_STRIDEN_THREAD_ALIGN, WorkBlock->beta, C, ldc);
}
}
inline
bool
MlasSgemmTryMultithread(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t M,
size_t N,
size_t K,
float alpha,
const float* A,
size_t lda,
const float* B,
size_t ldb,
float beta,
float* C,
size_t ldc,
void
MlasSgemmSchedule(
MLAS_SGEMM_WORK_BLOCK* WorkBlock,
MLAS_THREADPOOL* ThreadPool
)
/*++
Routine Description:
This routine attempts to launch a single precision matrix/matrix multiply
operation (SGEMM) across multiple threads.
This routine schedules the single precision matrix/matrix multiply
operation (SGEMM) across one or more threads.
Arguments:
TransA - Supplies the transpose operation for matrix A.
TransB - Supplies the transpose operation for matrix B.
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.
alpha - Supplies the scalar alpha multiplier (see SGEMM definition).
A - Supplies the address of matrix A.
lda - Supplies the first dimension of matrix A.
B - Supplies the address of matrix B.
ldb - Supplies the first dimension of matrix B.
beta - Supplies the scalar beta multiplier (see SGEMM definition).
C - Supplies the address of matrix C.
ldc - Supplies the first dimension of matrix C.
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:
Returns true if the operation was completed across multiple threads, else
false if the operation should fall back to a single thread.
None.
--*/
{
MLAS_SGEMM_WORK_BLOCK WorkBlock;
int32_t TargetThreadCount;
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);
const double Complexity = double(M) * double(N) * double(K);
int32_t TargetThreadCount;
if (Complexity < double(MLAS_SGEMM_THREAD_COMPLEXITY * MLAS_MAXIMUM_THREAD_COUNT)) {
TargetThreadCount = int32_t(Complexity / double(MLAS_SGEMM_THREAD_COMPLEXITY)) + 1;
@ -1193,90 +1340,36 @@ Return Value:
TargetThreadCount = MaximumThreadCount;
}
if (TargetThreadCount == 1) {
return false;
}
//
// Initialize the common fields of the work block.
//
WorkBlock.TransA = TransA;
WorkBlock.TransB = TransB;
WorkBlock.K = K;
WorkBlock.lda = lda;
WorkBlock.ldb = ldb;
WorkBlock.ldc = ldc;
WorkBlock.alpha = alpha;
WorkBlock.beta = beta;
//
// Segment the operation across multiple threads.
//
int32_t Index = 0;
// N.B. Currently, the operation is segmented as a 1D partition, which
// works okay for operations involving skinny matrices.
//
if (N > M) {
size_t StrideN = N / TargetThreadCount;
const size_t BlockedN = (N + MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1) /
MLAS_SGEMM_STRIDEN_THREAD_ALIGN;
if ((StrideN * TargetThreadCount) != N) {
StrideN++;
if (size_t(TargetThreadCount) > BlockedN) {
TargetThreadCount = int32_t(BlockedN);
}
StrideN =
(StrideN + MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1);
size_t pldb = (TransB == CblasNoTrans) ? 1 : ldb;
for (size_t CountN, n = 0; n < N; n += CountN) {
CountN = StrideN;
if (CountN > (N - n)) {
CountN = N - n;
}
WorkBlock.Segments[Index].M = M;
WorkBlock.Segments[Index].N = CountN;
WorkBlock.Segments[Index].A = A;
WorkBlock.Segments[Index].B = B + n * pldb;
WorkBlock.Segments[Index].C = C + n;
Index++;
}
WorkBlock->ThreadCountM = 1;
WorkBlock->ThreadCountN = TargetThreadCount;
} else {
size_t StrideM = M / TargetThreadCount;
if ((StrideM * TargetThreadCount) != M) {
StrideM++;
if (size_t(TargetThreadCount) > M) {
TargetThreadCount = int32_t(M);
}
size_t plda = (TransA == CblasNoTrans) ? lda : 1;
for (size_t CountM, m = 0; m < M; m += CountM) {
CountM = StrideM;
if (CountM > (M - m)) {
CountM = M - m;
}
WorkBlock.Segments[Index].M = CountM;
WorkBlock.Segments[Index].N = N;
WorkBlock.Segments[Index].A = A + m * plda;
WorkBlock.Segments[Index].B = B;
WorkBlock.Segments[Index].C = C + m * ldc;
Index++;
}
WorkBlock->ThreadCountM = TargetThreadCount;
WorkBlock->ThreadCountN = 1;
}
MlasExecuteThreaded(MlasSgemmOperationThreaded, &WorkBlock, Index, ThreadPool);
return true;
MlasExecuteThreaded(MlasSgemmThreaded, WorkBlock, TargetThreadCount, ThreadPool);
}
void
@ -1342,12 +1435,216 @@ Return Value:
--*/
{
MLAS_SGEMM_WORK_BLOCK WorkBlock;
//
// Try to run the operation across multiple threads or fall back to a
// single thread based on the GEMM parameters and system configuration.
// Capture the GEMM parameters to the work block.
//
if (!MlasSgemmTryMultithread(TransA, TransB, M, N, K, alpha, A, lda, B, ldb, beta, C, ldc, ThreadPool)) {
MlasSgemmOperation(TransA, TransB, M, N, K, alpha, A, lda, B, ldb, beta, C, ldc);
memset(&WorkBlock, 0, sizeof(MLAS_SGEMM_WORK_BLOCK));
WorkBlock.TransA = TransA;
WorkBlock.TransB = TransB;
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.alpha = alpha;
WorkBlock.beta = beta;
//
// Schedule the operation across a set of worker threads.
//
MlasSgemmSchedule(&WorkBlock, ThreadPool);
}
void
MLASCALL
MlasGemm(
CBLAS_TRANSPOSE TransA,
size_t M,
size_t N,
size_t K,
float alpha,
const float* A,
size_t lda,
const void* PackedB,
float beta,
float* C,
size_t ldc,
MLAS_THREADPOOL* ThreadPool
)
/*++
Routine Description:
This routine implements the single precision matrix/matrix multiply
operation (SGEMM).
Arguments:
TransA - Supplies the transpose operation for matrix A.
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.
alpha - Supplies the scalar alpha multiplier (see SGEMM definition).
A - Supplies the address of matrix A.
lda - Supplies the first dimension of matrix A.
PackedB - Supplies the address of packed matrix B.
beta - Supplies the scalar beta multiplier (see SGEMM definition).
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_SGEMM_WORK_BLOCK WorkBlock;
//
// Capture the GEMM parameters to the work block.
//
memset(&WorkBlock, 0, sizeof(MLAS_SGEMM_WORK_BLOCK));
WorkBlock.TransA = TransA;
WorkBlock.M = M;
WorkBlock.N = N;
WorkBlock.K = K;
WorkBlock.A = A;
WorkBlock.lda = lda;
WorkBlock.PackedB = PackedB;
WorkBlock.C = C;
WorkBlock.ldc = ldc;
WorkBlock.alpha = alpha;
WorkBlock.beta = beta;
//
// Schedule the operation across a set of worker threads.
//
MlasSgemmSchedule(&WorkBlock, ThreadPool);
}
size_t
MLASCALL
MlasGemmPackBSize(
size_t N,
size_t K
)
/*++
Routine Description:
This routine computes the length in bytes for the packed matrix B buffer.
Arguments:
N - Supplies the number of columns of matrix B.
K - Supplies the number of rows of matrix B.
Return Value:
Returns the size in bytes for the packed matrix B buffer.
--*/
{
//
// Compute the number of bytes required to hold the packed buffer.
//
const size_t AlignedN =
(N + MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1);
const size_t BytesRequired = AlignedN * K * sizeof(float);
const size_t BufferAlignment = MlasGetPreferredBufferAlignment();
const size_t AlignedBytesRequired = (BytesRequired + BufferAlignment - 1) &
~(BufferAlignment - 1);
return AlignedBytesRequired;
}
void
MLASCALL
MlasGemmPackB(
CBLAS_TRANSPOSE TransB,
size_t N,
size_t K,
const float* B,
size_t ldb,
void* PackedB
)
/*++
Routine Description:
This routine packs the contents of matrix B to the destination buffer. The
destination buffer should be sized based on MlasGemmPackBSize(). For best
performance, the destination buffer should be aligned to the value returned
from MlasGetPreferredBufferAlignment().
Arguments:
TransB - Supplies the transpose operation for matrix B.
N - Supplies the number of columns of matrix B.
K - Supplies the number of rows of matrix B.
B - Supplies the address of matrix B.
ldb - Supplies the first dimension of matrix B.
PackedB - Supplies the address of packed matrix B.
Return Value:
None.
--*/
{
const size_t AlignedN =
(N + MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1);
//
// Step through each slice of matrix B along the K dimension.
//
size_t CountK;
for (size_t k = 0; k < K; k += CountK) {
CountK = std::min(K - k, size_t(MLAS_SGEMM_PACKED_STRIDEK));
if (TransB == CblasNoTrans) {
MlasSgemmCopyPackB((float*)PackedB, B + k * ldb, ldb, N, CountK);
} else {
MlasSgemmTransposePackB((float*)PackedB, B + k, ldb, N, CountK);
}
PackedB = (float*)PackedB + AlignedN * CountK;
}
}

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

@ -5,6 +5,7 @@
#include "core/providers/cpu/math/matmul_helper.h"
#include "core/util/math.h"
#include "core/util/math_cpuonly.h"
#include "core/mlas/inc/mlas.h"
namespace onnxruntime {
@ -16,6 +17,23 @@ class MatMul final : public OpKernel {
Status Compute(OpKernelContext* context) const override;
};
#if !defined(USE_MKLML_FOR_BLAS)
template <>
class MatMul<float> final : public OpKernel {
public:
MatMul(const OpKernelInfo& info) : OpKernel(info) {}
Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override;
Status Compute(OpKernelContext* context) const override;
private:
TensorShape b_shape_;
BufferUniquePtr packed_b_;
};
#endif
ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(
MatMul,
1, 8,
@ -98,4 +116,86 @@ Status MatMul<T>::Compute(OpKernelContext* ctx) const {
return Status::OK();
}
#if !defined(USE_MKLML_FOR_BLAS)
Status MatMul<float>::PrePack(const Tensor& tensor, int input_idx, bool& is_packed) {
is_packed = false;
// only pack Matrix B
if (input_idx == 1) {
// Only handle the common case of a 2D weight matrix. Additional matrices
// could be handled by stacking the packed buffers.
b_shape_ = tensor.Shape();
if (b_shape_.NumDimensions() != 2) {
return Status::OK();
}
const size_t K = static_cast<size_t>(b_shape_[0]);
const size_t N = static_cast<size_t>(b_shape_[1]);
const size_t packed_b_size = MlasGemmPackBSize(N, K);
if (packed_b_size == 0) {
return Status::OK();
}
auto alloc = Info().GetAllocator(0, OrtMemTypeDefault);
auto* packed_b_data = alloc->Alloc(packed_b_size);
packed_b_ = BufferUniquePtr(packed_b_data, BufferDeleter(alloc));
MlasGemmPackB(CblasNoTrans, N, K, tensor.Data<float>(), N, packed_b_data);
is_packed = true;
}
return Status::OK();
}
Status MatMul<float>::Compute(OpKernelContext* ctx) const {
concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool();
const Tensor* a = ctx->Input<Tensor>(0);
const Tensor* b = packed_b_ ? nullptr : ctx->Input<Tensor>(1);
MatMulComputeHelper helper;
ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b ? b->Shape() : b_shape_));
Tensor* y = ctx->Output(0, helper.OutputShape());
// Bail out early if the output is going to be empty
if (y->Shape().Size() == 0)
return Status::OK();
const auto* a_data = a->Data<float>();
const auto* b_data = b ? b->Data<float>() : nullptr;
auto* y_data = y->MutableData<float>();
// TODO: replace it with GemmBatch for performance, it's OK for now as GemmBatch unrolls as well
size_t max_len = helper.OutputOffsets().size();
for (size_t i = 0; i < max_len; i++) {
if (packed_b_) {
MlasGemm(CblasNoTrans,
static_cast<size_t>(helper.M()),
static_cast<size_t>(helper.N()),
static_cast<size_t>(helper.K()),
1.0f,
a_data + helper.LeftOffsets()[i],
static_cast<size_t>(helper.K()),
packed_b_.get(),
0.0f,
y_data + helper.OutputOffsets()[i],
static_cast<size_t>(helper.N()),
thread_pool);
} else {
math::MatMul(
static_cast<int>(helper.M()),
static_cast<int>(helper.N()),
static_cast<int>(helper.K()),
a_data + helper.LeftOffsets()[i],
b_data + helper.RightOffsets()[i],
y_data + helper.OutputOffsets()[i],
thread_pool);
}
}
return Status::OK();
}
#endif
} // namespace onnxruntime

34
onnxruntime/test/contrib_ops/attention_op_test.cc
View file

@ -19,6 +19,7 @@ enum MaskIndexType {
static void RunAttentionTest(
const std::vector<float>& input_data, // input: [batch_size, sequence_length, hidden_size]
const std::vector<float>& weights_data, // weights: [hidden_size, 3 * hidden_size]
bool is_weights_constant,
const std::vector<float>& bias_data, // bias: [3 * hidden_size]
const std::vector<int32_t>& mask_index_data, // mask_index: [batch_size] or [batch_size, past_sequence_length + sequence_length] or empty
const std::vector<float>& output_data, // output: [batch_size, sequence_length, hidden_size]
@ -35,7 +36,7 @@ static void RunAttentionTest(
MaskIndexType mask_index_type = kMaskIndexEnd) {
int min_cuda_architecture = use_float16 ? 530 : 0;
bool enable_cuda = HasCudaEnvironment(min_cuda_architecture);
bool enable_cuda = HasCudaEnvironment(min_cuda_architecture) && !is_weights_constant;
bool enable_cpu = (nullptr != DefaultCpuExecutionProvider().get()) && !use_float16;
int head_size = hidden_size / number_of_heads;
if (enable_cpu || enable_cuda) {
@ -76,12 +77,12 @@ static void RunAttentionTest(
if (use_float16) {
tester.AddInput<MLFloat16>("input", input_dims, ToFloat16(input_data));
tester.AddInput<MLFloat16>("weight", weights_dims, ToFloat16(weights_data));
tester.AddInput<MLFloat16>("weight", weights_dims, ToFloat16(weights_data), is_weights_constant);
tester.AddInput<MLFloat16>("bias", bias_dims, ToFloat16(bias_data));
tester.AddOutput<MLFloat16>("output", output_dims, ToFloat16(output_data));
} else {
tester.AddInput<float>("input", input_dims, input_data);
tester.AddInput<float>("weight", weights_dims, weights_data);
tester.AddInput<float>("weight", weights_dims, weights_data, is_weights_constant);
tester.AddInput<float>("bias", bias_dims, bias_data);
tester.AddOutput<float>("output", output_dims, output_data);
}
@ -120,6 +121,33 @@ static void RunAttentionTest(
}
}
static void RunAttentionTest(
const std::vector<float>& input_data, // input: [batch_size, sequence_length, hidden_size]
const std::vector<float>& weights_data, // weights: [hidden_size, 3 * hidden_size]
const std::vector<float>& bias_data, // bias: [3 * hidden_size]
const std::vector<int32_t>& mask_index_data, // mask_index: [batch_size] or [batch_size, past_sequence_length + sequence_length] or empty
const std::vector<float>& output_data, // output: [batch_size, sequence_length, hidden_size]
int batch_size,
int sequence_length,
int hidden_size,
int number_of_heads,
bool use_float16 = false,
bool is_unidirectional = false,
bool use_past_state = false,
int past_sequence_length = 0,
const std::vector<float>* past_data = nullptr,
const std::vector<float>* present_data = nullptr,
MaskIndexType mask_index_type = kMaskIndexEnd) {
RunAttentionTest(input_data, weights_data, false, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length,
past_data, present_data, mask_index_type);
RunAttentionTest(input_data, weights_data, true, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length,
past_data, present_data, mask_index_type);
}
TEST(AttentionTest, AttentionBatch1) {
int batch_size = 1;
int sequence_length = 2;

99
onnxruntime/test/mlas/unittest.cpp
View file

@ -67,7 +67,7 @@ public:
ReleaseBuffer();
}
T* GetBuffer(size_t Elements)
T* GetBuffer(size_t Elements, bool ZeroFill = false)
{
//
// Check if the internal buffer needs to be reallocated.
@ -125,20 +125,27 @@ public:
T* GuardAddress = _GuardAddress;
T* buffer = GuardAddress - Elements;
const int MinimumFillValue = -23;
const int MaximumFillValue = 23;
if (ZeroFill) {
int FillValue = MinimumFillValue;
T* FillAddress = buffer;
std::fill_n(buffer, Elements, T(0));
while (FillAddress < GuardAddress) {
} else {
*FillAddress++ = (T)FillValue;
const int MinimumFillValue = -23;
const int MaximumFillValue = 23;
FillValue++;
int FillValue = MinimumFillValue;
T* FillAddress = buffer;
if (FillValue > MaximumFillValue) {
FillValue = MinimumFillValue;
while (FillAddress < GuardAddress) {
*FillAddress++ = (T)FillValue;
FillValue++;
if (FillValue > MaximumFillValue) {
FillValue = MinimumFillValue;
}
}
}
@ -206,8 +213,67 @@ public:
}
};
template<typename T, bool Packed>
class MlasFgemmTestBase;
template<typename T>
class MlasFgemmTest : public MlasTestBase
class MlasFgemmTestBase<T, false> : public MlasTestBase
{
public:
void
TestGemm(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t M,
size_t N,
size_t K,
float alpha,
const T* A,
size_t lda,
const T* B,
size_t ldb,
float beta,
T* C,
size_t ldc
)
{
MlasGemm(TransA, TransB, M, N, K, T(alpha), A, lda, B, ldb, T(beta), C, ldc, threadpool);
}
};
template<typename T>
class MlasFgemmTestBase<T, true> : public MlasTestBase
{
public:
void
TestGemm(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t M,
size_t N,
size_t K,
float alpha,
const T* A,
size_t lda,
const T* B,
size_t ldb,
float beta,
T* C,
size_t ldc
)
{
size_t PackedBSize = MlasGemmPackBSize(N, K);
void* PackedB = BufferBPacked.GetBuffer(PackedBSize, true);
MlasGemmPackB(TransB, N, K, B, ldb, PackedB);
MlasGemm(TransA, M, N, K, T(alpha), A, lda, PackedB, T(beta), C, ldc, threadpool);
}
private:
MatrixGuardBuffer<uint8_t> BufferBPacked;
};
template<typename T, bool Packed>
class MlasFgemmTest : public MlasFgemmTestBase<T, Packed>
{
private:
void
@ -251,7 +317,7 @@ private:
std::fill_n(C, M * N, -0.5f);
std::fill_n(CReference, M * N, -0.5f);
MlasGemm(TransA, TransB, M, N, K, T(alpha), A, lda, B, ldb, T(beta), C, ldc, threadpool);
this->TestGemm(TransA, TransB, M, N, K, alpha, A, lda, B, ldb, beta, C, ldc);
ReferenceGemm(TransA, TransB, M, N, K, alpha, A, lda, B, ldb, beta, CReference, ldc);
for (size_t f = 0; f < M * N; f++) {
@ -392,6 +458,9 @@ public:
for (size_t b = 256; b < 320; b += 32) {
Test(b, b, b, 1.0f, 0.0f);
}
Test(128, 3072, 768, 1.0f, 0.0f);
Test(128, 768, 3072, 1.0f, 0.0f);
}
void
@ -2705,10 +2774,12 @@ RunThreadedTests(
)
{
printf("SGEMM tests.\n");
onnxruntime::make_unique<MlasFgemmTest<float>>()->ExecuteShort();
onnxruntime::make_unique<MlasFgemmTest<float, false>>()->ExecuteShort();
printf("SGEMM packed tests.\n");
onnxruntime::make_unique<MlasFgemmTest<float, true>>()->ExecuteShort();
#ifdef MLAS_HAS_DGEMM
printf("DGEMM tests.\n");
onnxruntime::make_unique<MlasFgemmTest<double>>()->ExecuteShort();
onnxruntime::make_unique<MlasFgemmTest<double, false>>()->ExecuteShort();
#endif
#ifdef MLAS_HAS_QGEMM_U8X8

15
onnxruntime/test/providers/cpu/math/matmul_test.cc
View file

@ -94,7 +94,7 @@ std::vector<MatMulTestData<T>> GenerateTestCases() {
}
template <typename T>
void RunMatMulTest(int32_t opset_version = 7) {
void RunMatMulTest(int32_t opset_version, bool is_b_constant = false) {
std::vector<T> common_input_vals{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11};
for (auto t : GenerateTestCases<T>()) {
OpTester test("MatMul", opset_version);
@ -105,17 +105,24 @@ void RunMatMulTest(int32_t opset_version = 7) {
int64_t size1 = TensorShape::ReinterpretBaseType(t.input1_dims).SizeHelper(0, t.input1_dims.size());
std::vector<T> input1_vals(common_input_vals.cbegin(), common_input_vals.cbegin() + size1);
test.AddInput<T>("B", t.input1_dims, input1_vals);
test.AddInput<T>("B", t.input1_dims, input1_vals, is_b_constant);
test.AddOutput<T>("Y", t.expected_dims, t.expected_vals);
// OpenVINO EP: Disabled temporarily matmul broadcasting not fully supported
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); // Disable TensorRT because of unsupported data type
// Disable TensorRT because of unsupported data type
std::unordered_set<std::string> excluded_providers{kTensorrtExecutionProvider, kOpenVINOExecutionProvider};
if (is_b_constant) {
// NNAPI: currently fails for the "test 2D empty input" case
excluded_providers.insert(kNnapiExecutionProvider);
}
test.Run(OpTester::ExpectResult::kExpectSuccess, "", excluded_providers);
}
}
TEST(MathOpTest, MatMulFloatType) {
RunMatMulTest<float>(7);
RunMatMulTest<float>(7, false);
RunMatMulTest<float>(7, true);
}
TEST(MathOpTest, MatMulDoubleType) {