mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Add quantization support of GEMM directly with QGemm (#8447)
QGemm takes in quantized A, B, C, and quantization parameters of output Y, in which C and quantization parameters of Y are optional. Its output can be quantized or full precision, which depends on whether quantization parameters of Y exists or not. If quant params of Y are provided, the output will be requantized or is full precision. Comparing with QLinearMatMul and MatMulInteger, QGemm supports transpose, apha and beta attribute. The formula for quantized GEMM is: Y = alpha * scale_a * scale_b * ((A_int8 - zp_a) * (B_int8 - zp_b) + C_int32), in which, C_int32 is quantized with formula: C_int32 = (beta * C) / (alpha * scale_a * scale_b)
This commit is contained in:
parent
0f46b08646
commit
ceeb1a65d6
20 changed files with 743 additions and 68 deletions
|
|
@ -40,6 +40,7 @@ Do not modify directly.*
|
|||
* <a href="#com.microsoft.OptionalHasElement">com.microsoft.OptionalHasElement</a>
|
||||
* <a href="#com.microsoft.Pad">com.microsoft.Pad</a>
|
||||
* <a href="#com.microsoft.QAttention">com.microsoft.QAttention</a>
|
||||
* <a href="#com.microsoft.QGemm">com.microsoft.QGemm</a>
|
||||
* <a href="#com.microsoft.QLinearAdd">com.microsoft.QLinearAdd</a>
|
||||
* <a href="#com.microsoft.QLinearAveragePool">com.microsoft.QLinearAveragePool</a>
|
||||
* <a href="#com.microsoft.QLinearConcat">com.microsoft.QLinearConcat</a>
|
||||
|
|
@ -1894,6 +1895,73 @@ This version of the operator has been available since version 1 of the 'com.micr
|
|||
</dl>
|
||||
|
||||
|
||||
### <a name="com.microsoft.QGemm"></a><a name="com.microsoft.qgemm">**com.microsoft.QGemm**</a>
|
||||
|
||||
Quantized Gemm
|
||||
|
||||
#### Version
|
||||
|
||||
This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
|
||||
|
||||
#### Attributes
|
||||
|
||||
<dl>
|
||||
<dt><tt>alpha</tt> : float</dt>
|
||||
<dd>Scalar multiplier for the product of input tensors A * B.</dd>
|
||||
<dt><tt>transA</tt> : int</dt>
|
||||
<dd>Whether A should be transposed</dd>
|
||||
<dt><tt>transB</tt> : int</dt>
|
||||
<dd>Whether B should be transposed</dd>
|
||||
</dl>
|
||||
|
||||
#### Inputs (6 - 9)
|
||||
|
||||
<dl>
|
||||
<dt><tt>A</tt> : TA</dt>
|
||||
<dd>Input tensor A. The shape of A should be (M, K) if transA is 0, or (K, M) if transA is non-zero.</dd>
|
||||
<dt><tt>a_scale</tt> : T</dt>
|
||||
<dd>Scale of quantized input 'A'. It is a scalar,which means a per-tensor quantization.</dd>
|
||||
<dt><tt>a_zero_point</tt> : TA</dt>
|
||||
<dd>Zero point tensor for input 'A'. It is a scalar.</dd>
|
||||
<dt><tt>B</tt> : TB</dt>
|
||||
<dd>Input tensor B. The shape of B should be (K, N) if transB is 0, or (N, K) if transB is non-zero.</dd>
|
||||
<dt><tt>b_scale</tt> : T</dt>
|
||||
<dd>Scale of quantized input 'B'. It could be a scalar or a 1-D tensor, which means a per-tensor or per-column quantization. If it's a 1-D tensor, its number of elements should be equal to the number of columns of input 'B'.</dd>
|
||||
<dt><tt>b_zero_point</tt> : TB</dt>
|
||||
<dd>Zero point tensor for input 'B'. It's optional and default value is 0. It could be a scalar or a 1-D tensor, which means a per-tensor or per-column quantization. If it's a 1-D tensor, its number of elements should be equal to the number of columns of input 'B'.</dd>
|
||||
<dt><tt>C</tt> (optional) : TC</dt>
|
||||
<dd>Optional input tensor C. If not specified, the computation is done as if C is a scalar 0. The shape of C should be unidirectional broadcastable to (M, N). Its type is int32_t and must be quantized with zero_point = 0 and scale = alpha / beta * a_scale * b_scale.</dd>
|
||||
<dt><tt>y_scale</tt> (optional) : T</dt>
|
||||
<dd>Scale of output 'Y'. It is a scalar, which means a per-tensor quantization. It is optional. The output is full precision(float32) if it is not provided. Or the output is quantized.</dd>
|
||||
<dt><tt>y_zero_point</tt> (optional) : TYZ</dt>
|
||||
<dd>Zero point tensor for output 'Y'. It is a scalar, which means a per-tensor quantization. It is optional. The output is full precision(float32) if it is not provided. Or the output is quantized.</dd>
|
||||
</dl>
|
||||
|
||||
#### Outputs
|
||||
|
||||
<dl>
|
||||
<dt><tt>Y</tt> : TY</dt>
|
||||
<dd>Output tensor of shape (M, N).</dd>
|
||||
</dl>
|
||||
|
||||
#### Type Constraints
|
||||
|
||||
<dl>
|
||||
<dt><tt>T</tt> : tensor(float)</dt>
|
||||
<dd>Constrain scale types to float tensors.</dd>
|
||||
<dt><tt>TA</tt> : tensor(uint8), tensor(int8)</dt>
|
||||
<dd>Constrain input A and its zero point types to 8 bit tensors.</dd>
|
||||
<dt><tt>TB</tt> : tensor(uint8), tensor(int8)</dt>
|
||||
<dd>Constrain input B and its zero point types to 8 bit tensors.</dd>
|
||||
<dt><tt>TC</tt> : tensor(int32)</dt>
|
||||
<dd>Constrain input C to 32 bit integer tensors.</dd>
|
||||
<dt><tt>TYZ</tt> : tensor(uint8), tensor(int8)</dt>
|
||||
<dd>Constrain output zero point types to 8 bit tensors.</dd>
|
||||
<dt><tt>TY</tt> : tensor(float), tensor(uint8), tensor(int8)</dt>
|
||||
<dd>Constrain output type to float32 or 8 bit tensors.</dd>
|
||||
</dl>
|
||||
|
||||
|
||||
### <a name="com.microsoft.QLinearAdd"></a><a name="com.microsoft.qlinearadd">**com.microsoft.QLinearAdd**</a>
|
||||
|
||||
Performs element-wise binary addition on 8 bit data types (with Numpy-style broadcasting support).
|
||||
|
|
|
|||
|
|
@ -393,6 +393,7 @@ Do not modify directly.*
|
|||
|Pad|*in* data:**T**<br> *in* pads:**tensor(int64)**<br> *in* value:**T**<br> *out* output:**T**|1+|**T** = tensor(float)|
|
||||
|QAttention|*in* input:**T1**<br> *in* weight:**T2**<br> *in* bias:**T3**<br> *in* input_scale:**T3**<br> *in* weight_scale:**T3**<br> *in* mask_index:**T4**<br> *in* input_zero_point:**T1**<br> *in* weight_zero_point:**T2**<br> *in* past:**T3**<br> *out* output:**T3**<br> *out* present:**T3**|1+|**T1** = tensor(uint8)<br/> **T2** = tensor(int8), tensor(uint8)<br/> **T3** = tensor(float)<br/> **T4** = tensor(int32)|
|
||||
|QEmbedLayerNormalization|*in* input_ids:**T1**<br> *in* segment_ids:**T1**<br> *in* word_embedding_quant:**T2**<br> *in* position_embedding_quant:**T2**<br> *in* segment_embedding:**T2**<br> *in* gamma_quant:**T2**<br> *in* beta_quant:**T2**<br> *in* mask:**T1**<br> *in* word_embedding_scale:**T**<br> *in* position_embedding_scale:**T**<br> *in* segment_embedding_scale:**T**<br> *in* gamma_scale:**T**<br> *in* beta_scale:**T**<br> *in* word_embedding_zero_point:**T2**<br> *in* position_embedding_zero_point:**T2**<br> *in* segment_embedding_zero_point:**T2**<br> *in* gamma_zero_point:**T2**<br> *in* beta_zero_point:**T2**<br> *out* layernorm_out:**T**<br> *out* mask_index_out:**T1**|1+|**T** = tensor(float)|
|
||||
|QGemm|*in* A:**TA**<br> *in* a_scale:**T**<br> *in* a_zero_point:**TA**<br> *in* B:**TB**<br> *in* b_scale:**T**<br> *in* b_zero_point:**TB**<br> *in* C:**TC**<br> *in* y_scale:**T**<br> *in* y_zero_point:**TYZ**<br> *out* Y:**TY**|1+|**T** = tensor(float)<br/> **TA** = tensor(uint8)<br/> **TB** = tensor(int8), tensor(uint8)<br/> **TC** = tensor(int32)<br/> **TY** = tensor(float), tensor(uint8)<br/> **TYZ** = tensor(uint8)|
|
||||
|QLinearAdd|*in* A:**T**<br> *in* A_scale:**tensor(float)**<br> *in* A_zero_point:**T**<br> *in* B:**T**<br> *in* B_scale:**tensor(float)**<br> *in* B_zero_point:**T**<br> *in* C_scale:**tensor(float)**<br> *in* C_zero_point:**T**<br> *out* C:**T**|1+|**T** = tensor(int8), tensor(uint8)|
|
||||
|QLinearConv|*in* x:**T1**<br> *in* x_scale:**tensor(float)**<br> *in* x_zero_point:**T1**<br> *in* w:**T2**<br> *in* w_scale:**tensor(float)**<br> *in* w_zero_point:**T2**<br> *in* y_scale:**tensor(float)**<br> *in* y_zero_point:**T3**<br> *in* B:**T4**<br> *out* y:**T3**|1+|**T1** = tensor(uint8)<br/> **T2** = tensor(int8), tensor(uint8)<br/> **T3** = tensor(uint8)<br/> **T4** = tensor(int32)|
|
||||
|QLinearLeakyRelu|*in* X:**T**<br> *in* X_scale:**tensor(float)**<br> *in* X_zero_point:**T**<br> *in* Y_scale:**tensor(float)**<br> *in* Y_zero_point:**T**<br> *out* Y:**T**|1+|**T** = tensor(int8), tensor(uint8)|
|
||||
|
|
|
|||
|
|
@ -70,6 +70,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1,
|
|||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, QLinearConv);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, NhwcMaxPool);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, QEmbedLayerNormalization);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, uint8_t, QGemm);
|
||||
// ******** End: Quantization ******************* //
|
||||
|
||||
// This section includes all op kernel declarations for former experimental ops which have now been removed from onnx.
|
||||
|
|
@ -157,6 +158,7 @@ Status RegisterQuantizationKernels(KernelRegistry& kernel_registry) {
|
|||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, QLinearConv)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, NhwcMaxPool)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, QEmbedLayerNormalization)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, uint8_t, QGemm)>,
|
||||
};
|
||||
|
||||
for (auto& function_table_entry : function_table) {
|
||||
|
|
|
|||
|
|
@ -165,7 +165,7 @@ class DynamicQuantizeMatMul final : public MatMulIntegerToFloatBase {
|
|||
};
|
||||
|
||||
protected:
|
||||
int GetBIdx() override { return IN_B; }
|
||||
int GetBIdx() const override { return IN_B; }
|
||||
};
|
||||
|
||||
class MatMulIntegerToFloat final : public MatMulIntegerToFloatBase {
|
||||
|
|
@ -185,7 +185,7 @@ class MatMulIntegerToFloat final : public MatMulIntegerToFloatBase {
|
|||
};
|
||||
|
||||
protected:
|
||||
int GetBIdx() override { return IN_B; }
|
||||
int GetBIdx() const override { return IN_B; }
|
||||
|
||||
private:
|
||||
// a scale and b scale may be switched in fusion stage because of lack of shape information.
|
||||
|
|
|
|||
221
onnxruntime/contrib_ops/cpu/quantization/quant_gemm.cc
Normal file
221
onnxruntime/contrib_ops/cpu/quantization/quant_gemm.cc
Normal file
|
|
@ -0,0 +1,221 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "core/common/safeint.h"
|
||||
#include "core/providers/cpu/math/gemm_base.h"
|
||||
#include "core/providers/cpu/math/gemm_helper.h"
|
||||
#include "core/providers/cpu/math/matmul_integer_base.h"
|
||||
#include "core/quantization/quant_util.h"
|
||||
#include "core/util/math_cpuonly.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace contrib {
|
||||
|
||||
class QGemm : protected GemmBase, public MatMulIntegerBase {
|
||||
public:
|
||||
QGemm(const OpKernelInfo& info) : GemmBase(info), MatMulIntegerBase(info) {
|
||||
}
|
||||
|
||||
Status Compute(OpKernelContext* context) const override {
|
||||
const auto* a = context->Input<Tensor>(IN_A);
|
||||
const auto* b = packed_b_ ? nullptr : context->Input<Tensor>(IN_B);
|
||||
const auto& b_shape = b ? b->Shape() : b_shape_;
|
||||
|
||||
const auto* c = context->Input<Tensor>(IN_C);
|
||||
GemmHelper helper(a->Shape(), trans_A_ != CblasNoTrans,
|
||||
b_shape, trans_B_ != CblasNoTrans,
|
||||
c != nullptr ? c->Shape() : TensorShape({}));
|
||||
if (!helper.State().IsOK())
|
||||
return helper.State();
|
||||
|
||||
size_t M = SafeInt<size_t>(helper.M());
|
||||
size_t N = SafeInt<size_t>(helper.N());
|
||||
size_t K = SafeInt<size_t>(helper.K());
|
||||
|
||||
//validate scales and zero points
|
||||
const auto* a_zp = context->Input<Tensor>(IN_A_ZERO_POINT);
|
||||
const auto* b_zp = context->Input<Tensor>(IN_B_ZERO_POINT);
|
||||
const auto* y_zp = context->Input<Tensor>(IN_Y_ZERO_POINT);
|
||||
const auto* a_scale = context->Input<Tensor>(IN_A_SCALE);
|
||||
const auto* b_scale = context->Input<Tensor>(IN_B_SCALE);
|
||||
const auto* y_scale = context->Input<Tensor>(IN_Y_SCALE);
|
||||
CheckInputs(a_zp, b_zp, y_zp, a_scale, b_scale, y_scale, helper);
|
||||
|
||||
AllocatorPtr allocator;
|
||||
ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator));
|
||||
|
||||
BufferUniquePtr a_trans_buffer;
|
||||
const uint8_t* a_data = a->template Data<uint8_t>();
|
||||
if (trans_A_ == CblasTrans) {
|
||||
a_data = quantization::TransPoseInputData(a_data, a_trans_buffer, allocator, K, M);
|
||||
}
|
||||
|
||||
bool b_is_signed;
|
||||
const uint8_t* b_data = nullptr;
|
||||
BufferUniquePtr b_trans_buffer;
|
||||
if (nullptr == b) {
|
||||
b_data = static_cast<const uint8_t*>(packed_b_.get());
|
||||
b_is_signed = b_is_signed_;
|
||||
} else {
|
||||
b_data = static_cast<const uint8_t*>(b->DataRaw());
|
||||
b_is_signed = b->IsDataType<int8_t>();
|
||||
if (trans_B_ == CblasTrans) {
|
||||
b_data = quantization::TransPoseInputData(b_data, b_trans_buffer, allocator, N, K);
|
||||
}
|
||||
}
|
||||
|
||||
auto y = context->Output(OUT_Y, {SafeInt<int64_t>(M), SafeInt<int64_t>(N)});
|
||||
if (M == 0 || N == 0) return Status::OK();
|
||||
|
||||
// prepare output buffer of GEMM
|
||||
int32_t* gemm_output_data = nullptr;
|
||||
BufferUniquePtr gemm_output_buffer;
|
||||
bool need_requant = y_scale != nullptr;
|
||||
if (need_requant) {
|
||||
gemm_output_data = static_cast<int32_t*>(allocator->Alloc(SafeInt<size_t>(M * N) * sizeof(int32_t)));
|
||||
gemm_output_buffer.reset(gemm_output_data);
|
||||
} else {
|
||||
gemm_output_data = static_cast<int32_t*>(y->MutableDataRaw());
|
||||
}
|
||||
|
||||
if (c != nullptr) {
|
||||
GemmBroadcastBias(M, N, 1.f, c->template Data<int32_t>(), &(c->Shape()), gemm_output_data);
|
||||
}
|
||||
|
||||
MLAS_GEMM_U8X8_SHAPE_PARAMS gemm_shape{M, N, K, b_is_signed, c != nullptr};
|
||||
MLAS_GEMM_U8X8_DATA_PARAMS gemm_param;
|
||||
|
||||
gemm_param.A = a_data;
|
||||
gemm_param.lda = gemm_shape.K;
|
||||
gemm_param.ZeroPointA = *(a_zp->template Data<uint8_t>());
|
||||
|
||||
gemm_param.B = b_data;
|
||||
gemm_param.ldb = gemm_shape.N;
|
||||
gemm_param.BIsPacked = bool(packed_b_);
|
||||
gemm_param.ZeroPointB = static_cast<const uint8_t*>(b_zp->DataRaw());
|
||||
|
||||
gemm_param.C = gemm_output_data;
|
||||
gemm_param.ldc = gemm_shape.N;
|
||||
|
||||
gemm_param.PerColumnZeroPoints = !IsScalarOr1ElementVector(b_zp);
|
||||
|
||||
std::vector<float> output_scales = ComputeOutputScale(a_scale, b_scale, y_scale);
|
||||
std::unique_ptr<MLAS_QGEMM_SCALE_BIAS_OUTPUT_PROCESSOR> scale_bias_proc_ptr;
|
||||
std::unique_ptr<MLAS_QGEMM_REQUANT_OUTPUT_PROCESSOR> requant_proc_ptr;
|
||||
SetPostProcessor(y_zp, N, output_scales, y, gemm_param, scale_bias_proc_ptr, requant_proc_ptr);
|
||||
|
||||
MlasGemmBatch(gemm_shape, &gemm_param, 1, context->GetOperatorThreadPool());
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
protected:
|
||||
int GetBIdx() const override {
|
||||
return IN_B;
|
||||
}
|
||||
|
||||
virtual bool IsBTransposed() const override {
|
||||
return trans_B_ == CblasTrans;
|
||||
}
|
||||
|
||||
private:
|
||||
enum InputTensors : int {
|
||||
IN_A = 0,
|
||||
IN_A_SCALE = 1,
|
||||
IN_A_ZERO_POINT = 2,
|
||||
IN_B = 3,
|
||||
IN_B_SCALE = 4,
|
||||
IN_B_ZERO_POINT = 5,
|
||||
IN_C = 6,
|
||||
IN_Y_SCALE = 7,
|
||||
IN_Y_ZERO_POINT = 8
|
||||
};
|
||||
|
||||
enum OutputTensors : int {
|
||||
OUT_Y = 0
|
||||
};
|
||||
|
||||
static void CheckInputs(const Tensor* a_zp, const Tensor* b_zp, const Tensor* y_zp,
|
||||
const Tensor* a_scale, const Tensor* b_scale, const Tensor* y_scale, const GemmHelper& helper) {
|
||||
ORT_ENFORCE(IsScalarOr1ElementVector(a_scale),
|
||||
"QGemm : scale of input a must be a scalar or 1D tensor of size 1");
|
||||
ORT_ENFORCE(IsScalarOr1ElementVector(a_zp),
|
||||
"QGemm : zero point of input a must be a scalar or 1D tensor of size 1");
|
||||
|
||||
const auto& b_zp_shape = b_zp->Shape();
|
||||
const auto& b_scale_shape = b_scale->Shape();
|
||||
ORT_ENFORCE(b_zp_shape.NumDimensions() == 0 ||
|
||||
(b_zp_shape.NumDimensions() == 1 && (b_zp_shape[0] == 1 || b_zp_shape[0] == helper.N())),
|
||||
"QGemm : zero point of input b must be a scalar or 1D tensor of size 1 or N");
|
||||
ORT_ENFORCE(b_scale_shape.NumDimensions() == 0 ||
|
||||
(b_scale_shape.NumDimensions() == 1 && (b_scale_shape[0] == 1 || b_scale_shape[0] == helper.N())),
|
||||
"QGemm : scale of input b must be a scalar or 1D tensor of size 1 or N");
|
||||
ORT_ENFORCE(b_scale_shape.NumDimensions() == b_zp_shape.NumDimensions() &&
|
||||
(b_scale_shape.NumDimensions() == 0 || (b_scale_shape[0] == b_zp_shape[0])),
|
||||
"QGemm : zero point and scale of input b should have same shape size");
|
||||
|
||||
ORT_ENFORCE(y_zp == nullptr || IsScalarOr1ElementVector(y_zp),
|
||||
"QGemm : zero point of y must be null or a scalar or 1D tensor of size 1");
|
||||
ORT_ENFORCE(y_scale == nullptr || IsScalarOr1ElementVector(y_scale),
|
||||
"QGemm : scale of y must be null or a scalar or 1D tensor of size 1");
|
||||
}
|
||||
|
||||
std::vector<float> ComputeOutputScale(const Tensor* a_scale, const Tensor* b_scale, const Tensor* y_scale) const {
|
||||
const int64_t output_scale_size = b_scale->Shape().Size();
|
||||
std::vector<float> output_scales(output_scale_size);
|
||||
auto a_scale_value = *(a_scale->template Data<float>());
|
||||
const auto* b_scale_data = b_scale->template Data<float>();
|
||||
for (int64_t i = 0; i < output_scale_size; i++) {
|
||||
output_scales[i] = (alpha_ * a_scale_value * b_scale_data[i]);
|
||||
if (nullptr != y_scale) {
|
||||
output_scales[i] /= *(y_scale->template Data<float>());
|
||||
}
|
||||
}
|
||||
return output_scales;
|
||||
}
|
||||
|
||||
static void SetPostProcessor(const Tensor* y_zp,
|
||||
size_t out_lda,
|
||||
const std::vector<float>& output_scales,
|
||||
Tensor* y,
|
||||
MLAS_GEMM_U8X8_DATA_PARAMS& gemm_param,
|
||||
std::unique_ptr<MLAS_QGEMM_SCALE_BIAS_OUTPUT_PROCESSOR>& scale_bias_proc_ptr,
|
||||
std::unique_ptr<MLAS_QGEMM_REQUANT_OUTPUT_PROCESSOR>& requant_proc_ptr) {
|
||||
if (nullptr != y_zp) {
|
||||
requant_proc_ptr = std::make_unique<MLAS_QGEMM_REQUANT_OUTPUT_PROCESSOR>(
|
||||
static_cast<uint8_t*>(y->MutableDataRaw()),
|
||||
out_lda,
|
||||
nullptr,
|
||||
output_scales.data(),
|
||||
output_scales.size() > 1,
|
||||
*y_zp->template Data<uint8_t>());
|
||||
gemm_param.OutputProcessor = requant_proc_ptr.get();
|
||||
} else {
|
||||
scale_bias_proc_ptr = std::make_unique<MLAS_QGEMM_SCALE_BIAS_OUTPUT_PROCESSOR>(
|
||||
static_cast<float*>(y->MutableDataRaw()),
|
||||
out_lda,
|
||||
output_scales.data(),
|
||||
nullptr,
|
||||
MLAS_QGEMM_OUTPUT_MODE::ZeroMode,
|
||||
output_scales.size() > 1 ? MLAS_QUANTIZATION_GRANULARITY::PerColumn : MLAS_QUANTIZATION_GRANULARITY::PerMatrix);
|
||||
gemm_param.OutputProcessor = scale_bias_proc_ptr.get();
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
ONNX_OPERATOR_TYPED_KERNEL_EX(
|
||||
QGemm,
|
||||
kMSDomain,
|
||||
1,
|
||||
uint8_t,
|
||||
kCpuExecutionProvider,
|
||||
KernelDefBuilder()
|
||||
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>())
|
||||
.TypeConstraint("TA", DataTypeImpl::GetTensorType<uint8_t>())
|
||||
.TypeConstraint("TB", {DataTypeImpl::GetTensorType<uint8_t>(), DataTypeImpl::GetTensorType<int8_t>()})
|
||||
.TypeConstraint("TC", DataTypeImpl::GetTensorType<int32_t>())
|
||||
.TypeConstraint("TYZ", DataTypeImpl::GetTensorType<uint8_t>())
|
||||
.TypeConstraint("TY", {DataTypeImpl::GetTensorType<float>(), DataTypeImpl::GetTensorType<uint8_t>()}),
|
||||
QGemm);
|
||||
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -862,6 +862,130 @@ Wwhere the function `Sigmoid(x) = 1 / (1 + exp(-x))` )DOC";
|
|||
output_shape->mutable_dim(axis)->set_dim_value(total_length);
|
||||
}
|
||||
});
|
||||
|
||||
ONNX_CONTRIB_OPERATOR_SCHEMA(QGemm)
|
||||
.SetDomain(kMSDomain)
|
||||
.SinceVersion(1)
|
||||
.SetDoc("Quantized Gemm")
|
||||
.Input(0,
|
||||
"A",
|
||||
"Input tensor A. "
|
||||
"The shape of A should be (M, K) if transA is 0, "
|
||||
"or (K, M) if transA is non-zero.",
|
||||
"TA")
|
||||
.Input(1,
|
||||
"a_scale",
|
||||
"Scale of quantized input 'A'. "
|
||||
"It is a scalar,which means a per-tensor quantization.",
|
||||
"T")
|
||||
.Input(2,
|
||||
"a_zero_point",
|
||||
"Zero point tensor for input 'A'. It is a scalar.",
|
||||
"TA")
|
||||
.Input(3,
|
||||
"B",
|
||||
"Input tensor B. "
|
||||
"The shape of B should be (K, N) if transB is 0, "
|
||||
"or (N, K) if transB is non-zero.",
|
||||
"TB")
|
||||
.Input(4,
|
||||
"b_scale",
|
||||
"Scale of quantized input 'B'. It could be a scalar or a 1-D tensor, "
|
||||
"which means a per-tensor or per-column quantization. If it's a 1-D tensor, its number "
|
||||
"of elements should be equal to the number of columns of input 'B'.",
|
||||
"T")
|
||||
.Input(5,
|
||||
"b_zero_point",
|
||||
"Zero point tensor for input 'B'. It's optional and default value is 0. It could be a scalar or a 1-D tensor, "
|
||||
"which means a per-tensor or per-column quantization. If it's a 1-D tensor, its number "
|
||||
"of elements should be equal to the number of columns of input 'B'.",
|
||||
"TB")
|
||||
.Input(6,
|
||||
"C",
|
||||
"Optional input tensor C. "
|
||||
"If not specified, the computation is done as if C is a scalar 0. "
|
||||
"The shape of C should be unidirectional broadcastable to (M, N). "
|
||||
"Its type is int32_t and must be quantized with zero_point = 0 and "
|
||||
"scale = alpha / beta * a_scale * b_scale.",
|
||||
"TC",
|
||||
OpSchema::Optional)
|
||||
.Input(7,
|
||||
"y_scale",
|
||||
"Scale of output 'Y'. It is a scalar, which means a per-tensor quantization. "
|
||||
"It is optional. The output is full precision(float32) if it is not provided. "
|
||||
"Or the output is quantized.",
|
||||
"T",
|
||||
OpSchema::Optional)
|
||||
.Input(8,
|
||||
"y_zero_point",
|
||||
"Zero point tensor for output 'Y'. It is a scalar, which means a per-tensor quantization. "
|
||||
"It is optional. The output is full precision(float32) if it is not provided. "
|
||||
"Or the output is quantized.",
|
||||
"TYZ",
|
||||
OpSchema::Optional)
|
||||
.Output(0,
|
||||
"Y",
|
||||
"Output tensor of shape (M, N).",
|
||||
"TY")
|
||||
.Attr("transA",
|
||||
"Whether A should be transposed",
|
||||
AttributeProto::INT,
|
||||
static_cast<int64_t>(0))
|
||||
.Attr("transB",
|
||||
"Whether B should be transposed",
|
||||
AttributeProto::INT,
|
||||
static_cast<int64_t>(0))
|
||||
.Attr("alpha",
|
||||
"Scalar multiplier for the product of input tensors A * B.",
|
||||
AttributeProto::FLOAT,
|
||||
1.0f)
|
||||
.TypeConstraint("T",
|
||||
{"tensor(float)"},
|
||||
"Constrain scale types to float tensors.")
|
||||
.TypeConstraint("TA",
|
||||
{"tensor(uint8)", "tensor(int8)"},
|
||||
"Constrain input A and its zero point types to 8 bit tensors.")
|
||||
.TypeConstraint("TB",
|
||||
{"tensor(uint8)", "tensor(int8)"},
|
||||
"Constrain input B and its zero point types to 8 bit tensors.")
|
||||
.TypeConstraint("TC",
|
||||
{"tensor(int32)"},
|
||||
"Constrain input C to 32 bit integer tensors.")
|
||||
.TypeConstraint("TYZ",
|
||||
{"tensor(uint8)", "tensor(int8)"},
|
||||
"Constrain output zero point types to 8 bit tensors.")
|
||||
.TypeConstraint("TY",
|
||||
{"tensor(float)", "tensor(uint8)", "tensor(int8)"},
|
||||
"Constrain output type to float32 or 8 bit tensors.")
|
||||
.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
|
||||
if (ctx.getNumInputs() == 9 && nullptr != ctx.getInputType(8)) {
|
||||
propagateElemTypeFromInputToOutput(ctx, 8, 0);
|
||||
} else {
|
||||
updateOutputElemType(ctx, 0, ONNX_NAMESPACE::TensorProto::FLOAT);
|
||||
}
|
||||
|
||||
if (hasInputShape(ctx, 0) && hasInputShape(ctx, 3)) {
|
||||
auto transAAttr = ctx.getAttribute("transA");
|
||||
bool transA =
|
||||
transAAttr ? static_cast<int>(transAAttr->i()) != 0 : false;
|
||||
auto transBAttr = ctx.getAttribute("transB");
|
||||
bool transB =
|
||||
transBAttr ? static_cast<int>(transBAttr->i()) != 0 : false;
|
||||
auto& first_input_shape = getInputShape(ctx, 0);
|
||||
auto& second_input_shape = getInputShape(ctx, 3);
|
||||
if (first_input_shape.dim_size() != 2) {
|
||||
fail_shape_inference("First input does not have rank 2");
|
||||
}
|
||||
if (second_input_shape.dim_size() != 2) {
|
||||
fail_shape_inference("Second input does not have rank 2");
|
||||
}
|
||||
updateOutputShape(
|
||||
ctx,
|
||||
0,
|
||||
{first_input_shape.dim(transA ? 1 : 0),
|
||||
second_input_shape.dim(transB ? 0 : 1)});
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
} // namespace contrib
|
||||
|
|
|
|||
|
|
@ -527,6 +527,7 @@ struct MLAS_GEMM_U8X8_SHAPE_PARAMS {
|
|||
size_t N = 0;
|
||||
size_t K = 0;
|
||||
bool BIsSigned = false;
|
||||
bool IsAccumulateMode = false;
|
||||
};
|
||||
|
||||
struct MLAS_GEMM_U8X8_DATA_PARAMS {
|
||||
|
|
|
|||
|
|
@ -234,6 +234,7 @@ Return Value:
|
|||
int32_t* C = Data->C + RangeStartM * ldc + RangeStartN;
|
||||
const uint8_t* PackedZeroPointB = Data->PerColumnZeroPoints ?
|
||||
Data->ZeroPointB + RangeStartN : nullptr;
|
||||
bool IsAccumulateMode = Shape->IsAccumulateMode;
|
||||
|
||||
int32_t ZeroPointA = Data->ZeroPointA;
|
||||
int32_t ZeroPointB = typename KernelType::OffsetBType(*Data->ZeroPointB);
|
||||
|
|
@ -362,7 +363,7 @@ Return Value:
|
|||
int32_t* RowSums = RowSumBuffer;
|
||||
size_t RowsRemaining = CountM;
|
||||
|
||||
bool ZeroMode = (k == 0);
|
||||
bool ZeroMode = (k == 0) && !IsAccumulateMode;
|
||||
bool PostProcess = (k + CountK == K);
|
||||
|
||||
while (RowsRemaining > 0) {
|
||||
|
|
@ -459,6 +460,7 @@ Return Value:
|
|||
int32_t* C = Data->C + RangeStartM * ldc + RangeStartN;
|
||||
const uint8_t* PackedZeroPointB = Data->PerColumnZeroPoints ?
|
||||
Data->ZeroPointB + RangeStartN : nullptr;
|
||||
bool IsAccumulateMode = Shape->IsAccumulateMode;
|
||||
|
||||
int32_t ZeroPointA = Data->ZeroPointA;
|
||||
int32_t ZeroPointB = typename KernelType::OffsetBType(*Data->ZeroPointB);
|
||||
|
|
@ -581,7 +583,7 @@ Return Value:
|
|||
int32_t* RowSums = RowSumBuffer;
|
||||
size_t RowsRemaining = CountM;
|
||||
|
||||
bool ZeroMode = (k == 0);
|
||||
bool ZeroMode = (k == 0) && !IsAccumulateMode;
|
||||
bool PostProcess = (k + CountK == K);
|
||||
|
||||
while (RowsRemaining > 0) {
|
||||
|
|
|
|||
|
|
@ -108,30 +108,6 @@ bool GemmPackBFp32(AllocatorPtr& alloc,
|
|||
return true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void GemmBroadcastBias(int64_t M, int64_t N, float beta,
|
||||
const T* c_data, const TensorShape* c_shape,
|
||||
T* y_data) {
|
||||
// Broadcast the bias as needed if bias is given
|
||||
if (beta != 0 && c_data != nullptr) {
|
||||
ORT_ENFORCE(c_shape != nullptr, "c_shape is required if c_data is provided");
|
||||
auto output_mat = EigenMatrixMapRowMajor<T>(y_data, M, N);
|
||||
if (c_shape->Size() == 1) {
|
||||
// C is (), (1,) or (1, 1), set the scalar
|
||||
output_mat.setConstant(*c_data);
|
||||
} else if (c_shape->NumDimensions() == 1 || (*c_shape)[0] == 1) {
|
||||
// C is (N,) or (1, N)
|
||||
output_mat.rowwise() = ConstEigenVectorMap<T>(c_data, N).transpose();
|
||||
} else if ((*c_shape)[1] == 1) {
|
||||
// C is (M, 1)
|
||||
output_mat.colwise() = ConstEigenVectorMap<T>(c_data, M);
|
||||
} else {
|
||||
// C is (M, N), no broadcast needed.
|
||||
output_mat = ConstEigenMatrixMapRowMajor<T>(c_data, M, N);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void Gemm<T>::ComputeGemm(CBLAS_TRANSPOSE trans_a, CBLAS_TRANSPOSE trans_b,
|
||||
int64_t M, int64_t N, int64_t K,
|
||||
|
|
|
|||
|
|
@ -3,6 +3,8 @@
|
|||
|
||||
#pragma once
|
||||
|
||||
#include "gemm_base.h"
|
||||
|
||||
#include "core/framework/op_kernel.h"
|
||||
#include "core/common/common.h"
|
||||
#include "core/util/math.h"
|
||||
|
|
@ -11,18 +13,9 @@
|
|||
namespace onnxruntime {
|
||||
|
||||
template <typename T>
|
||||
class Gemm : public OpKernel {
|
||||
class Gemm : protected GemmBase, public OpKernel {
|
||||
public:
|
||||
Gemm(const OpKernelInfo& info) : OpKernel(info) {
|
||||
int64_t temp;
|
||||
ORT_ENFORCE(info.GetAttr<int64_t>("transA", &temp).IsOK());
|
||||
trans_A_ = temp == 0 ? CblasNoTrans : CblasTrans;
|
||||
|
||||
ORT_ENFORCE(info.GetAttr<int64_t>("transB", &temp).IsOK());
|
||||
trans_B_ = temp == 0 ? CblasNoTrans : CblasTrans;
|
||||
|
||||
ORT_ENFORCE(info.GetAttr<float>("alpha", &alpha_).IsOK());
|
||||
ORT_ENFORCE(info.GetAttr<float>("beta", &beta_).IsOK());
|
||||
Gemm(const OpKernelInfo& info) : GemmBase(info), OpKernel(info) {
|
||||
}
|
||||
|
||||
Status Compute(OpKernelContext* context) const override;
|
||||
|
|
@ -44,12 +37,6 @@ class Gemm : public OpKernel {
|
|||
T* y_data,
|
||||
concurrency::ThreadPool* thread_pool);
|
||||
|
||||
private:
|
||||
CBLAS_TRANSPOSE trans_A_;
|
||||
CBLAS_TRANSPOSE trans_B_;
|
||||
float alpha_;
|
||||
float beta_;
|
||||
|
||||
protected:
|
||||
TensorShape b_shape_;
|
||||
BufferUniquePtr packed_b_;
|
||||
|
|
|
|||
32
onnxruntime/core/providers/cpu/math/gemm_base.h
Normal file
32
onnxruntime/core/providers/cpu/math/gemm_base.h
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "core/framework/op_kernel.h"
|
||||
#include "core/common/common.h"
|
||||
#include "core/util/math.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
class GemmBase {
|
||||
protected:
|
||||
GemmBase(const OpKernelInfo& info) {
|
||||
int64_t temp;
|
||||
ORT_ENFORCE(info.GetAttr<int64_t>("transA", &temp).IsOK());
|
||||
trans_A_ = temp == 0 ? CblasNoTrans : CblasTrans;
|
||||
|
||||
ORT_ENFORCE(info.GetAttr<int64_t>("transB", &temp).IsOK());
|
||||
trans_B_ = temp == 0 ? CblasNoTrans : CblasTrans;
|
||||
|
||||
ORT_ENFORCE(info.GetAttr<float>("alpha", &alpha_).IsOK());
|
||||
info.GetAttrOrDefault<float>("beta", &beta_, 1.f);
|
||||
}
|
||||
|
||||
CBLAS_TRANSPOSE trans_A_;
|
||||
CBLAS_TRANSPOSE trans_B_;
|
||||
float alpha_;
|
||||
float beta_;
|
||||
};
|
||||
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -3,6 +3,7 @@
|
|||
|
||||
#pragma once
|
||||
#include "core/common/common.h"
|
||||
#include "core/util/math_cpuonly.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
|
|
@ -71,4 +72,32 @@ class GemmHelper {
|
|||
Status status_;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
void GemmBroadcastBias(int64_t M, int64_t N, float beta,
|
||||
const T* c_data, const TensorShape* c_shape,
|
||||
T* y_data)
|
||||
{
|
||||
// Broadcast the bias as needed if bias is given
|
||||
if (beta != 0 && c_data != nullptr) {
|
||||
ORT_ENFORCE(c_shape != nullptr, "c_shape is required if c_data is provided");
|
||||
auto output_mat = EigenMatrixMapRowMajor<T>(y_data, M, N);
|
||||
if (c_shape->Size() == 1) {
|
||||
// C is (), (1,) or (1, 1), set the scalar
|
||||
output_mat.setConstant(*c_data);
|
||||
}
|
||||
else if (c_shape->NumDimensions() == 1 || (*c_shape)[0] == 1) {
|
||||
// C is (N,) or (1, N)
|
||||
output_mat.rowwise() = ConstEigenVectorMap<T>(c_data, N).transpose();
|
||||
}
|
||||
else if ((*c_shape)[1] == 1) {
|
||||
// C is (M, 1)
|
||||
output_mat.colwise() = ConstEigenVectorMap<T>(c_data, M);
|
||||
}
|
||||
else {
|
||||
// C is (M, N), no broadcast needed.
|
||||
output_mat = ConstEigenMatrixMapRowMajor<T>(c_data, M, N);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ class MatMulInteger final : public MatMulIntegerBase {
|
|||
enum OutputTensors : int { OUT_Y = 0 };
|
||||
|
||||
protected:
|
||||
int GetBIdx() override { return IN_B; }
|
||||
int GetBIdx() const override { return IN_B; }
|
||||
};
|
||||
|
||||
ONNX_OPERATOR_TYPED_KERNEL_EX(
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
#include "core/framework/op_kernel.h"
|
||||
#include "core/mlas/inc/mlas.h"
|
||||
#include "core/providers/common.h"
|
||||
#include "core/quantization/quant_util.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
|
|
@ -27,11 +28,16 @@ class MatMulIntegerBase : public OpKernel {
|
|||
|
||||
b_is_signed_ = tensor.IsDataType<int8_t>();
|
||||
|
||||
const size_t K = static_cast<size_t>(b_shape_[0]);
|
||||
const size_t N = static_cast<size_t>(b_shape_[1]);
|
||||
size_t K = static_cast<size_t>(b_shape_[0]);
|
||||
size_t N = static_cast<size_t>(b_shape_[1]);
|
||||
|
||||
const auto* b_data = static_cast<const uint8_t*>(tensor.DataRaw());
|
||||
|
||||
BufferUniquePtr b_trans_buffer;
|
||||
if (IsBTransposed()) {
|
||||
std::swap(K, N);
|
||||
b_data = quantization::TransPoseInputData(b_data, b_trans_buffer, alloc, N, K);
|
||||
}
|
||||
const size_t packed_b_size = MlasGemmPackBSize(N, K, b_is_signed_);
|
||||
if (packed_b_size == 0) {
|
||||
return Status::OK();
|
||||
|
|
@ -75,7 +81,11 @@ class MatMulIntegerBase : public OpKernel {
|
|||
/**
|
||||
* @return input index of Matrix B, the weight tensor
|
||||
*/
|
||||
virtual int GetBIdx() = 0;
|
||||
virtual int GetBIdx() const = 0;
|
||||
|
||||
virtual bool IsBTransposed() const {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Check if quantization parameter of B is supported.
|
||||
// It should be in one of the formats below:
|
||||
|
|
|
|||
|
|
@ -33,7 +33,7 @@ class QLinearMatMul : public MatMulIntegerBase {
|
|||
};
|
||||
|
||||
protected:
|
||||
int GetBIdx() override {
|
||||
int GetBIdx() const override {
|
||||
return IN_B;
|
||||
}
|
||||
};
|
||||
|
|
|
|||
22
onnxruntime/core/quantization/quant_util.h
Normal file
22
onnxruntime/core/quantization/quant_util.h
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "core/common/common.h"
|
||||
#include "core/mlas/inc/mlas.h"
|
||||
#include "core/framework/allocator.h"
|
||||
#include "core/framework/buffer_deleter.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace quantization {
|
||||
// Transpose the input and store it to a new allocated buffer
|
||||
inline uint8_t* TransPoseInputData(const uint8_t* input, BufferUniquePtr& buffer_holder, AllocatorPtr& allocator, size_t M, size_t N) {
|
||||
uint8_t* output = static_cast<uint8_t*>(allocator->Alloc(M * N * sizeof(uint8_t)));
|
||||
MlasTranspose(input, output, M, N);
|
||||
buffer_holder.reset(output);
|
||||
return output;
|
||||
}
|
||||
|
||||
} // namespace quantization
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -42,6 +42,20 @@ inline std::vector<T> QuantizeLinearTestVector(
|
|||
return result;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
std::vector<T> ToVector(const int* value, int size) {
|
||||
std::vector<T> data(size);
|
||||
for (int i = 0; i < size; i++) {
|
||||
data[i] = static_cast<T>(value[i]);
|
||||
}
|
||||
return data;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T GetMiddle(const std::vector<T>& v) {
|
||||
const auto min_max_pair = std::minmax_element(v.begin(), v.end());
|
||||
return (*(min_max_pair.first) + *(min_max_pair.second)) / 2;
|
||||
}
|
||||
|
||||
} // namespace test
|
||||
} // namespace onnxruntime
|
||||
|
||||
|
|
|
|||
196
onnxruntime/test/contrib_ops/quant_gemm_test.cc
Normal file
196
onnxruntime/test/contrib_ops/quant_gemm_test.cc
Normal file
|
|
@ -0,0 +1,196 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
#include "test/common/cuda_op_test_utils.h"
|
||||
#include "test/common/quantization_test_utils.h"
|
||||
#include "test/providers/provider_test_utils.h"
|
||||
|
||||
#include "core/common/common.h"
|
||||
#include "core/common/safeint.h"
|
||||
#include "core/framework/op_kernel.h"
|
||||
#include "core/quantization/quantization.h"
|
||||
#include "core/util/math_cpuonly.h"
|
||||
#include "core/util/qmath.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <random>
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace test {
|
||||
|
||||
template <typename ScalarB, typename ScalarOutput>
|
||||
void RunQuantGemmU8X8Test(const int M,
|
||||
const int N,
|
||||
const int K,
|
||||
bool is_A_trans,
|
||||
bool is_B_trans,
|
||||
bool has_C,
|
||||
bool B_is_initializer,
|
||||
bool per_column = false) {
|
||||
static std::default_random_engine e(123);
|
||||
static std::uniform_int_distribution<int> n_unsigned(0, 127);
|
||||
static std::uniform_int_distribution<int> n_xint8(std::numeric_limits<ScalarB>::min(), std::numeric_limits<ScalarB>::max());
|
||||
static std::uniform_real_distribution<float> n_apha(1.0f, 2.0f);
|
||||
static std::uniform_real_distribution<float> n_scale(0.003f, 0.004f);
|
||||
|
||||
Eigen::MatrixXi matrix_a = Eigen::MatrixXi::Random(K, M)
|
||||
.unaryExpr([](int) { return n_unsigned(e); });
|
||||
std::vector<uint8_t> matrix_a_data;
|
||||
if (is_A_trans) {
|
||||
Eigen::MatrixXi matrix_a_trans = matrix_a.transpose().eval();
|
||||
matrix_a_data = ToVector<uint8_t>(matrix_a_trans.data(), M * K);
|
||||
} else {
|
||||
matrix_a_data = ToVector<uint8_t>(matrix_a.data(), M * K);
|
||||
}
|
||||
uint8_t a_zero_point = GetMiddle(matrix_a_data);
|
||||
Eigen::MatrixXi matrix_a_offset = matrix_a - a_zero_point * Eigen::MatrixXi::Ones(K, M);
|
||||
float a_scale = n_scale(e);
|
||||
|
||||
Eigen::MatrixXi matrix_b = Eigen::MatrixXi::Random(N, K)
|
||||
.unaryExpr([](int) { return n_xint8(e); });
|
||||
std::vector<ScalarB> matrix_b_data;
|
||||
if (is_B_trans) {
|
||||
Eigen::MatrixXi matrix_b_trans = matrix_b.transpose().eval();
|
||||
matrix_b_data = ToVector<ScalarB>(matrix_b_trans.data(), N * K);
|
||||
} else {
|
||||
matrix_b_data = ToVector<ScalarB>(matrix_b.data(), N * K);
|
||||
}
|
||||
ScalarB b_zero_point = GetMiddle(matrix_b_data);
|
||||
std::vector<float> b_scale({n_scale(e)});
|
||||
std::vector<ScalarB> b_zp_per_column({b_zero_point});
|
||||
Eigen::MatrixXi b_zp_matrix = b_zero_point * Eigen::MatrixXi::Ones(N, K);
|
||||
Eigen::MatrixXf b_scale_matrix = b_scale[0] * Eigen::MatrixXf::Ones(N, M);
|
||||
if (per_column) {
|
||||
b_zp_per_column.resize(N);
|
||||
b_scale.resize(N);
|
||||
for (int i = 0; i < N; i++) {
|
||||
b_zp_per_column[i] = b_zero_point + i % 2 == 0 ? 1 : -1;
|
||||
b_zp_matrix.row(i).setConstant(b_zp_per_column[i]);
|
||||
b_scale[i] = n_scale(e);
|
||||
b_scale_matrix.row(i).setConstant(b_scale[i]);
|
||||
}
|
||||
}
|
||||
float alpha = n_apha(e);
|
||||
|
||||
Eigen::MatrixXi matrix_c = Eigen::MatrixXi::Random(N, M)
|
||||
.unaryExpr([](int) { return n_xint8(e); });
|
||||
|
||||
Eigen::MatrixXi matrix_int32 = (matrix_b - b_zp_matrix) * matrix_a_offset;
|
||||
if (has_C) {
|
||||
matrix_int32 = matrix_int32 + matrix_c;
|
||||
}
|
||||
Eigen::MatrixXf matrix_output = alpha * a_scale * (b_scale_matrix.cwiseProduct((matrix_int32.eval().cast<float>())));
|
||||
|
||||
OpTester test("QGemm", 1, onnxruntime::kMSDomain);
|
||||
test.AddAttribute<int64_t>("transA", is_A_trans ? 1 : 0);
|
||||
test.AddAttribute<int64_t>("transB", is_B_trans ? 1 : 0);
|
||||
test.AddAttribute<float>("alpha", alpha);
|
||||
test.AddInput<uint8_t>("A", is_A_trans ? std::vector<int64_t>({K, M}) : std::vector<int64_t>({M, K}), std::move(matrix_a_data));
|
||||
test.AddInput<float>("a_scale", {}, {a_scale});
|
||||
test.AddInput<uint8_t>("a_zero_point", {}, {a_zero_point});
|
||||
test.AddInput<ScalarB>("B", is_B_trans ? std::vector<int64_t>({N, K}) : std::vector<int64_t>({K, N}), std::move(matrix_b_data), B_is_initializer);
|
||||
test.AddInput<float>("b_scale", {SafeInt<int64_t>(b_scale.size())}, b_scale);
|
||||
test.AddInput<ScalarB>("b_zero_point", {SafeInt<int64_t>(b_zp_per_column.size())}, b_zp_per_column);
|
||||
|
||||
if (has_C) {
|
||||
test.AddInput<int32_t>("C", {M, N}, ToVector<int32_t>(matrix_c.data(), M * N));
|
||||
} else {
|
||||
test.AddOptionalInputEdge<int32_t>();
|
||||
}
|
||||
|
||||
if constexpr (std::is_same_v<ScalarOutput, float>) {
|
||||
test.AddOptionalInputEdge<float>();
|
||||
test.AddOptionalInputEdge<uint8_t>();
|
||||
test.AddOutput<float>("Y", {M, N}, std::vector<float>(matrix_output.data(), matrix_output.data() + M * N));
|
||||
} else {
|
||||
std::vector<uint8_t> quant_output(M * N);
|
||||
quantization::Params<uint8_t> quant_param = quantization::QuantizeLinear(matrix_output.data(), quant_output.data(), M * N);
|
||||
test.AddInput<float>("y_scale", {}, {quant_param.scale});
|
||||
test.AddInput<uint8_t>("y_zero_point", {}, {quant_param.zero_point});
|
||||
test.AddOutput<uint8_t>("Y", {M, N}, quant_output);
|
||||
}
|
||||
|
||||
test.Run();
|
||||
}
|
||||
|
||||
void RunQuantGemmTest(const int M,
|
||||
const int N,
|
||||
const int K,
|
||||
bool is_A_trans,
|
||||
bool is_B_trans,
|
||||
bool has_C,
|
||||
bool B_is_initializer,
|
||||
bool per_column = false) {
|
||||
RunQuantGemmU8X8Test<int8_t, float>(M, N, K, is_A_trans, is_B_trans, has_C, B_is_initializer, per_column);
|
||||
RunQuantGemmU8X8Test<int8_t, uint8_t>(M, N, K, is_A_trans, is_B_trans, has_C, B_is_initializer, per_column);
|
||||
RunQuantGemmU8X8Test<uint8_t, float>(M, N, K, is_A_trans, is_B_trans, has_C, B_is_initializer, per_column);
|
||||
RunQuantGemmU8X8Test<uint8_t, uint8_t>(M, N, K, is_A_trans, is_B_trans, has_C, B_is_initializer, per_column);
|
||||
}
|
||||
|
||||
void RunQuantGemmTestBatch(const int M, const int N, const int K) {
|
||||
// No Trans
|
||||
RunQuantGemmTest(M, N, K,
|
||||
false /*is_A_trans*/, false /*is_B_trans*/,
|
||||
false /*has_C*/, false /*B_is_initializer*/,
|
||||
false /*per_column*/);
|
||||
RunQuantGemmTest(M, N, K,
|
||||
false /*is_A_trans*/, false /*is_B_trans*/,
|
||||
true /*has_C*/, true /*B_is_initializer*/,
|
||||
true /*per_column*/);
|
||||
|
||||
// A Trans
|
||||
RunQuantGemmTest(M, N, K,
|
||||
true /*is_A_trans*/, false /*is_B_trans*/,
|
||||
false /*has_C*/, true /*B_is_initializer*/,
|
||||
false /*per_column*/);
|
||||
RunQuantGemmTest(M, N, K,
|
||||
true /*is_A_trans*/, false /*is_B_trans*/,
|
||||
true /*has_C*/, false /*B_is_initializer*/,
|
||||
true /*per_column*/);
|
||||
|
||||
// B Trans
|
||||
RunQuantGemmTest(M, N, K,
|
||||
false /*is_A_trans*/, true /*is_B_trans*/,
|
||||
false /*has_C*/, false /*B_is_initializer*/,
|
||||
false /*per_column*/);
|
||||
RunQuantGemmTest(M, N, K,
|
||||
false /*is_A_trans*/, true /*is_B_trans*/,
|
||||
true /*has_C*/, true /*B_is_initializer*/, // B uses prepacking
|
||||
true /*per_column*/);
|
||||
|
||||
// A and B Trans
|
||||
RunQuantGemmTest(M, N, K,
|
||||
true /*is_A_trans*/, true /*is_B_trans*/,
|
||||
false /*has_C*/, true /*B_is_initializer*/,
|
||||
true /*per_column*/);
|
||||
RunQuantGemmTest(M, N, K,
|
||||
true /*is_A_trans*/, true /*is_B_trans*/,
|
||||
true /*has_C*/, false /*B_is_initializer*/,
|
||||
false /*per_column*/);
|
||||
}
|
||||
|
||||
TEST(QuantGemmTest, Scalar) {
|
||||
RunQuantGemmTestBatch(1, 1, 32);
|
||||
RunQuantGemmTestBatch(1, 1, 260);
|
||||
RunQuantGemmTestBatch(1, 1, 288);
|
||||
}
|
||||
|
||||
TEST(QuantGemmTest, GEMV) {
|
||||
RunQuantGemmTestBatch(1, 2, 16);
|
||||
RunQuantGemmTestBatch(1, 2, 64);
|
||||
RunQuantGemmTestBatch(1, 8, 36);
|
||||
RunQuantGemmTestBatch(1, 8, 68);
|
||||
RunQuantGemmTestBatch(1, 8, 400);
|
||||
RunQuantGemmTestBatch(1, 512, 1024);
|
||||
}
|
||||
|
||||
TEST(QuantGemmTest, GEMM) {
|
||||
RunQuantGemmTestBatch(2, 2, 40);
|
||||
RunQuantGemmTestBatch(2, 48, 33);
|
||||
RunQuantGemmTestBatch(2, 51, 40);
|
||||
RunQuantGemmTestBatch(4, 8, 68);
|
||||
}
|
||||
|
||||
} // namespace test
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -3,6 +3,7 @@
|
|||
|
||||
#include "gtest/gtest.h"
|
||||
#include "test/common/cuda_op_test_utils.h"
|
||||
#include "test/common/quantization_test_utils.h"
|
||||
#include "test/providers/provider_test_utils.h"
|
||||
|
||||
#include "core/common/common.h"
|
||||
|
|
@ -293,21 +294,6 @@ TEST(MatmulIntegerOpTest, MatMulInteger_PerColumn_ND) {
|
|||
test.Run();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
std::vector<T> ToVector(const int* value, int size) {
|
||||
std::vector<T> data(size);
|
||||
for (int i = 0; i < size; i++) {
|
||||
data[i] = static_cast<T>(value[i]);
|
||||
}
|
||||
return data;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T GetMiddle(const std::vector<T>& v) {
|
||||
const auto min_max_pair = std::minmax_element(v.begin(), v.end());
|
||||
return (*(min_max_pair.first) + *(min_max_pair.second)) / 2;
|
||||
}
|
||||
|
||||
// [M x N] = [M x K] x [K x N] = [batch_seq x input_dim] x [input_dim x embed_dim]
|
||||
template <typename ScalarB>
|
||||
void RunMatMulIntegerU8X8Test(const int M, const int N, const int K, bool non_zero_zp, bool B_is_initializer, bool per_column_zp = false) {
|
||||
|
|
|
|||
|
|
@ -260,7 +260,11 @@
|
|||
1734858160766311432
|
||||
],
|
||||
[
|
||||
"QEmbedLayerNormalization com.microsoft CPUExecutionProvider",
|
||||
9235385557940152248
|
||||
"QEmbedLayerNormalization com.microsoft CPUExecutionProvider",
|
||||
9235385557940152248
|
||||
],
|
||||
[
|
||||
"QGemm com.microsoft CPUExecutionProvider",
|
||||
13737193491843065240
|
||||
]
|
||||
]
|
||||
|
|
|
|||
Loading…
Reference in a new issue