From f30581ed2c61c716ffe1f3108c92950e54c25f2e Mon Sep 17 00:00:00 2001
From: Jing Fang <126209182+fajin-corp@users.noreply.github.com>
Date: Fri, 9 Aug 2024 12:15:11 -0700
Subject: [PATCH] [CPU EP] Add block quantized Gather contrib op (#21630)
### Description
Add a gather that supports block-quantized input data.
### Motivation and Context
To support Web inference scenario with quantized vocabulary embeddings.
---
docs/ContribOperators.md | 59 +++
docs/OperatorKernels.md | 1 +
.../contrib_ops/cpu/cpu_contrib_kernels.cc | 9 +
.../quantization/gather_block_quantized.cc | 282 +++++++++++
.../core/graph/contrib_ops/contrib_defs.cc | 118 +++++
.../gather_block_quantized_op_test.cc | 468 ++++++++++++++++++
6 files changed, 937 insertions(+)
create mode 100644 onnxruntime/contrib_ops/cpu/quantization/gather_block_quantized.cc
create mode 100644 onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc
diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index ed9e2a0567..c60b25f341 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -37,6 +37,7 @@ Do not modify directly.*
* com.microsoft.FusedMatMul
* com.microsoft.FusedMatMulActivation
* com.microsoft.GatedRelativePositionBias
+ * com.microsoft.GatherBlockQuantized
* com.microsoft.GatherND
* com.microsoft.Gelu
* com.microsoft.GemmFastGelu
@@ -2030,6 +2031,64 @@ This version of the operator has been available since version 1 of the 'com.micr
+### **com.microsoft.GatherBlockQuantized**
+
+ GatherBlockQuantized is a Gather with data quantized. It is similar to Gather (https://github.com/onnx/onnx/blob/main/docs/Operators.md#gather) with differences:
+ 1. Input `data` is a constant. It is quantized block-wise along attribute `quantize_axis` with block size specified by attribute `block_size`.
+ `block_size must` be a power of 2 and not smaller than 16, like 16, 32, 64, 128, ..
+ 2. Input `data`'s scale and zero point are specified by input `scales` and `zero_points`. `scales` and `zero_points` are also constants.
+ If `zero_points` is not provided, 0 is the zero point.
+ 3. During the op execution, `data` and `indices` are first used to generate the quantized output. Then, `scales` and `zero_points` are used
+ to dequantize the output.
+ 4. The `output` and `scales` have the same type. The `data` and `zero_points` have the same type.
+
+#### Version
+
+This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
+
+#### Attributes
+
+
+- block_size : int
+- (Optional) block size used for weight quantization. It needs to be a power of 2 and not smaller than 16.
+- gather_axis : int
+- (Optional) Which axis to gather on. Negative value means counting dimensions from the back. Accepted range is [-r, r-1] where r = rank(data).
+- quantize_axis : int
+- (Optional) Which axis to block-wise quantize. Negative value means counting dimensions from the back. Accepted range is [-r, r-1] where r = rank(data).
+
+
+#### Inputs (3 - 4)
+
+
+- data : T1
+- Tensor of rank r >= 1. Block-wise quantized.
+- indices : Tind
+- Tensor of int32/int64 indices, of any rank q. All index values are expected to be within bounds [-s, s-1] along axis of size s. It is an error if any of the index values are out of bounds.
+- scales : T2
+- quantization scale
+- zero_points (optional) : T1
+- quantization zero points
+
+
+#### Outputs
+
+
+- output : T2
+- Dequantized output tensor of rank q + (r - 1).
+
+
+#### Type Constraints
+
+
+- T1 : tensor(int4), tensor(uint4)
+- Constrain quantized types.
+- T2 : tensor(float), tensor(float16), tensor(bfloat16)
+- Constrain dequantized types.
+- Tind : tensor(int32), tensor(int64)
+- Constrain indices to integer types.
+
+
+
### **com.microsoft.GatherND**
Given `data` tensor of rank r >= 1, and `indices` tensor of rank q >= 1, gather
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index 529c676321..f0aa332ff3 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -477,6 +477,7 @@ Do not modify directly.*
|FusedConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*in* Z:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
|FusedGemm|*in* A:**T**
*in* B:**T**
*in* C:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
|FusedMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
+|GatherBlockQuantized|*in* data:**T1**
*in* indices:**Tind**
*in* scales:**T2**
*in* zero_points:**T1**
*out* output:**T2**|1+|**T1** = tensor(int4), tensor(uint4)
**T2** = tensor(float), tensor(float16)
**Tind** = tensor(int32), tensor(int64)|
|GatherND|*in* data:**T**
*in* indices:**Tind**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**Tind** = tensor(int32), tensor(int64)|
|Gelu|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
|GreedySearch|*in* input_ids:**I**
*in* max_length:**I**
*in* min_length:**I**
*in* repetition_penalty:**T**
*in* vocab_mask:**I**
*in* prefix_vocab_mask:**I**
*in* attention_mask:**I**
*out* sequences:**I**|1+|**T** = tensor(float)|
diff --git a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
index 84f9ca88ec..e9c1b4c434 100644
--- a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
+++ b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
@@ -3,6 +3,7 @@
#include "contrib_ops/cpu/cpu_contrib_kernels.h"
#include "core/graph/constants.h"
+#include "core/framework/int4.h"
#include "core/mlas/inc/mlas.h"
namespace onnxruntime {
@@ -33,6 +34,10 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Trans
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FusedMatMul);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, 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);
+class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Int4x2, int32_t, GatherBlockQuantized);
+class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Int4x2, int64_t, GatherBlockQuantized);
#ifndef ORT_MINIMAL_BUILD
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulFpQ4);
#endif
@@ -298,6 +303,10 @@ Status RegisterCpuContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
#ifndef ORT_MINIMAL_BUILD
BuildKernelCreateInfo,
#endif
diff --git a/onnxruntime/contrib_ops/cpu/quantization/gather_block_quantized.cc b/onnxruntime/contrib_ops/cpu/quantization/gather_block_quantized.cc
new file mode 100644
index 0000000000..5935663f11
--- /dev/null
+++ b/onnxruntime/contrib_ops/cpu/quantization/gather_block_quantized.cc
@@ -0,0 +1,282 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include
+#include
+
+#include "core/common/common.h"
+#include "core/common/narrow.h"
+#include "core/common/safeint.h"
+#include "core/framework/float16.h"
+#include "core/framework/int4.h"
+#include "core/framework/op_kernel.h"
+#include "core/platform/threadpool.h"
+#include "core/providers/common.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+template
+class GatherBlockQuantized : public OpKernel {
+ public:
+ GatherBlockQuantized(const OpKernelInfo& info) : OpKernel(info) {
+ if (!info.GetAttr("gather_axis", &gather_axis_).IsOK()) {
+ gather_axis_ = 0;
+ }
+
+ if (!info.GetAttr("quantize_axis", &quantize_axis_).IsOK()) {
+ quantize_axis_ = 1;
+ }
+
+ if (!info.GetAttr("block_size", &block_size_).IsOK()) {
+ block_size_ = 128;
+ }
+
+ ORT_ENFORCE(block_size_ >= 16 && ((block_size_ - 1) & block_size_) == 0,
+ "'block_size' must be 2's power and not less than 16.");
+ }
+
+ Status Compute(OpKernelContext* context) const override;
+
+ protected:
+ struct Prepare {
+ const Tensor* data_tensor;
+ const Tensor* indices_tensor;
+ const Tensor* scales_tensor;
+ const Tensor* zero_points_tensor;
+ Tensor* output_tensor;
+ int64_t gather_axis;
+ int64_t quantize_axis;
+ };
+
+ Status PrepareForCompute(OpKernelContext* context, Prepare& args) const;
+
+ template
+ Status CopyDataAndDequantize(const T1* data_ptr,
+ const Tind* indices_ptr,
+ const T2* scales_ptr,
+ const T1* zero_points_ptr,
+ T2* output_ptr,
+ const int64_t gather_M,
+ const int64_t gather_N,
+ const int64_t gather_axis_dim,
+ const int64_t gather_block,
+ const int64_t quantize_axis_dim,
+ const int64_t quantize_N,
+ concurrency::ThreadPool* tp) const;
+
+ private:
+ int64_t gather_axis_;
+ int64_t quantize_axis_;
+ int64_t block_size_;
+};
+
+template
+Status GatherBlockQuantized::PrepareForCompute(OpKernelContext* context, Prepare& p) const {
+ p.data_tensor = context->Input(0);
+ p.indices_tensor = context->Input(1);
+ p.scales_tensor = context->Input(2);
+ p.zero_points_tensor = context->Input(3);
+
+ const auto& data_shape = p.data_tensor->Shape();
+ const auto& indices_shape = p.indices_tensor->Shape();
+ const auto data_rank = data_shape.NumDimensions();
+ p.gather_axis = HandleNegativeAxis(gather_axis_, narrow(data_rank));
+ p.quantize_axis = HandleNegativeAxis(quantize_axis_, narrow(data_rank));
+
+ std::vector shape;
+ shape.reserve(data_rank - 1 + indices_shape.NumDimensions());
+
+ // get output tensor
+ // replace the dimension for p.gather_axis with the shape from the indices
+ for (int64_t i = 0; i < p.gather_axis; ++i)
+ shape.push_back(data_shape[narrow(i)]);
+
+ for (const auto dim : indices_shape.GetDims())
+ shape.push_back(dim);
+
+ for (int64_t i = p.gather_axis + 1; i < static_cast(data_rank); ++i)
+ shape.push_back(data_shape[narrow(i)]);
+
+ p.output_tensor = context->Output(0, TensorShape(std::move(shape)));
+
+ // validate quantization parameters
+ const auto& scales_shape = p.scales_tensor->Shape();
+ ORT_RETURN_IF_NOT(data_shape.NumDimensions() == scales_shape.NumDimensions(),
+ "data and scales must have the same rank.");
+ for (size_t i = 0; i < data_shape.NumDimensions(); ++i) {
+ ORT_RETURN_IF_NOT(i == static_cast(p.quantize_axis)
+ ? (data_shape[i] + block_size_ - 1) / block_size_ == scales_shape[i]
+ : data_shape[i] == scales_shape[i],
+ "data and scales do not match shapes.");
+ }
+
+ if (p.zero_points_tensor) {
+ const auto& zero_points_shape = p.zero_points_tensor->Shape();
+ ORT_RETURN_IF_NOT(scales_shape.NumDimensions() == zero_points_shape.NumDimensions(),
+ "scales and zero_points must have the same rank.");
+ for (size_t i = 0; i < scales_shape.NumDimensions(); ++i) {
+ ORT_RETURN_IF_NOT(scales_shape[i] == zero_points_shape[i],
+ "scales and zero_points must have the same shape.");
+ }
+ }
+
+ return Status::OK();
+}
+
+template
+template
+Status GatherBlockQuantized::CopyDataAndDequantize(const T1* data_ptr,
+ const Tind* indices_ptr,
+ const T2* scales_ptr,
+ const T1* zero_points_ptr,
+ T2* output_ptr,
+ const int64_t gather_M,
+ const int64_t gather_N,
+ const int64_t gather_axis_dim,
+ const int64_t gather_block,
+ const int64_t quantize_axis_dim,
+ const int64_t quantize_N,
+ concurrency::ThreadPool* tp) const {
+ auto data_full_block = gather_axis_dim * gather_block;
+ auto quantize_full_block = quantize_axis_dim * quantize_N;
+ auto scale_full_block = (quantize_axis_dim + block_size_ - 1) / block_size_ * quantize_N;
+
+ auto lambda = [&](int64_t gather_MN_idx, std::unordered_map& cache) {
+ int64_t gather_M_idx = gather_MN_idx / gather_N;
+ int64_t gather_N_idx = gather_MN_idx % gather_N;
+
+ int64_t indices_val = static_cast(indices_ptr[gather_N_idx]);
+ ORT_ENFORCE(indices_val >= -gather_axis_dim && indices_val < gather_axis_dim,
+ "indices element out of data bounds, idx=", indices_val,
+ " must be within the inclusive range [", -gather_axis_dim, ",", gather_axis_dim - 1, "]");
+
+ indices_val = indices_val < 0 ? indices_val + gather_axis_dim : indices_val;
+ int64_t output_idx_base = gather_MN_idx * gather_block;
+ int64_t data_idx_base = gather_M_idx * data_full_block + indices_val * gather_block;
+
+ if (auto it = cache.find(data_idx_base); it != cache.end()) {
+ int64_t output_src_idx = it->second;
+ memcpy(output_ptr + output_idx_base, output_ptr + output_src_idx, narrow(gather_block * sizeof(T2)));
+ return;
+ }
+
+ // TODO(fajin): use SIMD
+ int64_t output_idx = output_idx_base;
+ int64_t data_idx = data_idx_base;
+ for (int64_t i = 0; i < gather_block; ++i, ++output_idx, ++data_idx) {
+ auto data_val = static_cast(data_ptr[data_idx >> 1].GetElem(narrow(data_idx & 1)));
+
+ int64_t x = data_idx / quantize_full_block;
+ int64_t y = data_idx % quantize_full_block / quantize_N;
+ int64_t z = data_idx % quantize_N;
+ int64_t scale_idx = x * scale_full_block + y / block_size_ * quantize_N + z;
+ auto scale_val = static_cast(scales_ptr[scale_idx]);
+ auto zp_val = static_cast(zero_points_ptr
+ ? zero_points_ptr[scale_idx >> 1].GetElem(narrow(scale_idx & 1))
+ : 0);
+
+ output_ptr[output_idx] = static_cast(static_cast(data_val - zp_val) * scale_val);
+ }
+
+ cache[data_idx_base] = output_idx_base;
+ };
+
+ concurrency::ThreadPool::TryParallelFor(
+ tp,
+ SafeInt(gather_M) * gather_N,
+ static_cast(gather_block * 3),
+ [&lambda](ptrdiff_t first, ptrdiff_t last) {
+ // cache dequantized gather_block. Key is data_idx_base. Value is the output_idx_base.
+ // cache is per thread to avoid contention.
+ std::unordered_map cache;
+
+ for (auto index = static_cast(first), end = static_cast(last);
+ index < end;
+ ++index) {
+ lambda(index, cache);
+ }
+ });
+
+ return Status::OK();
+}
+
+template
+Status GatherBlockQuantized::Compute(OpKernelContext* context) const {
+ Prepare p;
+ ORT_RETURN_IF_ERROR(PrepareForCompute(context, p));
+
+ const auto& data_shape = p.data_tensor->Shape();
+ // re-shape the data tensor to [gather_M, gather_axis_dim, gather_block]
+ // re-shape the indices tensor to [gather_N]
+ // re-shape the output tensor to [gather_M, gather_N, gather_block]
+ // For an index i in the output tensor:
+ // 1> the output block index is blk_i = i / gather_block, block element index is blk_ele_i = i % gather_block,
+ // 2> block is picked from data based on value from indices: axis_i = indices[blk_i % gather_N],
+ // 3> get the corresponding block in data tensor: data_blk = data[blk_i / gather_N, axis_i, :],
+ // 4> pick the element from the block: value_i = data_blk[blk_ele_i]
+ const int64_t gather_block = data_shape.SizeFromDimension(SafeInt(p.gather_axis) + 1);
+ const int64_t gather_axis_dim = data_shape[narrow(p.gather_axis)];
+ const int64_t gather_M = data_shape.SizeToDimension(narrow(p.gather_axis));
+ const int64_t gather_N = p.indices_tensor->Shape().Size();
+ // re-shape the data tensor to [quantize_M, quantize_axis_dim, quantize_N]
+ // For an index i in the output tensor:
+ // 1> based on previous comment, corresponding data index is (blk_i / gather_N, axis_i, blk_ele_i)
+ // 2> flatten the data index:
+ // data_i = blk_i / gather_N * gather_axis_dim * gather_block + axis_i * gather_block + blk_ele_i
+ // 3> map data_i to quantize shape: (x, y, z) =
+ // (data_i / (quantize_axis_dim * quantize_N),
+ // data_i % (quantize_axis_dim * quantize_N) / quantize_N,
+ // data_i % quantize_N)
+ // 4> get scale index: (x, y / block_size_, z)
+ const int64_t quantize_axis_dim = data_shape[narrow(p.quantize_axis)];
+ const int64_t quantize_N = data_shape.SizeFromDimension(SafeInt(p.quantize_axis) + 1);
+
+ concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
+ const auto* data_ptr = p.data_tensor->template Data();
+ const auto* indices_ptr = p.indices_tensor->template Data();
+ const auto* zero_points_ptr = p.zero_points_tensor ? p.zero_points_tensor->template Data() : nullptr;
+ const auto dequantized_type = p.scales_tensor->GetElementType();
+
+ if (dequantized_type == ONNX_NAMESPACE::TensorProto::FLOAT) {
+ const auto* scales_ptr = p.scales_tensor->template Data();
+ auto* output_ptr = p.output_tensor->template MutableData();
+
+ return CopyDataAndDequantize(data_ptr, indices_ptr, scales_ptr, zero_points_ptr,
+ output_ptr, gather_M, gather_N, gather_axis_dim, gather_block,
+ quantize_axis_dim, quantize_N,
+ tp);
+ } else if (dequantized_type == ONNX_NAMESPACE::TensorProto::FLOAT16) {
+ const auto* scales_ptr = p.scales_tensor->template Data();
+ auto* output_ptr = p.output_tensor->template MutableData();
+
+ return CopyDataAndDequantize(data_ptr, indices_ptr, scales_ptr, zero_points_ptr,
+ output_ptr, gather_M, gather_N, gather_axis_dim, gather_block,
+ quantize_axis_dim, quantize_N,
+ tp);
+ } else if (dequantized_type == ONNX_NAMESPACE::TensorProto::BFLOAT16) {
+ ORT_THROW("DequantizeLinear into BFLOAT16 is not implemented yet.");
+ } else {
+ ORT_THROW("Unsupported dequantized type: ", dequantized_type);
+ }
+}
+
+#define REGISTER_GATHERBLOCKQUANTIZED(T1, Tind) \
+ ONNX_OPERATOR_TWO_TYPED_KERNEL_EX( \
+ GatherBlockQuantized, \
+ kMSDomain, 1, \
+ T1, Tind, \
+ kCpuExecutionProvider, \
+ KernelDefBuilder() \
+ .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \
+ .TypeConstraint("T2", {DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}) \
+ .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \
+ GatherBlockQuantized);
+
+REGISTER_GATHERBLOCKQUANTIZED(UInt4x2, int32_t);
+REGISTER_GATHERBLOCKQUANTIZED(UInt4x2, int64_t);
+REGISTER_GATHERBLOCKQUANTIZED(Int4x2, int32_t);
+REGISTER_GATHERBLOCKQUANTIZED(Int4x2, int64_t);
+
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
index 2d51658953..aebe726afe 100644
--- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
@@ -3544,6 +3544,124 @@ MatMulBnb4 is a MatMul with weight quantized with 4 bits using either FP4 or NF4
MatmulWithQuantWeightShapeInference(ctx, in_features, out_features, transB);
});
+ static const char* GatherBlockQuantized_ver1_doc = R"DOC(
+GatherBlockQuantized is a Gather with data quantized. It is similar to Gather (https://github.com/onnx/onnx/blob/main/docs/Operators.md#gather) with differences:
+ 1. Input `data` is a constant. It is quantized block-wise along attribute `quantize_axis` with block size specified by attribute `block_size`.
+ `block_size must` be a power of 2 and not smaller than 16, like 16, 32, 64, 128, ..
+ 2. Input `data`'s scale and zero point are specified by input `scales` and `zero_points`. `scales` and `zero_points` are also constants.
+ If `zero_points` is not provided, 0 is the zero point.
+ 3. During the op execution, `data` and `indices` are first used to generate the quantized output. Then, `scales` and `zero_points` are used
+ to dequantize the output.
+ 4. The `output` and `scales` have the same type. The `data` and `zero_points` have the same type.
+)DOC";
+
+ ONNX_CONTRIB_OPERATOR_SCHEMA(GatherBlockQuantized)
+ .SetDomain(kMSDomain)
+ .SinceVersion(1)
+ .SetDoc(GatherBlockQuantized_ver1_doc)
+ .Attr("gather_axis",
+ "(Optional) Which axis to gather on. Negative value means "
+ "counting dimensions from the back. Accepted range is [-r, r-1] where r = rank(data).",
+ AttributeProto::INT, static_cast(0))
+ .Attr("quantize_axis",
+ "(Optional) Which axis to block-wise quantize. Negative value means "
+ "counting dimensions from the back. Accepted range is [-r, r-1] where r = rank(data).",
+ AttributeProto::INT, static_cast(1))
+ .Attr("block_size",
+ "(Optional) block size used for weight quantization. It needs to be a power of 2 and not smaller than 16.",
+ AttributeProto::INT,
+ static_cast(128))
+ .Input(0, "data", "Tensor of rank r >= 1. Block-wise quantized.", "T1")
+ .Input(1,
+ "indices",
+ "Tensor of int32/int64 indices, of any rank q. All index values are expected to be within bounds [-s, s-1] "
+ "along axis of size s. It is an error if any of the index values are out of bounds.",
+ "Tind")
+ .Input(2, "scales", "quantization scale", "T2")
+ .Input(3, "zero_points", "quantization zero points", "T1", OpSchema::Optional)
+ .Output(0, "output", "Dequantized output tensor of rank q + (r - 1).", "T2")
+ .TypeConstraint("T1", {"tensor(int4)", "tensor(uint4)"}, "Constrain quantized types.")
+ .TypeConstraint("T2", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain dequantized types.")
+ .TypeConstraint("Tind", {"tensor(int32)", "tensor(int64)"}, "Constrain indices to integer types.")
+ .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
+ // Type inference
+ propagateElemTypeFromInputToOutput(ctx, 2, 0);
+ // Shape inference
+ if (!hasNInputShapes(ctx, 3)) {
+ return;
+ }
+ const TensorShapeProto& data_shape = ctx.getInputType(0)->tensor_type().shape();
+ const TensorShapeProto& indices_shape = ctx.getInputType(1)->tensor_type().shape();
+ const TensorShapeProto& scales_shape = ctx.getInputType(2)->tensor_type().shape();
+ int r = data_shape.dim_size();
+
+ if (r < 1) {
+ fail_shape_inference("data tensor must have rank >= 1");
+ }
+
+ int gather_axis = static_cast(getAttribute(ctx, "gather_axis", 0));
+ int quantize_axis = static_cast(getAttribute(ctx, "quantize_axis", 1));
+ auto block_size = getAttribute(ctx, "block_size", 128);
+ if (gather_axis < -r || gather_axis >= r) {
+ fail_shape_inference("gather_axis must be in [-r, r-1]");
+ }
+ if (quantize_axis < -r || quantize_axis >= r) {
+ fail_shape_inference("quantize_axis must be in [-r, r-1]");
+ }
+ if (block_size < 0) {
+ fail_shape_inference("block_size must be non-negative");
+ }
+
+ gather_axis = (gather_axis + r) % r;
+ quantize_axis = (quantize_axis + r) % r;
+
+ if (scales_shape.dim_size() != r) {
+ fail_shape_inference("scales must have the same rank as data");
+ }
+
+ for (int i = 0; i < r; ++i) {
+ if (!data_shape.dim(i).has_dim_value() ||
+ !scales_shape.dim(i).has_dim_value() ||
+ (i == quantize_axis && (data_shape.dim(i).dim_value() + block_size - 1) / block_size != scales_shape.dim(i).dim_value()) ||
+ (i != quantize_axis && data_shape.dim(i).dim_value() != scales_shape.dim(i).dim_value())) {
+ fail_shape_inference("data shape and scales shape do not match");
+ }
+ }
+
+ // validate zero point shape
+ if (ctx.hasInput(3)) {
+ if (!hasInputShape(ctx, 3)) {
+ fail_shape_inference("zero_points shape must be known");
+ }
+
+ const auto& zp_shape = getInputShape(ctx, 3);
+ if (zp_shape.dim_size() != r) {
+ fail_shape_inference("zero points must have the same rank as data");
+ }
+
+ for (int i = 0; i < r; ++i) {
+ if (!zp_shape.dim(i).has_dim_value() ||
+ zp_shape.dim(i).dim_value() != scales_shape.dim(i).dim_value()) {
+ fail_shape_inference("zero points shape and scales shape do not match");
+ }
+ }
+ }
+
+ int q = indices_shape.dim_size();
+ int out_rank = q + r - 1;
+ if (out_rank == 0) {
+ ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
+ }
+ for (int i = 0; i < out_rank; ++i) {
+ *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape()->add_dim() =
+ (i < gather_axis)
+ ? data_shape.dim(i)
+ : (i >= gather_axis && i < gather_axis + q)
+ ? indices_shape.dim(i - gather_axis)
+ : data_shape.dim(i - q + 1);
+ }
+ });
+
#ifdef ENABLE_ATEN
ONNX_CONTRIB_OPERATOR_SCHEMA(ATen)
.SetDomain(kPytorchAtenDomain)
diff --git a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc
new file mode 100644
index 0000000000..c4536fc56a
--- /dev/null
+++ b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc
@@ -0,0 +1,468 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include
+#include
+#include
+#include
+
+#include "core/common/common.h"
+#include "core/framework/execution_provider.h"
+#include "gtest/gtest.h"
+#include "test/providers/provider_test_utils.h"
+#include "test/util/include/default_providers.h"
+
+namespace onnxruntime {
+namespace test {
+
+// Combinations: types, gather_axis, quantize_axis, block_size, indices, scale shape vs data shape
+template
+void RunGatherBlockQuantized(const std::vector& data,
+ const std::vector& data_shape,
+ const std::vector& indices,
+ const std::vector& indices_shape,
+ const std::vector& scales,
+ const std::vector& scales_shape,
+ const std::vector& zero_points,
+ const int64_t gather_axis,
+ const int64_t quantize_axis,
+ const int64_t block_size,
+ const std::vector& output,
+ const std::vector& output_shape,
+ OpTester::ExpectResult expect_result = OpTester::ExpectResult::kExpectSuccess) {
+ auto run_test = [&](bool indices_is_initializer) {
+ OpTester test("GatherBlockQuantized", 1, kMSDomain);
+
+ test.AddAttribute("gather_axis", gather_axis);
+ test.AddAttribute("quantize_axis", quantize_axis);
+ test.AddAttribute("block_size", block_size);
+
+ test.AddInput("data", data_shape, data);
+ test.AddInput("indices", indices_shape, indices, indices_is_initializer);
+ test.AddInput("scales", scales_shape, scales);
+ if (!zero_points.empty()) {
+ test.AddInput("zero_points", scales_shape, zero_points);
+ }
+
+ test.AddOutput("output", output_shape, output);
+
+ std::vector> eps;
+ eps.push_back(DefaultCpuExecutionProvider());
+ test.Run(expect_result, "", {}, nullptr, &eps);
+ };
+
+ run_test(false);
+ run_test(true);
+}
+
+template
+typename std::enable_if<
+ (boost::mp11::mp_contains, T1>::value && std::is_same::value) ||
+ (std::is_integral::value && std::is_same::value),
+ std::vector>::type
+ToType(const std::vector& vec) {
+ std::vector result;
+ for (auto v : vec) {
+ result.push_back(static_cast(v));
+ }
+
+ return result;
+}
+
+template
+typename std::enable_if, T>::value, std::vector>::type
+ToType(const std::vector& vec) {
+ std::vector result;
+ size_t i = 0;
+ constexpr int offset = std::is_same::value ? 0 : 8;
+ for (i = 0; i + 1 < vec.size(); i += 2) {
+ result.push_back(T(vec[i] + offset, vec[i + 1] + offset));
+ }
+ if (i < vec.size()) {
+ result.push_back(T(vec[i] + offset, 0 + offset));
+ }
+
+ return result;
+}
+
+template
+void Test_Fail_WithZeroPoints(int64_t gather_axis,
+ int64_t quantize_axis,
+ int64_t block_size) {
+ std::vector data = {-8, -7, -6, -5,
+ -4, -3, -2, -1,
+ 0, 1, 2, 3,
+ 4, 5, 6, 7,
+ 4, 5, 6, 7,
+ -4, -3, -2, -1};
+ std::vector data_shape = {2, 3, 4};
+ std::vector indices = {1};
+ std::vector indices_shape = {1};
+ std::vector scales = {1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f};
+ std::vector scales_shape = {2, 3, 1};
+ std::vector zero_points = {-1, 1, 0, 0, 1, -1};
+ std::vector output = {8.f, 10.f, 12.f, 14.f,
+ 3.f, 4.f, 5.f, 6.f,
+ -6.f, -4.f, -2.f, 0.f};
+ std::vector output_shape = {1, 3, 4};
+
+ RunGatherBlockQuantized(ToType(data),
+ data_shape,
+ ToType(indices),
+ indices_shape,
+ ToType(scales),
+ scales_shape,
+ ToType(zero_points),
+ gather_axis,
+ quantize_axis,
+ block_size,
+ ToType(output),
+ output_shape,
+ OpTester::ExpectResult::kExpectFailure);
+}
+
+TEST(GatherBlockQuantizedOpTest, UnsupportedTypes) {
+ Test_Fail_WithZeroPoints(0, 2, 16);
+ Test_Fail_WithZeroPoints(0, 2, 16);
+ Test_Fail_WithZeroPoints(0, 2, 16);
+ Test_Fail_WithZeroPoints(0, 2, 16);
+ Test_Fail_WithZeroPoints(0, 2, 16);
+ Test_Fail_WithZeroPoints(0, 2, 16);
+ Test_Fail_WithZeroPoints(0, 2, 16);
+ Test_Fail_WithZeroPoints(0, 2, 16);
+ Test_Fail_WithZeroPoints(0, 2, 16);
+ Test_Fail_WithZeroPoints(0, 2, 16);
+ Test_Fail_WithZeroPoints(0, 2, 16);
+ Test_Fail_WithZeroPoints(0, 2, 16);
+}
+
+TEST(GatherBlockQuantizedOpTest, InvalidBlockSize) {
+ Test_Fail_WithZeroPoints(0, 2, 8);
+ Test_Fail_WithZeroPoints(0, 2, 17);
+}
+
+TEST(GatherBlockQuantizedOpTest, InvalidGatherAxis) {
+ Test_Fail_WithZeroPoints(3, 2, 16);
+ Test_Fail_WithZeroPoints(-4, 2, 16);
+}
+
+TEST(GatherBlockQuantizedOpTest, InvalidQuantizeAxis) {
+ Test_Fail_WithZeroPoints(0, 3, 16);
+ Test_Fail_WithZeroPoints(0, -4, 16);
+}
+
+template
+void Test_ShapeMismatch_WithZeroPoints() {
+ std::vector data = {-8, -7, -6, -5,
+ -4, -3, -2, -1,
+ 0, 1, 2, 3,
+ 4, 5, 6, 7,
+ 4, 5, 6, 7,
+ -4, -3, -2, -1};
+ std::vector data_shape = {2, 3, 4};
+ std::vector indices = {1};
+ std::vector indices_shape = {1};
+ std::vector scales = {1.0f, 2.0f, 1.0f, 2.0f};
+ std::vector scales_shape = {2, 2, 1};
+ std::vector zero_points = {-1, 1, 0, 0};
+ std::vector output = {8.f, 10.f, 12.f, 14.f,
+ 3.f, 4.f, 5.f, 6.f,
+ -6.f, -4.f, -2.f, 0.f};
+ std::vector output_shape = {1, 3, 4};
+
+ RunGatherBlockQuantized(ToType(data),
+ data_shape,
+ ToType(indices),
+ indices_shape,
+ ToType(scales),
+ scales_shape,
+ ToType(zero_points),
+ 0,
+ 2,
+ 16,
+ ToType(output),
+ output_shape,
+ OpTester::ExpectResult::kExpectFailure);
+}
+
+TEST(GatherBlockQuantizedOpTest, ShapeMismatch) {
+ Test_ShapeMismatch_WithZeroPoints();
+ Test_ShapeMismatch_WithZeroPoints();
+}
+
+template
+void Test_InvalidIndices_WithZeroPoints() {
+ std::vector data = {-8, -7, -6, -5,
+ -4, -3, -2, -1,
+ 0, 1, 2, 3,
+ 4, 5, 6, 7,
+ 4, 5, 6, 7,
+ -4, -3, -2, -1};
+ std::vector data_shape = {2, 3, 4};
+ std::vector indices = {2};
+ std::vector indices_shape = {1};
+ std::vector scales = {1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f};
+ std::vector scales_shape = {2, 3, 1};
+ std::vector zero_points = {-1, 1, 0, 0, 1, -1};
+ std::vector output = {8.f, 10.f, 12.f, 14.f,
+ 3.f, 4.f, 5.f, 6.f,
+ -6.f, -4.f, -2.f, 0.f};
+ std::vector output_shape = {1, 3, 4};
+
+ RunGatherBlockQuantized(ToType(data),
+ data_shape,
+ ToType(indices),
+ indices_shape,
+ ToType(scales),
+ scales_shape,
+ ToType(zero_points),
+ 0,
+ 2,
+ 16,
+ ToType(output),
+ output_shape,
+ OpTester::ExpectResult::kExpectFailure);
+}
+
+TEST(GatherBlockQuantizedOpTest, InvalidIndices) {
+ Test_InvalidIndices_WithZeroPoints();
+ Test_InvalidIndices_WithZeroPoints();
+}
+
+template
+void Test_GatherAxis0_WithZeroPoints() {
+ std::vector data = {-8, -7, -6, -5, -8, -7, -6, -5, -8, -7, -6, -5, -8, -7, -6, -5, -8,
+ -4, -3, -2, -1, -4, -3, -2, -1, -4, -3, -2, -1, -4, -3, -2, -1, -4,
+ 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0,
+ 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4,
+ 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4,
+ -4, -3, -2, -1, -4, -3, -2, -1, -4, -3, -2, -1, -4, -3, -2, -1, -4};
+ std::vector data_shape = {2, 3, 17};
+ std::vector indices = {1};
+ std::vector indices_shape = {1};
+ std::vector scales = {1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f,
+ 2.0f, 2.0f, 1.0f, 1.0f, 2.0f, 1.0f};
+ std::vector scales_shape = {2, 3, 2};
+ std::vector zero_points = {-1, 1, 0, 0, 1, -1,
+ 1, -1, 1, 0, -1, 1};
+ std::vector output = {6, 8, 10, 12, 6, 8, 10, 12, 6, 8, 10, 12, 6, 8, 10, 12, 10,
+ 3, 4, 5, 6, 3, 4, 5, 6, 3, 4, 5, 6, 3, 4, 5, 6, 4,
+ -6, -4, -2, 0, -6, -4, -2, 0, -6, -4, -2, 0, -6, -4, -2, 0, -5};
+ std::vector output_shape = {1, 3, 17};
+
+ RunGatherBlockQuantized(ToType(data),
+ data_shape,
+ ToType(indices),
+ indices_shape,
+ ToType(scales),
+ scales_shape,
+ ToType(zero_points),
+ 0,
+ 2,
+ 16,
+ ToType(output),
+ output_shape,
+ OpTester::ExpectResult::kExpectSuccess);
+ RunGatherBlockQuantized(ToType(data),
+ data_shape,
+ ToType(indices),
+ indices_shape,
+ ToType(scales),
+ scales_shape,
+ ToType(zero_points),
+ -3,
+ -1,
+ 16,
+ ToType(output),
+ output_shape,
+ OpTester::ExpectResult::kExpectSuccess);
+}
+
+TEST(GatherBlockQuantizedOpTest, GatherAxis0WithZeroPoints) {
+ Test_GatherAxis0_WithZeroPoints();
+ Test_GatherAxis0_WithZeroPoints();
+ Test_GatherAxis0_WithZeroPoints();
+ Test_GatherAxis0_WithZeroPoints();
+ Test_GatherAxis0_WithZeroPoints();
+ Test_GatherAxis0_WithZeroPoints();
+ Test_GatherAxis0_WithZeroPoints();
+ Test_GatherAxis0_WithZeroPoints();
+}
+
+template
+void Test_GatherAxis0_NoZeroPoints() {
+ std::vector data = {-8, -7, -6, -5,
+ -4, -3, -2, -1,
+ 0, 1, 2, 3,
+ 4, 5, 6, 7,
+ 4, 5, 6, 7,
+ -4, -3, -2, -1};
+ std::vector data_shape = {2, 3, 4};
+ std::vector indices = {1};
+ std::vector indices_shape = {1};
+ std::vector scales = {1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f};
+ std::vector scales_shape = {2, 3, 1};
+ std::vector output = {8.f, 10.f, 12.f, 14.f,
+ 4.f, 5.f, 6.f, 7.f,
+ -8.f, -6.f, -4.f, -2.f};
+ std::vector output_shape = {1, 3, 4};
+
+ RunGatherBlockQuantized(ToType(data),
+ data_shape,
+ ToType(indices),
+ indices_shape,
+ ToType(scales),
+ scales_shape,
+ {},
+ 0,
+ 2,
+ 16,
+ ToType(output),
+ output_shape,
+ OpTester::ExpectResult::kExpectSuccess);
+ RunGatherBlockQuantized(ToType(data),
+ data_shape,
+ ToType(indices),
+ indices_shape,
+ ToType(scales),
+ scales_shape,
+ {},
+ -3,
+ -1,
+ 16,
+ ToType(output),
+ output_shape,
+ OpTester::ExpectResult::kExpectSuccess);
+}
+
+TEST(GatherBlockQuantizedOpTest, GatherAxis0NoZeroPoints) {
+ Test_GatherAxis0_NoZeroPoints();
+ Test_GatherAxis0_NoZeroPoints();
+ Test_GatherAxis0_NoZeroPoints();
+ Test_GatherAxis0_NoZeroPoints();
+}
+
+template
+void Test_GatherAxis1_WithZeroPoints() {
+ std::vector data = {-8, -7, -6, -5,
+ -4, -3, -2, -1,
+ 0, 1, 2, 3,
+ 4, 5, 6, 7,
+ 4, 5, 6, 7,
+ -4, -3, -2, -1};
+ std::vector data_shape = {2, 3, 4};
+ std::vector indices = {2, -3, 2};
+ std::vector indices_shape = {1, 3};
+ std::vector scales = {1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f};
+ std::vector scales_shape = {2, 1, 4};
+ std::vector zero_points = {-1, 1, 0, 0, 1, -1, 0, 0};
+ std::vector output = {1.f, 0.f, 2.f, 6.f,
+ -7.f, -16.f, -6.f, -10.f,
+ 1.f, 0.f, 2.f, 6.f,
+ -5.f, -4.f, -2.f, -2.f,
+ 3.f, 12.f, 6.f, 14.f,
+ -5.f, -4.f, -2.f, -2.f};
+ std::vector output_shape = {2, 1, 3, 4};
+
+ RunGatherBlockQuantized(ToType(data),
+ data_shape,
+ ToType(indices),
+ indices_shape,
+ ToType(scales),
+ scales_shape,
+ ToType(zero_points),
+ 1,
+ 1,
+ 16,
+ ToType(output),
+ output_shape,
+ OpTester::ExpectResult::kExpectSuccess);
+ RunGatherBlockQuantized(ToType(data),
+ data_shape,
+ ToType(indices),
+ indices_shape,
+ ToType(scales),
+ scales_shape,
+ ToType(zero_points),
+ -2,
+ -2,
+ 16,
+ ToType(output),
+ output_shape,
+ OpTester::ExpectResult::kExpectSuccess);
+}
+
+TEST(GatherBlockQuantizedOpTest, GatherAxis1) {
+ Test_GatherAxis1_WithZeroPoints();
+ Test_GatherAxis1_WithZeroPoints();
+ Test_GatherAxis1_WithZeroPoints();
+ Test_GatherAxis1_WithZeroPoints();
+ Test_GatherAxis1_WithZeroPoints();
+ Test_GatherAxis1_WithZeroPoints();
+ Test_GatherAxis1_WithZeroPoints();
+ Test_GatherAxis1_WithZeroPoints