diff --git a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc index a6db615c89..193eb78a3a 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc @@ -64,7 +64,10 @@ Status MatMulIntegerToFloatBase::ComputeCommon(OpKernelContext* ctx, const Tensor* b_zp_tensor, const Tensor* bias_tensor) const { MatMulComputeHelper helper; - ORT_RETURN_IF_ERROR(helper.Compute(a_shape, b_tensor ? b_tensor->Shape() : b_shape_)); + ORT_RETURN_IF_ERROR(helper.Compute(a_shape, + b_tensor ? b_tensor->Shape() : b_shape_, + b_scale_tensor ? &b_scale_tensor->Shape() : nullptr, + b_zp_tensor ? &b_zp_tensor->Shape() : nullptr)); Tensor* y = ctx->Output(OUT_Y, helper.OutputShape()); // Bail out early if the output is going to be empty @@ -122,12 +125,9 @@ Status MatMulIntegerToFloatBase::ComputeCommon(OpKernelContext* ctx, std::vector gemm_data_vec(num_gemms); for (size_t gemm_idx = 0; gemm_idx < num_gemms; gemm_idx++) { - int64_t scale_zp_offset = helper.RightOffsets()[gemm_idx] / helper.K(); - int64_t scale_offset = is_b_scale_per_column ? scale_zp_offset : 0; - int64_t zp_offset = is_b_zp_per_column ? scale_zp_offset : 0; gemm_scale_procs.emplace_back(y_data + helper.OutputOffsets()[gemm_idx], gemm_shape.N, - b_scale_data + scale_offset, + b_scale_data + helper.RightScaleOffsets()[gemm_idx], bias_data, MLAS_QGEMM_OUTPUT_MODE::ZeroMode, is_b_scale_per_column ? MLAS_QUANTIZATION_GRANULARITY::PerColumn : MLAS_QUANTIZATION_GRANULARITY::PerMatrix); @@ -139,7 +139,7 @@ Status MatMulIntegerToFloatBase::ComputeCommon(OpKernelContext* ctx, params.BIsPacked = bool(packed_b_); params.B = b_tensor ? static_cast(b_tensor->DataRaw()) + helper.RightOffsets()[gemm_idx] : packed_b_.get(); params.ldb = gemm_shape.N; - params.ZeroPointB = b_zp_ptr + zp_offset; + params.ZeroPointB = b_zp_ptr + helper.RightZeroPointOffsets()[gemm_idx]; params.PerColumnZeroPoints = is_b_zp_per_column; params.C = reinterpret_cast(y_data + helper.OutputOffsets()[gemm_idx]); params.ldc = gemm_shape.N; diff --git a/onnxruntime/core/providers/cpu/math/matmul_helper.h b/onnxruntime/core/providers/cpu/math/matmul_helper.h index d4d1b8f8bb..05e5281c92 100644 --- a/onnxruntime/core/providers/cpu/math/matmul_helper.h +++ b/onnxruntime/core/providers/cpu/math/matmul_helper.h @@ -149,6 +149,33 @@ class MatMulComputeHelper { return Status::OK(); } + Status Compute(const TensorShape& left_shape, const TensorShape& right_shape, + const TensorShape* right_scale_shape, const TensorShape* right_zp_shape, + bool transa = false, bool transb = false) { + ORT_RETURN_IF_ERROR(Compute(left_shape, right_shape, transa, transb)); + right_zp_offsets_.clear(); + right_scale_offsets_.clear(); + right_zp_offsets_.resize(right_offsets_.size()); + right_scale_offsets_.resize(right_offsets_.size()); + + auto SetRightMatrixQuantParam = [this, &right_shape](const TensorShape* param_shape, std::vector& quant_param_offsets) { + if (nullptr != param_shape && param_shape->NumDimensions() > 1) { + ORT_RETURN_IF_NOT(param_shape->NumDimensions() == right_shape.NumDimensions() && param_shape->Size() * K_ == right_shape.Size(), + "Per-column quantization parameter of batched matrix should have same dimension as the matrix," + "and its size by K should be equal to the matrix's size."); + for (size_t batch_id = 0; batch_id < quant_param_offsets.size(); batch_id++) { + quant_param_offsets[batch_id] = right_offsets_[batch_id] / K_; + } + } + return Status::OK(); + }; + + ORT_RETURN_IF_ERROR(SetRightMatrixQuantParam(right_zp_shape, right_zp_offsets_)); + ORT_RETURN_IF_ERROR(SetRightMatrixQuantParam(right_scale_shape, right_scale_offsets_)); + + return Status::OK(); + } + private: void ComputeBroadcastOffsets() { num_broadcasted_dims_ = left_padded_dims_.size() - 2; @@ -228,6 +255,9 @@ class MatMulComputeHelper { std::vector right_offsets_; std::vector output_offsets_; + std::vector right_zp_offsets_; + std::vector right_scale_offsets_; + public: // output shape const TensorShape& OutputShape() const { @@ -264,6 +294,16 @@ class MatMulComputeHelper { return output_offsets_; } + // Batched Scale Offset for right matrices + const std::vector& RightScaleOffsets() const { + return right_scale_offsets_; + } + + // Batched Zero Point Offset for right matrices + const std::vector& RightZeroPointOffsets() const { + return right_zp_offsets_; + } + template static void OffsetToArrays(T* p, const std::vector& offsets, gsl::span arrays) { auto len = offsets.size(); diff --git a/onnxruntime/core/providers/cpu/math/matmul_integer.cc b/onnxruntime/core/providers/cpu/math/matmul_integer.cc index f7bc00cb9b..b8ab9f08e2 100644 --- a/onnxruntime/core/providers/cpu/math/matmul_integer.cc +++ b/onnxruntime/core/providers/cpu/math/matmul_integer.cc @@ -41,27 +41,9 @@ ONNX_OPERATOR_TYPED_KERNEL_EX( MatMulInteger); Status MatMulInteger::Compute(OpKernelContext* ctx) const { - MatMulComputeHelper helper; const auto* a = ctx->Input(IN_A); const auto* b = packed_b_ ? nullptr : ctx->Input(IN_B); - const uint8_t* b_data; - bool b_is_signed; - if (nullptr != b) { - ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b->Shape())); - b_data = static_cast(b->DataRaw()); - b_is_signed = b->IsDataType(); - } else { - ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape_)); - b_data = static_cast(packed_b_.get()); - b_is_signed = b_is_signed_; - } - - Tensor* y = ctx->Output(OUT_Y, helper.OutputShape()); - // Bail out early if the output is going to be empty - if (y->Shape().Size() == 0) - return Status::OK(); - // validate zero points uint8_t a_offset = 0; const auto* a_zero_point = ctx->Input(IN_A_ZERO_POINT); @@ -82,6 +64,24 @@ Status MatMulInteger::Compute(OpKernelContext* ctx) const { b_offset_ptr = static_cast(b_zero_point->DataRaw()); } + MatMulComputeHelper helper; + const uint8_t* b_data; + bool b_is_signed; + if (nullptr != b) { + ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b->Shape(), nullptr, b_zero_point ? &b_zero_point->Shape() : nullptr)); + b_data = static_cast(b->DataRaw()); + b_is_signed = b->IsDataType(); + } else { + ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape_, nullptr, b_zero_point ? &b_zero_point->Shape() : nullptr)); + b_data = static_cast(packed_b_.get()); + b_is_signed = b_is_signed_; + } + + Tensor* y = ctx->Output(OUT_Y, helper.OutputShape()); + // Bail out early if the output is going to be empty + if (y->Shape().Size() == 0) + return Status::OK(); + const auto* a_data = a->template Data(); auto* y_data = y->template MutableData(); @@ -95,12 +95,11 @@ Status MatMulInteger::Compute(OpKernelContext* ctx) const { std::vector gemm_data_vec(batch_size); for (size_t batch = 0; batch < batch_size; batch++) { - int64_t zp_offset = is_b_zp_per_column ? helper.RightOffsets()[batch] / helper.K() : 0; auto& gemm_params = gemm_data_vec[batch]; gemm_params.lda = gemm_shape.K; gemm_params.ZeroPointA = a_offset; gemm_params.ldb = gemm_shape.N; - gemm_params.ZeroPointB = b_offset_ptr + zp_offset; + gemm_params.ZeroPointB = b_offset_ptr + helper.RightZeroPointOffsets()[batch]; gemm_params.PerColumnZeroPoints = is_b_zp_per_column; gemm_params.ldc = gemm_shape.N; gemm_params.BIsPacked = bool(packed_b_); diff --git a/onnxruntime/core/providers/cpu/math/quantize_linear_matmul.cc b/onnxruntime/core/providers/cpu/math/quantize_linear_matmul.cc index 50a41cfa4c..60068885b6 100644 --- a/onnxruntime/core/providers/cpu/math/quantize_linear_matmul.cc +++ b/onnxruntime/core/providers/cpu/math/quantize_linear_matmul.cc @@ -25,32 +25,8 @@ ONNX_OPERATOR_KERNEL_EX( QLinearMatMul); Status QLinearMatMul::Compute(OpKernelContext* ctx) const { - MatMulComputeHelper helper; const auto* a = ctx->Input(IN_A); - - const uint8_t* b_data; - bool b_is_signed; // can't modify b_is_signed_, this is a const method - if (packed_b_) { - ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape_)); - b_data = static_cast(packed_b_.get()); - b_is_signed = b_is_signed_; - } else { - const Tensor* b = ctx->Input(IN_B); - if (b == nullptr) { - // For required input, the framework has checks to ensure this won't happen, - // dead code to quiet the compiler warning. - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Required input B can not be null!"); - } - ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b->Shape())); - b_data = static_cast(b->DataRaw()); - b_is_signed = b->IsDataType(); - } - - Tensor* y = ctx->Output(OUT_Y, helper.OutputShape()); - // Bail out early if the output is going to be empty - if (y->Shape().Size() == 0) - return Status::OK(); + const auto* b = packed_b_ ? nullptr : ctx->Input(IN_B); // validate offsets const auto* a_offset = ctx->Input(IN_A_ZERO_POINT); @@ -58,8 +34,8 @@ Status QLinearMatMul::Compute(OpKernelContext* ctx) const { const auto* y_offset = ctx->Input(IN_Y_ZERO_POINT); ORT_ENFORCE(IsScalarOr1ElementVector(a_offset), "QLinearMatmul : input zero point must be a scalar or 1D tensor of size 1"); - ORT_ENFORCE(IsScalarOr1ElementVector(b_offset), - "QLinearMatmul : weight zero point must be a scalar or 1D tensor of size 1"); + ORT_ENFORCE(IsBQuantParamSupported(b_offset->Shape(), b ? b->Shape() : b_shape_), + "QLinearMatmul : weight zero point must be a scalar, 1D tensor of size 1, or last to second dimension is 1"); ORT_ENFORCE(IsScalarOr1ElementVector(y_offset), "QLinearMatmul : result zero point must be a scalar or 1D tensor of size 1"); @@ -69,16 +45,38 @@ Status QLinearMatMul::Compute(OpKernelContext* ctx) const { const auto* y_scale = ctx->Input(IN_Y_SCALE); ORT_ENFORCE(IsScalarOr1ElementVector(a_scale), "QLinearMatmul : input scale must be a scalar or 1D tensor of size 1"); - ORT_ENFORCE(IsScalarOr1ElementVector(b_scale), - "QLinearMatmul : weight scale must be a scalar or 1D tensor of size 1"); + ORT_ENFORCE(IsBQuantParamSupported(b_scale->Shape(), b ? b->Shape() : b_shape_), + "QLinearMatmul : weight scale must be a scalar, 1D tensor of size 1, or last to second dimension is 1"); ORT_ENFORCE(IsScalarOr1ElementVector(y_scale), "QLinearMatmul : result scale must be a scalar or 1D tensor of size 1"); + MatMulComputeHelper helper; + const uint8_t* b_data; + bool b_is_signed; + if (nullptr != b) { + ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b->Shape(), &b_scale->Shape(), &b_offset->Shape())); + b_data = static_cast(b->DataRaw()); + b_is_signed = b->IsDataType(); + } else { + ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape_, &b_scale->Shape(), &b_offset->Shape())); + b_data = static_cast(packed_b_.get()); + b_is_signed = b_is_signed_; + } + + Tensor* y = ctx->Output(OUT_Y, helper.OutputShape()); + // Bail out early if the output is going to be empty + if (y->Shape().Size() == 0) + return Status::OK(); + + const auto* b_scale_data = b_scale->template Data(); auto a_scale_data = *(a_scale->template Data()); - auto b_scale_data = *(b_scale->template Data()); auto y_scale_data = *(y_scale->template Data()); - const float real_multiplier = (a_scale_data * b_scale_data) / y_scale_data; + const int64_t output_scale_size = b_scale->Shape().Size(); + std::vector output_scales(output_scale_size); + for (int64_t i = 0; i < output_scale_size; i++) { + output_scales[i] = (a_scale_data * b_scale_data[i] / y_scale_data); + } AllocatorPtr alloc; ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&alloc)); @@ -97,14 +95,16 @@ Status QLinearMatMul::Compute(OpKernelContext* ctx) const { gemm_params.lda = gemm_shape.K; gemm_params.ZeroPointA = *a_offset->template Data(); gemm_params.ldb = gemm_shape.N; - gemm_params.ZeroPointB = static_cast(b_offset->DataRaw()); gemm_params.C = gemm_output; gemm_params.ldc = gemm_shape.N; gemm_params.BIsPacked = bool(packed_b_); + gemm_params.PerColumnZeroPoints = !IsScalarOr1ElementVector(b_offset); + auto b_zp_data = static_cast(b_offset->DataRaw()); for (size_t i = 0; i < helper.OutputOffsets().size(); i++) { gemm_params.A = a->template Data() + helper.LeftOffsets()[i]; gemm_params.B = b_data + helper.RightOffsets()[i]; + gemm_params.ZeroPointB = b_zp_data + helper.RightZeroPointOffsets()[i]; MlasGemm(gemm_shape, gemm_params, ctx->GetOperatorThreadPool()); @@ -114,8 +114,8 @@ Status QLinearMatMul::Compute(OpKernelContext* ctx) const { nullptr, static_cast(helper.M()), static_cast(helper.N()), - &real_multiplier, - false, + output_scales.data() + helper.RightScaleOffsets()[i], + output_scales.size() > 1, *y_offset->template Data()); } diff --git a/onnxruntime/test/providers/cpu/math/quantize_linear_matmul_test.cc b/onnxruntime/test/providers/cpu/math/quantize_linear_matmul_test.cc index dee74464b5..24a3bc9a65 100644 --- a/onnxruntime/test/providers/cpu/math/quantize_linear_matmul_test.cc +++ b/onnxruntime/test/providers/cpu/math/quantize_linear_matmul_test.cc @@ -194,6 +194,79 @@ TEST(QuantizeLinearMatmulOpTest, QLinearMatMulAllInputExceptT1AreInitializers) { QLinearMatMul2DTest(true); } +TEST(QuantizeLinearMatmulOpTest, PerColumn_2D) { + OpTester test("QLinearMatMul", 10); + test.AddInput("a", + {2, 4}, + {125, 135, 133, 122, + 132, 123, 136, 135}); + test.AddInput("a_scale", {}, {0.1f}); + test.AddInput("a_zero_point", {}, {133}); + test.AddInput("b", + {4, 4}, + {0, -8, 2, 3, + -11, -13, -8, 1, + 2, 4, 4, -10, + 3, 2, -11, 2}); + test.AddInput("b_scale", {4}, + {0.1f, 0.2f, 0.3f, 0.4f}); + test.AddInput("b_zero_point", + {1, 4}, + {1, -2, 2, -1}); + test.AddInput("y_scale", {}, {0.2f}); + test.AddInput("y_zero_point", {}, {130}); + + test.AddOutput("y", + {2, 4}, + {128, 128, 148, 118, + 136, 144, 142, 121}); + + test.Run(); +} + +TEST(QuantizeLinearMatmulOpTest, PerColumn_ND) { + OpTester test("QLinearMatMul", 10); + test.AddInput("a", + {2, 2, 4}, + {125, 135, 133, 122, + 132, 123, 136, 135, + + 125, 135, 133, 122, + 132, 123, 136, 135}); + test.AddInput("a_scale", {}, {0.1f}); + test.AddInput("a_zero_point", {}, {133}); + test.AddInput("b", + {2, 4, 4}, + {0, -8, 2, 3, + -11, -13, -8, 1, + 2, 4, 4, -10, + 3, 2, -11, 2, + + 0, -8, 2, 3, + -11, -13, -8, 1, + 2, 4, 4, -10, + 3, 2, -11, 2}); + test.AddInput("b_scale", {2, 1, 4}, + {0.1f, 0.2f, 0.3f, 0.4f, + 0.4f, 0.3f, 0.2f, 0.1f}); + test.AddInput("b_zero_point", + {2, 1, 4}, + {1, -2, 2, -1, + 2, -4, -1, 0}); + test.AddInput("y_scale", {}, {0.2f}); + test.AddInput("y_zero_point", {}, {130}); + + test.AddOutput("y", + {2, 2, 4}, + {128, 128, 148, 118, + 136, 144, 142, 121, + + 126, 122, 137, 128, + 157, 150, 136, 128}); + + test.Run(); +} + /** * @brief Extend QLinearMatMul for verifying prepacking behavior */