diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index eb66c33cc7..44d6c8ae4b 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -40,6 +40,7 @@ Do not modify directly.* * com.microsoft.OptionalHasElement * com.microsoft.Pad * com.microsoft.QAttention + * com.microsoft.QGemm * com.microsoft.QLinearAdd * com.microsoft.QLinearAveragePool * com.microsoft.QLinearConcat @@ -1894,6 +1895,73 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.QGemm** + + Quantized Gemm + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
alpha : float
+
Scalar multiplier for the product of input tensors A * B.
+
transA : int
+
Whether A should be transposed
+
transB : int
+
Whether B should be transposed
+
+ +#### Inputs (6 - 9) + +
+
A : TA
+
Input tensor A. The shape of A should be (M, K) if transA is 0, or (K, M) if transA is non-zero.
+
a_scale : T
+
Scale of quantized input 'A'. It is a scalar,which means a per-tensor quantization.
+
a_zero_point : TA
+
Zero point tensor for input 'A'. It is a scalar.
+
B : TB
+
Input tensor B. The shape of B should be (K, N) if transB is 0, or (N, K) if transB is non-zero.
+
b_scale : T
+
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'.
+
b_zero_point : TB
+
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'.
+
C (optional) : TC
+
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.
+
y_scale (optional) : T
+
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.
+
y_zero_point (optional) : TYZ
+
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.
+
+ +#### Outputs + +
+
Y : TY
+
Output tensor of shape (M, N).
+
+ +#### Type Constraints + +
+
T : tensor(float)
+
Constrain scale types to float tensors.
+
TA : tensor(uint8), tensor(int8)
+
Constrain input A and its zero point types to 8 bit tensors.
+
TB : tensor(uint8), tensor(int8)
+
Constrain input B and its zero point types to 8 bit tensors.
+
TC : tensor(int32)
+
Constrain input C to 32 bit integer tensors.
+
TYZ : tensor(uint8), tensor(int8)
+
Constrain output zero point types to 8 bit tensors.
+
TY : tensor(float), tensor(uint8), tensor(int8)
+
Constrain output type to float32 or 8 bit tensors.
+
+ + ### **com.microsoft.QLinearAdd** Performs element-wise binary addition on 8 bit data types (with Numpy-style broadcasting support). diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 76aaa7047c..e5d4ff1e1e 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -393,6 +393,7 @@ Do not modify directly.* |Pad|*in* data:**T**
*in* pads:**tensor(int64)**
*in* value:**T**
*out* output:**T**|1+|**T** = tensor(float)| |QAttention|*in* input:**T1**
*in* weight:**T2**
*in* bias:**T3**
*in* input_scale:**T3**
*in* weight_scale:**T3**
*in* mask_index:**T4**
*in* input_zero_point:**T1**
*in* weight_zero_point:**T2**
*in* past:**T3**
*out* output:**T3**
*out* present:**T3**|1+|**T1** = tensor(uint8)
**T2** = tensor(int8), tensor(uint8)
**T3** = tensor(float)
**T4** = tensor(int32)| |QEmbedLayerNormalization|*in* input_ids:**T1**
*in* segment_ids:**T1**
*in* word_embedding_quant:**T2**
*in* position_embedding_quant:**T2**
*in* segment_embedding:**T2**
*in* gamma_quant:**T2**
*in* beta_quant:**T2**
*in* mask:**T1**
*in* word_embedding_scale:**T**
*in* position_embedding_scale:**T**
*in* segment_embedding_scale:**T**
*in* gamma_scale:**T**
*in* beta_scale:**T**
*in* word_embedding_zero_point:**T2**
*in* position_embedding_zero_point:**T2**
*in* segment_embedding_zero_point:**T2**
*in* gamma_zero_point:**T2**
*in* beta_zero_point:**T2**
*out* layernorm_out:**T**
*out* mask_index_out:**T1**|1+|**T** = tensor(float)| +|QGemm|*in* A:**TA**
*in* a_scale:**T**
*in* a_zero_point:**TA**
*in* B:**TB**
*in* b_scale:**T**
*in* b_zero_point:**TB**
*in* C:**TC**
*in* y_scale:**T**
*in* y_zero_point:**TYZ**
*out* Y:**TY**|1+|**T** = tensor(float)
**TA** = tensor(uint8)
**TB** = tensor(int8), tensor(uint8)
**TC** = tensor(int32)
**TY** = tensor(float), tensor(uint8)
**TYZ** = tensor(uint8)| |QLinearAdd|*in* A:**T**
*in* A_scale:**tensor(float)**
*in* A_zero_point:**T**
*in* B:**T**
*in* B_scale:**tensor(float)**
*in* B_zero_point:**T**
*in* C_scale:**tensor(float)**
*in* C_zero_point:**T**
*out* C:**T**|1+|**T** = tensor(int8), tensor(uint8)| |QLinearConv|*in* x:**T1**
*in* x_scale:**tensor(float)**
*in* x_zero_point:**T1**
*in* w:**T2**
*in* w_scale:**tensor(float)**
*in* w_zero_point:**T2**
*in* y_scale:**tensor(float)**
*in* y_zero_point:**T3**
*in* B:**T4**
*out* y:**T3**|1+|**T1** = tensor(uint8)
**T2** = tensor(int8), tensor(uint8)
**T3** = tensor(uint8)
**T4** = tensor(int32)| |QLinearLeakyRelu|*in* X:**T**
*in* X_scale:**tensor(float)**
*in* X_zero_point:**T**
*in* Y_scale:**tensor(float)**
*in* Y_zero_point:**T**
*out* Y:**T**|1+|**T** = tensor(int8), tensor(uint8)| diff --git a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc index 73b8488f04..6d88f5afeb 100644 --- a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc @@ -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, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, }; for (auto& function_table_entry : function_table) { diff --git a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc index 193eb78a3a..e4fb9428d1 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc @@ -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. diff --git a/onnxruntime/contrib_ops/cpu/quantization/quant_gemm.cc b/onnxruntime/contrib_ops/cpu/quantization/quant_gemm.cc new file mode 100644 index 0000000000..91aca25ede --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/quantization/quant_gemm.cc @@ -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(IN_A); + const auto* b = packed_b_ ? nullptr : context->Input(IN_B); + const auto& b_shape = b ? b->Shape() : b_shape_; + + const auto* c = context->Input(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(helper.M()); + size_t N = SafeInt(helper.N()); + size_t K = SafeInt(helper.K()); + + //validate scales and zero points + const auto* a_zp = context->Input(IN_A_ZERO_POINT); + const auto* b_zp = context->Input(IN_B_ZERO_POINT); + const auto* y_zp = context->Input(IN_Y_ZERO_POINT); + const auto* a_scale = context->Input(IN_A_SCALE); + const auto* b_scale = context->Input(IN_B_SCALE); + const auto* y_scale = context->Input(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(); + 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(packed_b_.get()); + b_is_signed = b_is_signed_; + } else { + b_data = static_cast(b->DataRaw()); + b_is_signed = b->IsDataType(); + if (trans_B_ == CblasTrans) { + b_data = quantization::TransPoseInputData(b_data, b_trans_buffer, allocator, N, K); + } + } + + auto y = context->Output(OUT_Y, {SafeInt(M), SafeInt(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(allocator->Alloc(SafeInt(M * N) * sizeof(int32_t))); + gemm_output_buffer.reset(gemm_output_data); + } else { + gemm_output_data = static_cast(y->MutableDataRaw()); + } + + if (c != nullptr) { + GemmBroadcastBias(M, N, 1.f, c->template Data(), &(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()); + + gemm_param.B = b_data; + gemm_param.ldb = gemm_shape.N; + gemm_param.BIsPacked = bool(packed_b_); + gemm_param.ZeroPointB = static_cast(b_zp->DataRaw()); + + gemm_param.C = gemm_output_data; + gemm_param.ldc = gemm_shape.N; + + gemm_param.PerColumnZeroPoints = !IsScalarOr1ElementVector(b_zp); + + std::vector output_scales = ComputeOutputScale(a_scale, b_scale, y_scale); + std::unique_ptr scale_bias_proc_ptr; + std::unique_ptr 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 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 output_scales(output_scale_size); + auto a_scale_value = *(a_scale->template Data()); + const auto* b_scale_data = b_scale->template Data(); + 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()); + } + } + return output_scales; + } + + static void SetPostProcessor(const Tensor* y_zp, + size_t out_lda, + const std::vector& output_scales, + Tensor* y, + MLAS_GEMM_U8X8_DATA_PARAMS& gemm_param, + std::unique_ptr& scale_bias_proc_ptr, + std::unique_ptr& requant_proc_ptr) { + if (nullptr != y_zp) { + requant_proc_ptr = std::make_unique( + static_cast(y->MutableDataRaw()), + out_lda, + nullptr, + output_scales.data(), + output_scales.size() > 1, + *y_zp->template Data()); + gemm_param.OutputProcessor = requant_proc_ptr.get(); + } else { + scale_bias_proc_ptr = std::make_unique( + static_cast(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()) + .TypeConstraint("TA", DataTypeImpl::GetTensorType()) + .TypeConstraint("TB", {DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}) + .TypeConstraint("TC", DataTypeImpl::GetTensorType()) + .TypeConstraint("TYZ", DataTypeImpl::GetTensorType()) + .TypeConstraint("TY", {DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}), + QGemm); + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/contrib_ops/quantization_defs.cc b/onnxruntime/core/graph/contrib_ops/quantization_defs.cc index 2202b640fc..335b7efe8b 100644 --- a/onnxruntime/core/graph/contrib_ops/quantization_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/quantization_defs.cc @@ -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(0)) + .Attr("transB", + "Whether B should be transposed", + AttributeProto::INT, + static_cast(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(transAAttr->i()) != 0 : false; + auto transBAttr = ctx.getAttribute("transB"); + bool transB = + transBAttr ? static_cast(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 diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index bec4cdf479..2d46605035 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -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 { diff --git a/onnxruntime/core/mlas/lib/qgemm.h b/onnxruntime/core/mlas/lib/qgemm.h index 95d1702059..51b21da37d 100644 --- a/onnxruntime/core/mlas/lib/qgemm.h +++ b/onnxruntime/core/mlas/lib/qgemm.h @@ -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) { diff --git a/onnxruntime/core/providers/cpu/math/gemm.cc b/onnxruntime/core/providers/cpu/math/gemm.cc index c23398ee51..e770d59b8b 100644 --- a/onnxruntime/core/providers/cpu/math/gemm.cc +++ b/onnxruntime/core/providers/cpu/math/gemm.cc @@ -108,30 +108,6 @@ bool GemmPackBFp32(AllocatorPtr& alloc, return true; } -template -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(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(c_data, N).transpose(); - } else if ((*c_shape)[1] == 1) { - // C is (M, 1) - output_mat.colwise() = ConstEigenVectorMap(c_data, M); - } else { - // C is (M, N), no broadcast needed. - output_mat = ConstEigenMatrixMapRowMajor(c_data, M, N); - } - } -} - template void Gemm::ComputeGemm(CBLAS_TRANSPOSE trans_a, CBLAS_TRANSPOSE trans_b, int64_t M, int64_t N, int64_t K, diff --git a/onnxruntime/core/providers/cpu/math/gemm.h b/onnxruntime/core/providers/cpu/math/gemm.h index d0cf142cd0..b8cac2f15a 100644 --- a/onnxruntime/core/providers/cpu/math/gemm.h +++ b/onnxruntime/core/providers/cpu/math/gemm.h @@ -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 -class Gemm : public OpKernel { +class Gemm : protected GemmBase, public OpKernel { public: - Gemm(const OpKernelInfo& info) : OpKernel(info) { - int64_t temp; - ORT_ENFORCE(info.GetAttr("transA", &temp).IsOK()); - trans_A_ = temp == 0 ? CblasNoTrans : CblasTrans; - - ORT_ENFORCE(info.GetAttr("transB", &temp).IsOK()); - trans_B_ = temp == 0 ? CblasNoTrans : CblasTrans; - - ORT_ENFORCE(info.GetAttr("alpha", &alpha_).IsOK()); - ORT_ENFORCE(info.GetAttr("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_; diff --git a/onnxruntime/core/providers/cpu/math/gemm_base.h b/onnxruntime/core/providers/cpu/math/gemm_base.h new file mode 100644 index 0000000000..02662f0a0f --- /dev/null +++ b/onnxruntime/core/providers/cpu/math/gemm_base.h @@ -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("transA", &temp).IsOK()); + trans_A_ = temp == 0 ? CblasNoTrans : CblasTrans; + + ORT_ENFORCE(info.GetAttr("transB", &temp).IsOK()); + trans_B_ = temp == 0 ? CblasNoTrans : CblasTrans; + + ORT_ENFORCE(info.GetAttr("alpha", &alpha_).IsOK()); + info.GetAttrOrDefault("beta", &beta_, 1.f); + } + + CBLAS_TRANSPOSE trans_A_; + CBLAS_TRANSPOSE trans_B_; + float alpha_; + float beta_; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/math/gemm_helper.h b/onnxruntime/core/providers/cpu/math/gemm_helper.h index 7b7240217d..5e58f10a71 100644 --- a/onnxruntime/core/providers/cpu/math/gemm_helper.h +++ b/onnxruntime/core/providers/cpu/math/gemm_helper.h @@ -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 +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(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(c_data, N).transpose(); + } + else if ((*c_shape)[1] == 1) { + // C is (M, 1) + output_mat.colwise() = ConstEigenVectorMap(c_data, M); + } + else { + // C is (M, N), no broadcast needed. + output_mat = ConstEigenMatrixMapRowMajor(c_data, M, N); + } + } +} + } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/math/matmul_integer.cc b/onnxruntime/core/providers/cpu/math/matmul_integer.cc index b8ab9f08e2..1b80fa2d21 100644 --- a/onnxruntime/core/providers/cpu/math/matmul_integer.cc +++ b/onnxruntime/core/providers/cpu/math/matmul_integer.cc @@ -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( diff --git a/onnxruntime/core/providers/cpu/math/matmul_integer_base.h b/onnxruntime/core/providers/cpu/math/matmul_integer_base.h index 920f7cac71..507bb00242 100644 --- a/onnxruntime/core/providers/cpu/math/matmul_integer_base.h +++ b/onnxruntime/core/providers/cpu/math/matmul_integer_base.h @@ -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(); - const size_t K = static_cast(b_shape_[0]); - const size_t N = static_cast(b_shape_[1]); + size_t K = static_cast(b_shape_[0]); + size_t N = static_cast(b_shape_[1]); const auto* b_data = static_cast(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: diff --git a/onnxruntime/core/providers/cpu/math/quantize_linear_matmul.h b/onnxruntime/core/providers/cpu/math/quantize_linear_matmul.h index 9769230d07..569b334c09 100644 --- a/onnxruntime/core/providers/cpu/math/quantize_linear_matmul.h +++ b/onnxruntime/core/providers/cpu/math/quantize_linear_matmul.h @@ -33,7 +33,7 @@ class QLinearMatMul : public MatMulIntegerBase { }; protected: - int GetBIdx() override { + int GetBIdx() const override { return IN_B; } }; diff --git a/onnxruntime/core/quantization/quant_util.h b/onnxruntime/core/quantization/quant_util.h new file mode 100644 index 0000000000..7020fe2daa --- /dev/null +++ b/onnxruntime/core/quantization/quant_util.h @@ -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(allocator->Alloc(M * N * sizeof(uint8_t))); + MlasTranspose(input, output, M, N); + buffer_holder.reset(output); + return output; +} + +} // namespace quantization +} // namespace onnxruntime diff --git a/onnxruntime/test/common/quantization_test_utils.h b/onnxruntime/test/common/quantization_test_utils.h index 9abdc445fa..c2ba1a5847 100644 --- a/onnxruntime/test/common/quantization_test_utils.h +++ b/onnxruntime/test/common/quantization_test_utils.h @@ -42,6 +42,20 @@ inline std::vector QuantizeLinearTestVector( return result; } +template +std::vector ToVector(const int* value, int size) { + std::vector data(size); + for (int i = 0; i < size; i++) { + data[i] = static_cast(value[i]); + } + return data; +} + +template +T GetMiddle(const std::vector& 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 - diff --git a/onnxruntime/test/contrib_ops/quant_gemm_test.cc b/onnxruntime/test/contrib_ops/quant_gemm_test.cc new file mode 100644 index 0000000000..897320b30e --- /dev/null +++ b/onnxruntime/test/contrib_ops/quant_gemm_test.cc @@ -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 +#include + +namespace onnxruntime { +namespace test { + +template +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 n_unsigned(0, 127); + static std::uniform_int_distribution n_xint8(std::numeric_limits::min(), std::numeric_limits::max()); + static std::uniform_real_distribution n_apha(1.0f, 2.0f); + static std::uniform_real_distribution n_scale(0.003f, 0.004f); + + Eigen::MatrixXi matrix_a = Eigen::MatrixXi::Random(K, M) + .unaryExpr([](int) { return n_unsigned(e); }); + std::vector matrix_a_data; + if (is_A_trans) { + Eigen::MatrixXi matrix_a_trans = matrix_a.transpose().eval(); + matrix_a_data = ToVector(matrix_a_trans.data(), M * K); + } else { + matrix_a_data = ToVector(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 matrix_b_data; + if (is_B_trans) { + Eigen::MatrixXi matrix_b_trans = matrix_b.transpose().eval(); + matrix_b_data = ToVector(matrix_b_trans.data(), N * K); + } else { + matrix_b_data = ToVector(matrix_b.data(), N * K); + } + ScalarB b_zero_point = GetMiddle(matrix_b_data); + std::vector b_scale({n_scale(e)}); + std::vector 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()))); + + OpTester test("QGemm", 1, onnxruntime::kMSDomain); + test.AddAttribute("transA", is_A_trans ? 1 : 0); + test.AddAttribute("transB", is_B_trans ? 1 : 0); + test.AddAttribute("alpha", alpha); + test.AddInput("A", is_A_trans ? std::vector({K, M}) : std::vector({M, K}), std::move(matrix_a_data)); + test.AddInput("a_scale", {}, {a_scale}); + test.AddInput("a_zero_point", {}, {a_zero_point}); + test.AddInput("B", is_B_trans ? std::vector({N, K}) : std::vector({K, N}), std::move(matrix_b_data), B_is_initializer); + test.AddInput("b_scale", {SafeInt(b_scale.size())}, b_scale); + test.AddInput("b_zero_point", {SafeInt(b_zp_per_column.size())}, b_zp_per_column); + + if (has_C) { + test.AddInput("C", {M, N}, ToVector(matrix_c.data(), M * N)); + } else { + test.AddOptionalInputEdge(); + } + + if constexpr (std::is_same_v) { + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + test.AddOutput("Y", {M, N}, std::vector(matrix_output.data(), matrix_output.data() + M * N)); + } else { + std::vector quant_output(M * N); + quantization::Params quant_param = quantization::QuantizeLinear(matrix_output.data(), quant_output.data(), M * N); + test.AddInput("y_scale", {}, {quant_param.scale}); + test.AddInput("y_zero_point", {}, {quant_param.zero_point}); + test.AddOutput("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(M, N, K, is_A_trans, is_B_trans, has_C, B_is_initializer, per_column); + RunQuantGemmU8X8Test(M, N, K, is_A_trans, is_B_trans, has_C, B_is_initializer, per_column); + RunQuantGemmU8X8Test(M, N, K, is_A_trans, is_B_trans, has_C, B_is_initializer, per_column); + RunQuantGemmU8X8Test(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 diff --git a/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc b/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc index 57722628ef..da52ac0f5d 100644 --- a/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc +++ b/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc @@ -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 -std::vector ToVector(const int* value, int size) { - std::vector data(size); - for (int i = 0; i < size; i++) { - data[i] = static_cast(value[i]); - } - return data; -} - -template -T GetMiddle(const std::vector& 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 void RunMatMulIntegerU8X8Test(const int M, const int N, const int K, bool non_zero_zp, bool B_is_initializer, bool per_column_zp = false) { diff --git a/onnxruntime/test/testdata/kernel_def_hashes/contrib.cpu.json b/onnxruntime/test/testdata/kernel_def_hashes/contrib.cpu.json index 6158c8ae1e..0b69923ca2 100644 --- a/onnxruntime/test/testdata/kernel_def_hashes/contrib.cpu.json +++ b/onnxruntime/test/testdata/kernel_def_hashes/contrib.cpu.json @@ -260,7 +260,11 @@ 1734858160766311432 ], [ - "QEmbedLayerNormalization com.microsoft CPUExecutionProvider", - 9235385557940152248 + "QEmbedLayerNormalization com.microsoft CPUExecutionProvider", + 9235385557940152248 + ], + [ + "QGemm com.microsoft CPUExecutionProvider", + 13737193491843065240 ] ]