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

MLAS: support prepacking APIs for quantized GEMM (#4433)

Browse source 
Add support for prepacking matrix B for use in the quantized GEMMs.
This commit is contained in:
Tracy Sharpe 2020-07-06 15:20:10 -07:00 committed by GitHub
parent dd73e8c016
commit 3ef449816c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 1186 additions and 259 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

77
onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc
View file

@ -22,6 +22,14 @@ class QAttention : public OpKernel, public AttentionCPUBase {
QAttention(const OpKernelInfo& info);
Status Compute(OpKernelContext* context) const override;
private:
void TryPackWeights(const OpKernelInfo& info);
#ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8
BufferUniquePtr packed_weights_;
size_t packed_weights_size_;
#endif
};
// These ops are internal-only, so register outside of onnx
@ -40,6 +48,53 @@ ONNX_OPERATOR_TYPED_KERNEL_EX(
template <typename T>
QAttention<T>::QAttention(const OpKernelInfo& info) : OpKernel(info), AttentionCPUBase(info) {
TryPackWeights(info);
}
template <typename T>
void QAttention<T>::TryPackWeights(const OpKernelInfo& info) {
#ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8
// Check if the weights tensor is constant.
const Tensor* weights;
if (!info.TryGetConstantInput(1, &weights)) {
return;
}
const auto& weights_dims = weights->Shape().GetDims();
if (weights_dims.size() != 2) {
return;
}
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;
}
const auto* weights_data = static_cast<const uint8_t*>(weights->DataRaw());
const bool weights_is_signed = weights->IsDataType<int8_t>();
packed_weights_size_ = MlasGemmPackBSize(head_size, hidden_size, weights_is_signed);
if (packed_weights_size_ == 0) {
return;
}
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(head_size, hidden_size, weights_data, hidden_size_x3, weights_is_signed, packed_weights_data);
packed_weights_data += packed_weights_size_;
weights_data += head_size;
}
#else
ORT_UNUSED_PARAMETER(info);
#endif
}
template <typename T>
@ -140,6 +195,28 @@ Status QAttention<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
#ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8
if (packed_weights_) {
const auto* packed_weight =
static_cast<const uint8_t*>(packed_weights_.get()) + packed_weights_size_ * (weights_offset / head_size);
MlasGemm(
sequence_length, // M = S
head_size, // N = H
hidden_size, // K = NH
input_data + input_offset, // A
hidden_size, // lda = NH
input_zero_point, // input zero point
packed_weight, // B
weight_zero_point, // weight zero point
weights_is_signed, // weight data type
qkv_dest + qkv_offset, // C
head_size, // ldc
&dequant_scale, // output scale
bias_data + weights_offset, // bias
nullptr); // use single-thread
continue;
}
#endif
QGemm(sequence_length, // M = S
head_size, // N = H
hidden_size, // K = NH

213
onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc
View file

@ -14,69 +14,71 @@
namespace onnxruntime {
namespace contrib {
class DynamicQuantizeMatMul final : public OpKernel {
class MatMulIntegerToFloatBase : public OpKernel {
public:
DynamicQuantizeMatMul(const OpKernelInfo& info) : OpKernel(info) {}
MatMulIntegerToFloatBase(const OpKernelInfo& info) : OpKernel(info) {
TryPackWeights(info);
}
Status Compute(OpKernelContext* context) const override;
protected:
void TryPackWeights(const OpKernelInfo& info);
Status ComputeCommon(OpKernelContext* ctx,
const uint8_t* a_data,
const TensorShape& a_shape,
uint8_t a_zero_point,
const Tensor* b,
uint8_t b_zero_point,
float multiplier,
const Tensor* bias_tensor) const;
#ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8
BufferUniquePtr packed_b_;
#endif
};
ONNX_OPERATOR_TYPED_KERNEL_EX(
DynamicQuantizeMatMul,
kMSDomain,
1,
float,
kCpuExecutionProvider,
KernelDefBuilder()
.TypeConstraint("T1", DataTypeImpl::GetTensorType<float>())
.TypeConstraint("T2", {DataTypeImpl::GetTensorType<uint8_t>(), DataTypeImpl::GetTensorType<int8_t>()}),
DynamicQuantizeMatMul);
void MatMulIntegerToFloatBase::TryPackWeights(const OpKernelInfo& info) {
#ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8
// Check if the weights tensor is constant.
const Tensor* b;
if (!info.TryGetConstantInput(1, &b)) {
return;
}
class MatMulIntegerToFloat final : public OpKernel {
public:
MatMulIntegerToFloat(const OpKernelInfo& info) : OpKernel(info) {}
// Only handle the common case of a 2D weight matrix. Additional matrices
// could be handled by stacking the packed buffers.
const auto& b_shape = b->Shape();
if (b_shape.NumDimensions() != 2) {
return;
}
Status Compute(OpKernelContext* context) const override;
};
const size_t K = static_cast<size_t>(b_shape[0]);
const size_t N = static_cast<size_t>(b_shape[1]);
ONNX_OPERATOR_TYPED_KERNEL_EX(
MatMulIntegerToFloat,
kMSDomain,
1,
uint8_t,
kCpuExecutionProvider,
KernelDefBuilder()
.TypeConstraint("T1", DataTypeImpl::GetTensorType<uint8_t>())
.TypeConstraint("T2", {DataTypeImpl::GetTensorType<uint8_t>(), DataTypeImpl::GetTensorType<int8_t>()})
.TypeConstraint("T3", DataTypeImpl::GetTensorType<float>()),
MatMulIntegerToFloat);
const auto* b_data = static_cast<const uint8_t*>(b->DataRaw());
const bool b_is_signed = b->IsDataType<int8_t>();
static void GetQuantizationParameter(const float* data, int64_t num_of_elements, float& scale, uint8_t& zp) {
// find input range min and max
float min, max;
MlasFindMinMaxElement(data, &min, &max, num_of_elements);
const size_t packed_b_size = MlasGemmPackBSize(N, K, b_is_signed);
if (packed_b_size == 0) {
return;
}
// ensure the input range includes zero
min = std::min(min, 0.0f);
max = std::max(max, 0.0f);
// find scale and zero point
uint8_t qmin = 0;
uint8_t qmax = 255;
scale = max == min ? 1.0f : (max - min) / (qmax - qmin);
float initial_zero_point = qmin - min / scale;
zp = static_cast<uint8_t>(RoundHalfToEven(std::max(float(qmin), std::min(float(qmax), initial_zero_point))));
auto alloc = info.GetAllocator(0, OrtMemTypeDefault);
auto* packed_b_data = alloc->Alloc(packed_b_size);
packed_b_ = BufferUniquePtr(packed_b_data, BufferDeleter(alloc));
MlasGemmPackB(N, K, b_data, N, b_is_signed, packed_b_data);
#else
ORT_UNUSED_PARAMETER(info);
#endif
}
static Status MatMulIntegerToFloatCommon(OpKernelContext* ctx,
const uint8_t* a_data,
const TensorShape& a_shape,
uint8_t a_zero_point,
const Tensor* b,
uint8_t b_zero_point,
float multiplier,
const Tensor* bias_tensor) {
Status MatMulIntegerToFloatBase::ComputeCommon(OpKernelContext* ctx,
const uint8_t* a_data,
const TensorShape& a_shape,
uint8_t a_zero_point,
const Tensor* b,
uint8_t b_zero_point,
float multiplier,
const Tensor* bias_tensor) const {
MatMulComputeHelper helper;
ORT_RETURN_IF_ERROR(helper.Compute(a_shape, b->Shape()));
Tensor* y = ctx->Output(0, helper.OutputShape());
@ -89,6 +91,25 @@ static Status MatMulIntegerToFloatCommon(OpKernelContext* ctx,
concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool();
for (size_t i = 0; i < helper.OutputOffsets().size(); i++) {
#ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8
if (packed_b_) {
MlasGemm(static_cast<size_t>(helper.M()),
static_cast<size_t>(helper.N()),
static_cast<size_t>(helper.K()),
a_data + helper.LeftOffsets()[i],
static_cast<size_t>(helper.K()),
a_zero_point,
packed_b_.get(),
b_zero_point,
b_is_signed,
y_data + helper.OutputOffsets()[i],
static_cast<size_t>(helper.N()),
&multiplier,
nullptr,
thread_pool);
continue;
}
#endif
QGemm(static_cast<int>(helper.M()),
static_cast<int>(helper.N()),
static_cast<int>(helper.K()),
@ -109,6 +130,38 @@ static Status MatMulIntegerToFloatCommon(OpKernelContext* ctx,
return Status::OK();
}
class DynamicQuantizeMatMul final : public MatMulIntegerToFloatBase {
public:
DynamicQuantizeMatMul(const OpKernelInfo& info) : MatMulIntegerToFloatBase(info) {}
Status Compute(OpKernelContext* context) const override;
};
class MatMulIntegerToFloat final : public MatMulIntegerToFloatBase {
public:
MatMulIntegerToFloat(const OpKernelInfo& info) : MatMulIntegerToFloatBase(info) {}
Status Compute(OpKernelContext* context) const override;
};
static void GetQuantizationParameter(const float* data, int64_t num_of_elements, float& scale, uint8_t& zp) {
// find input range min and max
float min, max;
MlasFindMinMaxElement(data, &min, &max, num_of_elements);
// ensure the input range includes zero
min = std::min(min, 0.0f);
max = std::max(max, 0.0f);
// find scale and zero point
uint8_t qmin = 0;
uint8_t qmax = 255;
scale = max == min ? 1.0f : (max - min) / (qmax - qmin);
float initial_zero_point = qmin - min / scale;
zp = static_cast<uint8_t>(RoundHalfToEven(std::max(float(qmin), std::min(float(qmax), initial_zero_point))));
}
Status DynamicQuantizeMatMul::Compute(OpKernelContext* ctx) const {
const auto* a = ctx->Input<Tensor>(0);
const auto* b = ctx->Input<Tensor>(1);
@ -116,7 +169,6 @@ Status DynamicQuantizeMatMul::Compute(OpKernelContext* ctx) const {
const auto* b_scale_tensor = ctx->Input<Tensor>(2);
ORT_ENFORCE(IsScalarOr1ElementVector(b_scale_tensor),
"DynamicQuantizeMatMul : input B scale must be a scalar or 1D tensor of size 1. Per-Channel is not supported yet.");
float b_scale = *b_scale_tensor->template Data<float>();
const auto* b_zero_point_tensor = ctx->Input<Tensor>(3);
@ -142,14 +194,14 @@ Status DynamicQuantizeMatMul::Compute(OpKernelContext* ctx) const {
// quantize the data
MlasQuantizeLinear(a_data, a_data_quant, num_of_elements, a_scale, a_zero_point);
return MatMulIntegerToFloatCommon(ctx,
a_data_quant,
a->Shape(),
a_zero_point,
b,
b_zero_point,
a_scale * b_scale,
ctx->Input<Tensor>(4));
return ComputeCommon(ctx,
a_data_quant,
a->Shape(),
a_zero_point,
b,
b_zero_point,
a_scale * b_scale,
ctx->Input<Tensor>(4));
}
Status MatMulIntegerToFloat::Compute(OpKernelContext* ctx) const {
@ -183,15 +235,38 @@ Status MatMulIntegerToFloat::Compute(OpKernelContext* ctx) const {
b_zero_point = *static_cast<const uint8_t*>(b_zero_point_tensor->DataRaw());
}
return MatMulIntegerToFloatCommon(ctx,
a->Data<uint8_t>(),
a->Shape(),
a_zero_point,
b,
b_zero_point,
a_scale * b_scale,
ctx->Input<Tensor>(6));
return ComputeCommon(ctx,
a->Data<uint8_t>(),
a->Shape(),
a_zero_point,
b,
b_zero_point,
a_scale * b_scale,
ctx->Input<Tensor>(6));
}
ONNX_OPERATOR_TYPED_KERNEL_EX(
DynamicQuantizeMatMul,
kMSDomain,
1,
float,
kCpuExecutionProvider,
KernelDefBuilder()
.TypeConstraint("T1", DataTypeImpl::GetTensorType<float>())
.TypeConstraint("T2", {DataTypeImpl::GetTensorType<uint8_t>(), DataTypeImpl::GetTensorType<int8_t>()}),
DynamicQuantizeMatMul);
ONNX_OPERATOR_TYPED_KERNEL_EX(
MatMulIntegerToFloat,
kMSDomain,
1,
uint8_t,
kCpuExecutionProvider,
KernelDefBuilder()
.TypeConstraint("T1", DataTypeImpl::GetTensorType<uint8_t>())
.TypeConstraint("T2", {DataTypeImpl::GetTensorType<uint8_t>(), DataTypeImpl::GetTensorType<int8_t>()})
.TypeConstraint("T3", DataTypeImpl::GetTensorType<float>()),
MatMulIntegerToFloat);
} // namespace contrib
} // namespace onnxruntime

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

@ -185,6 +185,65 @@ MlasGemm(
MLAS_THREADPOOL* ThreadPool
);
void
MLASCALL
MlasGemm(
size_t M,
size_t N,
size_t K,
const uint8_t* A,
size_t lda,
uint8_t offa,
const void* PackedB,
uint8_t offb,
bool BIsSigned,
int32_t* C,
size_t ldc,
MLAS_THREADPOOL* ThreadPool
);
void
MLASCALL
MlasGemm(
size_t M,
size_t N,
size_t K,
const uint8_t* A,
size_t lda,
uint8_t offa,
const void* PackedB,
uint8_t offb,
bool BIsSigned,
float* C,
size_t ldc,
const float* Scale,
const float* Bias,
MLAS_THREADPOOL* ThreadPool
);
//
// Buffer packing routines.
//
size_t
MLASCALL
MlasGemmPackBSize(
size_t N,
size_t K,
bool BIsSigned
);
void
MLASCALL
MlasGemmPackB(
size_t N,
size_t K,
const uint8_t* B,
size_t ldb,
bool BIsSigned,
void* PackedB
);
//
// Convolution routines.
//

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

@ -543,10 +543,6 @@ extern "C" {
MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE MlasSgemmTransposePackB16x4Avx;
#endif
#if defined(MLAS_TARGET_AMD64_IX86)
MLAS_GEMM_U8X8_OPERATION MlasGemmU8X8OperationSse;
MLAS_GEMM_U8X8_OPERATION MlasGemmU8S8OperationAvx2;
MLAS_GEMM_U8X8_OPERATION MlasGemmU8U8OperationAvx2;
#if defined(MLAS_TARGET_AMD64)
MLAS_GEMM_U8S8_KERNEL MlasGemmU8S8KernelAvx2;
MLAS_GEMV_U8S8_KERNEL MlasGemvU8S8KernelAvx2;
@ -557,7 +553,6 @@ extern "C" {
MLAS_GEMM_U8U8_KERNEL MlasGemmU8U8KernelAvx2;
MLAS_GEMM_U8U8_KERNEL MlasGemmU8U8KernelAvx512Core;
#endif
#endif
#if defined(MLAS_TARGET_AMD64)
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelSse;
@ -674,6 +669,28 @@ MlasSgemmOperation(
size_t ldc
);
//
// Quantized integer matrix/matrix multiply operation.
//
struct MLAS_GEMM_U8X8_KERNEL_SSE;
struct MLAS_GEMM_U8S8_KERNEL_AVX2;
struct MLAS_GEMM_U8U8_KERNEL_AVX2;
template<typename KernelType>
void
MLASCALL
MlasGemmU8X8Operation(
const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock
);
template<typename KernelType>
void
MLASCALL
MlasGemmU8X8PackedOperation(
const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock
);
//
// Environment information class.
//
@ -684,8 +701,6 @@ struct MLAS_PLATFORM {
#if defined(MLAS_TARGET_AMD64_IX86)
PMLAS_GEMM_FLOAT_KERNEL GemmFloatKernel;
PMLAS_GEMM_U8X8_OPERATION GemmU8S8Operation;
PMLAS_GEMM_U8X8_OPERATION GemmU8U8Operation;
#endif
#if defined(MLAS_TARGET_AMD64)
@ -693,8 +708,12 @@ struct MLAS_PLATFORM {
PMLAS_SGEMM_KERNEL_M1_ROUTINE KernelM1TransposeBRoutine;
PMLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE TransposePackB16x4Routine;
PMLAS_GEMM_DOUBLE_KERNEL GemmDoubleKernel;
PMLAS_GEMM_U8X8_OPERATION GemmU8S8Operation;
PMLAS_GEMM_U8X8_OPERATION GemmU8S8PackedOperation;
PMLAS_GEMM_U8S8_KERNEL GemmU8S8Kernel;
PMLAS_GEMV_U8S8_KERNEL GemvU8S8Kernel;
PMLAS_GEMM_U8X8_OPERATION GemmU8U8Operation;
PMLAS_GEMM_U8X8_OPERATION GemmU8U8PackedOperation;
PMLAS_GEMM_U8U8_KERNEL GemmU8U8Kernel;
PMLAS_CONV_FLOAT_KERNEL ConvNchwFloatKernel;
PMLAS_CONV_FLOAT_KERNEL ConvNchwcFloatKernel;

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

@ -115,13 +115,13 @@ Return Value:
//
this->GemmFloatKernel = MlasGemmFloatKernelSse;
this->GemmU8S8Operation = MlasGemmU8X8OperationSse;
this->GemmU8U8Operation = MlasGemmU8X8OperationSse;
#if defined(MLAS_TARGET_AMD64)
this->TransposePackB16x4Routine = MlasSgemmTransposePackB16x4Sse;
this->GemmDoubleKernel = MlasGemmDoubleKernelSse;
this->GemmU8S8Operation = MlasGemmU8X8Operation<MLAS_GEMM_U8X8_KERNEL_SSE>;
this->GemmU8U8Operation = MlasGemmU8X8Operation<MLAS_GEMM_U8X8_KERNEL_SSE>;
this->ConvNchwFloatKernel = MlasConvNchwFloatKernelSse;
this->ConvNchwcFloatKernel = MlasConvNchwcFloatKernelSse;
this->ConvDepthwiseFloatKernel = MlasConvDepthwiseFloatKernelSse;
@ -200,10 +200,12 @@ Return Value:
if (((Cpuid1[2] & 0x1000) != 0) && ((Cpuid7[1] & 0x20) != 0)) {
this->GemmU8S8Operation = MlasGemmU8S8OperationAvx2;
this->GemmU8S8Operation = MlasGemmU8X8Operation<MLAS_GEMM_U8S8_KERNEL_AVX2>;
this->GemmU8S8PackedOperation = MlasGemmU8X8PackedOperation<MLAS_GEMM_U8S8_KERNEL_AVX2>;
this->GemmU8S8Kernel = MlasGemmU8S8KernelAvx2;
this->GemvU8S8Kernel = MlasGemvU8S8KernelAvx2;
this->GemmU8U8Operation = MlasGemmU8U8OperationAvx2;
this->GemmU8U8Operation = MlasGemmU8X8Operation<MLAS_GEMM_U8U8_KERNEL_AVX2>;
this->GemmU8U8PackedOperation = MlasGemmU8X8PackedOperation<MLAS_GEMM_U8U8_KERNEL_AVX2>;
this->GemmU8U8Kernel = MlasGemmU8U8KernelAvx2;
this->GemmFloatKernel = MlasGemmFloatKernelFma3;
@ -262,7 +264,8 @@ Return Value:
if ((Cpuid7[2] & 0x800) != 0) {
this->GemmU8U8Operation = MlasGemmU8S8OperationAvx2;
this->GemmU8U8Operation = MlasGemmU8X8Operation<MLAS_GEMM_U8S8_KERNEL_AVX2>;
this->GemmU8U8PackedOperation = MlasGemmU8X8PackedOperation<MLAS_GEMM_U8S8_KERNEL_AVX2>;
this->GemmU8S8Kernel = MlasGemmU8S8KernelAvx512Vnni;
this->GemvU8S8Kernel = MlasGemvU8S8KernelAvx512Vnni;
}

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

File diff suppressed because it is too large Load diff

4
onnxruntime/core/util/qmath.h
View file

@ -12,6 +12,10 @@
#define MLAS_SUPPORTS_GEMM_U8X8
#endif
#if defined(_M_AMD64) || defined(__x86_64__)
#define MLAS_SUPPORTS_PACKED_GEMM_U8X8
#endif
namespace onnxruntime {
void QGemm(

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

@ -45,6 +45,10 @@ Abstract:
#define MLAS_HAS_QGEMM_U8X8
#endif
#if defined(_M_AMD64) || defined(__x86_64__)
#define MLAS_HAS_PACKED_QGEMM_U8X8
#endif
MLAS_THREADPOOL* threadpool = nullptr;
template<typename T>
@ -468,11 +472,129 @@ public:
#ifdef MLAS_HAS_QGEMM_U8X8
template<typename xint8_t, typename OutputType>
template<bool Packed>
class MlasQgemmU8X8U8X8TestBase;
template<>
class MlasQgemmU8X8U8X8TestBase<false> : public MlasTestBase
{
protected:
void
TestGemm(
size_t M,
size_t N,
size_t K,
const uint8_t* A,
size_t lda,
uint8_t offa,
const uint8_t* B,
size_t ldb,
uint8_t offb,
bool BIsSigned,
int32_t* C,
size_t ldc
)
{
MlasGemm(M, N, K, A, lda, offa, B, ldb, offb, BIsSigned, C, ldc, threadpool);
}
void
TestGemm(
size_t M,
size_t N,
size_t K,
const uint8_t* A,
size_t lda,
uint8_t offa,
const uint8_t* B,
size_t ldb,
uint8_t offb,
bool BIsSigned,
float* C,
size_t ldc,
float CScale,
const float* Bias
)
{
MlasGemm(M, N, K, A, lda, offa, B, ldb, offb, BIsSigned, C, ldc, &CScale, Bias, threadpool);
}
};
#ifdef MLAS_HAS_PACKED_QGEMM_U8X8
template<>
class MlasQgemmU8X8U8X8TestBase<true> : public MlasTestBase
{
private:
void*
PackB(
size_t N,
size_t K,
const uint8_t* B,
size_t ldb,
bool BIsSigned
)
{
size_t PackedBSize = MlasGemmPackBSize(N, K, BIsSigned);
void* PackedB = BufferBPacked.GetBuffer(PackedBSize);
MlasGemmPackB(N, K, B, ldb, BIsSigned, PackedB);
return PackedB;
}
protected:
void
TestGemm(
size_t M,
size_t N,
size_t K,
const uint8_t* A,
size_t lda,
uint8_t offa,
const uint8_t* B,
size_t ldb,
uint8_t offb,
bool BIsSigned,
int32_t* C,
size_t ldc
)
{
const void* PackedB = PackB(N, K, B, ldb, BIsSigned);
MlasGemm(M, N, K, A, lda, offa, PackedB, offb, BIsSigned, C, ldc, threadpool);
}
void
TestGemm(
size_t M,
size_t N,
size_t K,
const uint8_t* A,
size_t lda,
uint8_t offa,
const uint8_t* B,
size_t ldb,
uint8_t offb,
bool BIsSigned,
float* C,
size_t ldc,
float CScale,
const float* Bias
)
{
const void* PackedB = PackB(N, K, B, ldb, BIsSigned);
MlasGemm(M, N, K, A, lda, offa, PackedB, offb, BIsSigned, C, ldc, &CScale, Bias, threadpool);
}
private:
MatrixGuardBuffer<uint8_t> BufferBPacked;
};
#endif
template<typename xint8_t, typename OutputType, bool Packed>
class MlasQgemmU8X8Test;
template<typename xint8_t>
class MlasQgemmU8X8Test<xint8_t, int32_t> : public MlasTestBase
template<typename xint8_t, bool Packed>
class MlasQgemmU8X8Test<xint8_t, int32_t, Packed> : public MlasQgemmU8X8U8X8TestBase<Packed>
{
private:
void
@ -511,7 +633,7 @@ private:
std::fill_n(C, M * N, -1);
std::fill_n(CReference, M * N, -1);
MlasGemm(M, N, K, A, lda, offa, B, ldb, offb, BIsSigned, C, ldc, threadpool);
this->TestGemm(M, N, K, A, lda, offa, B, ldb, offb, BIsSigned, C, ldc);
ReferenceQgemm(M, N, K, A, lda, offa, (const xint8_t*)B, ldb, (xint8_t)offb, CReference, ldc);
for (size_t f = 0; f < M * N; f++) {
@ -569,6 +691,9 @@ public:
void
) override
{
for (size_t b = 1; b < 16; b++) {
Test(b, b, b, 14, 211);
}
for (size_t b = 1; b < 16; b++) {
Test(b, b, b, 14, 211);
}
@ -583,7 +708,8 @@ public:
Test(1, 32, b, 0, 0);
Test(1, b, b, 0, 0);
}
Test(1024, 1024, 1024, 5, 8);
Test(43, 500, 401, 183, 223);
Test(1023, 1023, 1023, 5, 8);
}
void
@ -651,8 +777,8 @@ public:
}
};
template<typename xint8_t>
class MlasQgemmU8X8Test<xint8_t, float> : public MlasTestBase
template<typename xint8_t, bool Packed>
class MlasQgemmU8X8Test<xint8_t, float, Packed> : public MlasQgemmU8X8U8X8TestBase<Packed>
{
private:
void
@ -714,7 +840,7 @@ private:
}
}
MlasGemm(M, N, K, A, lda, offa, B, ldb, offb, BIsSigned, C, ldc, &CScale, Bias, threadpool);
this->TestGemm(M, N, K, A, lda, offa, B, ldb, offb, BIsSigned, C, ldc, CScale, Bias);
for (size_t f = 0; f < M * N; f++) {
// Sensitive to comparing positive/negative zero.
@ -767,6 +893,7 @@ public:
for (size_t b = 1; b < 96; b++) {
Test(1, b, 32, 0, 0);
}
Test(43, 503, 401, 183, 223);
Test(1024, 1024, 256, 13, 15);
}
};
@ -2435,25 +2562,40 @@ RunThreadedTests(
#ifdef MLAS_HAS_QGEMM_U8X8
printf("QGEMM U8S8=int32_t tests.\n");
onnxruntime::make_unique<MlasQgemmU8X8Test<int8_t, int32_t>>()->ExecuteShort();
printf("QGEMM U8U8=int32_t tests.\n");
onnxruntime::make_unique<MlasQgemmU8X8Test<uint8_t, int32_t>>()->ExecuteShort();
onnxruntime::make_unique<MlasQgemmU8X8Test<int8_t, int32_t, false>>()->ExecuteShort();
printf("QGEMM U8S8=float tests.\n");
onnxruntime::make_unique<MlasQgemmU8X8Test<int8_t, float>>()->ExecuteShort();
onnxruntime::make_unique<MlasQgemmU8X8Test<int8_t, float, false>>()->ExecuteShort();
printf("QGEMM U8U8=int32_t tests.\n");
onnxruntime::make_unique<MlasQgemmU8X8Test<uint8_t, int32_t, false>>()->ExecuteShort();
printf("QGEMM U8U8=float tests.\n");
onnxruntime::make_unique<MlasQgemmU8X8Test<uint8_t, float>>()->ExecuteShort();
onnxruntime::make_unique<MlasQgemmU8X8Test<uint8_t, float, false>>()->ExecuteShort();
#endif
#ifdef MLAS_HAS_PACKED_QGEMM_U8X8
if (MlasGemmPackBSize(128, 128, true) > 0) {
printf("QGEMM U8S8=int32_t packed tests.\n");
onnxruntime::make_unique<MlasQgemmU8X8Test<int8_t, int32_t, true>>()->ExecuteShort();
printf("QGEMM U8S8=float packed tests.\n");
onnxruntime::make_unique<MlasQgemmU8X8Test<int8_t, float, true>>()->ExecuteShort();
}
if (MlasGemmPackBSize(128, 128, false) > 0) {
printf("QGEMM U8U8=int32_t packed tests.\n");
onnxruntime::make_unique<MlasQgemmU8X8Test<uint8_t, int32_t, true>>()->ExecuteShort();
printf("QGEMM U8U8=float packed tests.\n");
onnxruntime::make_unique<MlasQgemmU8X8Test<uint8_t, float, true>>()->ExecuteShort();
}
#endif
printf("Conv2D tests.\n");
onnxruntime::make_unique<MlasConv2DTest>()->ExecuteShort();
if (MlasNchwcGetBlockSize() > 1) {
onnxruntime::make_unique<MlasNchwcConv2DTest>()->ExecuteShort();
onnxruntime::make_unique<MlasNchwcConv2DTest>()->ExecuteShort();
}
printf("Pool2D tests.\n");
onnxruntime::make_unique<MlasPool2DTest>()->ExecuteShort();
if (MlasNchwcGetBlockSize() > 1) {
onnxruntime::make_unique<MlasNchwcPool2DTest>()->ExecuteShort();
onnxruntime::make_unique<MlasNchwcPool2DTest>()->ExecuteShort();
}
printf("Pool3D tests.\n");