saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-07 17:15:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add per-column support for QLinearMatMul (#7729)

Browse source 
* Add per-column support for QLinearMatMul
This commit is contained in:
Yufeng Li 2021-05-17 18:57:42 -07:00 committed by GitHub
parent d3c4b70ede
commit 9075488368
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 172 additions and 60 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_matmul.cc
View file

@ -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<MLAS_GEMM_U8X8_DATA_PARAMS> 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<const uint8_t*>(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<int32_t*>(y_data + helper.OutputOffsets()[gemm_idx]);
params.ldc = gemm_shape.N;

40
onnxruntime/core/providers/cpu/math/matmul_helper.h
View file

@ -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<size_t>& 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<size_t> right_offsets_;
std::vector<size_t> output_offsets_;
std::vector<size_t> right_zp_offsets_;
std::vector<size_t> 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<size_t>& RightScaleOffsets() const {
return right_scale_offsets_;
}
// Batched Zero Point Offset for right matrices
const std::vector<size_t>& RightZeroPointOffsets() const {
return right_zp_offsets_;
}
template <typename T>
static void OffsetToArrays(T* p, const std::vector<size_t>& offsets, gsl::span<T*> arrays) {
auto len = offsets.size();

39
onnxruntime/core/providers/cpu/math/matmul_integer.cc
View file

@ -41,27 +41,9 @@ ONNX_OPERATOR_TYPED_KERNEL_EX(
MatMulInteger);
Status MatMulInteger::Compute(OpKernelContext* ctx) const {
MatMulComputeHelper helper;
const auto* a = ctx->Input<Tensor>(IN_A);
const auto* b = packed_b_ ? nullptr : ctx->Input<Tensor>(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<const uint8_t*>(b->DataRaw());
b_is_signed = b->IsDataType<int8_t>();
} else {
ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape_));
b_data = static_cast<const uint8_t*>(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<Tensor>(IN_A_ZERO_POINT);
@ -82,6 +64,24 @@ Status MatMulInteger::Compute(OpKernelContext* ctx) const {
b_offset_ptr = static_cast<const uint8_t*>(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<const uint8_t*>(b->DataRaw());
b_is_signed = b->IsDataType<int8_t>();
} else {
ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape_, nullptr, b_zero_point ? &b_zero_point->Shape() : nullptr));
b_data = static_cast<const uint8_t*>(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<uint8_t>();
auto* y_data = y->template MutableData<int32_t>();
@ -95,12 +95,11 @@ Status MatMulInteger::Compute(OpKernelContext* ctx) const {
std::vector<MLAS_GEMM_U8X8_DATA_PARAMS> 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_);

68
onnxruntime/core/providers/cpu/math/quantize_linear_matmul.cc
View file

@ -25,32 +25,8 @@ ONNX_OPERATOR_KERNEL_EX(
QLinearMatMul);
Status QLinearMatMul::Compute(OpKernelContext* ctx) const {
MatMulComputeHelper helper;
const auto* a = ctx->Input<Tensor>(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<const uint8_t*>(packed_b_.get());
b_is_signed = b_is_signed_;
} else {
const Tensor* b = ctx->Input<Tensor>(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<const uint8_t*>(b->DataRaw());
b_is_signed = b->IsDataType<int8_t>();
}
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<Tensor>(IN_B);
// validate offsets
const auto* a_offset = ctx->Input<Tensor>(IN_A_ZERO_POINT);
@ -58,8 +34,8 @@ Status QLinearMatMul::Compute(OpKernelContext* ctx) const {
const auto* y_offset = ctx->Input<Tensor>(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<Tensor>(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<const uint8_t*>(b->DataRaw());
b_is_signed = b->IsDataType<int8_t>();
} else {
ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape_, &b_scale->Shape(), &b_offset->Shape()));
b_data = static_cast<const uint8_t*>(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<float>();
auto a_scale_data = *(a_scale->template Data<float>());
auto b_scale_data = *(b_scale->template Data<float>());
auto y_scale_data = *(y_scale->template Data<float>());
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<float> 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<uint8_t>();
gemm_params.ldb = gemm_shape.N;
gemm_params.ZeroPointB = static_cast<const uint8_t*>(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<const uint8_t*>(b_offset->DataRaw());
for (size_t i = 0; i < helper.OutputOffsets().size(); i++) {
gemm_params.A = a->template Data<uint8_t>() + 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<size_t>(helper.M()),
static_cast<size_t>(helper.N()),
&real_multiplier,
false,
output_scales.data() + helper.RightScaleOffsets()[i],
output_scales.size() > 1,
*y_offset->template Data<uint8_t>());
}

73
onnxruntime/test/providers/cpu/math/quantize_linear_matmul_test.cc
View file

@ -194,6 +194,79 @@ TEST(QuantizeLinearMatmulOpTest, QLinearMatMulAllInputExceptT1AreInitializers) {
QLinearMatMul2DTest(true);
}
TEST(QuantizeLinearMatmulOpTest, PerColumn_2D) {
OpTester test("QLinearMatMul", 10);
test.AddInput<uint8_t>("a",
{2, 4},
{125, 135, 133, 122,
132, 123, 136, 135});
test.AddInput<float>("a_scale", {}, {0.1f});
test.AddInput<uint8_t>("a_zero_point", {}, {133});
test.AddInput<int8_t>("b",
{4, 4},
{0, -8, 2, 3,
-11, -13, -8, 1,
2, 4, 4, -10,
3, 2, -11, 2});
test.AddInput<float>("b_scale", {4},
{0.1f, 0.2f, 0.3f, 0.4f});
test.AddInput<int8_t>("b_zero_point",
{1, 4},
{1, -2, 2, -1});
test.AddInput<float>("y_scale", {}, {0.2f});
test.AddInput<uint8_t>("y_zero_point", {}, {130});
test.AddOutput<uint8_t>("y",
{2, 4},
{128, 128, 148, 118,
136, 144, 142, 121});
test.Run();
}
TEST(QuantizeLinearMatmulOpTest, PerColumn_ND) {
OpTester test("QLinearMatMul", 10);
test.AddInput<uint8_t>("a",
{2, 2, 4},
{125, 135, 133, 122,
132, 123, 136, 135,
125, 135, 133, 122,
132, 123, 136, 135});
test.AddInput<float>("a_scale", {}, {0.1f});
test.AddInput<uint8_t>("a_zero_point", {}, {133});
test.AddInput<int8_t>("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<float>("b_scale", {2, 1, 4},
{0.1f, 0.2f, 0.3f, 0.4f,
0.4f, 0.3f, 0.2f, 0.1f});
test.AddInput<int8_t>("b_zero_point",
{2, 1, 4},
{1, -2, 2, -1,
2, -4, -1, 0});
test.AddInput<float>("y_scale", {}, {0.2f});
test.AddInput<uint8_t>("y_zero_point", {}, {130});
test.AddOutput<uint8_t>("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
*/