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

[CPU EP] Refactor MatMulNBits to decouple type implementation (#22140)

Browse source 
### Description
Decouple implementation for different A types to improve readability and
maintainability.

### Motivation and Context
As more types are added, the implementation can differ a lot between
types. Besides, different hardware may require different
implementations.
This PR creates an abstraction boundary where different implemetation
can plug in easily.
This commit is contained in:
Jing Fang 2024-09-19 17:57:35 -07:00 committed by GitHub
parent c270fe6dd3
commit b0ef1f3923
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 431 additions and 281 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
View file

@ -32,7 +32,8 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, WordC
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GatherND);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, TransposeMatMul); // backward compatibility
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FusedMatMul);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulNBits);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, MatMulNBits);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MLFloat16, MatMulNBits);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulBnb4);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, UInt4x2, int32_t, GatherBlockQuantized);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, UInt4x2, int64_t, GatherBlockQuantized);
@ -301,7 +302,8 @@ Status RegisterCpuContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MurmurHash3)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, TransposeMatMul)>, // backward compatibility
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FusedMatMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulNBits)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, MatMulNBits)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MLFloat16, MatMulNBits)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulBnb4)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, UInt4x2, int32_t, GatherBlockQuantized)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, UInt4x2, int64_t, GatherBlockQuantized)>,

706
onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc
View file

@ -79,6 +79,9 @@ bool GetType(const NodeArg& node_arg, int32_t& type) {
return true;
}
// T1 is the type of the input matrix A, scales and biases.
// Use class level template to facilitate specialization for different types.
template <typename T1>
class MatMulNBits final : public OpKernel {
public:
MatMulNBits(const OpKernelInfo& info)
@ -89,10 +92,10 @@ class MatMulNBits final : public OpKernel {
nbits_{narrow<size_t>(info.GetAttr<int64_t>("bits"))},
accuracy_level_{GetAccuracyLevel(nbits_, block_size_, info.GetAttr<int64_t>("accuracy_level"))},
has_g_idx_{info.GetInputCount() > InputIndex::g_idx && info.node().InputDefs()[InputIndex::g_idx]->Exists()},
has_bias_{info.GetInputCount() > InputIndex::bias && info.node().InputDefs()[InputIndex::bias]->Exists()} {
has_bias_{info.GetInputCount() > InputIndex::bias && info.node().InputDefs()[InputIndex::bias]->Exists()},
compute_type_{static_cast<MLAS_SQNBIT_GEMM_COMPUTE_TYPE>(accuracy_level_)} {
const auto& node = info.node();
auto input_defs = node.InputDefs();
const NodeArg* zero_point_arg =
(info.GetInputCount() > InputIndex::zero_points && input_defs[InputIndex::zero_points]->Exists())
? input_defs[3]
@ -134,6 +137,7 @@ class MatMulNBits final : public OpKernel {
const int64_t accuracy_level_;
const bool has_g_idx_;
const bool has_bias_;
const MLAS_SQNBIT_GEMM_COMPUTE_TYPE compute_type_;
bool has_unquantized_zero_point_{false};
const bool column_wise_quant_{true};
IAllocatorUniquePtr<void> packed_b_{};
@ -147,17 +151,58 @@ class MatMulNBits final : public OpKernel {
#endif // defined(ORT_NEURAL_SPEED)
template <typename AType>
Status ComputeTyped(OpKernelContext* ctx) const;
// dequantize B first and then compute float gemm
Status ComputeBUnpacked(const Tensor* a,
const Tensor* b,
const Tensor* scales,
const Tensor* zero_points,
const Tensor* reorder_idx,
const Tensor* bias,
Tensor* y,
AllocatorPtr& allocator,
concurrency::ThreadPool* thread_pool,
const MatMulComputeHelper& helper) const {
ORT_THROW("ComputeBUnpacked is not supported for T1 type.");
}
Status ComputeBPacked(const Tensor* a,
const Tensor* scales,
const Tensor* zero_points,
const Tensor* bias,
Tensor* y,
AllocatorPtr& allocator,
concurrency::ThreadPool* thread_pool,
const MatMulComputeHelper& helper) const {
ORT_THROW("ComputeBPacked is not supported for T1 type.");
}
void PackScale(const Tensor& tensor) {
ORT_THROW("PackScale is not supported for T1 type.");
}
};
bool IsATypeFloat16(const Tensor& tensor) {
return tensor.GetElementType() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16;
#ifdef MLAS_TARGET_AMD64_IX86
template <>
void MatMulNBits<MLFloat16>::PackScale(const Tensor& tensor) {
auto sptr = tensor.Data<MLFloat16>();
std::vector<float> scales_v(static_cast<unsigned int>(tensor.Shape().Size()));
MlasConvertHalfToFloatBuffer(sptr, &scales_v[0], scales_v.size());
MlasSQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, compute_type_, nullptr, packed_b_.get(), &scales_v[0],
has_zp_input_, nullptr, nullptr);
}
Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ AllocatorPtr alloc,
/*out*/ bool& is_packed,
/*out*/ PrePackedWeights* prepacked_weights) {
template <>
void MatMulNBits<float>::PackScale(const Tensor& tensor) {
auto sptr = tensor.Data<float>();
MlasSQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, compute_type_, nullptr, packed_b_.get(), sptr,
has_zp_input_, nullptr, nullptr);
}
#endif
template <typename T1>
Status MatMulNBits<T1>::PrePack(const Tensor& tensor, int input_idx, /*out*/ AllocatorPtr alloc,
/*out*/ bool& is_packed,
/*out*/ PrePackedWeights* prepacked_weights) {
is_packed = false;
if (has_g_idx_ || has_unquantized_zero_point_) {
return Status::OK();
@ -178,16 +223,15 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ Allocat
MLAS_THREADPOOL* pool = nullptr;
auto comp_type = static_cast<NS_SQNBIT_COMPUTE_TYPE>(accuracy_level_);
auto nbits = static_cast<int>(nbits_);
if (input_idx == InputIndex::B) {
packed_b_size_ = NSNBitsGemmPackBSize(N_, K_, block_size_, nbits, is_asym_, comp_type);
packed_b_size_ = NSNBitsGemmPackBSize(N_, K_, block_size_, nbits, is_asym_, compute_type_);
if (packed_b_size_ == 0) return Status::OK();
auto qptr = tensor.Data<uint8_t>();
packed_b_ = IAllocator::MakeUniquePtr<void>(alloc, packed_b_size_, true);
std::memset(packed_b_.get(), 0, packed_b_size_);
NSNBitsGemmPackB(packed_b_.get(), qptr, nullptr, nullptr, N_, K_, K_, block_size_, nbits, is_asym_, false,
comp_type, pool);
compute_type_, pool);
if (prepacked_weights) {
prepacked_weights->buffers_.push_back(std::move(packed_b_));
prepacked_weights->buffer_sizes_.push_back(packed_b_size_);
@ -197,7 +241,7 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ Allocat
if (input_idx == InputIndex::scales && packed_b_ != nullptr) {
auto sptr = tensor.Data<float>();
NSNBitsGemmPackB(packed_b_.get(), nullptr, sptr, nullptr, N_, K_, K_, block_size_, nbits, is_asym_, !is_asym_,
comp_type, pool);
compute_type_, pool);
if (prepacked_weights) {
prepacked_weights->buffers_.push_back(std::move(packed_b_));
prepacked_weights->buffer_sizes_.push_back(packed_b_size_);
@ -207,7 +251,7 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ Allocat
if (input_idx == InputIndex::zero_points && packed_b_ != nullptr) {
auto zptr = tensor.Data<uint8_t>();
NSNBitsGemmPackB(packed_b_.get(), nullptr, nullptr, zptr, N_, K_, K_, block_size_, nbits, is_asym_, is_asym_,
comp_type, pool);
compute_type_, pool);
if (prepacked_weights) {
prepacked_weights->buffers_.push_back(std::move(packed_b_));
prepacked_weights->buffer_sizes_.push_back(packed_b_size_);
@ -217,35 +261,26 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ Allocat
#else // defined(ORT_NEURAL_SPEED)
ORT_UNUSED_PARAMETER(prepacked_weights);
const auto compute_type = static_cast<MLAS_SQNBIT_GEMM_COMPUTE_TYPE>(accuracy_level_);
if (!MlasIsSQNBitGemmAvailable(nbits_, block_size_, compute_type)) {
if (!MlasIsSQNBitGemmAvailable(nbits_, block_size_, compute_type_)) {
return Status::OK();
}
if (input_idx == InputIndex::B) {
packed_b_size_ = MlasSQNBitGemmPackQuantBDataSize(N_, K_, nbits_, block_size_, compute_type);
packed_b_size_ = MlasSQNBitGemmPackQuantBDataSize(N_, K_, nbits_, block_size_, compute_type_);
if (packed_b_size_ == 0) {
return Status::OK();
}
auto qptr = tensor.DataRaw();
packed_b_ = IAllocator::MakeUniquePtr<void>(alloc, packed_b_size_, true);
MlasSQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, compute_type, qptr, packed_b_.get(), nullptr, has_zp_input_, nullptr, nullptr);
MlasSQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, compute_type_, qptr, packed_b_.get(), nullptr, has_zp_input_, nullptr, nullptr);
is_packed = true;
} else if (compute_type == CompInt8) {
} else if (compute_type_ == CompInt8) {
#ifdef MLAS_TARGET_AMD64_IX86
if (input_idx == InputIndex::scales && packed_b_ != nullptr) {
if (IsATypeFloat16(tensor)) {
auto sptr = tensor.Data<MLFloat16>();
std::vector<float> scales_v(static_cast<unsigned int>(tensor.Shape().Size()));
MlasConvertHalfToFloatBuffer(sptr, &scales_v[0], scales_v.size());
MlasSQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, compute_type, nullptr, packed_b_.get(), &scales_v[0], has_zp_input_, nullptr, nullptr);
} else {
auto sptr = tensor.Data<float>();
MlasSQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, compute_type, nullptr, packed_b_.get(), sptr, has_zp_input_, nullptr, nullptr);
}
PackScale(tensor);
is_packed = false;
} else if (input_idx == InputIndex::zero_points && packed_b_ != nullptr) {
auto zptr = tensor.Data<uint8_t>();
MlasSQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, compute_type, nullptr, packed_b_.get(), nullptr, has_zp_input_, zptr, nullptr);
MlasSQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, compute_type_, nullptr, packed_b_.get(), nullptr, has_zp_input_, zptr, nullptr);
is_packed = false;
}
#endif
@ -255,8 +290,9 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ Allocat
return Status::OK();
}
Status MatMulNBits::UseSharedPrePackedBuffers(std::vector<BufferUniquePtr>& prepacked_buffers, int input_idx,
/*out*/ bool& used_shared_buffers) {
template <typename T1>
Status MatMulNBits<T1>::UseSharedPrePackedBuffers(std::vector<BufferUniquePtr>& prepacked_buffers, int input_idx,
/*out*/ bool& used_shared_buffers) {
used_shared_buffers = false;
#if defined(ORT_NEURAL_SPEED)
@ -287,34 +323,20 @@ Status MatMulNBits::UseSharedPrePackedBuffers(std::vector<BufferUniquePtr>& prep
return Status::OK();
}
Status MatMulNBits::Compute(OpKernelContext* ctx) const {
const Tensor* a = ctx->Input<Tensor>(InputIndex::A);
if (IsATypeFloat16(*a)) {
return ComputeTyped<MLFloat16>(ctx);
} else {
return ComputeTyped<float>(ctx);
}
}
template <typename AType>
Status MatMulNBits::ComputeTyped(OpKernelContext* ctx) const {
concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool();
const Tensor* a = ctx->Input<Tensor>(InputIndex::A);
const auto* a_data = a->Data<AType>();
TensorShape b_shape({static_cast<int64_t>(N_), static_cast<int64_t>(K_)});
MatMulComputeHelper helper;
ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape, false, true));
Tensor* y = ctx->Output(0, helper.OutputShape());
// Bail out early if the output is going to be empty
if (y->Shape().Size() == 0) {
return Status::OK();
}
auto* y_data = y->MutableData<AType>();
template <>
Status MatMulNBits<float>::ComputeBPacked(const Tensor* a,
const Tensor* scales,
const Tensor* zero_points,
const Tensor* bias,
Tensor* y,
AllocatorPtr& allocator,
concurrency::ThreadPool* thread_pool,
const MatMulComputeHelper& helper) const {
const auto* a_data = a->Data<float>();
const auto* scales_data = scales->Data<float>();
const auto* zero_points_data = zero_points == nullptr ? nullptr : zero_points->DataRaw();
const auto* bias_data = bias == nullptr ? nullptr : bias->Data<float>();
auto* y_data = y->MutableData<float>();
const size_t batch_count = helper.OutputOffsets().size();
const size_t M = static_cast<size_t>(helper.M());
@ -322,152 +344,239 @@ Status MatMulNBits::ComputeTyped(OpKernelContext* ctx) const {
const size_t K = static_cast<size_t>(helper.K());
const size_t lda = helper.Lda(false);
// clang-format off
const bool has_single_b_matrix = std::all_of(
helper.RightOffsets().begin(),
helper.RightOffsets().end(),
[](size_t offset) { return offset == 0; });
// clang-format on
#if defined(ORT_NEURAL_SPEED)
if (has_single_b_matrix &&
packed_b_) {
InlinedVector<NS_SQNBITS_GEMM_DATA_PACKED_PARAMS> gemm_params(batch_count);
AllocatorPtr allocator;
ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&allocator));
for (size_t i = 0; i < batch_count; i++) {
gemm_params[i].A = a_data + helper.LeftOffsets()[i];
gemm_params[i].lda = lda;
gemm_params[i].B = packed_b_.get();
gemm_params[i].C = y_data + helper.OutputOffsets()[i];
gemm_params[i].ldc = N;
}
auto ws_size = NSSQNBitsGemmBatchWorkspaceSize(M, N, K, batch_count, gemm_params.data());
// workspace for activation process(dynamic quantization and others)
auto ws_ptr = IAllocator::MakeUniquePtr<int8_t>(allocator, ws_size);
NSSQNBitsGemmBatchPackedB(M, N, K, batch_count, gemm_params.data(), ws_ptr.get(), thread_pool);
return Status::OK();
IAllocatorUniquePtr<std::byte> workspace{};
const size_t workspace_size = MlasSQNBitGemmBatchWorkspaceSize(
M, N, K, batch_count, nbits_, block_size_, compute_type_);
if (workspace_size > 0) {
workspace = IAllocator::MakeUniquePtr<std::byte>(allocator, workspace_size);
}
#else // defined(ORT_NEURAL_SPEED)
if (has_single_b_matrix &&
packed_b_) { // Assume that MlasSQNBitGemmBatch() always requires packed B.
// If this changes, i.e., if MlasIsSQNBitGemmAvailable() can return true while
// MlasSQNBitGemmPackQuantBDataSize() returns 0, we can consider calling MlasSQNBitGemmBatch()
// with B directly too.
const auto compute_type = static_cast<MLAS_SQNBIT_GEMM_COMPUTE_TYPE>(accuracy_level_);
if (MlasIsSQNBitGemmAvailable(nbits_, block_size_, compute_type)) {
const Tensor* scales = ctx->Input<Tensor>(InputIndex::scales);
const Tensor* zero_points = ctx->Input<Tensor>(InputIndex::zero_points);
const Tensor* bias = ctx->Input<Tensor>(InputIndex::bias);
const auto* scales_data = scales->Data<AType>();
const auto* zero_points_data = zero_points == nullptr ? nullptr : zero_points->DataRaw();
const auto* bias_data = bias == nullptr ? nullptr : bias->Data<AType>();
IAllocatorUniquePtr<std::byte> workspace{};
const size_t workspace_size = MlasSQNBitGemmBatchWorkspaceSize(
M, N, K, batch_count, nbits_, block_size_, compute_type);
if (workspace_size > 0) {
AllocatorPtr allocator;
ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&allocator));
workspace = IAllocator::MakeUniquePtr<std::byte>(allocator, workspace_size);
}
if constexpr (std::is_same<AType, MLFloat16>::value) {
InlinedVector<MLAS_SQNBIT_GEMM_DATA_PARAMS> data(batch_count);
AllocatorPtr allocator;
ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&allocator));
auto tmp_a_data_ptr = IAllocator::MakeUniquePtr<float>(allocator, (size_t)(a->Shape().Size()));
MlasConvertHalfToFloatBuffer(a_data, tmp_a_data_ptr.get(), static_cast<size_t>(a->Shape().Size()));
auto tmp_scales_data_ptr = IAllocator::MakeUniquePtr<float>(allocator, (size_t)(scales->Shape().Size()));
MlasConvertHalfToFloatBuffer(scales_data, tmp_scales_data_ptr.get(), static_cast<size_t>(scales->Shape().Size()));
std::vector<float> bias_data_v;
if (bias_data != nullptr) {
bias_data_v.resize((const unsigned int)(bias->Shape().Size()));
MlasConvertHalfToFloatBuffer(bias_data, &bias_data_v[0], bias_data_v.size());
}
std::vector<float> C_v((const unsigned int)(y->Shape().Size()));
for (size_t i = 0; i < batch_count; ++i) {
data[i].A = tmp_a_data_ptr.get() + helper.LeftOffsets()[i];
data[i].lda = lda;
InlinedVector<MLAS_SQNBIT_GEMM_DATA_PARAMS> data(batch_count);
for (size_t i = 0; i < batch_count; ++i) {
data[i].A = a_data + helper.LeftOffsets()[i];
data[i].lda = lda;
#ifdef MLAS_TARGET_AMD64_IX86
if (compute_type == CompInt8) {
data[i].QuantBDataWorkspace = packed_b_.get();
}
#endif
data[i].PackedQuantBData = static_cast<std::byte*>(packed_b_.get());
data[i].QuantBScale = tmp_scales_data_ptr.get();
data[i].QuantBZeroPoint = zero_points_data;
data[i].Bias = bias_data != nullptr ? &bias_data_v[0] : nullptr;
data[i].C = &C_v[0] + helper.OutputOffsets()[i];
data[i].ldc = N;
}
MlasSQNBitGemmBatch(M, N, K, batch_count, nbits_, block_size_, compute_type, data.data(), workspace.get(),
thread_pool);
MlasConvertFloatToHalfBuffer(&C_v[0], y_data, C_v.size());
return Status::OK();
} else {
InlinedVector<MLAS_SQNBIT_GEMM_DATA_PARAMS> data(batch_count);
for (size_t i = 0; i < batch_count; ++i) {
data[i].A = a_data + helper.LeftOffsets()[i];
data[i].lda = lda;
#ifdef MLAS_TARGET_AMD64_IX86
if (compute_type == CompInt8) {
data[i].QuantBDataWorkspace = packed_b_.get();
}
#endif
data[i].PackedQuantBData = static_cast<std::byte*>(packed_b_.get());
data[i].QuantBScale = scales_data;
data[i].QuantBZeroPoint = zero_points_data;
data[i].Bias = bias_data;
data[i].C = y_data + helper.OutputOffsets()[i];
data[i].ldc = N;
}
MlasSQNBitGemmBatch(M, N, K, batch_count, nbits_, block_size_, compute_type, data.data(), workspace.get(),
thread_pool);
return Status::OK();
}
if (compute_type_ == CompInt8) {
data[i].QuantBDataWorkspace = packed_b_.get();
}
#endif
data[i].PackedQuantBData = static_cast<std::byte*>(packed_b_.get());
data[i].QuantBScale = scales_data;
data[i].QuantBZeroPoint = zero_points_data;
data[i].Bias = bias_data;
data[i].C = y_data + helper.OutputOffsets()[i];
data[i].ldc = N;
}
MlasSQNBitGemmBatch(M, N, K, batch_count, nbits_, block_size_, compute_type_, data.data(), workspace.get(),
thread_pool);
return Status::OK();
}
template <>
Status MatMulNBits<MLFloat16>::ComputeBPacked(const Tensor* a,
const Tensor* scales,
const Tensor* zero_points,
const Tensor* bias,
Tensor* y,
AllocatorPtr& allocator,
concurrency::ThreadPool* thread_pool,
const MatMulComputeHelper& helper) const {
const auto* a_data = a->Data<MLFloat16>();
const auto* scales_data = scales->Data<MLFloat16>();
const auto* zero_points_data = zero_points == nullptr ? nullptr : zero_points->DataRaw();
const auto* bias_data = bias == nullptr ? nullptr : bias->Data<MLFloat16>();
auto* y_data = y->MutableData<MLFloat16>();
const size_t batch_count = helper.OutputOffsets().size();
const size_t M = static_cast<size_t>(helper.M());
const size_t N = static_cast<size_t>(helper.N());
const size_t K = static_cast<size_t>(helper.K());
const size_t lda = helper.Lda(false);
IAllocatorUniquePtr<std::byte> workspace{};
const size_t workspace_size = MlasSQNBitGemmBatchWorkspaceSize(
M, N, K, batch_count, nbits_, block_size_, compute_type_);
if (workspace_size > 0) {
workspace = IAllocator::MakeUniquePtr<std::byte>(allocator, workspace_size);
}
#endif // !defined(ORT_NEURAL_SPEED)
auto tmp_a_data_ptr = IAllocator::MakeUniquePtr<float>(allocator, (size_t)(a->Shape().Size()));
MlasConvertHalfToFloatBuffer(a_data, tmp_a_data_ptr.get(), static_cast<size_t>(a->Shape().Size()));
// fallback implementation - dequantize B first and then compute float gemm
auto tmp_scales_data_ptr = IAllocator::MakeUniquePtr<float>(allocator, (size_t)(scales->Shape().Size()));
MlasConvertHalfToFloatBuffer(scales_data, tmp_scales_data_ptr.get(), static_cast<size_t>(scales->Shape().Size()));
const Tensor* scales = ctx->Input<Tensor>(InputIndex::scales);
const Tensor* zero_points = ctx->Input<Tensor>(InputIndex::zero_points);
const Tensor* reorder_idx = ctx->Input<Tensor>(InputIndex::g_idx);
const auto* scales_data = scales->Data<AType>();
const float* scales_data_;
std::vector<float> scales_data_v;
if constexpr (std::is_same<AType, MLFloat16>::value) {
scales_data_v.resize((const unsigned int)scales->Shape().Size());
MlasConvertHalfToFloatBuffer(scales_data, &scales_data_v[0], scales_data_v.size());
scales_data_ = &scales_data_v[0];
} else {
scales_data_ = scales_data;
std::vector<float> bias_data_v;
if (bias_data != nullptr) {
bias_data_v.resize(static_cast<size_t>(bias->Shape().Size()));
MlasConvertHalfToFloatBuffer(bias_data, &bias_data_v[0], bias_data_v.size());
}
std::vector<float> C_v(static_cast<size_t>(y->Shape().Size()));
InlinedVector<MLAS_SQNBIT_GEMM_DATA_PARAMS> data(batch_count);
for (size_t i = 0; i < batch_count; ++i) {
data[i].A = tmp_a_data_ptr.get() + helper.LeftOffsets()[i];
data[i].lda = lda;
#ifdef MLAS_TARGET_AMD64_IX86
if (compute_type_ == CompInt8) {
data[i].QuantBDataWorkspace = packed_b_.get();
}
#endif
data[i].PackedQuantBData = static_cast<std::byte*>(packed_b_.get());
data[i].QuantBScale = tmp_scales_data_ptr.get();
data[i].QuantBZeroPoint = zero_points_data;
data[i].Bias = bias_data != nullptr ? &bias_data_v[0] : nullptr;
data[i].C = &C_v[0] + helper.OutputOffsets()[i];
data[i].ldc = N;
}
MlasSQNBitGemmBatch(M, N, K, batch_count, nbits_, block_size_, compute_type_, data.data(), workspace.get(),
thread_pool);
MlasConvertFloatToHalfBuffer(&C_v[0], y_data, C_v.size());
return Status::OK();
}
template <>
Status MatMulNBits<float>::ComputeBUnpacked(const Tensor* a,
const Tensor* b,
const Tensor* scales,
const Tensor* zero_points,
const Tensor* reorder_idx,
const Tensor* bias,
Tensor* y,
AllocatorPtr& allocator,
concurrency::ThreadPool* thread_pool,
const MatMulComputeHelper& helper) const {
const auto* a_data = a->Data<float>();
const uint8_t* b_data = b->Data<uint8_t>();
const auto* scales_data = scales->Data<float>();
const auto* zero_points_data = zero_points == nullptr ? nullptr : zero_points->DataRaw();
const auto* reorder_idx_data = reorder_idx == nullptr ? nullptr : reorder_idx->Data<int32_t>();
auto* y_data = y->MutableData<float>();
const Tensor* b = ctx->Input<Tensor>(InputIndex::B);
const uint8_t* b_data = b->Data<uint8_t>();
const size_t batch_count = helper.OutputOffsets().size();
const size_t M = static_cast<size_t>(helper.M());
const size_t N = static_cast<size_t>(helper.N());
const size_t K = static_cast<size_t>(helper.K());
const size_t lda = helper.Lda(false);
const size_t ldb = helper.Ldb(true);
AllocatorPtr allocator;
ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&allocator));
auto tmp_b_data_ptr = IAllocator::MakeUniquePtr<float>(allocator, SafeInt<size_t>(K_) * N_);
if ((reorder_idx_data == nullptr) && (!zero_points || !zero_points->IsDataType<AType>())) {
if ((reorder_idx_data == nullptr) && (!zero_points || !zero_points->IsDataType<float>())) {
// dequantize b, only 4b quantization is supported for now
MlasDequantizeBlockwise<float, 4>(
tmp_b_data_ptr.get(), // dequantized output
b_data, // quantized input
scales_data, // quantization scales
static_cast<const uint8_t*>(zero_points_data), // quantization zero points
static_cast<int32_t>(block_size_), // quantization block size
column_wise_quant_, // columnwise quantization or row-wise
static_cast<int32_t>(K_), // number of rows in quantized input
static_cast<int32_t>(N_), // number of columns in quantized input
thread_pool);
} else {
ORT_ENFORCE(column_wise_quant_, "Row-wise quantization is not supported for now");
// !!!!!!!!!!!!!! naive implementation, need to be optimized !!!!!!!!!!!!!!
if (zero_points && zero_points->IsDataType<float>()) {
DequantizeBlockwise<float, float>(
tmp_b_data_ptr.get(), // dequantized output
b_data, // quantized input
scales_data, // quantization scales
static_cast<const float*>(zero_points_data), // quantization zero points
reorder_idx_data,
static_cast<int32_t>(block_size_), // quantization block size
column_wise_quant_, // columnwise quantization or row-wise
static_cast<int32_t>(K_), // number of rows in quantized input
static_cast<int32_t>(N_), // number of columns in quantized input
thread_pool);
} else {
DequantizeBlockwise<float, uint8_t>(
tmp_b_data_ptr.get(), // dequantized output
b_data, // quantized input
scales_data, // quantization scales
static_cast<const uint8_t*>(zero_points_data), // quantization zero points
reorder_idx_data,
static_cast<int32_t>(block_size_), // quantization block size
column_wise_quant_, // columnwise quantization or row-wise
static_cast<int32_t>(K_), // number of rows in quantized input
static_cast<int32_t>(N_), // number of columns in quantized input
thread_pool);
}
}
#if 0 // for debug
auto tm_b_data_ptr_trans = IAllocator::MakeUniquePtr<float>(allocator, SafeInt<size_t>(K_) * N_);
MlasTranspose(tmp_b_data_ptr.get(), tm_b_data_ptr_trans.get(), N_, K_);
#endif
std::vector<MLAS_SGEMM_DATA_PARAMS> data(batch_count);
for (size_t i = 0; i < batch_count; i++) {
data[i].BIsPacked = false;
data[i].A = a_data + helper.LeftOffsets()[i];
data[i].lda = lda;
data[i].B = tmp_b_data_ptr.get() + helper.RightOffsets()[i];
data[i].ldb = ldb;
data[i].C = y_data + helper.OutputOffsets()[i];
data[i].ldc = N;
data[i].alpha = 1.f;
data[i].beta = 0.0f;
}
// if there is a bias input, copy bias values into C and set beta to 1.0f
if (bias) {
gsl::span<const float> bias_span = bias->DataAsSpan<float>();
for (size_t i = 0; i < batch_count; ++i) {
float* C_row = data[i].C;
const size_t ldc = data[i].ldc;
for (size_t m = 0; m < M; ++m) {
memcpy(C_row, bias_span.data(), bias_span.size_bytes());
C_row += ldc;
}
data[i].beta = 1.0f;
}
}
MlasGemmBatch(CblasNoTrans, CblasTrans,
M, N, K, data.data(), batch_count, thread_pool);
return Status::OK();
}
template <>
Status MatMulNBits<MLFloat16>::ComputeBUnpacked(const Tensor* a,
const Tensor* b,
const Tensor* scales,
const Tensor* zero_points,
const Tensor* reorder_idx,
const Tensor* bias,
Tensor* y,
AllocatorPtr& allocator,
concurrency::ThreadPool* thread_pool,
const MatMulComputeHelper& helper) const {
const auto* a_data = a->Data<MLFloat16>();
const uint8_t* b_data = b->Data<uint8_t>();
const auto* scales_data = scales->Data<MLFloat16>();
const auto* zero_points_data = zero_points == nullptr ? nullptr : zero_points->DataRaw();
const auto* reorder_idx_data = reorder_idx == nullptr ? nullptr : reorder_idx->Data<int32_t>();
auto* y_data = y->MutableData<MLFloat16>();
const float* scales_data_;
std::vector<float> scales_data_v;
scales_data_v.resize(static_cast<size_t>(scales->Shape().Size()));
MlasConvertHalfToFloatBuffer(scales_data, &scales_data_v[0], scales_data_v.size());
scales_data_ = &scales_data_v[0];
const size_t batch_count = helper.OutputOffsets().size();
const size_t M = static_cast<size_t>(helper.M());
const size_t N = static_cast<size_t>(helper.N());
const size_t K = static_cast<size_t>(helper.K());
const size_t lda = helper.Lda(false);
const size_t ldb = helper.Ldb(true);
auto tmp_b_data_ptr = IAllocator::MakeUniquePtr<float>(allocator, SafeInt<size_t>(K_) * N_);
if ((reorder_idx_data == nullptr) && (!zero_points || !zero_points->IsDataType<MLFloat16>())) {
// dequantize b, only 4b quantization is supported for now
MlasDequantizeBlockwise<float, 4>(
tmp_b_data_ptr.get(), // dequantized output
@ -482,12 +591,12 @@ Status MatMulNBits::ComputeTyped(OpKernelContext* ctx) const {
} else {
ORT_ENFORCE(column_wise_quant_, "Row-wise quantization is not supported for now");
// !!!!!!!!!!!!!! naive implementation, need to be optimized !!!!!!!!!!!!!!
if ((zero_points && zero_points->IsDataType<AType>())) {
DequantizeBlockwise<float, AType>(
tmp_b_data_ptr.get(), // dequantized output
b_data, // quantized input
scales_data_, // quantization scales
static_cast<const AType*>(zero_points_data), // quantization zero points
if (zero_points && zero_points->IsDataType<MLFloat16>()) {
DequantizeBlockwise<float, MLFloat16>(
tmp_b_data_ptr.get(), // dequantized output
b_data, // quantized input
scales_data_, // quantization scales
static_cast<const MLFloat16*>(zero_points_data), // quantization zero points
reorder_idx_data,
static_cast<int32_t>(block_size_), // quantization block size
column_wise_quant_, // columnwise quantization or row-wise
@ -512,93 +621,132 @@ Status MatMulNBits::ComputeTyped(OpKernelContext* ctx) const {
auto tm_b_data_ptr_trans = IAllocator::MakeUniquePtr<float>(allocator, SafeInt<size_t>(K_) * N_);
MlasTranspose(tmp_b_data_ptr.get(), tm_b_data_ptr_trans.get(), N_, K_);
#endif
if constexpr (std::is_same<AType, MLFloat16>::value) {
std::vector<MLAS_SGEMM_DATA_PARAMS> data(batch_count);
auto tmp_a_data_ptr = IAllocator::MakeUniquePtr<float>(allocator, (size_t)(a->Shape().Size()));
MlasConvertHalfToFloatBuffer(a_data, tmp_a_data_ptr.get(), static_cast<size_t>(a->Shape().Size()));
auto tmp_c_ptr = IAllocator::MakeUniquePtr<float>(allocator, (size_t)(y->Shape().Size()));
for (size_t i = 0; i < batch_count; i++) {
data[i].BIsPacked = false;
data[i].A = tmp_a_data_ptr.get() + helper.LeftOffsets()[i];
data[i].lda = lda;
data[i].B = tmp_b_data_ptr.get() + helper.RightOffsets()[i];
data[i].ldb = ldb;
data[i].C = tmp_c_ptr.get() + helper.OutputOffsets()[i];
data[i].ldc = N;
data[i].alpha = 1.f;
data[i].beta = 0.0f;
}
// if there is a bias input, copy bias values into C and set beta to 1.0f
if (const Tensor* bias = ctx->Input<Tensor>(InputIndex::bias);
bias != nullptr) {
auto tmp_bias_data_ptr = IAllocator::MakeUniquePtr<float>(allocator, (size_t)(bias->Shape().Size()));
MlasConvertHalfToFloatBuffer(bias->Data<AType>(), tmp_bias_data_ptr.get(), static_cast<size_t>(bias->Shape().Size()));
for (size_t i = 0; i < batch_count; ++i) {
float* C_row = data[i].C;
const size_t ldc = data[i].ldc;
for (size_t m = 0; m < M; ++m) {
std::copy(tmp_bias_data_ptr.get(), tmp_bias_data_ptr.get() + bias->Shape().Size(), C_row);
C_row += ldc;
}
data[i].beta = 1.0f;
}
}
MlasGemmBatch(CblasNoTrans, CblasTrans,
M, N, K, data.data(), batch_count, thread_pool);
MlasConvertFloatToHalfBuffer(tmp_c_ptr.get(), y_data, static_cast<size_t>(y->Shape().Size()));
return Status::OK();
} else {
std::vector<MLAS_SGEMM_DATA_PARAMS> data(batch_count);
for (size_t i = 0; i < batch_count; i++) {
data[i].BIsPacked = false;
data[i].A = a_data + helper.LeftOffsets()[i];
data[i].lda = lda;
data[i].B = tmp_b_data_ptr.get() + helper.RightOffsets()[i];
data[i].ldb = ldb;
data[i].C = y_data + helper.OutputOffsets()[i];
data[i].ldc = N;
data[i].alpha = 1.f;
data[i].beta = 0.0f;
}
// if there is a bias input, copy bias values into C and set beta to 1.0f
if (const Tensor* bias = ctx->Input<Tensor>(InputIndex::bias);
bias != nullptr) {
gsl::span<const float> bias_span = bias->DataAsSpan<float>();
for (size_t i = 0; i < batch_count; ++i) {
float* C_row = data[i].C;
const size_t ldc = data[i].ldc;
for (size_t m = 0; m < M; ++m) {
memcpy(C_row, bias_span.data(), bias_span.size_bytes());
C_row += ldc;
}
data[i].beta = 1.0f;
}
}
MlasGemmBatch(CblasNoTrans, CblasTrans,
M, N, K, data.data(), batch_count, thread_pool);
return Status::OK();
std::vector<MLAS_SGEMM_DATA_PARAMS> data(batch_count);
auto tmp_a_data_ptr = IAllocator::MakeUniquePtr<float>(allocator, (size_t)(a->Shape().Size()));
MlasConvertHalfToFloatBuffer(a_data, tmp_a_data_ptr.get(), static_cast<size_t>(a->Shape().Size()));
auto tmp_c_ptr = IAllocator::MakeUniquePtr<float>(allocator, (size_t)(y->Shape().Size()));
for (size_t i = 0; i < batch_count; i++) {
data[i].BIsPacked = false;
data[i].A = tmp_a_data_ptr.get() + helper.LeftOffsets()[i];
data[i].lda = lda;
data[i].B = tmp_b_data_ptr.get() + helper.RightOffsets()[i];
data[i].ldb = ldb;
data[i].C = tmp_c_ptr.get() + helper.OutputOffsets()[i];
data[i].ldc = N;
data[i].alpha = 1.f;
data[i].beta = 0.0f;
}
// if there is a bias input, copy bias values into C and set beta to 1.0f
if (bias) {
auto tmp_bias_data_ptr = IAllocator::MakeUniquePtr<float>(allocator, (size_t)(bias->Shape().Size()));
MlasConvertHalfToFloatBuffer(bias->Data<MLFloat16>(),
tmp_bias_data_ptr.get(),
static_cast<size_t>(bias->Shape().Size()));
for (size_t i = 0; i < batch_count; ++i) {
float* C_row = data[i].C;
const size_t ldc = data[i].ldc;
for (size_t m = 0; m < M; ++m) {
std::copy(tmp_bias_data_ptr.get(), tmp_bias_data_ptr.get() + bias->Shape().Size(), C_row);
C_row += ldc;
}
data[i].beta = 1.0f;
}
}
MlasGemmBatch(CblasNoTrans, CblasTrans,
M, N, K, data.data(), batch_count, thread_pool);
MlasConvertFloatToHalfBuffer(tmp_c_ptr.get(), y_data, static_cast<size_t>(y->Shape().Size()));
return Status::OK();
}
ONNX_OPERATOR_KERNEL_EX(
MatMulNBits,
kMSDomain,
1,
kCpuExecutionProvider,
KernelDefBuilder()
.TypeConstraint("T1", {DataTypeImpl::GetTensorType<float>(), DataTypeImpl::GetTensorType<MLFloat16>()})
.TypeConstraint("T2", DataTypeImpl::GetTensorType<uint8_t>())
.TypeConstraint("T3", {DataTypeImpl::GetTensorType<uint8_t>(), DataTypeImpl::GetTensorType<float>(), DataTypeImpl::GetTensorType<MLFloat16>()})
.TypeConstraint("T4", DataTypeImpl::GetTensorType<int32_t>()),
MatMulNBits);
template <typename T1>
Status MatMulNBits<T1>::Compute(OpKernelContext* ctx) const {
concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool();
const Tensor* a = ctx->Input<Tensor>(InputIndex::A);
const Tensor* scales = ctx->Input<Tensor>(InputIndex::scales);
const Tensor* zero_points = ctx->Input<Tensor>(InputIndex::zero_points);
const Tensor* reorder_idx = ctx->Input<Tensor>(InputIndex::g_idx);
const Tensor* bias = ctx->Input<Tensor>(InputIndex::bias);
TensorShape b_shape({static_cast<int64_t>(N_), static_cast<int64_t>(K_)});
MatMulComputeHelper helper;
ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape, false, true));
Tensor* y = ctx->Output(0, helper.OutputShape());
// Bail out early if the output is going to be empty
if (y->Shape().Size() == 0) {
return Status::OK();
}
AllocatorPtr allocator;
ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&allocator));
// clang-format off
const bool has_single_b_matrix = std::all_of(
helper.RightOffsets().begin(),
helper.RightOffsets().end(),
[](size_t offset) { return offset == 0; });
// clang-format on
if (has_single_b_matrix &&
packed_b_) { // Assume that MlasSQNBitGemmBatch() always requires packed B.
// If this changes, i.e., if MlasIsSQNBitGemmAvailable() can return true while
// MlasSQNBitGemmPackQuantBDataSize() returns 0, we can consider calling MlasSQNBitGemmBatch()
// with B directly too.
#if defined(ORT_NEURAL_SPEED)
const auto* a_data = a->Data<T1>();
auto* y_data = y->MutableData<T1>();
const size_t batch_count = helper.OutputOffsets().size();
const size_t M = static_cast<size_t>(helper.M());
const size_t N = static_cast<size_t>(helper.N());
const size_t K = static_cast<size_t>(helper.K());
const size_t lda = helper.Lda(false);
InlinedVector<NS_SQNBITS_GEMM_DATA_PACKED_PARAMS> gemm_params(batch_count);
for (size_t i = 0; i < batch_count; i++) {
gemm_params[i].A = a_data + helper.LeftOffsets()[i];
gemm_params[i].lda = lda;
gemm_params[i].B = packed_b_.get();
gemm_params[i].C = y_data + helper.OutputOffsets()[i];
gemm_params[i].ldc = N;
}
auto ws_size = NSSQNBitsGemmBatchWorkspaceSize(M, N, K, batch_count, gemm_params.data());
// workspace for activation process(dynamic quantization and others)
auto ws_ptr = IAllocator::MakeUniquePtr<int8_t>(allocator, ws_size);
NSSQNBitsGemmBatchPackedB(M, N, K, batch_count, gemm_params.data(), ws_ptr.get(), thread_pool);
return Status::OK();
#else // defined(ORT_NEURAL_SPEED)
if (MlasIsSQNBitGemmAvailable(nbits_, block_size_, compute_type_)) {
return ComputeBPacked(a, scales, zero_points, bias, y, allocator, thread_pool, helper);
}
#endif // !defined(ORT_NEURAL_SPEED)
}
// If B is prepacked, B would have been removed from the context
const Tensor* b = ctx->Input<Tensor>(InputIndex::B);
return ComputeBUnpacked(a, b, scales, zero_points, reorder_idx, bias, y, allocator, thread_pool, helper);
}
#define REGISTER_MatMulNBits(T1) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
MatMulNBits, \
kMSDomain, \
1, \
T1, \
kCpuExecutionProvider, \
KernelDefBuilder() \
.TypeConstraint("T1", DataTypeImpl::GetTensorType<T1>()) \
.TypeConstraint("T2", DataTypeImpl::GetTensorType<uint8_t>()) \
.TypeConstraint("T3", {DataTypeImpl::GetTensorType<uint8_t>(), \
DataTypeImpl::GetTensorType<float>(), \
DataTypeImpl::GetTensorType<MLFloat16>()}) \
.TypeConstraint("T4", DataTypeImpl::GetTensorType<int32_t>()), \
MatMulNBits<T1>);
REGISTER_MatMulNBits(float);
REGISTER_MatMulNBits(MLFloat16);
} // namespace contrib
} // namespace onnxruntime