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:
Yufeng Li 2021-07-27 21:21:49 -07:00 committed by GitHub
parent 0f46b08646
commit ceeb1a65d6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
20 changed files with 743 additions and 68 deletions

View file

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

View file

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

View file

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

View file

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

View 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

View file

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

View file

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

View file

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

View file

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

View file

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

View 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

View file

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

View file

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

View file

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

View file

@ -33,7 +33,7 @@ class QLinearMatMul : public MatMulIntegerBase {
};
protected:
int GetBIdx() override {
int GetBIdx() const override {
return IN_B;
}
};

View 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

View file

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

View 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

View file

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

View file

@ -260,7 +260,11 @@
1734858160766311432
],
[
"QEmbedLayerNormalization com.microsoft CPUExecutionProvider",
9235385557940152248
"QEmbedLayerNormalization com.microsoft CPUExecutionProvider",
9235385557940152248
],
[
"QGemm com.microsoft CPUExecutionProvider",
13737193491843065240
]
]