From 3ef449816c2c01703d6aa9d99f15bd77f1325515 Mon Sep 17 00:00:00 2001 From: Tracy Sharpe <42477615+tracysh@users.noreply.github.com> Date: Mon, 6 Jul 2020 15:20:10 -0700 Subject: [PATCH] MLAS: support prepacking APIs for quantized GEMM (#4433) Add support for prepacking matrix B for use in the quantized GEMMs. --- .../cpu/quantization/attention_quant.cc | 77 ++ .../quantization/dynamic_quantize_matmul.cc | 213 +++-- onnxruntime/core/mlas/inc/mlas.h | 59 ++ onnxruntime/core/mlas/lib/mlasi.h | 33 +- onnxruntime/core/mlas/lib/platform.cpp | 13 +- onnxruntime/core/mlas/lib/qgemm.cpp | 874 ++++++++++++++---- onnxruntime/core/util/qmath.h | 4 + onnxruntime/test/mlas/unittest.cpp | 172 +++- 8 files changed, 1186 insertions(+), 259 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc b/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc index 3721665a5e..ebec77f365 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc @@ -22,6 +22,14 @@ class QAttention : public OpKernel, public AttentionCPUBase { QAttention(const OpKernelInfo& info); Status Compute(OpKernelContext* context) const override; + + private: + void TryPackWeights(const OpKernelInfo& info); + +#ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8 + BufferUniquePtr packed_weights_; + size_t packed_weights_size_; +#endif }; // These ops are internal-only, so register outside of onnx @@ -40,6 +48,53 @@ ONNX_OPERATOR_TYPED_KERNEL_EX( template QAttention::QAttention(const OpKernelInfo& info) : OpKernel(info), AttentionCPUBase(info) { + TryPackWeights(info); +} + +template +void QAttention::TryPackWeights(const OpKernelInfo& info) { +#ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8 + // Check if the weights tensor is constant. + const Tensor* weights; + if (!info.TryGetConstantInput(1, &weights)) { + return; + } + + const auto& weights_dims = weights->Shape().GetDims(); + if (weights_dims.size() != 2) { + return; + } + + const size_t hidden_size = static_cast(weights_dims[0]); + const size_t hidden_size_x3 = static_cast(weights_dims[1]); + const size_t head_size = hidden_size / num_heads_; + + // Bail out if the weights shape has an expected shape. + if ((hidden_size == 0) || ((hidden_size % num_heads_) != 0) || (hidden_size_x3 != 3 * hidden_size)) { + return; + } + + const auto* weights_data = static_cast(weights->DataRaw()); + const bool weights_is_signed = weights->IsDataType(); + + packed_weights_size_ = MlasGemmPackBSize(head_size, hidden_size, weights_is_signed); + if (packed_weights_size_ == 0) { + return; + } + + const size_t loop_len = 3 * num_heads_; + auto alloc = info.GetAllocator(0, OrtMemTypeDefault); + auto* packed_weights_data = static_cast(alloc->Alloc(packed_weights_size_ * loop_len)); + packed_weights_ = BufferUniquePtr(packed_weights_data, BufferDeleter(alloc)); + + for (size_t i = 0; i < loop_len; i++) { + MlasGemmPackB(head_size, hidden_size, weights_data, hidden_size_x3, weights_is_signed, packed_weights_data); + packed_weights_data += packed_weights_size_; + weights_data += head_size; + } +#else + ORT_UNUSED_PARAMETER(info); +#endif } template @@ -140,6 +195,28 @@ Status QAttention::Compute(OpKernelContext* context) const { // A: input (BxSxNxH) (B.)S x NH S x NH // B: weights (NxHx3xNxH) NH x (3.N.)H NH x H // C: QKV[qkv_index] (3xBxNxSxH) (3.B.N.)S x H S x H +#ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8 + if (packed_weights_) { + const auto* packed_weight = + static_cast(packed_weights_.get()) + packed_weights_size_ * (weights_offset / head_size); + MlasGemm( + sequence_length, // M = S + head_size, // N = H + hidden_size, // K = NH + input_data + input_offset, // A + hidden_size, // lda = NH + input_zero_point, // input zero point + packed_weight, // B + weight_zero_point, // weight zero point + weights_is_signed, // weight data type + qkv_dest + qkv_offset, // C + head_size, // ldc + &dequant_scale, // output scale + bias_data + weights_offset, // bias + nullptr); // use single-thread + continue; + } +#endif QGemm(sequence_length, // M = S head_size, // N = H hidden_size, // K = NH diff --git a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc index e29c977032..983cf72814 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc @@ -14,69 +14,71 @@ namespace onnxruntime { namespace contrib { -class DynamicQuantizeMatMul final : public OpKernel { +class MatMulIntegerToFloatBase : public OpKernel { public: - DynamicQuantizeMatMul(const OpKernelInfo& info) : OpKernel(info) {} + MatMulIntegerToFloatBase(const OpKernelInfo& info) : OpKernel(info) { + TryPackWeights(info); + } - Status Compute(OpKernelContext* context) const override; + protected: + void TryPackWeights(const OpKernelInfo& info); + Status ComputeCommon(OpKernelContext* ctx, + const uint8_t* a_data, + const TensorShape& a_shape, + uint8_t a_zero_point, + const Tensor* b, + uint8_t b_zero_point, + float multiplier, + const Tensor* bias_tensor) const; + +#ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8 + BufferUniquePtr packed_b_; +#endif }; -ONNX_OPERATOR_TYPED_KERNEL_EX( - DynamicQuantizeMatMul, - kMSDomain, - 1, - float, - kCpuExecutionProvider, - KernelDefBuilder() - .TypeConstraint("T1", DataTypeImpl::GetTensorType()) - .TypeConstraint("T2", {DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}), - DynamicQuantizeMatMul); +void MatMulIntegerToFloatBase::TryPackWeights(const OpKernelInfo& info) { +#ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8 + // Check if the weights tensor is constant. + const Tensor* b; + if (!info.TryGetConstantInput(1, &b)) { + return; + } -class MatMulIntegerToFloat final : public OpKernel { - public: - MatMulIntegerToFloat(const OpKernelInfo& info) : OpKernel(info) {} + // Only handle the common case of a 2D weight matrix. Additional matrices + // could be handled by stacking the packed buffers. + const auto& b_shape = b->Shape(); + if (b_shape.NumDimensions() != 2) { + return; + } - Status Compute(OpKernelContext* context) const override; -}; + const size_t K = static_cast(b_shape[0]); + const size_t N = static_cast(b_shape[1]); -ONNX_OPERATOR_TYPED_KERNEL_EX( - MatMulIntegerToFloat, - kMSDomain, - 1, - uint8_t, - kCpuExecutionProvider, - KernelDefBuilder() - .TypeConstraint("T1", DataTypeImpl::GetTensorType()) - .TypeConstraint("T2", {DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}) - .TypeConstraint("T3", DataTypeImpl::GetTensorType()), - MatMulIntegerToFloat); + const auto* b_data = static_cast(b->DataRaw()); + const bool b_is_signed = b->IsDataType(); -static void GetQuantizationParameter(const float* data, int64_t num_of_elements, float& scale, uint8_t& zp) { - // find input range min and max - float min, max; - MlasFindMinMaxElement(data, &min, &max, num_of_elements); + const size_t packed_b_size = MlasGemmPackBSize(N, K, b_is_signed); + if (packed_b_size == 0) { + return; + } - // ensure the input range includes zero - min = std::min(min, 0.0f); - max = std::max(max, 0.0f); - - // find scale and zero point - uint8_t qmin = 0; - uint8_t qmax = 255; - scale = max == min ? 1.0f : (max - min) / (qmax - qmin); - - float initial_zero_point = qmin - min / scale; - zp = static_cast(RoundHalfToEven(std::max(float(qmin), std::min(float(qmax), initial_zero_point)))); + auto alloc = info.GetAllocator(0, OrtMemTypeDefault); + auto* packed_b_data = alloc->Alloc(packed_b_size); + packed_b_ = BufferUniquePtr(packed_b_data, BufferDeleter(alloc)); + MlasGemmPackB(N, K, b_data, N, b_is_signed, packed_b_data); +#else + ORT_UNUSED_PARAMETER(info); +#endif } -static Status MatMulIntegerToFloatCommon(OpKernelContext* ctx, - const uint8_t* a_data, - const TensorShape& a_shape, - uint8_t a_zero_point, - const Tensor* b, - uint8_t b_zero_point, - float multiplier, - const Tensor* bias_tensor) { +Status MatMulIntegerToFloatBase::ComputeCommon(OpKernelContext* ctx, + const uint8_t* a_data, + const TensorShape& a_shape, + uint8_t a_zero_point, + const Tensor* b, + uint8_t b_zero_point, + float multiplier, + const Tensor* bias_tensor) const { MatMulComputeHelper helper; ORT_RETURN_IF_ERROR(helper.Compute(a_shape, b->Shape())); Tensor* y = ctx->Output(0, helper.OutputShape()); @@ -89,6 +91,25 @@ static Status MatMulIntegerToFloatCommon(OpKernelContext* ctx, concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool(); for (size_t i = 0; i < helper.OutputOffsets().size(); i++) { +#ifdef MLAS_SUPPORTS_PACKED_GEMM_U8X8 + if (packed_b_) { + MlasGemm(static_cast(helper.M()), + static_cast(helper.N()), + static_cast(helper.K()), + a_data + helper.LeftOffsets()[i], + static_cast(helper.K()), + a_zero_point, + packed_b_.get(), + b_zero_point, + b_is_signed, + y_data + helper.OutputOffsets()[i], + static_cast(helper.N()), + &multiplier, + nullptr, + thread_pool); + continue; + } +#endif QGemm(static_cast(helper.M()), static_cast(helper.N()), static_cast(helper.K()), @@ -109,6 +130,38 @@ static Status MatMulIntegerToFloatCommon(OpKernelContext* ctx, return Status::OK(); } +class DynamicQuantizeMatMul final : public MatMulIntegerToFloatBase { + public: + DynamicQuantizeMatMul(const OpKernelInfo& info) : MatMulIntegerToFloatBase(info) {} + + Status Compute(OpKernelContext* context) const override; +}; + +class MatMulIntegerToFloat final : public MatMulIntegerToFloatBase { + public: + MatMulIntegerToFloat(const OpKernelInfo& info) : MatMulIntegerToFloatBase(info) {} + + Status Compute(OpKernelContext* context) const override; +}; + +static void GetQuantizationParameter(const float* data, int64_t num_of_elements, float& scale, uint8_t& zp) { + // find input range min and max + float min, max; + MlasFindMinMaxElement(data, &min, &max, num_of_elements); + + // ensure the input range includes zero + min = std::min(min, 0.0f); + max = std::max(max, 0.0f); + + // find scale and zero point + uint8_t qmin = 0; + uint8_t qmax = 255; + scale = max == min ? 1.0f : (max - min) / (qmax - qmin); + + float initial_zero_point = qmin - min / scale; + zp = static_cast(RoundHalfToEven(std::max(float(qmin), std::min(float(qmax), initial_zero_point)))); +} + Status DynamicQuantizeMatMul::Compute(OpKernelContext* ctx) const { const auto* a = ctx->Input(0); const auto* b = ctx->Input(1); @@ -116,7 +169,6 @@ Status DynamicQuantizeMatMul::Compute(OpKernelContext* ctx) const { const auto* b_scale_tensor = ctx->Input(2); ORT_ENFORCE(IsScalarOr1ElementVector(b_scale_tensor), "DynamicQuantizeMatMul : input B scale must be a scalar or 1D tensor of size 1. Per-Channel is not supported yet."); - float b_scale = *b_scale_tensor->template Data(); const auto* b_zero_point_tensor = ctx->Input(3); @@ -142,14 +194,14 @@ Status DynamicQuantizeMatMul::Compute(OpKernelContext* ctx) const { // quantize the data MlasQuantizeLinear(a_data, a_data_quant, num_of_elements, a_scale, a_zero_point); - return MatMulIntegerToFloatCommon(ctx, - a_data_quant, - a->Shape(), - a_zero_point, - b, - b_zero_point, - a_scale * b_scale, - ctx->Input(4)); + return ComputeCommon(ctx, + a_data_quant, + a->Shape(), + a_zero_point, + b, + b_zero_point, + a_scale * b_scale, + ctx->Input(4)); } Status MatMulIntegerToFloat::Compute(OpKernelContext* ctx) const { @@ -183,15 +235,38 @@ Status MatMulIntegerToFloat::Compute(OpKernelContext* ctx) const { b_zero_point = *static_cast(b_zero_point_tensor->DataRaw()); } - return MatMulIntegerToFloatCommon(ctx, - a->Data(), - a->Shape(), - a_zero_point, - b, - b_zero_point, - a_scale * b_scale, - ctx->Input(6)); + return ComputeCommon(ctx, + a->Data(), + a->Shape(), + a_zero_point, + b, + b_zero_point, + a_scale * b_scale, + ctx->Input(6)); } +ONNX_OPERATOR_TYPED_KERNEL_EX( + DynamicQuantizeMatMul, + kMSDomain, + 1, + float, + kCpuExecutionProvider, + KernelDefBuilder() + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", {DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}), + DynamicQuantizeMatMul); + +ONNX_OPERATOR_TYPED_KERNEL_EX( + MatMulIntegerToFloat, + kMSDomain, + 1, + uint8_t, + kCpuExecutionProvider, + KernelDefBuilder() + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", {DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}) + .TypeConstraint("T3", DataTypeImpl::GetTensorType()), + MatMulIntegerToFloat); + } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index 77a1a46682..9345667bc4 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -185,6 +185,65 @@ MlasGemm( MLAS_THREADPOOL* ThreadPool ); +void +MLASCALL +MlasGemm( + size_t M, + size_t N, + size_t K, + const uint8_t* A, + size_t lda, + uint8_t offa, + const void* PackedB, + uint8_t offb, + bool BIsSigned, + int32_t* C, + size_t ldc, + MLAS_THREADPOOL* ThreadPool + ); + +void +MLASCALL +MlasGemm( + size_t M, + size_t N, + size_t K, + const uint8_t* A, + size_t lda, + uint8_t offa, + const void* PackedB, + uint8_t offb, + bool BIsSigned, + float* C, + size_t ldc, + const float* Scale, + const float* Bias, + MLAS_THREADPOOL* ThreadPool + ); + +// +// Buffer packing routines. +// + +size_t +MLASCALL +MlasGemmPackBSize( + size_t N, + size_t K, + bool BIsSigned + ); + +void +MLASCALL +MlasGemmPackB( + size_t N, + size_t K, + const uint8_t* B, + size_t ldb, + bool BIsSigned, + void* PackedB + ); + // // Convolution routines. // diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index 21106a7dff..d44f696aea 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -543,10 +543,6 @@ extern "C" { MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE MlasSgemmTransposePackB16x4Avx; #endif -#if defined(MLAS_TARGET_AMD64_IX86) - MLAS_GEMM_U8X8_OPERATION MlasGemmU8X8OperationSse; - MLAS_GEMM_U8X8_OPERATION MlasGemmU8S8OperationAvx2; - MLAS_GEMM_U8X8_OPERATION MlasGemmU8U8OperationAvx2; #if defined(MLAS_TARGET_AMD64) MLAS_GEMM_U8S8_KERNEL MlasGemmU8S8KernelAvx2; MLAS_GEMV_U8S8_KERNEL MlasGemvU8S8KernelAvx2; @@ -557,7 +553,6 @@ extern "C" { MLAS_GEMM_U8U8_KERNEL MlasGemmU8U8KernelAvx2; MLAS_GEMM_U8U8_KERNEL MlasGemmU8U8KernelAvx512Core; #endif -#endif #if defined(MLAS_TARGET_AMD64) MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelSse; @@ -674,6 +669,28 @@ MlasSgemmOperation( size_t ldc ); +// +// Quantized integer matrix/matrix multiply operation. +// + +struct MLAS_GEMM_U8X8_KERNEL_SSE; +struct MLAS_GEMM_U8S8_KERNEL_AVX2; +struct MLAS_GEMM_U8U8_KERNEL_AVX2; + +template +void +MLASCALL +MlasGemmU8X8Operation( + const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock + ); + +template +void +MLASCALL +MlasGemmU8X8PackedOperation( + const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock + ); + // // Environment information class. // @@ -684,8 +701,6 @@ struct MLAS_PLATFORM { #if defined(MLAS_TARGET_AMD64_IX86) PMLAS_GEMM_FLOAT_KERNEL GemmFloatKernel; - PMLAS_GEMM_U8X8_OPERATION GemmU8S8Operation; - PMLAS_GEMM_U8X8_OPERATION GemmU8U8Operation; #endif #if defined(MLAS_TARGET_AMD64) @@ -693,8 +708,12 @@ struct MLAS_PLATFORM { PMLAS_SGEMM_KERNEL_M1_ROUTINE KernelM1TransposeBRoutine; PMLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE TransposePackB16x4Routine; PMLAS_GEMM_DOUBLE_KERNEL GemmDoubleKernel; + PMLAS_GEMM_U8X8_OPERATION GemmU8S8Operation; + PMLAS_GEMM_U8X8_OPERATION GemmU8S8PackedOperation; PMLAS_GEMM_U8S8_KERNEL GemmU8S8Kernel; PMLAS_GEMV_U8S8_KERNEL GemvU8S8Kernel; + PMLAS_GEMM_U8X8_OPERATION GemmU8U8Operation; + PMLAS_GEMM_U8X8_OPERATION GemmU8U8PackedOperation; PMLAS_GEMM_U8U8_KERNEL GemmU8U8Kernel; PMLAS_CONV_FLOAT_KERNEL ConvNchwFloatKernel; PMLAS_CONV_FLOAT_KERNEL ConvNchwcFloatKernel; diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp index c31459e9e8..270010df43 100644 --- a/onnxruntime/core/mlas/lib/platform.cpp +++ b/onnxruntime/core/mlas/lib/platform.cpp @@ -115,13 +115,13 @@ Return Value: // this->GemmFloatKernel = MlasGemmFloatKernelSse; - this->GemmU8S8Operation = MlasGemmU8X8OperationSse; - this->GemmU8U8Operation = MlasGemmU8X8OperationSse; #if defined(MLAS_TARGET_AMD64) this->TransposePackB16x4Routine = MlasSgemmTransposePackB16x4Sse; this->GemmDoubleKernel = MlasGemmDoubleKernelSse; + this->GemmU8S8Operation = MlasGemmU8X8Operation; + this->GemmU8U8Operation = MlasGemmU8X8Operation; this->ConvNchwFloatKernel = MlasConvNchwFloatKernelSse; this->ConvNchwcFloatKernel = MlasConvNchwcFloatKernelSse; this->ConvDepthwiseFloatKernel = MlasConvDepthwiseFloatKernelSse; @@ -200,10 +200,12 @@ Return Value: if (((Cpuid1[2] & 0x1000) != 0) && ((Cpuid7[1] & 0x20) != 0)) { - this->GemmU8S8Operation = MlasGemmU8S8OperationAvx2; + this->GemmU8S8Operation = MlasGemmU8X8Operation; + this->GemmU8S8PackedOperation = MlasGemmU8X8PackedOperation; this->GemmU8S8Kernel = MlasGemmU8S8KernelAvx2; this->GemvU8S8Kernel = MlasGemvU8S8KernelAvx2; - this->GemmU8U8Operation = MlasGemmU8U8OperationAvx2; + this->GemmU8U8Operation = MlasGemmU8X8Operation; + this->GemmU8U8PackedOperation = MlasGemmU8X8PackedOperation; this->GemmU8U8Kernel = MlasGemmU8U8KernelAvx2; this->GemmFloatKernel = MlasGemmFloatKernelFma3; @@ -262,7 +264,8 @@ Return Value: if ((Cpuid7[2] & 0x800) != 0) { - this->GemmU8U8Operation = MlasGemmU8S8OperationAvx2; + this->GemmU8U8Operation = MlasGemmU8X8Operation; + this->GemmU8U8PackedOperation = MlasGemmU8X8PackedOperation; this->GemmU8S8Kernel = MlasGemmU8S8KernelAvx512Vnni; this->GemvU8S8Kernel = MlasGemvU8S8KernelAvx512Vnni; } diff --git a/onnxruntime/core/mlas/lib/qgemm.cpp b/onnxruntime/core/mlas/lib/qgemm.cpp index cd40cebd81..04337ce1b3 100644 --- a/onnxruntime/core/mlas/lib/qgemm.cpp +++ b/onnxruntime/core/mlas/lib/qgemm.cpp @@ -25,12 +25,16 @@ Abstract: struct MLAS_GEMM_U8X8_WORK_BLOCK { int32_t ThreadCountM; int32_t ThreadCountN; + size_t RangeStartM; + size_t RangeStartN; + size_t RangeCountM; + size_t RangeCountN; size_t M; size_t N; size_t K; const uint8_t* A; size_t lda; - const uint8_t* B; + const void* B; size_t ldb; int32_t* C; size_t ldc; @@ -38,20 +42,32 @@ struct MLAS_GEMM_U8X8_WORK_BLOCK { const float* BiasFloat; uint8_t offa; uint8_t offb; + bool BIsPacked; bool BIsSigned; bool CIsFloat; }; +// +// Define the default striding parameters used for the quantized integer +// matrix/matrix multiply operation. +// + +struct MLAS_GEMM_U8X8_STRIDES { + size_t M; + size_t N; + size_t K; +}; + void MlasGemmU8X8ScaleSumBuffer( - int32_t* D, - const int32_t* S, + int32_t* Output, + const int32_t* Input, size_t N, int32_t Scale ) { for (size_t n = 0; n < N; n++) { - D[n] = S[n] * Scale; + Output[n] = Input[n] * Scale; } } @@ -67,8 +83,8 @@ MlasGemmU8X8ScaleSumBuffer( } template -MLAS_FORCEINLINE void +MLASCALL MlasGemmU8X8Operation( const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock ) @@ -76,7 +92,7 @@ MlasGemmU8X8Operation( Routine Description: - This module implements the quantized integer matrix/matrix multiply + This routine implements the quantized integer matrix/matrix multiply operation (QGEMM). Arguments: @@ -89,23 +105,39 @@ Return Value: --*/ { - MLAS_DECLSPEC_ALIGN(typename KernelType::PackedAType PanelA[KernelType::StrideM * KernelType::StrideK], 64); - MLAS_DECLSPEC_ALIGN(typename KernelType::PackedBType PanelB[KernelType::StrideN * KernelType::StrideK], 64); + constexpr MLAS_GEMM_U8X8_STRIDES Strides = KernelType::Strides; - MLAS_DECLSPEC_ALIGN(int32_t RowSumBuffer[KernelType::StrideM], 64); - MLAS_DECLSPEC_ALIGN(int32_t ColumnSumBuffer[KernelType::StrideN], 64); + MLAS_DECLSPEC_ALIGN(typename KernelType::PackedAType PanelA[Strides.M * Strides.K], 64); + MLAS_DECLSPEC_ALIGN(typename KernelType::PackedBType PanelB[Strides.N * Strides.K], 64); - const uint8_t* A = WorkBlock->A; - const uint8_t* B = WorkBlock->B; - int32_t* C = WorkBlock->C; + MLAS_DECLSPEC_ALIGN(int32_t RowSumBuffer[Strides.M], 64); + MLAS_DECLSPEC_ALIGN(int32_t ColumnSumBuffer[Strides.N], 64); + + const size_t M = WorkBlock->RangeCountM; + const size_t N = WorkBlock->RangeCountN; + const size_t K = WorkBlock->K; const size_t lda = WorkBlock->lda; const size_t ldb = WorkBlock->ldb; const size_t ldc = WorkBlock->ldc; + const uint8_t* A = WorkBlock->A + WorkBlock->RangeStartM * lda; + const uint8_t* B = (const uint8_t*)WorkBlock->B + WorkBlock->RangeStartN; + int32_t* C = WorkBlock->C + WorkBlock->RangeStartM * ldc + WorkBlock->RangeStartN; + int32_t offa = WorkBlock->offa; int32_t offb = typename KernelType::OffsetBType(WorkBlock->offb); + // + // Try to use a GEMV kernel if supported by this kernel type. + // + + if ((M == 1) && (offa == 0) && (offb == 0) && !WorkBlock->CIsFloat) { + if (KernelType::TryGemvKernel(A, B, ldb, C, K, N, WorkBlock->BIsSigned)) { + return; + } + } + // // Flip the sign bit of the zero point offset of matrix B if the kernel uses // signed types and the matrix B data is unsigned. @@ -121,14 +153,11 @@ Return Value: // Step through each slice of matrix B along the K dimension. // - const size_t M = WorkBlock->M; - const size_t N = WorkBlock->N; - const size_t K = WorkBlock->K; size_t CountK; for (size_t k = 0; k < K; k += CountK) { - CountK = (std::min)(K - k, KernelType::StrideK); + CountK = (std::min)(K - k, Strides.K); // // Step through each slice of matrix B along the N dimension. @@ -138,7 +167,7 @@ Return Value: for (size_t n = 0; n < N; n += CountN) { - CountN = (std::min)(N - n, KernelType::StrideN); + CountN = (std::min)(N - n, Strides.N); // // Copy a panel of matrix B to a local packed buffer. @@ -162,7 +191,7 @@ Return Value: for (size_t m = 0; m < M; m += CountM) { - CountM = (std::min)(M - m, KernelType::StrideM); + CountM = (std::min)(M - m, Strides.M); // // Copy a panel of matrix A to a local packed buffer. @@ -188,7 +217,7 @@ Return Value: size_t RowsHandled; - RowsHandled = KernelType::Kernel(pa, PanelB, c, PackedCountK, + RowsHandled = KernelType::GemmKernel(pa, PanelB, c, PackedCountK, RowsRemaining, CountN, ldc, RowSums, ColumnSumBuffer, DepthValue, ZeroMode); @@ -209,6 +238,162 @@ Return Value: } } +template +void +MLASCALL +MlasGemmU8X8PackedOperation( + const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock + ) +/*++ + +Routine Description: + + This routine implements the quantized integer matrix/matrix multiply + operation (QGEMM). + +Arguments: + + WorkBlock - Supplies the structure containing the GEMM parameters. + +Return Value: + + None. + +--*/ +{ + constexpr MLAS_GEMM_U8X8_STRIDES Strides = KernelType::PackedStrides; + + MLAS_DECLSPEC_ALIGN(typename KernelType::PackedAType PanelA[Strides.M * Strides.K], 64); + + MLAS_DECLSPEC_ALIGN(int32_t RowSumBuffer[Strides.M], 64); + MLAS_DECLSPEC_ALIGN(int32_t ColumnSumBuffer[Strides.N], 64); + + const size_t M = WorkBlock->RangeCountM; + const size_t N = WorkBlock->RangeCountN; + const size_t K = WorkBlock->K; + + const size_t lda = WorkBlock->lda; + const size_t ldc = WorkBlock->ldc; + + const uint8_t* A = WorkBlock->A + WorkBlock->RangeStartM * lda; + const uint8_t* PackedB = (const uint8_t*)WorkBlock->B; + int32_t* C = WorkBlock->C + WorkBlock->RangeStartM * ldc + WorkBlock->RangeStartN; + + int32_t offa = WorkBlock->offa; + int32_t offb = typename KernelType::OffsetBType(WorkBlock->offb); + + // + // Flip the sign bit of the zero point offset of matrix B if the kernel uses + // signed types and the matrix B data is unsigned. + // + + if (std::is_signed::value) { + if (!WorkBlock->BIsSigned) { + offb = typename KernelType::OffsetBType(offb ^ 0x80); + } + } + + // + // Extract the pointer to the column sum buffer from the packed matrix. + // + + const size_t AlignedN = + (WorkBlock->N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1); + const int32_t* PackedColumnSumBuffer = (const int32_t*)PackedB; + PackedB = (const uint8_t*)(PackedColumnSumBuffer + AlignedN); + PackedColumnSumBuffer += WorkBlock->RangeStartN; + + // + // Step through each slice of matrix B along the K dimension. + // + + size_t CountK; + + for (size_t k = 0; k < K; k += CountK) { + + CountK = (std::min)(K - k, Strides.K); + + const size_t PackedCountK = (CountK + KernelType::PackedK - 1) / + KernelType::PackedK; + + if (k > 0) { + std::fill_n(ColumnSumBuffer, Strides.N, 0); + } + + // + // Step through each slice of matrix B along the N dimension. + // + + size_t CountN; + + for (size_t n = 0; n < N; n += CountN) { + + CountN = (std::min)(N - n, Strides.N); + + if (k == 0) { + MlasGemmU8X8ScaleSumBuffer(ColumnSumBuffer, PackedColumnSumBuffer + n, + CountN, -offa); + } + + // + // Step through each slice of matrix A along the M dimension. + // + + const int32_t DepthValue = int32_t(CountK) * offa * offb; + const uint8_t* b = PackedB + (WorkBlock->RangeStartN + n) * + KernelType::PackedK * PackedCountK; + int32_t* c = C + n; + size_t CountM; + + for (size_t m = 0; m < M; m += CountM) { + + CountM = (std::min)(M - m, Strides.M); + + // + // Copy a panel of matrix A to a local packed buffer. + // + + KernelType::CopyPackA(PanelA, A + m * lda, lda, CountM, CountK, + RowSumBuffer); + + MlasGemmU8X8ScaleSumBuffer(RowSumBuffer, CountM, -offb); + + // + // Step through the rows of the local packed buffer. + // + + typename KernelType::PackedAType* pa = PanelA; + int32_t* RowSums = RowSumBuffer; + size_t RowsRemaining = CountM; + + bool ZeroMode = (k == 0); + bool PostProcess = (k + CountK == K); + + while (RowsRemaining > 0) { + + size_t RowsHandled; + + RowsHandled = KernelType::GemmKernel(pa, b, c, PackedCountK, + RowsRemaining, CountN, ldc, RowSums, ColumnSumBuffer, + DepthValue, ZeroMode); + + if (PostProcess && WorkBlock->CIsFloat) { + KernelType::OutputFloat(WorkBlock, c, n, RowsHandled, CountN); + } + + c += ldc * RowsHandled; + pa += KernelType::PackedK * PackedCountK * RowsHandled; + RowSums += RowsHandled; + RowsRemaining -= RowsHandled; + } + } + } + + A += CountK; + PackedB = (const uint8_t*)PackedB + AlignedN * CountK; + } +} + #ifdef MLAS_TARGET_AMD64_IX86 void @@ -775,7 +960,7 @@ Return Value: if (BiasFloat != nullptr) { - BiasFloat += StartN; + BiasFloat += WorkBlock->RangeStartN + StartN; while (CountM-- > 0) { @@ -842,9 +1027,31 @@ struct MLAS_GEMM_U8X8_KERNEL_SSE typedef int8_t OffsetBType; static constexpr size_t PackedK = 2; - static constexpr size_t StrideM = 12; - static constexpr size_t StrideN = 128; - static constexpr size_t StrideK = 128; + static constexpr MLAS_GEMM_U8X8_STRIDES Strides{12, 128, 128}; + + MLAS_FORCEINLINE + static + bool + TryGemvKernel( + const uint8_t* A, + const uint8_t* B, + size_t ldb, + int32_t* C, + size_t CountK, + size_t CountN, + bool BIsSigned + ) + { + MLAS_UNREFERENCED_PARAMETER(A); + MLAS_UNREFERENCED_PARAMETER(B); + MLAS_UNREFERENCED_PARAMETER(ldb); + MLAS_UNREFERENCED_PARAMETER(C); + MLAS_UNREFERENCED_PARAMETER(CountK); + MLAS_UNREFERENCED_PARAMETER(CountN); + MLAS_UNREFERENCED_PARAMETER(BIsSigned); + + return false; + } MLAS_FORCEINLINE static @@ -881,7 +1088,7 @@ struct MLAS_GEMM_U8X8_KERNEL_SSE MLAS_FORCEINLINE static size_t - Kernel( + GemmKernel( const PackedAType* A, const PackedBType* B, int32_t* C, @@ -920,36 +1127,14 @@ struct MLAS_GEMM_U8X8_KERNEL_SSE }; constexpr size_t MLAS_GEMM_U8X8_KERNEL_SSE::PackedK; -constexpr size_t MLAS_GEMM_U8X8_KERNEL_SSE::StrideM; -constexpr size_t MLAS_GEMM_U8X8_KERNEL_SSE::StrideN; -constexpr size_t MLAS_GEMM_U8X8_KERNEL_SSE::StrideK; +constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8X8_KERNEL_SSE::Strides; +template void MLASCALL -MlasGemmU8X8OperationSse( +MlasGemmU8X8Operation( const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock - ) -/*++ - -Routine Description: - - This module implements the quantized integer matrix/matrix multiply - operation (QGEMM). - - This implementation supports SSE2 U8S8/U8U8. - -Arguments: - - WorkBlock - Supplies the structure containing the GEMM parameters. - -Return Value: - - None. - ---*/ -{ - return MlasGemmU8X8Operation(WorkBlock); -} + ); #endif @@ -1021,9 +1206,29 @@ struct MLAS_GEMM_U8S8_KERNEL_AVX2 typedef int8_t OffsetBType; static constexpr size_t PackedK = 4; - static constexpr size_t StrideM = 24; - static constexpr size_t StrideN = 256; - static constexpr size_t StrideK = 128; + static constexpr MLAS_GEMM_U8X8_STRIDES Strides{24, 256, 128}; + static constexpr MLAS_GEMM_U8X8_STRIDES PackedStrides{48, 256, 384}; + + MLAS_FORCEINLINE + static + bool + TryGemvKernel( + const uint8_t* A, + const uint8_t* B, + size_t ldb, + int32_t* C, + size_t CountK, + size_t CountN, + bool BIsSigned + ) + { + if (BIsSigned) { + MlasPlatform.GemvU8S8Kernel(A, B, C, CountK, CountN, ldb); + return true; + } + + return false; + } MLAS_FORCEINLINE static @@ -1060,7 +1265,7 @@ struct MLAS_GEMM_U8S8_KERNEL_AVX2 MLAS_FORCEINLINE static size_t - Kernel( + GemmKernel( const PackedAType* A, const PackedBType* B, int32_t* C, @@ -1094,46 +1299,20 @@ struct MLAS_GEMM_U8S8_KERNEL_AVX2 }; constexpr size_t MLAS_GEMM_U8S8_KERNEL_AVX2::PackedK; -constexpr size_t MLAS_GEMM_U8S8_KERNEL_AVX2::StrideM; -constexpr size_t MLAS_GEMM_U8S8_KERNEL_AVX2::StrideN; -constexpr size_t MLAS_GEMM_U8S8_KERNEL_AVX2::StrideK; +constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8S8_KERNEL_AVX2::Strides; +constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8S8_KERNEL_AVX2::PackedStrides; +template void -MLASCALL -MlasGemmU8S8OperationAvx2( +MlasGemmU8X8Operation( const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock - ) -/*++ + ); -Routine Description: - - This module implements the quantized integer matrix/matrix multiply - operation (QGEMM). - - This implementation supports AVX2/AVX512 U8S8 and AVX512VNNI U8S8/U8U8. - -Arguments: - - WorkBlock - Supplies the structure containing the GEMM parameters. - -Return Value: - - None. - ---*/ -{ - if ((WorkBlock->M == 1) && WorkBlock->BIsSigned && !WorkBlock->CIsFloat && - (WorkBlock->offa == 0) && (WorkBlock->offb == 0)) { - - if (MlasPlatform.GemvU8S8Kernel != nullptr) { - MlasPlatform.GemvU8S8Kernel(WorkBlock->A, WorkBlock->B, WorkBlock->C, - WorkBlock->K, WorkBlock->N, WorkBlock->ldb); - return; - } - } - - return MlasGemmU8X8Operation(WorkBlock); -} +template +void +MlasGemmU8X8PackedOperation( + const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock + ); struct MLAS_GEMM_U8U8_KERNEL_AVX2 { @@ -1142,9 +1321,32 @@ struct MLAS_GEMM_U8U8_KERNEL_AVX2 typedef uint8_t OffsetBType; static constexpr size_t PackedK = 2; - static constexpr size_t StrideM = 24; - static constexpr size_t StrideN = 256; - static constexpr size_t StrideK = 128; + static constexpr MLAS_GEMM_U8X8_STRIDES Strides{24, 256, 128}; + static constexpr MLAS_GEMM_U8X8_STRIDES PackedStrides{48, 256, 384}; + + MLAS_FORCEINLINE + static + bool + TryGemvKernel( + const uint8_t* A, + const uint8_t* B, + size_t ldb, + int32_t* C, + size_t CountK, + size_t CountN, + bool BIsSigned + ) + { + MLAS_UNREFERENCED_PARAMETER(A); + MLAS_UNREFERENCED_PARAMETER(B); + MLAS_UNREFERENCED_PARAMETER(ldb); + MLAS_UNREFERENCED_PARAMETER(C); + MLAS_UNREFERENCED_PARAMETER(CountK); + MLAS_UNREFERENCED_PARAMETER(CountN); + MLAS_UNREFERENCED_PARAMETER(BIsSigned); + + return false; + } MLAS_FORCEINLINE static @@ -1182,7 +1384,7 @@ struct MLAS_GEMM_U8U8_KERNEL_AVX2 MLAS_FORCEINLINE static size_t - Kernel( + GemmKernel( const PackedAType* A, const PackedBType* B, int32_t* C, @@ -1216,36 +1418,21 @@ struct MLAS_GEMM_U8U8_KERNEL_AVX2 }; constexpr size_t MLAS_GEMM_U8U8_KERNEL_AVX2::PackedK; -constexpr size_t MLAS_GEMM_U8U8_KERNEL_AVX2::StrideM; -constexpr size_t MLAS_GEMM_U8U8_KERNEL_AVX2::StrideN; -constexpr size_t MLAS_GEMM_U8U8_KERNEL_AVX2::StrideK; +constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8U8_KERNEL_AVX2::Strides; +constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8U8_KERNEL_AVX2::PackedStrides; +template void MLASCALL -MlasGemmU8U8OperationAvx2( +MlasGemmU8X8Operation( const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock - ) -/*++ + ); -Routine Description: - - This module implements the quantized integer matrix/matrix multiply - operation (QGEMM). - - This implementation supports AVX2/AVX512 U8U8. - -Arguments: - - WorkBlock - Supplies the structure containing the GEMM parameters. - -Return Value: - - None. - ---*/ -{ - return MlasGemmU8X8Operation(WorkBlock); -} +template +void +MlasGemmU8X8PackedOperation( + const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock + ); #endif @@ -1275,70 +1462,54 @@ Return Value: --*/ { - const auto* WorkBlock = (MLAS_GEMM_U8X8_WORK_BLOCK*)Context; + MLAS_GEMM_U8X8_WORK_BLOCK WorkBlock; - const int32_t ThreadCountM = WorkBlock->ThreadCountM; - const int32_t ThreadCountN = WorkBlock->ThreadCountN; + memcpy(&WorkBlock, Context, sizeof(MLAS_GEMM_U8X8_WORK_BLOCK)); - const int32_t ThreadIdM = ThreadId / ThreadCountN; - const int32_t ThreadIdN = ThreadId % ThreadCountN; + const int32_t ThreadIdM = ThreadId / WorkBlock.ThreadCountN; + const int32_t ThreadIdN = ThreadId % WorkBlock.ThreadCountN; // // Partition the operation along the M dimension. // - size_t M = WorkBlock->M; - size_t m; - size_t CountM; - - MlasPartitionWork(ThreadIdM, ThreadCountM, M, &m, &CountM); + MlasPartitionWork(ThreadIdM, WorkBlock.ThreadCountM, WorkBlock.M, + &WorkBlock.RangeStartM, &WorkBlock.RangeCountM); // // Partition the operation along the N dimension. // - size_t N = WorkBlock->N; - size_t n; - size_t CountN; - - const size_t BlockedN = (N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) / + const size_t BlockedN = (WorkBlock.N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) / MLAS_QGEMM_STRIDEN_THREAD_ALIGN; - MlasPartitionWork(ThreadIdN, ThreadCountN, BlockedN, &n, &CountN); + MlasPartitionWork(ThreadIdN, WorkBlock.ThreadCountN, BlockedN, + &WorkBlock.RangeStartN, &WorkBlock.RangeCountN); - n *= MLAS_QGEMM_STRIDEN_THREAD_ALIGN; - CountN *= MLAS_QGEMM_STRIDEN_THREAD_ALIGN; + WorkBlock.RangeStartN *= MLAS_QGEMM_STRIDEN_THREAD_ALIGN; + WorkBlock.RangeCountN *= MLAS_QGEMM_STRIDEN_THREAD_ALIGN; - if (CountN > N - n) { - CountN = N - n; - } + WorkBlock.RangeCountN = std::min(WorkBlock.N - WorkBlock.RangeStartN, + WorkBlock.RangeCountN); // // Dispatch the partitioned operation. // - MLAS_GEMM_U8X8_WORK_BLOCK LocalWorkBlock; - - memcpy(&LocalWorkBlock, WorkBlock, sizeof(MLAS_GEMM_U8X8_WORK_BLOCK)); - - LocalWorkBlock.M = CountM; - LocalWorkBlock.N = CountN; - LocalWorkBlock.A += m * LocalWorkBlock.lda; - LocalWorkBlock.B += n; - LocalWorkBlock.C += m * LocalWorkBlock.ldc + n; - - if (LocalWorkBlock.BiasFloat != nullptr) { - LocalWorkBlock.BiasFloat += n; - } - #if defined(MLAS_TARGET_AMD64) - if (LocalWorkBlock.BIsSigned) { - MlasPlatform.GemmU8S8Operation(&LocalWorkBlock); + PMLAS_GEMM_U8X8_OPERATION GemmU8X8Operation; + + if (WorkBlock.BIsSigned) { + GemmU8X8Operation = WorkBlock.BIsPacked ? + MlasPlatform.GemmU8S8PackedOperation : MlasPlatform.GemmU8S8Operation; } else { - MlasPlatform.GemmU8U8Operation(&LocalWorkBlock); + GemmU8X8Operation = WorkBlock.BIsPacked ? + MlasPlatform.GemmU8U8PackedOperation : MlasPlatform.GemmU8U8Operation; } + + GemmU8X8Operation(&WorkBlock); #else - MlasGemmU8X8OperationSse(&LocalWorkBlock); + MlasGemmU8X8Operation(&WorkBlock); #endif } @@ -1351,8 +1522,8 @@ MlasGemmU8X8Schedule( Routine Description: - This module implements the quantized integer matrix/matrix multiply - operation (QGEMM). + This routine schedules the quantized integer matrix/matrix multiply + operation (QGEMM) across one or more threads. Arguments: @@ -1445,7 +1616,7 @@ MlasGemm( Routine Description: - This module implements the quantized integer matrix/matrix multiply + This routine implements the quantized integer matrix/matrix multiply operation (QGEMM). Arguments: @@ -1536,7 +1707,7 @@ MlasGemm( Routine Description: - This module implements the quantized integer matrix/matrix multiply + This routine implements the quantized integer matrix/matrix multiply operation (QGEMM). Arguments: @@ -1615,3 +1786,380 @@ Return Value: } #endif + +#ifdef MLAS_TARGET_AMD64 + +void +MLASCALL +MlasGemm( + size_t M, + size_t N, + size_t K, + const uint8_t* A, + size_t lda, + uint8_t offa, + const void* PackedB, + uint8_t offb, + bool BIsSigned, + int32_t* C, + size_t ldc, + MLAS_THREADPOOL* ThreadPool + ) +/*++ + +Routine Description: + + This routine implements the quantized integer matrix/matrix multiply + operation (QGEMM). + +Arguments: + + M - Supplies the number of rows of matrix A and matrix C. + + N - Supplies the number of columns of matrix B and matrix C. + + K - Supplies the number of columns of matrix A and the number of rows of + matrix B. + + A - Supplies the address of matrix A. + + lda - Supplies the first dimension of matrix A. + + offa - Supplies the zero point offset of matrix A. + + PackedB - Supplies the address of packed matrix B. + + offb - Supplies the zero point offset of matrix B. + + BIsSigned - Supplies true if matrix B is signed data, else false if matrix + B is unsigned data. + + C - Supplies the address of matrix C. + + ldc - Supplies the first dimension of matrix C. + + ThreadPool - Supplies the thread pool object to use, else nullptr if the + base library threading support should be used. + +Return Value: + + None. + +--*/ +{ + MLAS_GEMM_U8X8_WORK_BLOCK WorkBlock; + + // + // Capture the GEMM parameters to the work block. + // + + memset(&WorkBlock, 0, sizeof(MLAS_GEMM_U8X8_WORK_BLOCK)); + + WorkBlock.M = M; + WorkBlock.N = N; + WorkBlock.K = K; + WorkBlock.A = A; + WorkBlock.lda = lda; + WorkBlock.B = PackedB; + WorkBlock.C = C; + WorkBlock.ldc = ldc; + WorkBlock.offa = offa; + WorkBlock.offb = offb; + WorkBlock.BIsPacked = true; + WorkBlock.BIsSigned = BIsSigned; + + // + // Schedule the operation across a set of worker threads. + // + + MlasGemmU8X8Schedule(&WorkBlock, ThreadPool); +} + +void +MLASCALL +MlasGemm( + size_t M, + size_t N, + size_t K, + const uint8_t* A, + size_t lda, + uint8_t offa, + const void* PackedB, + uint8_t offb, + bool BIsSigned, + float* C, + size_t ldc, + const float* Scale, + const float* Bias, + MLAS_THREADPOOL* ThreadPool + ) +/*++ + +Routine Description: + + This routine implements the quantized integer matrix/matrix multiply + operation (QGEMM). + +Arguments: + + M - Supplies the number of rows of matrix A and matrix C. + + N - Supplies the number of columns of matrix B and matrix C. + + K - Supplies the number of columns of matrix A and the number of rows of + matrix B. + + A - Supplies the address of matrix A. + + lda - Supplies the first dimension of matrix A. + + offa - Supplies the zero point offset of matrix A. + + PackedB - Supplies the address of packed matrix B. + + offb - Supplies the zero point offset of matrix B. + + BIsSigned - Supplies true if matrix B is signed data, else false if matrix + B is unsigned data. + + C - Supplies the address of matrix C. + + ldc - Supplies the first dimension of matrix C. + + Scale - Supplies the scale multiplier to apply to each element of matrix C. + Used to scale the integer output of the QGEMM back to a floating point + number. + + Bias - Supplies the bias vector to apply to element of matrix C. The vector + is of length N. + + ThreadPool - Supplies the thread pool object to use, else nullptr if the + base library threading support should be used. + +Return Value: + + None. + +--*/ +{ + MLAS_GEMM_U8X8_WORK_BLOCK WorkBlock; + + // + // Capture the GEMM parameters to the work block. + // + + memset(&WorkBlock, 0, sizeof(MLAS_GEMM_U8X8_WORK_BLOCK)); + + WorkBlock.M = M; + WorkBlock.N = N; + WorkBlock.K = K; + WorkBlock.A = A; + WorkBlock.lda = lda; + WorkBlock.B = PackedB; + WorkBlock.C = (int32_t*)C; + WorkBlock.ldc = ldc; + WorkBlock.Scale = Scale; + WorkBlock.BiasFloat = Bias; + WorkBlock.offa = offa; + WorkBlock.offb = offb; + WorkBlock.BIsPacked = true; + WorkBlock.BIsSigned = BIsSigned; + WorkBlock.CIsFloat = true; + + // + // Schedule the operation across a set of worker threads. + // + + MlasGemmU8X8Schedule(&WorkBlock, ThreadPool); +} + +size_t +MLASCALL +MlasGemmPackBSize( + size_t N, + size_t K, + bool BIsSigned + ) +/*++ + +Routine Description: + + This routine computes the number of bytes required to pack a matrix with + the supplied shape and type. + +Arguments: + + N - Supplies the number of columns of matrix B. + + K - Supplies the the number of rows of matrix B. + + BIsSigned - Supplies true if matrix B is signed data, else false if matrix + B is unsigned data. + +Return Value: + + Returns the number of bytes required to pack the matrix. + +--*/ +{ + // + // Retrieve the address of the packed operation function for the platform. + // + + PMLAS_GEMM_U8X8_OPERATION GemmU8X8Operation = BIsSigned ? + MlasPlatform.GemmU8S8PackedOperation : MlasPlatform.GemmU8U8PackedOperation; + + // + // Retrieve the packing parameters based on the packed operation function. + // + + size_t PackedK; + + if (GemmU8X8Operation == &MlasGemmU8X8PackedOperation) { + PackedK = MLAS_GEMM_U8S8_KERNEL_AVX2::PackedK; + } else if (GemmU8X8Operation == &MlasGemmU8X8PackedOperation) { + PackedK = MLAS_GEMM_U8U8_KERNEL_AVX2::PackedK; + } else { + return 0; + } + + // + // Compute the number of bytes required to hold the packed buffer. + // + + const size_t AlignedN = + (N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1); + const size_t AlignedK = (K + PackedK - 1) & ~(PackedK - 1); + + const size_t BytesRequired = + (AlignedN * sizeof(int32_t)) + (AlignedN * AlignedK * sizeof(uint8_t)); + const size_t BufferAlignment = MlasGetPreferredBufferAlignment(); + const size_t AlignedBytesRequired = (BytesRequired + BufferAlignment - 1) & + ~(BufferAlignment - 1); + + return AlignedBytesRequired; +} + +void +MLASCALL +MlasGemmPackB( + size_t N, + size_t K, + const uint8_t* B, + size_t ldb, + bool BIsSigned, + void* PackedB + ) +/*++ + +Routine Description: + + This routine packs the supplied matrix B to the supplied packed matrix B + buffer. The size of the packed buffer was obtained from MlasGemmPackBSize. + +Arguments: + + N - Supplies the number of columns of matrix B. + + K - Supplies the the number of rows of matrix B. + + B - Supplies the address of matrix B. + + ldb - Supplies the first dimension of matrix B. + + BIsSigned - Supplies true if matrix B is signed data, else false if matrix + B is unsigned data. + + PackedB - Supplies the address of packed matrix B. + +Return Value: + + None. + +--*/ +{ + // + // Retrieve the address of the packed operation function for the platform. + // + + PMLAS_GEMM_U8X8_OPERATION GemmU8X8Operation = BIsSigned ? + MlasPlatform.GemmU8S8PackedOperation : MlasPlatform.GemmU8U8PackedOperation; + + // + // Retrieve the packing parameters based on the packed operation function. + // + + size_t PackedK; + size_t StrideK; + + if (GemmU8X8Operation == &MlasGemmU8X8PackedOperation) { + PackedK = MLAS_GEMM_U8S8_KERNEL_AVX2::PackedK; + StrideK = MLAS_GEMM_U8S8_KERNEL_AVX2::PackedStrides.K; + } else if (GemmU8X8Operation == &MlasGemmU8X8PackedOperation) { + PackedK = MLAS_GEMM_U8U8_KERNEL_AVX2::PackedK; + StrideK = MLAS_GEMM_U8U8_KERNEL_AVX2::PackedStrides.K; + } else { + throw std::runtime_error("packing unavailable"); + } + + // + // Reserve and initialize storage for the column sum buffer to hold the sums + // of the elements along each of the columns. + // + + const size_t AlignedN = + (N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1); + + int32_t* PackedColumnSumBuffer = (int32_t*)PackedB; + std::fill_n(PackedColumnSumBuffer, AlignedN, 0); + PackedB = PackedColumnSumBuffer + AlignedN; + + // + // Step through each slice of matrix B along the K dimension. + // + + size_t CountK; + + for (size_t k = 0; k < K; k += CountK) { + + CountK = (std::min)(K - k, StrideK); + + // + // Step through each slice of matrix B along the N dimension. + // + + const size_t AlignedK = (CountK + PackedK - 1) & ~(PackedK - 1); + uint8_t* pb = (uint8_t*)PackedB; + size_t CountN; + + for (size_t n = 0; n < N; n += CountN) { + + constexpr size_t BatchedN = 128; + MLAS_DECLSPEC_ALIGN(int32_t ColumnSumBuffer[BatchedN], 64); + + CountN = (std::min)(N - n, BatchedN); + + if (GemmU8X8Operation == &MlasGemmU8X8PackedOperation) { + MLAS_GEMM_U8S8_KERNEL_AVX2::CopyPackB(pb, B + n, ldb, CountN, CountK, ColumnSumBuffer, BIsSigned); + } else { + MLAS_GEMM_U8U8_KERNEL_AVX2::CopyPackB(pb, B + n, ldb, CountN, CountK, ColumnSumBuffer, BIsSigned); + } + + // + // Accumulate this batch of the column sum buffer into the packed + // buffer accumulators. + // + + for (size_t nn = 0; nn < CountN; nn++) { + PackedColumnSumBuffer[n + nn] += ColumnSumBuffer[nn]; + } + + pb += CountN * AlignedK; + } + + PackedB = (uint8_t*)PackedB + AlignedN * AlignedK; + B += ldb * CountK; + } +} + +#endif diff --git a/onnxruntime/core/util/qmath.h b/onnxruntime/core/util/qmath.h index 9d2d03d1bd..4e9f4c61cd 100644 --- a/onnxruntime/core/util/qmath.h +++ b/onnxruntime/core/util/qmath.h @@ -12,6 +12,10 @@ #define MLAS_SUPPORTS_GEMM_U8X8 #endif +#if defined(_M_AMD64) || defined(__x86_64__) +#define MLAS_SUPPORTS_PACKED_GEMM_U8X8 +#endif + namespace onnxruntime { void QGemm( diff --git a/onnxruntime/test/mlas/unittest.cpp b/onnxruntime/test/mlas/unittest.cpp index 9c74291f95..bb3b0d6189 100644 --- a/onnxruntime/test/mlas/unittest.cpp +++ b/onnxruntime/test/mlas/unittest.cpp @@ -45,6 +45,10 @@ Abstract: #define MLAS_HAS_QGEMM_U8X8 #endif +#if defined(_M_AMD64) || defined(__x86_64__) +#define MLAS_HAS_PACKED_QGEMM_U8X8 +#endif + MLAS_THREADPOOL* threadpool = nullptr; template @@ -468,11 +472,129 @@ public: #ifdef MLAS_HAS_QGEMM_U8X8 -template +template +class MlasQgemmU8X8U8X8TestBase; + +template<> +class MlasQgemmU8X8U8X8TestBase : public MlasTestBase +{ +protected: + void + TestGemm( + size_t M, + size_t N, + size_t K, + const uint8_t* A, + size_t lda, + uint8_t offa, + const uint8_t* B, + size_t ldb, + uint8_t offb, + bool BIsSigned, + int32_t* C, + size_t ldc + ) + { + MlasGemm(M, N, K, A, lda, offa, B, ldb, offb, BIsSigned, C, ldc, threadpool); + } + + void + TestGemm( + size_t M, + size_t N, + size_t K, + const uint8_t* A, + size_t lda, + uint8_t offa, + const uint8_t* B, + size_t ldb, + uint8_t offb, + bool BIsSigned, + float* C, + size_t ldc, + float CScale, + const float* Bias + ) + { + MlasGemm(M, N, K, A, lda, offa, B, ldb, offb, BIsSigned, C, ldc, &CScale, Bias, threadpool); + } +}; + +#ifdef MLAS_HAS_PACKED_QGEMM_U8X8 + +template<> +class MlasQgemmU8X8U8X8TestBase : public MlasTestBase +{ +private: + void* + PackB( + size_t N, + size_t K, + const uint8_t* B, + size_t ldb, + bool BIsSigned + ) + { + size_t PackedBSize = MlasGemmPackBSize(N, K, BIsSigned); + void* PackedB = BufferBPacked.GetBuffer(PackedBSize); + MlasGemmPackB(N, K, B, ldb, BIsSigned, PackedB); + return PackedB; + } + +protected: + void + TestGemm( + size_t M, + size_t N, + size_t K, + const uint8_t* A, + size_t lda, + uint8_t offa, + const uint8_t* B, + size_t ldb, + uint8_t offb, + bool BIsSigned, + int32_t* C, + size_t ldc + ) + { + const void* PackedB = PackB(N, K, B, ldb, BIsSigned); + MlasGemm(M, N, K, A, lda, offa, PackedB, offb, BIsSigned, C, ldc, threadpool); + } + + void + TestGemm( + size_t M, + size_t N, + size_t K, + const uint8_t* A, + size_t lda, + uint8_t offa, + const uint8_t* B, + size_t ldb, + uint8_t offb, + bool BIsSigned, + float* C, + size_t ldc, + float CScale, + const float* Bias + ) + { + const void* PackedB = PackB(N, K, B, ldb, BIsSigned); + MlasGemm(M, N, K, A, lda, offa, PackedB, offb, BIsSigned, C, ldc, &CScale, Bias, threadpool); + } + +private: + MatrixGuardBuffer BufferBPacked; +}; + +#endif + +template class MlasQgemmU8X8Test; -template -class MlasQgemmU8X8Test : public MlasTestBase +template +class MlasQgemmU8X8Test : public MlasQgemmU8X8U8X8TestBase { private: void @@ -511,7 +633,7 @@ private: std::fill_n(C, M * N, -1); std::fill_n(CReference, M * N, -1); - MlasGemm(M, N, K, A, lda, offa, B, ldb, offb, BIsSigned, C, ldc, threadpool); + this->TestGemm(M, N, K, A, lda, offa, B, ldb, offb, BIsSigned, C, ldc); ReferenceQgemm(M, N, K, A, lda, offa, (const xint8_t*)B, ldb, (xint8_t)offb, CReference, ldc); for (size_t f = 0; f < M * N; f++) { @@ -569,6 +691,9 @@ public: void ) override { + for (size_t b = 1; b < 16; b++) { + Test(b, b, b, 14, 211); + } for (size_t b = 1; b < 16; b++) { Test(b, b, b, 14, 211); } @@ -583,7 +708,8 @@ public: Test(1, 32, b, 0, 0); Test(1, b, b, 0, 0); } - Test(1024, 1024, 1024, 5, 8); + Test(43, 500, 401, 183, 223); + Test(1023, 1023, 1023, 5, 8); } void @@ -651,8 +777,8 @@ public: } }; -template -class MlasQgemmU8X8Test : public MlasTestBase +template +class MlasQgemmU8X8Test : public MlasQgemmU8X8U8X8TestBase { private: void @@ -714,7 +840,7 @@ private: } } - MlasGemm(M, N, K, A, lda, offa, B, ldb, offb, BIsSigned, C, ldc, &CScale, Bias, threadpool); + this->TestGemm(M, N, K, A, lda, offa, B, ldb, offb, BIsSigned, C, ldc, CScale, Bias); for (size_t f = 0; f < M * N; f++) { // Sensitive to comparing positive/negative zero. @@ -767,6 +893,7 @@ public: for (size_t b = 1; b < 96; b++) { Test(1, b, 32, 0, 0); } + Test(43, 503, 401, 183, 223); Test(1024, 1024, 256, 13, 15); } }; @@ -2435,25 +2562,40 @@ RunThreadedTests( #ifdef MLAS_HAS_QGEMM_U8X8 printf("QGEMM U8S8=int32_t tests.\n"); - onnxruntime::make_unique>()->ExecuteShort(); - printf("QGEMM U8U8=int32_t tests.\n"); - onnxruntime::make_unique>()->ExecuteShort(); + onnxruntime::make_unique>()->ExecuteShort(); printf("QGEMM U8S8=float tests.\n"); - onnxruntime::make_unique>()->ExecuteShort(); + onnxruntime::make_unique>()->ExecuteShort(); + printf("QGEMM U8U8=int32_t tests.\n"); + onnxruntime::make_unique>()->ExecuteShort(); printf("QGEMM U8U8=float tests.\n"); - onnxruntime::make_unique>()->ExecuteShort(); + onnxruntime::make_unique>()->ExecuteShort(); +#endif + +#ifdef MLAS_HAS_PACKED_QGEMM_U8X8 + if (MlasGemmPackBSize(128, 128, true) > 0) { + printf("QGEMM U8S8=int32_t packed tests.\n"); + onnxruntime::make_unique>()->ExecuteShort(); + printf("QGEMM U8S8=float packed tests.\n"); + onnxruntime::make_unique>()->ExecuteShort(); + } + if (MlasGemmPackBSize(128, 128, false) > 0) { + printf("QGEMM U8U8=int32_t packed tests.\n"); + onnxruntime::make_unique>()->ExecuteShort(); + printf("QGEMM U8U8=float packed tests.\n"); + onnxruntime::make_unique>()->ExecuteShort(); + } #endif printf("Conv2D tests.\n"); onnxruntime::make_unique()->ExecuteShort(); if (MlasNchwcGetBlockSize() > 1) { - onnxruntime::make_unique()->ExecuteShort(); + onnxruntime::make_unique()->ExecuteShort(); } printf("Pool2D tests.\n"); onnxruntime::make_unique()->ExecuteShort(); if (MlasNchwcGetBlockSize() > 1) { - onnxruntime::make_unique()->ExecuteShort(); + onnxruntime::make_unique()->ExecuteShort(); } printf("Pool3D tests.\n");