diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index 31389e4451..0d782a96a4 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -58,6 +58,7 @@ Do not modify directly.*
* com.microsoft.QLinearReduceMean
* com.microsoft.QLinearSigmoid
* com.microsoft.QLinearSoftmax
+ * com.microsoft.QLinearWhere
* com.microsoft.QOrderedAttention
* com.microsoft.QOrderedGelu
* com.microsoft.QOrderedLayerNormalization
@@ -3000,6 +3001,56 @@ This version of the operator has been available since version 1 of the 'com.micr
+### **com.microsoft.QLinearWhere**
+
+ Return elements, either from X or Y, depending on condition.
+
+#### Version
+
+This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
+
+#### Inputs
+
+
+- condition : B
+- When True (nonzero), yield x, otherwise yield y
+- X : T
+- Y's zero point.
+- x_scale : TF
+- X's scale.
+- x_zero_point : T
+- X's zero point.
+- Y : T
+- Y's zero point.
+- y_scale : TF
+- Y's scale.
+- y_zero_point : T
+- Y's zero point.
+- z_scale : TF
+- Z's scale.
+- z_zero_point : T
+- Z's zero point.
+
+
+#### Outputs
+
+
+- Z : T
+- Tensor of shape equal to the broadcasted shape of condition, X, and Y
+
+
+#### Type Constraints
+
+
+- B : tensor(bool)
+- Constrain input and output types to 8 bit signed and unsigned tensors.
+- TF : tensor(float)
+- Constrain scale types to any float tensor type.
+- T : tensor(uint8), tensor(int8)
+- Constrain input and output types to 8 bit signed and unsigned tensors.
+
+
+
### **com.microsoft.QOrderedAttention**
Quantized version of simplified Multi-Head Self Attention(using int8 with specific matrix Layout).
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index 6860f270c8..9c345ee5f8 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -433,6 +433,7 @@ Do not modify directly.*
|QLinearMul|*in* A:**T**
*in* A_scale:**tensor(float)**
*in* A_zero_point:**T**
*in* B:**T**
*in* B_scale:**tensor(float)**
*in* B_zero_point:**T**
*in* C_scale:**tensor(float)**
*in* C_zero_point:**T**
*out* C:**T**|1+|**T** = tensor(int8), tensor(uint8)|
|QLinearSigmoid|*in* X:**T**
*in* X_scale:**tensor(float)**
*in* X_zero_point:**T**
*in* Y_scale:**tensor(float)**
*in* Y_zero_point:**T**
*out* Y:**T**|1+|**T** = tensor(int8), tensor(uint8)|
|QLinearSoftmax|*in* X:**T**
*in* X_scale:**tensor(float)**
*in* x_zero_point:**T**
*in* y_scale:**tensor(float)**
*in* y_zero_point:**T**
*out* Y:**T**|1+|**T** = tensor(int8), tensor(uint8)|
+|QLinearWhere|*in* condition:**B**
*in* X:**T**
*in* x_scale:**TF**
*in* x_zero_point:**T**
*in* Y:**T**
*in* y_scale:**TF**
*in* y_zero_point:**T**
*in* z_scale:**TF**
*in* z_zero_point:**T**
*out* Z:**T**|1+|**T** = tensor(int8), tensor(uint8)|
|QuantizeLinear|*in* x:**T1**
*in* y_scale:**T1**
*in* y_zero_point:**T2**
*out* y:**T2**|1+|**T1** = tensor(float)
**T2** = tensor(int8), tensor(uint8)|
|QuickGelu|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
|Range|*in* start:**T**
*in* limit:**T**
*in* delta:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64)|
diff --git a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
index f10fa6233c..3f459dff99 100644
--- a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
+++ b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
@@ -47,6 +47,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Quick
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulInteger16);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, QLinearGlobalAveragePool);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, QLinearConcat);
+class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, QLinearWhere);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, QLinearAveragePool);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, uint8_t, DequantizeLinear);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, int8_t, DequantizeLinear);
@@ -145,6 +146,7 @@ Status RegisterQuantizationKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
diff --git a/onnxruntime/contrib_ops/cpu/quantization/qlinear_concat.cc b/onnxruntime/contrib_ops/cpu/quantization/qlinear_concat.cc
index 0e868ae62a..2e314b5e5d 100644
--- a/onnxruntime/contrib_ops/cpu/quantization/qlinear_concat.cc
+++ b/onnxruntime/contrib_ops/cpu/quantization/qlinear_concat.cc
@@ -1,6 +1,7 @@
// Copyright (c Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
+#include "qlinear_util.h"
#include "qlinear_concat.h"
#include "qlinear_lookup_table.h"
@@ -12,24 +13,6 @@
namespace onnxruntime {
namespace contrib {
-constexpr int LOOKUP_TABLE_IS_FIXED = 1;
-constexpr int LOOKUP_TABLE_IS_COPY = 2;
-
-static inline bool has_same_scale(const Tensor* tensor_x_scale, const Tensor* tensor_y_scale) {
- return *(tensor_x_scale->Data()) == *(tensor_y_scale->Data());
-}
-
-static inline bool has_same_zero_point(bool is_signed, const Tensor* tensor_x_zero_point, const Tensor* tensor_y_zero_point) {
- if (is_signed) {
- const int8_t X_zero_point = (tensor_x_zero_point == nullptr) ? static_cast(0) : *(tensor_x_zero_point->Data());
- const int8_t Y_zero_point = (tensor_y_zero_point == nullptr) ? static_cast(0) : *(tensor_y_zero_point->Data());
- return X_zero_point == Y_zero_point;
- } else {
- const uint8_t X_zero_point = (tensor_x_zero_point == nullptr) ? static_cast(0) : *(tensor_x_zero_point->Data());
- const uint8_t Y_zero_point = (tensor_y_zero_point == nullptr) ? static_cast(0) : *(tensor_y_zero_point->Data());
- return X_zero_point == Y_zero_point;
- }
-}
QLinearConcat::QLinearConcat(const OpKernelInfo& info) : OpKernel(info), ConcatBase(info) {
size_t input_def_count = info.node().InputDefs().size();
diff --git a/onnxruntime/contrib_ops/cpu/quantization/qlinear_util.h b/onnxruntime/contrib_ops/cpu/quantization/qlinear_util.h
new file mode 100644
index 0000000000..0ff3a9073c
--- /dev/null
+++ b/onnxruntime/contrib_ops/cpu/quantization/qlinear_util.h
@@ -0,0 +1,32 @@
+// Copyright (c Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#ifndef ONNXRUNTIME_QLINEAR_UTIL_H
+#define ONNXRUNTIME_QLINEAR_UTIL_H
+#include
+#include "core/framework/tensor.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+constexpr int LOOKUP_TABLE_IS_FIXED = 1;
+constexpr int LOOKUP_TABLE_IS_COPY = 2;
+
+static inline bool has_same_scale(const Tensor* tensor_x_scale, const Tensor* tensor_y_scale) {
+ return *(tensor_x_scale->Data()) == *(tensor_y_scale->Data());
+}
+
+static inline bool has_same_zero_point(bool is_signed, const Tensor* tensor_x_zero_point, const Tensor* tensor_y_zero_point) {
+ if (is_signed) {
+ const int8_t X_zero_point = (tensor_x_zero_point == nullptr) ? static_cast(0) : *(tensor_x_zero_point->Data());
+ const int8_t Y_zero_point = (tensor_y_zero_point == nullptr) ? static_cast(0) : *(tensor_y_zero_point->Data());
+ return X_zero_point == Y_zero_point;
+ }
+ const uint8_t X_zero_point = (tensor_x_zero_point == nullptr) ? static_cast(0) : *(tensor_x_zero_point->Data());
+ const uint8_t Y_zero_point = (tensor_y_zero_point == nullptr) ? static_cast(0) : *(tensor_y_zero_point->Data());
+ return X_zero_point == Y_zero_point;
+
+}
+} // namespace contrib
+} // namespace onnxruntime
+#endif // ONNXRUNTIME_QLINEAR_UTIL_H
diff --git a/onnxruntime/contrib_ops/cpu/quantization/qlinear_where.cc b/onnxruntime/contrib_ops/cpu/quantization/qlinear_where.cc
new file mode 100644
index 0000000000..1a5e8d87d9
--- /dev/null
+++ b/onnxruntime/contrib_ops/cpu/quantization/qlinear_where.cc
@@ -0,0 +1,404 @@
+#include "qlinear_where.h"
+#include "qlinear_lookup_table.h"
+
+#include "core/providers/common.h"
+#include "core/mlas/inc/mlas.h"
+#include "core/platform/threadpool.h"
+
+#include "core/providers/cpu/math/element_wise_ops.h"
+
+namespace onnxruntime {
+namespace {
+
+template
+using EnableIfEigenScalar = typename std::enable_if::value, R>::type;
+
+template
+using EnableIfEigenNotScalar = typename std::enable_if::value, R>::type;
+
+template
+ProcessBroadcastSpanFuncs CreateScalarBroadcastFuncs() {
+ return ProcessBroadcastSpanFuncs{
+ // Scalar Condition + Eigen Input -> Eigen output
+ [](BroadcastHelper& per_iter_bh) {
+ auto* user_data = static_cast(per_iter_bh.GetUserData());
+ bool target = user_data[0] == 1;
+ bool is_copy = user_data[1] == 1;
+ bool condition = per_iter_bh.ScalarInput0();
+ auto value = per_iter_bh.EigenInput1();
+ auto output = per_iter_bh.OutputEigen();
+ if (condition == target) {
+ output = value;
+ } else {
+ output = EigenVectorMap::PlainObject::Constant(value.size(), T{});
+ }
+ // Transform the output to the correct value from LookupTable
+ if (!is_copy) {
+ auto* look_up_table = user_data + 2;
+ std::transform(value.cbegin(), value.cend(), output.begin(),
+ [condition, target, &look_up_table](const T& value_element) {
+ return condition == target ? look_up_table[value_element] : T{};
+ });
+ }
+ },
+ // Eigen Condition + Scalar Input -> Eigen Output
+ [](BroadcastHelper& per_iter_bh) {
+ auto* user_data = static_cast(per_iter_bh.GetUserData());
+ bool target = user_data[0] == 1;
+ bool is_copy = user_data[1] == 1;
+ auto condition = per_iter_bh.EigenInput0();
+ const T& value = per_iter_bh.ScalarInput1();
+ auto output = per_iter_bh.OutputEigen();
+ output = (condition.array() == target)
+ .select(value, EigenVectorMap::PlainObject::Constant(condition.size(), T{}));
+ // Transform the output to the correct value from LookupTable
+ if (!is_copy) {
+ auto* look_up_table = user_data + 2;
+ std::transform(condition.cbegin(), condition.cend(), output.begin(),
+ [target, &look_up_table, &value](bool condition_element) {
+ return condition_element == target ? look_up_table[value] : T{};
+ });
+ }
+ },
+ // Eigen Condition + Eigen Input -> Eigen Output
+ [](BroadcastHelper& per_iter_bh) {
+ auto* user_data = static_cast(per_iter_bh.GetUserData());
+ bool target = user_data[0] == 1;
+ bool is_copy = user_data[1] == 1;
+ auto condition = per_iter_bh.EigenInput0();
+ auto value = per_iter_bh.EigenInput1();
+ auto output = per_iter_bh.OutputEigen();
+ output = (condition.array() == target)
+ .select(value, EigenVectorMap::PlainObject::Constant(condition.size(), T{}));
+ // Transform the output to the correct value from LookupTable
+ if (!is_copy) {
+ auto* look_up_table = user_data + 2;
+ std::transform(condition.cbegin(), condition.cend(), value.cbegin(), output.begin(),
+ [target, &look_up_table](bool condition_element, const T& value_element) {
+ return (condition_element == target) ? look_up_table[value_element] : T{};
+ });
+ }
+ }};
+}
+
+template
+ProcessBroadcastSpanFuncs CreateNonScalarBroadcastFuncs() {
+ return ProcessBroadcastSpanFuncs{
+ [](BroadcastHelper& per_iter_bh) {
+ auto* user_data = static_cast(per_iter_bh.GetUserData());
+ bool target = user_data[0] == 1;
+ bool is_copy = user_data[1] == 1;
+ auto* look_up_table = user_data + 2;
+ bool condition = per_iter_bh.ScalarInput0();
+ auto value = per_iter_bh.SpanInput1();
+ auto output = per_iter_bh.OutputSpan();
+ if (condition == target) {
+ // Transform the output to the correct value from LookupTable
+ std::transform(value.cbegin(), value.cend(), output.begin(),
+ [condition, target, &look_up_table,is_copy](const T& value_element) {
+ return is_copy ? value_element : look_up_table[value_element];
+ });
+ } else {
+ std::fill(output.begin(), output.end(), T{});
+ }
+ },
+ [](BroadcastHelper& per_iter_bh) {
+ auto* user_data = static_cast(per_iter_bh.GetUserData());
+ bool target = user_data[0] == 1;
+ bool is_copy = user_data[1] == 1;
+ auto* look_up_table = user_data + 2;
+ auto condition = per_iter_bh.SpanInput0();
+ const T& value = per_iter_bh.ScalarInput1();
+ auto output = per_iter_bh.OutputSpan();
+ // Transform the output to the correct value from LookupTable
+ std::transform(condition.begin(), condition.end(), output.begin(),
+ [target, &value,&look_up_table,is_copy](bool condition_element) {
+ return condition_element == target ? is_copy ? value : look_up_table[value] : T{};
+ });
+ },
+ [](BroadcastHelper& per_iter_bh) {
+ auto* user_data = static_cast(per_iter_bh.GetUserData());
+ bool target = user_data[0] == 1;
+ bool is_copy = user_data[1] == 1;
+ auto* look_up_table = user_data + 2;
+ auto condition = per_iter_bh.SpanInput0();
+ auto value = per_iter_bh.SpanInput1();
+ auto output = per_iter_bh.OutputSpan();
+ // Transform the output to the correct value from LookupTable
+ std::transform(condition.begin(), condition.end(), value.cbegin(), output.begin(),
+ [target,&look_up_table,is_copy](bool condition_element, const T& value_element) {
+ return condition_element == target ? is_copy ? value_element : look_up_table[value_element] : T{};
+ });
+ }};
+}
+
+template
+EnableIfEigenScalar SelectBroadcastFuncs() {
+ // NOTE: Workaround a VS2017 bug by calling a separate function to create the broadcast funcs.
+ // If we create them directly here it doesn't bring in the definitions of the Eigen classes leading to
+ // a 'class has no constructors' error
+ return CreateScalarBroadcastFuncs();
+}
+
+template
+EnableIfEigenNotScalar SelectBroadcastFuncs() {
+ return CreateNonScalarBroadcastFuncs();
+}
+// function pointer to create typed tensor from type agnostic code whilst avoiding the overhead of std::function
+using AllocTensorFunc = std::unique_ptr (*)(const TensorAllocator& allocator, const TensorShape& shape);
+
+static std::unique_ptr UntypedSelect(OpKernelContext& context, std::vector& user_data, const ProcessBroadcastSpanFuncs& functors, const TensorAllocator& allocator, AllocTensorFunc allocate_tensor) {
+ const auto& condition = *context.Input(0);
+ // select the X input (input 1) for 'true', and Y input (input 2) for 'false'
+ bool target = user_data[0] == 1;
+ const auto& values = *context.Input(target ? 1 : 4);
+
+ InputBroadcaster input_broadcaster(condition, values);
+ std::unique_ptr selection_tensor = allocate_tensor(allocator, input_broadcaster.GetOutputShape());
+ OutputBroadcaster output_broadcaster(input_broadcaster.GetSpanSize(), *selection_tensor);
+ // store value of 'target' directly in void* for user_data so it's accessible in the state-less functors
+ BroadcastHelper broadcast_helper(input_broadcaster, output_broadcaster, reinterpret_cast(user_data.data()));
+ BroadcastLooper(broadcast_helper, functors);
+ return selection_tensor;
+}
+
+// Merging Functions
+template
+void MergeScalarAndVector(EigenVectorMap output, const T& scalar_value, ConstEigenVectorMap vector_value) {
+ if (scalar_value != T{}) {
+ output = EigenVectorMap::PlainObject::Constant(vector_value.size(), scalar_value);
+ } else {
+ output = vector_value;
+ }
+};
+
+template
+EnableIfEigenScalar MergeBroadcastFuncs() {
+ return ProcessBroadcastSpanFuncs{
+ [](BroadcastHelper& per_iter_bh) {
+ MergeScalarAndVector(per_iter_bh.OutputEigen(),
+ per_iter_bh.ScalarInput0(), // X selection
+ per_iter_bh.EigenInput1()); // Y selection
+ },
+ [](BroadcastHelper& per_iter_bh) {
+ MergeScalarAndVector(per_iter_bh.OutputEigen(),
+ per_iter_bh.ScalarInput1(), // Y selection
+ per_iter_bh.EigenInput0()); // X selection
+ },
+ [](BroadcastHelper& per_iter_bh) {
+ auto X_selection = per_iter_bh.EigenInput0();
+ auto Y_selection = per_iter_bh.EigenInput1();
+ per_iter_bh.OutputEigen() = X_selection.binaryExpr(Y_selection,
+ [](T x, T y) -> T {
+ return x != T{} ? x : y;
+ });
+ }};
+}
+
+template
+void MergeScalarAndVector(gsl::span output, const T& scalar_value, gsl::span vector_value) {
+ if (!scalar_value.empty()) {
+ std::fill(output.begin(), output.end(), scalar_value);
+ } else {
+ std::copy(vector_value.cbegin(), vector_value.cend(), output.begin());
+ }
+};
+
+template
+
+EnableIfEigenNotScalar MergeBroadcastFuncs() {
+ return ProcessBroadcastSpanFuncs{
+ [](BroadcastHelper& per_iter_bh) {
+ MergeScalarAndVector(per_iter_bh.OutputSpan(),
+ per_iter_bh.ScalarInput0(), // X selection
+ per_iter_bh.SpanInput1()); // Y selection
+ },
+ [](BroadcastHelper& per_iter_bh) {
+ MergeScalarAndVector(per_iter_bh.OutputSpan(),
+ per_iter_bh.ScalarInput1(), // Y selection
+ per_iter_bh.SpanInput0()); // X selection
+ },
+ [](BroadcastHelper& per_iter_bh) {
+ auto X_selection = per_iter_bh.SpanInput0();
+ auto Y_selection = per_iter_bh.SpanInput1();
+ auto output = per_iter_bh.OutputSpan();
+ std::transform(X_selection.cbegin(), X_selection.cend(), Y_selection.cbegin(), output.begin(),
+ [](const T& x, const T& y) { return !x.empty() ? x : y; });
+ }};
+}
+static void UntypedMerge(OpKernelContext& context,
+ const Tensor& X_selection_tensor, const Tensor& Y_selection_tensor,
+ const ProcessBroadcastSpanFuncs& functors) {
+ InputBroadcaster merge_broadcaster{X_selection_tensor, Y_selection_tensor};
+ Tensor& output = *context.Output(0, merge_broadcaster.GetOutputShape());
+
+ OutputBroadcaster output_broadcaster{merge_broadcaster.GetSpanSize(), output};
+ BroadcastHelper broadcast_helper(merge_broadcaster, output_broadcaster);
+
+ BroadcastLooper(broadcast_helper, functors);
+}
+} // namespace
+
+namespace contrib {
+
+QLinearWhere::QLinearWhere(const OpKernelInfo& info) : OpKernel(info) {
+ size_t input_def_count = info.node().InputDefs().size();
+ ORT_ENFORCE(input_def_count == kExpected_input_count,
+ "There must be ", kExpected_input_count, " inputs! (condition, x, x_scale, x_zero_point, y, y_scale, y_zero_point, z_scale, z_zero_point)");
+ const Tensor* tensor_x_scale = nullptr;
+ const Tensor* tensor_x_zero_point = nullptr;
+ const Tensor* tensor_y_scale = nullptr;
+ const Tensor* tensor_y_zero_point = nullptr;
+ const Tensor* tensor_z_scale = nullptr;
+ const Tensor* tensor_z_zero_point = nullptr;
+
+ bool get_x_scale = info.TryGetConstantInput(2, &tensor_x_scale);
+ bool get_x_zero_point = info.TryGetConstantInput(3, &tensor_x_zero_point);
+ bool get_y_scale = info.TryGetConstantInput(5, &tensor_y_scale);
+ bool get_y_zero_point = info.TryGetConstantInput(6, &tensor_y_zero_point);
+ bool get_z_scale = info.TryGetConstantInput(7, &tensor_z_scale);
+ bool get_z_zero_point = info.TryGetConstantInput(8, &tensor_z_zero_point);
+ if (!get_z_scale || !get_z_zero_point) {
+ // Can not build fix lookup table
+ return;
+ }
+ ORT_ENFORCE(
+ tensor_x_zero_point->GetElementType() == tensor_y_zero_point->GetElementType() &&
+ tensor_x_zero_point->GetElementType() == tensor_z_zero_point->GetElementType() &&
+ tensor_y_zero_point->GetElementType() == tensor_z_zero_point->GetElementType(),
+ "Wrong input type encountered for zero point input def of x, y, z");
+ bool is_signed_int8 = tensor_z_zero_point->IsDataType();
+ const auto identity_float = [](float v) -> float { return v; };
+
+ if (get_x_scale && get_x_zero_point) {
+ // Build fix lookup table for x
+ is_x_fixed_copy_ = has_same_scale(tensor_x_scale, tensor_z_scale) &&
+ has_same_zero_point(is_signed_int8, tensor_x_zero_point, tensor_z_zero_point);
+ if (!is_x_fixed_copy_) {
+ x_fixed_lookup_table_.resize(256);
+ if (is_signed_int8) {
+ QlinearBuildLookupTable(
+ x_fixed_lookup_table_.data(), tensor_x_scale, tensor_x_zero_point,
+ tensor_z_scale, tensor_z_zero_point, identity_float);
+ } else {
+ QlinearBuildLookupTable(
+ x_fixed_lookup_table_.data(), tensor_x_scale, tensor_x_zero_point,
+ tensor_z_scale, tensor_z_zero_point, identity_float);
+ }
+ }
+ is_x_dynamic_ = false;
+ }
+
+ if (get_y_scale && get_y_zero_point) {
+ // Build fix lookup table for y
+ is_y_fixed_copy_ = has_same_scale(tensor_y_scale, tensor_z_scale) &&
+ has_same_zero_point(is_signed_int8, tensor_y_zero_point, tensor_z_zero_point);
+ if (!is_y_fixed_copy_) {
+ y_fixed_lookup_table_.resize(256);
+ if (is_signed_int8) {
+ QlinearBuildLookupTable(
+ y_fixed_lookup_table_.data(), tensor_y_scale, tensor_y_zero_point,
+ tensor_z_scale, tensor_z_zero_point, identity_float);
+ } else {
+ QlinearBuildLookupTable(
+ y_fixed_lookup_table_.data(), tensor_y_scale, tensor_y_zero_point,
+ tensor_z_scale, tensor_z_zero_point, identity_float);
+ }
+ }
+ is_y_dynamic_ = false;
+ }
+}
+
+Status QLinearWhere::Compute(OpKernelContext* ctx) const {
+// const auto* tensor_condition = ctx->Input(0);
+// const auto* tensor_x_input = ctx->Input(1);
+ const auto* tensor_x_scale = ctx->Input(2);
+ const auto* tensor_x_zero_point = ctx->Input(3);
+// const auto* tensor_y_input = ctx->Input(4);
+ const auto* tensor_y_scale = ctx->Input(5);
+ const auto* tensor_y_zero_point = ctx->Input(6);
+ const auto* tensor_z_scale = ctx->Input(7);
+ const auto* tensor_z_zero_point = ctx->Input(8);
+// auto* tensor_output = ctx->Output(0, tensor_condition->Shape());
+ ORT_ENFORCE(tensor_x_scale->IsDataType(), "Input scale is not float for quantized input x @ 2");
+ ORT_ENFORCE(tensor_y_scale->IsDataType(), "Input scale is not float for quantized input y @ 5");
+ ORT_ENFORCE(tensor_z_scale->IsDataType(), "Input scale is not float for quantized output z @ 7");
+ ORT_ENFORCE(tensor_x_zero_point->GetElementType() == tensor_y_zero_point->GetElementType() &&
+ tensor_x_zero_point->GetElementType() == tensor_z_zero_point->GetElementType() &&
+ tensor_y_zero_point->GetElementType() == tensor_z_zero_point->GetElementType(),
+ "Wrong input type encountered for zero point of quantized input @", 3, 6, 8);
+ bool is_signed_int8 = tensor_z_zero_point->IsDataType();
+ const auto identity_float = [](float v) -> float { return v; };
+
+ std::vector x_dynamic_lookup_table;
+ bool is_x_copy = !is_x_dynamic_ ? is_x_fixed_copy_ : has_same_scale(tensor_x_scale, tensor_z_scale) && has_same_zero_point(is_signed_int8, tensor_x_zero_point, tensor_z_zero_point);
+ if (is_x_dynamic_ && !is_x_copy) {
+ x_dynamic_lookup_table.resize(256);
+ if (is_signed_int8) {
+ QlinearBuildLookupTable(
+ x_dynamic_lookup_table.data(), tensor_x_scale, tensor_x_zero_point,
+ tensor_z_scale, tensor_z_zero_point, identity_float);
+ } else {
+ QlinearBuildLookupTable(
+ x_dynamic_lookup_table.data(), tensor_x_scale, tensor_x_zero_point,
+ tensor_z_scale, tensor_z_zero_point, identity_float);
+ }
+ }
+
+ // Build dynamic lookup table for y
+ std::vector y_dynamic_lookup_table;
+ bool is_y_copy = !is_y_dynamic_ ? is_y_fixed_copy_ : has_same_scale(tensor_y_scale, tensor_z_scale) && has_same_zero_point(is_signed_int8, tensor_y_zero_point, tensor_z_zero_point);
+ if (is_y_dynamic_ && !is_y_copy) {
+ y_dynamic_lookup_table.resize(256);
+ if (is_signed_int8) {
+ QlinearBuildLookupTable(
+ y_dynamic_lookup_table.data(), tensor_y_scale, tensor_y_zero_point,
+ tensor_z_scale, tensor_z_zero_point, identity_float);
+ } else {
+ QlinearBuildLookupTable(
+ y_dynamic_lookup_table.data(), tensor_y_scale, tensor_y_zero_point,
+ tensor_z_scale, tensor_z_zero_point, identity_float);
+ }
+ }
+ // Each lookup table is 256 bytes
+ const auto& x_lookup_table = is_x_dynamic_ ? x_dynamic_lookup_table : x_fixed_lookup_table_;
+ const auto& y_lookup_table = is_y_dynamic_ ? y_dynamic_lookup_table : y_fixed_lookup_table_;
+
+ // Each user_data will sized at 2 + 256 bytes, and contain {is_x/y, is_x/y_copy, x/y_lookup_table} respectively
+ std::vector x_user_data(258);
+ std::vector y_user_data(258);
+ x_user_data[0] = 1;
+ y_user_data[0] = 0;
+ x_user_data[1] = is_x_copy ? 1 : 0;
+ y_user_data[1] = is_y_copy ? 1 : 0;
+ if (!is_x_copy) {
+ std::copy(x_lookup_table.begin(), x_lookup_table.end(), x_user_data.begin() + 2);
+ }
+ if (!is_y_copy) {
+ std::copy(y_lookup_table.begin(), y_lookup_table.end(), y_user_data.begin() + 2);
+ }
+ //Allocator, Allocation, and SelectBroadcastFuncs are the same implementation from where_op.cc
+ const auto typed_tensor_allocation = [](const TensorAllocator& allocator,
+ const TensorShape& shape) {
+ return allocator.Allocate(shape);
+ };
+ TensorAllocator tensor_allocator{*ctx};
+ ProcessBroadcastSpanFuncs funcs = SelectBroadcastFuncs();
+ // UntypedSelect is MODIFIED from where_op.cc
+ auto X_selection_tensor = UntypedSelect(*ctx, x_user_data, funcs, tensor_allocator, typed_tensor_allocation);
+ auto Y_selection_tensor = UntypedSelect(*ctx, y_user_data, funcs, tensor_allocator, typed_tensor_allocation);
+ // UntypedMerge is the same as from where_op.cc
+ UntypedMerge(*ctx, *X_selection_tensor, *Y_selection_tensor, MergeBroadcastFuncs());
+
+ return Status();
+}
+
+ONNX_CPU_OPERATOR_MS_KERNEL(
+ QLinearWhere,
+ 1,
+ KernelDefBuilder().TypeConstraint(
+ "T",
+ {DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}),
+ QLinearWhere)
+} // namespace contrib
+} // namespace onnxruntime
\ No newline at end of file
diff --git a/onnxruntime/contrib_ops/cpu/quantization/qlinear_where.h b/onnxruntime/contrib_ops/cpu/quantization/qlinear_where.h
new file mode 100644
index 0000000000..1c70538f0f
--- /dev/null
+++ b/onnxruntime/contrib_ops/cpu/quantization/qlinear_where.h
@@ -0,0 +1,32 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "qlinear_util.h"
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+#include "core/util/math_cpuonly.h"
+#include "core/providers/cpu/tensor/where_op.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+class QLinearWhere final : public OpKernel {
+ public:
+ QLinearWhere(const OpKernelInfo& info);
+
+ Status Compute(OpKernelContext* context) const override;
+
+ private:
+ constexpr static size_t kExpected_input_count = 9;
+ std::vector y_fixed_lookup_table_;
+ std::vector x_fixed_lookup_table_;
+ bool is_x_dynamic_ = true;
+ bool is_y_dynamic_ = true;
+ bool is_x_fixed_copy_ = false;
+ bool is_y_fixed_copy_ = false;
+};
+
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h
index 9446d30536..ff32e568e0 100644
--- a/onnxruntime/core/graph/contrib_ops/ms_opset.h
+++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h
@@ -26,6 +26,7 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QEmbedLayerNormalization
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QGemm);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QLinearAdd);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QLinearConcat);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QLinearWhere);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QLinearLeakyRelu);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QLinearMul);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QLinearReduceMean);
@@ -108,6 +109,7 @@ class OpSet_Microsoft_ver1 {
fn(GetOpSchema());
fn(GetOpSchema());
fn(GetOpSchema());
+ fn(GetOpSchema());
fn(GetOpSchema());
fn(GetOpSchema());
fn(GetOpSchema());
diff --git a/onnxruntime/core/graph/contrib_ops/quantization_defs.cc b/onnxruntime/core/graph/contrib_ops/quantization_defs.cc
index a2949498fb..60b0ef95cd 100644
--- a/onnxruntime/core/graph/contrib_ops/quantization_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/quantization_defs.cc
@@ -9,6 +9,7 @@
#include "core/graph/constants.h"
#include "core/graph/contrib_ops/contrib_defs.h"
#include "core/graph/contrib_ops/shape_inference_functions.h"
+#include "onnx/onnx-ml.pb.h" // ?
// Suppress a warning: global initializer calls a non-constexpr function 'symbol' which is from
// ONNX_OPERATOR_SET_SCHEMA_EX macro and only happens in debug build
@@ -817,10 +818,46 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
}
}
- if (all_lengths_known) {
- output_shape->mutable_dim(axis)->set_dim_value(total_length);
- }
- }));
+ if (all_lengths_known) {
+ output_shape->mutable_dim(axis)->set_dim_value(total_length);
+ }
+ }));
+
+ ONNX_MS_OPERATOR_SET_SCHEMA(QLinearWhere, 1, OpSchema()
+ .SetDoc("Return elements, either from X or Y, depending on condition.")
+ .Input(0, "condition", " When True (nonzero), yield x, otherwise yield y", "B")
+ .Input(1, "X", "Y's zero point.", "T")
+ .Input(2, "x_scale", "X's scale.", "TF")
+ .Input(3, "x_zero_point", "X's zero point.", "T")
+ .Input(4, "Y", "Y's zero point.", "T")
+ .Input(5, "y_scale", "Y's scale.", "TF")
+ .Input(6, "y_zero_point", "Y's zero point.", "T")
+ .Input(7, "z_scale", "Z's scale.", "TF")
+ .Input(8, "z_zero_point", "Z's zero point.", "T")
+ .Output(0, "Z", "Tensor of shape equal to the broadcasted shape of condition, X, and Y", "T")
+ .TypeConstraint(
+ "B",
+ {"tensor(bool)"},
+ "Constrain input and output types to 8 bit signed and unsigned tensors.")
+ .TypeConstraint(
+ "TF",
+ {"tensor(float)"},
+ "Constrain scale types to any float tensor type.")
+ .TypeConstraint(
+ "T",
+ {"tensor(uint8)", "tensor(int8)"},
+ "Constrain input and output types to 8 bit signed and unsigned tensors.")
+ .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+ propagateElemTypeFromInputToOutput(ctx, 1, 0);
+ if (hasNInputShapes(ctx, 9)) {
+ std::vector shapes;
+ shapes.push_back(&ctx.getInputType(0)->tensor_type().shape());
+ shapes.push_back(&ctx.getInputType(1)->tensor_type().shape());
+ shapes.push_back(&ctx.getInputType(4)->tensor_type().shape());
+ multidirectionalBroadcastShapeInference(
+ shapes, *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape());
+ }
+ }));
ONNX_MS_OPERATOR_SET_SCHEMA(
QGemm, 1,
diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc
index d7407fcf20..eca8183875 100644
--- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc
+++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc
@@ -72,7 +72,22 @@ std::vector ConvMoves() {
return moves;
}
+std::vector WhereMoves(){
+ NTO::NodeLocation dq_x{NTO::NodeType::kInput, 0};
+ NTO::NodeLocation dq_y{NTO::NodeType::kInput, 1};
+ NTO::NodeLocation target{NTO::NodeType::kTarget, 0};
+ NTO::NodeLocation q{NTO::NodeType::kOutput, 0};
+ std::vector moves{
+ MoveAndAppend(target,ArgType::kInput, 0, ArgType::kInput), // move the condition to the new node
+ MoveAll(dq_x, ArgType::kInput), // append all inputs from x
+ MoveAll(dq_y, ArgType::kInput), // append all inputs from x
+ MoveAndAppend(q, ArgType::kInput, 1, ArgType::kInput), // append scale (input 1) from q
+ MoveAndAppend(q, ArgType::kInput, 2, ArgType::kInput), // append zp (input 2) from q
+ MoveAll(q, ArgType::kOutput)
+ };
+ return moves;
+}
QDQReplaceWithNew SplitReplacer() {
NTO::NodeLocation dq{NTO::NodeType::kInput, 0};
NTO::NodeLocation q{NTO::NodeType::kOutput, 0};
@@ -224,7 +239,9 @@ VariadicReplaceWithQLinear::VariadicReplaceWithQLinear(std::string domain)
ConvReplaceWithQLinear::ConvReplaceWithQLinear()
: ReplaceWithQLinear(kOnnxDomain, ConvMoves()) {
}
-
+WhereReplaceWithQLinear::WhereReplaceWithQLinear()
+ : ReplaceWithQLinear(kMSDomain, WhereMoves()) {
+}
MatMulReplaceWithQLinear::MatMulReplaceWithQLinear()
: matmul_int_to_float_replacer_{MatMulIntToFloatReplacer()},
qlinear_matmul_replacer_{kOnnxDomain} {
diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.h b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.h
index 0337f11306..8179a03050 100644
--- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.h
+++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.h
@@ -59,7 +59,9 @@ struct VariadicReplaceWithQLinear : ReplaceWithQLinear {
struct ConvReplaceWithQLinear : ReplaceWithQLinear {
ConvReplaceWithQLinear();
};
-
+struct WhereReplaceWithQLinear : ReplaceWithQLinear {
+ WhereReplaceWithQLinear();
+};
struct SplitReplaceWithQuant : public Action {
Status Run(Graph&, const NodesToOptimize& selected_nodes) const override;
};
diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.cc
index 54ba592144..b3b73b66fb 100644
--- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.cc
+++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.cc
@@ -206,6 +206,26 @@ void GemmQDQRules(SelectorActionRegistry& qdq_selector_action_registry) {
#endif
}
+void WhereQDQRules (SelectorActionRegistry& qdq_selector_action_registry) {
+ // 3 nodes. 2 x DQ for inputs and 1X Q for output
+ // Compare to other BinaryOperators (Add, Mul), Where also have a special case that it has boolean input
+ // Where Replace with QLinearWhere
+ // Delete all original nodes.
+ const std::string action_name{"Where"};
+ std::unique_ptr action = std::make_unique();
+
+#if !defined(ORT_MINIMAL_BUILD)
+ std::unique_ptr selector = std::make_unique();
+ qdq_selector_action_registry.RegisterSelectorAndAction(action_name,
+ {{"Where", {}}},
+ std::move(selector),
+ std::move(action));
+
+#else
+ qdq_selector_action_registry.RegisterAction(action_name, std::move(action));
+#endif
+}
+
SelectorActionRegistry CreateSelectorActionRegistry(bool is_int8_allowed) {
SelectorActionRegistry qdq_selector_action_registry;
SplitQDQRules(qdq_selector_action_registry);
@@ -217,6 +237,7 @@ SelectorActionRegistry CreateSelectorActionRegistry(bool is_int8_allowed) {
ConvQDQRules(qdq_selector_action_registry, is_int8_allowed);
MatMulQDQRules(qdq_selector_action_registry, is_int8_allowed);
GemmQDQRules(qdq_selector_action_registry);
+ WhereQDQRules(qdq_selector_action_registry);
return qdq_selector_action_registry;
}
diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc
index 6b588314c4..838a03ca5f 100644
--- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc
+++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc
@@ -309,6 +309,22 @@ void GemmSelector::UpdateBuilder(NodesToOptimizeIndicesBuilder& builder) const {
builder.input_nodes.resize(3, NodesToOptimizeIndices::kEmptyNodeIndex);
}
+bool WhereNodeGroupSelector::Check(const GraphViewer &graph_viewer, const Node &node,
+ const std::vector &dq_nodes,
+ const std::vector &q_nodes) const {
+ // Where has 1 boolean input and 2 dq inputs
+ if (!CheckQDQNodes(graph_viewer, node, dq_nodes, q_nodes,2)) {
+ return false;
+ }
+
+ const int32_t dt_input_1 = dq_nodes[0]->InputDefs()[0]->TypeAsProto()->tensor_type().elem_type();
+ const int32_t dt_input_2 = dq_nodes[1]->InputDefs()[0]->TypeAsProto()->tensor_type().elem_type();
+ const int32_t dt_output = q_nodes[0]->OutputDefs()[0]->TypeAsProto()->tensor_type().elem_type();
+ return dt_input_1 == dt_input_2 &&
+ dt_input_1 == dt_output;
+
+}
+
} // namespace QDQ
} // namespace onnxruntime
diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h
index 4f52f8493b..ec6c16d6a7 100644
--- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h
+++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h
@@ -103,6 +103,17 @@ class ConvNodeGroupSelector : public NodeGroupSelector {
bool int8_allowed_;
};
+class WhereNodeGroupSelector : public NodeGroupSelector {
+public:
+ WhereNodeGroupSelector()= default;
+
+private:
+ bool Check(const GraphViewer& graph_viewer, const Node& node,
+ const std::vector& dq_nodes,
+ const std::vector& q_nodes) const override;
+
+};
+
// 2 DQ nodes for input -> node -> optional Q if QLinearMatMul, MatMulIntegerToFloat if not
// The lack of a trailing Q isn't really a QDQ node group, so we default support for that to off.
class MatMulNodeGroupSelector : public NodeGroupSelector {
@@ -200,7 +211,10 @@ class ConvSelector : public BaseSelector {
void UpdateBuilder(NodesToOptimizeIndicesBuilder&) const override;
};
-
+class WhereSelector : public BaseSelector {
+public:
+ WhereSelector() : BaseSelector(std::make_unique()) {}
+};
// 2 DQ nodes for input -> node -> optional Q if QLinearMatMul, MatMulIntegerToFloat if not
class MatMulSelector : public BaseSelector {
public:
diff --git a/onnxruntime/python/tools/quantization/onnx_quantizer.py b/onnxruntime/python/tools/quantization/onnx_quantizer.py
index 3a5cb95276..b1cbc88339 100644
--- a/onnxruntime/python/tools/quantization/onnx_quantizer.py
+++ b/onnxruntime/python/tools/quantization/onnx_quantizer.py
@@ -1057,7 +1057,6 @@ class ONNXQuantizer:
if node.input[0] not in self.tensors_range.keys() or node.output[0] not in self.tensors_range.keys():
continue
self.tensors_range[node.input[0]] = self.tensors_range[node.output[0]]
-
quantization_params = {}
for tensor_name in self.tensors_range.keys():
rmin, rmax = self.tensors_range[tensor_name]
diff --git a/onnxruntime/python/tools/quantization/operators/where.py b/onnxruntime/python/tools/quantization/operators/where.py
new file mode 100644
index 0000000000..3843af408e
--- /dev/null
+++ b/onnxruntime/python/tools/quantization/operators/where.py
@@ -0,0 +1,87 @@
+import onnx
+
+from ..quant_utils import TENSOR_NAME_QUANT_SUFFIX, QuantizedValue, QuantizedValueType, attribute_to_kwarg, ms_domain
+from .base_operator import QuantOperatorBase
+from .qdq_base_operator import QDQOperatorBase
+
+
+class QLinearWhere(QuantOperatorBase):
+ def should_quantize(self):
+ return True
+
+ def quantize(self):
+ node = self.node
+ assert node.op_type == "Where"
+ if not self.quantizer.force_quantize_no_input_check:
+ self.quantizer.new_nodes += [node]
+ return
+ (
+ data_found,
+ output_scale_name,
+ output_zp_name,
+ _,
+ _,
+ ) = self.quantizer._get_quantization_params(node.output[0])
+ (
+ q_input_names,
+ zero_point_names,
+ scale_names,
+ nodes,
+ ) = self.quantizer.quantize_activation(node, [1, 2])
+ if not data_found or q_input_names is None:
+ return super().quantize()
+ qlinear_output = node.output[0] + TENSOR_NAME_QUANT_SUFFIX
+ qlinear_output_name = node.name + "_quant" if node.name != "" else ""
+
+ q_output = QuantizedValue(
+ node.output[0],
+ qlinear_output,
+ output_scale_name,
+ output_zp_name,
+ QuantizedValueType.Input,
+ )
+ self.quantizer.quantized_value_map[node.output[0]] = q_output
+
+ kwargs = {}
+ for attribute in node.attribute:
+ kwargs.update(attribute_to_kwarg(attribute))
+ kwargs["domain"] = ms_domain
+
+ qlwhere_inputs = [
+ node.input[0],
+ q_input_names[0],
+ scale_names[0],
+ zero_point_names[0],
+ q_input_names[1],
+ scale_names[1],
+ zero_point_names[1],
+ output_scale_name,
+ output_zp_name,
+ ]
+ qlwhere_node = onnx.helper.make_node(
+ "QLinearWhere", qlwhere_inputs, [qlinear_output], qlinear_output_name, **kwargs
+ )
+
+ self.quantizer.new_nodes += nodes
+ self.quantizer.new_nodes += [qlwhere_node]
+
+
+class QDQWhere(QDQOperatorBase):
+ def quantize(self):
+ node = self.node
+ assert node.op_type == "Where"
+ if self.quantizer.force_quantize_no_input_check:
+ if not self.quantizer.is_tensor_quantized(node.input[1]):
+ self.quantizer.quantize_activation_tensor(node.input[1])
+ if not self.quantizer.is_tensor_quantized(node.input[2]):
+ self.quantizer.quantize_activation_tensor(node.input[2])
+ if not self.disable_qdq_for_node_output:
+ for output in node.output:
+ self.quantizer.quantize_activation_tensor(output)
+ elif (
+ self.quantizer.is_tensor_quantized(node.input[1])
+ and self.quantizer.is_tensor_quantized(node.input[2])
+ and not self.disable_qdq_for_node_output
+ ):
+ for output in node.output:
+ self.quantizer.quantize_activation_tensor(output)
diff --git a/onnxruntime/python/tools/quantization/registry.py b/onnxruntime/python/tools/quantization/registry.py
index 0227da5e24..d77f4a3f9a 100644
--- a/onnxruntime/python/tools/quantization/registry.py
+++ b/onnxruntime/python/tools/quantization/registry.py
@@ -19,6 +19,7 @@ from .operators.qdq_base_operator import QDQOperatorBase
from .operators.resize import QDQResize, QResize
from .operators.softmax import QDQSoftmax, QLinearSoftmax
from .operators.split import QDQSplit, QSplit
+from .operators.where import QDQWhere, QLinearWhere
from .quant_utils import QuantizationMode
CommonOpsRegistry = {
@@ -57,6 +58,7 @@ QLinearOpsRegistry = {
"AveragePool": QLinearPool,
"Concat": QLinearConcat,
"Softmax": QLinearSoftmax,
+ "Where": QLinearWhere,
}
QLinearOpsRegistry.update(CommonOpsRegistry)
@@ -76,6 +78,7 @@ QDQRegistry = {
"Split": QDQSplit,
"Gather": QDQGather,
"Softmax": QDQSoftmax,
+ "Where": QDQWhere,
}
diff --git a/onnxruntime/test/contrib_ops/qlinear_where_test.cc b/onnxruntime/test/contrib_ops/qlinear_where_test.cc
new file mode 100644
index 0000000000..0240abce24
--- /dev/null
+++ b/onnxruntime/test/contrib_ops/qlinear_where_test.cc
@@ -0,0 +1,126 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+#include "gtest/gtest.h"
+#include "test/providers/provider_test_utils.h"
+#include
+
+namespace onnxruntime {
+namespace test {
+
+template
+void RunQLinearWhere(
+ OpTester &test,
+ const std::vector &x,
+ const std::vector &y,
+ const std::vector &z,
+ float x_scale,
+ T x_zero_point,
+ float y_scale,
+ T y_zero_point,
+ float z_scale,
+ T z_zero_point,
+ const std::vector &x_shape,
+ const std::vector &y_shape,
+ const std::vector &z_shape) {
+
+ test.AddInput("X", x_shape, x);
+ test.AddInput("x_scale", {}, {x_scale}, true);
+ test.AddInput("x_zero_point", {}, {x_zero_point}, true);
+ test.AddInput("Y", y_shape, y);
+ test.AddInput("y_scale", {}, {y_scale}, true);
+ test.AddInput("y_zero_point", {}, {y_zero_point}, true);
+ test.AddInput("z_scale", {}, {z_scale});
+ test.AddInput("z_zero_point", {}, {z_zero_point});
+ test.AddOutput("Z", z_shape, z);
+ test.Run(OpTester::ExpectResult::kExpectSuccess);
+}
+template
+void QLinearWhereScalarAll() {
+ constexpr float scale = 0.039f;
+ constexpr uint8_t zp = 135;
+ OpTester test("QLinearWhere", 1, onnxruntime::kMSDomain);
+ test.AddInput("condition", {1}, {true}, true );
+ RunQLinearWhere(
+ test,
+ {1}, {2}, {1},// x ,y ,z
+ scale, zp, scale, zp, scale, zp, // x_scale, x_zp, y_scale, y_zp, z_scale, z_zp
+ {1}, {1}, {1});
+}
+
+TEST(QLinearWhereTest, QLinearWhereScalarAll) {
+ QLinearWhereScalarAll();
+ QLinearWhereScalarAll();
+}
+
+template
+void QLinearWhereVectorAll() {
+ constexpr float scale = 0.039f;
+ constexpr uint8_t zp = 135;
+ OpTester test("QLinearWhere", 1, onnxruntime::kMSDomain);
+ test.AddInput("condition", {4}, {true,false,false,true}, true );
+ RunQLinearWhere(
+ test,
+ {1,1,1,1}, {2,2,2,2}, {1,2,2,1},// x ,y ,z
+ scale, zp, scale, zp, scale, zp, // x_scale, x_zp, y_scale, y_zp, z_scale, z_zp
+ {4}, {4}, {4});
+}
+
+TEST(QLinearWhereTest,QLinearWhereVectorAll){
+ QLinearWhereVectorAll();
+ QLinearWhereVectorAll();
+}
+template
+void QLinearWhereMatrixAll() {
+ constexpr float scale = 0.039f;
+ constexpr uint8_t zp = 135;
+ OpTester test("QLinearWhere", 1, onnxruntime::kMSDomain);
+ test.AddInput("condition", {2,2}, {true,false,false,true}, true );
+ RunQLinearWhere(
+ test,
+ {1,1,1,1}, {2,2,2,2}, {1,2,2,1},// x ,y ,z
+ scale, zp, scale, zp, scale, zp, // x_scale, x_zp, y_scale, y_zp, z_scale, z_zp
+ {2,2}, {2,2}, {2,2});
+}
+
+TEST(QLinearWhereTest,QLinearWhereMatrixAll){
+ QLinearWhereMatrixAll();
+ QLinearWhereMatrixAll();
+}
+template
+void QLinearWhereScalarX_VectorY_MatrixCondition() {
+ constexpr float scale = 0.039f;
+ constexpr uint8_t zp = 135;
+ OpTester test("QLinearWhere", 1, onnxruntime::kMSDomain);
+ test.AddInput("condition", {2,2}, {true,false,false,true}, true );
+ RunQLinearWhere(
+ test,
+ {1}, {2,2}, {1,2,2,1},// x ,y ,z
+ scale, zp, scale, zp, scale, zp, // x_scale, x_zp, y_scale, y_zp, z_scale, z_zp
+ {1}, {2}, {2,2});
+}
+
+TEST(QLinearWhereTest,QLinearWhereScalarX_VectorY_MatrixCondition){
+ QLinearWhereScalarX_VectorY_MatrixCondition();
+ QLinearWhereScalarX_VectorY_MatrixCondition();
+}
+
+template
+void QLinearWhereVectorX_VectorY_MatrixCondition() {
+ constexpr float scale = 0.039f;
+ constexpr uint8_t zp = 135;
+ OpTester test("QLinearWhere", 1, onnxruntime::kMSDomain);
+ test.AddInput("condition", {2,2}, {true,false,false,true}, true );
+ RunQLinearWhere(
+ test,
+ {1,1}, {2,2}, {1,2,2,1},// x ,y ,z
+ scale, zp, scale, zp, scale, zp, // x_scale, x_zp, y_scale, y_zp, z_scale, z_zp
+ {2}, {2}, {2,2});
+}
+
+TEST(QLinearWhereTest,QLinearWhereVectorX_VectorY_MatrixCondition){
+ QLinearWhereVectorX_VectorY_MatrixCondition();
+ QLinearWhereVectorX_VectorY_MatrixCondition();
+}
+
+} // namespace test
+} // namespace onnxruntime
\ No newline at end of file
diff --git a/onnxruntime/test/optimizer/qdq_test_utils.h b/onnxruntime/test/optimizer/qdq_test_utils.h
index 9cdb4972cf..f9a21135fa 100644
--- a/onnxruntime/test/optimizer/qdq_test_utils.h
+++ b/onnxruntime/test/optimizer/qdq_test_utils.h
@@ -389,6 +389,37 @@ GetQDQTestCaseFn BuildQDQSplitTestCase(
builder.AddQuantizeLinearNode(split_output_3, .003f, q_zp, q_split_output_3);
};
}
+template
+GetQDQTestCaseFn BuildQDQWhereTestCase(
+ const std::vector& cond_shape,
+ const std::vector& x_shape,
+ const std::vector& y_shape
+ ) {
+ return [cond_shape,x_shape,y_shape](ModelTestBuilder& builder)
+ {
+ auto* input_cond_arg = builder.MakeInputBool(cond_shape);
+ auto* input_x_arg = builder.MakeInput(x_shape,
+ std::numeric_limits::min(),
+ std::numeric_limits::max());
+ auto* input_y_arg = builder.MakeInput(y_shape,
+ std::numeric_limits::min(),
+ std::numeric_limits::max());
+ InputType zp = std::numeric_limits::max() / 2;
+ constexpr float scale = 0.003f;
+ auto* dq_x_output = builder.MakeIntermediate();
+ auto* dq_y_output = builder.MakeIntermediate();
+ builder.AddDequantizeLinearNode(input_x_arg, scale, zp, dq_x_output);
+ builder.AddDequantizeLinearNode(input_y_arg, scale, zp, dq_y_output);
+ // add Where
+
+ auto* where_output = builder.MakeIntermediate();
+ builder.AddNode("Where", {input_cond_arg,dq_x_output,dq_y_output}, {where_output});
+
+ // add Q
+ auto* q_where_output = builder.MakeOutput();
+ builder.AddQuantizeLinearNode(where_output, scale, zp, q_where_output); // Model input (node_token_1)
+ };
+}
template
GetQDQTestCaseFn BuildQDQTransposeTestCase(
diff --git a/onnxruntime/test/optimizer/qdq_transformer_test.cc b/onnxruntime/test/optimizer/qdq_transformer_test.cc
index e746b2ae34..bd3a2eae44 100644
--- a/onnxruntime/test/optimizer/qdq_transformer_test.cc
+++ b/onnxruntime/test/optimizer/qdq_transformer_test.cc
@@ -779,6 +779,21 @@ TEST(QDQTransformerTests, Split) {
};
test_case({6, 18, 54}, 0);
}
+TEST(QDQTransformerTests, Where) {
+ auto test_case = [&](const std::vector& cond_shape, const std::vector& x_shape,const std::vector& y_shape) {
+ auto check_graph = [&](InferenceSessionWrapper& session) {
+ auto op_to_count = CountOpsInGraph(session.GetGraph());
+ EXPECT_EQ(op_to_count["com.microsoft.QLinearWhere"], 1);
+ EXPECT_EQ(op_to_count["QuantizeLinear"], 0);
+ EXPECT_EQ(op_to_count["DequantizeLinear"], 0);
+ };
+ TransformerTester(BuildQDQWhereTestCase(cond_shape,x_shape,y_shape),
+ check_graph,
+ TransformerLevel::Level1,
+ TransformerLevel::Level2);
+ };
+ test_case({1},{1},{1});
+}
TEST(QDQTransformerTests, Transpose) {
auto test_case = [&](const std::vector& input_shape, const std::vector& perms) {
diff --git a/onnxruntime/test/python/quantization/test_op_where.py b/onnxruntime/test/python/quantization/test_op_where.py
new file mode 100644
index 0000000000..43d6fe4fd4
--- /dev/null
+++ b/onnxruntime/test/python/quantization/test_op_where.py
@@ -0,0 +1,155 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+
+import unittest
+
+import numpy as np
+from onnx import TensorProto, helper, save
+from op_test_utils import TestDataFeeds, check_model_correctness, check_op_type_count, check_qtype_by_node_type
+
+from onnxruntime.quantization import QuantFormat, QuantType, quantize_static
+
+
+class TestWhereModel(unittest.TestCase):
+ @staticmethod
+ def input_feeds_for_where(n, name2shape):
+ input_data_list = []
+ for i in range(n):
+ inputs = {}
+ for name, shape in name2shape.items():
+ if name == "condition":
+ inputs.update({name: np.ones(shape).astype(np.bool_)})
+ else:
+ inputs.update({name: np.ones(shape).astype(np.float32)})
+ input_data_list.extend([inputs])
+
+ dr = TestDataFeeds(input_data_list)
+ return dr
+
+ @staticmethod
+ def construct_model(model_path, input_shape):
+ initializers = []
+ input_condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, input_shape)
+ input_x = helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape)
+ input_y = helper.make_tensor_value_info("y", TensorProto.FLOAT, input_shape)
+ out_put = helper.make_tensor_value_info("z", TensorProto.FLOAT, input_shape)
+ node = helper.make_node(
+ "Where",
+ inputs=["condition", "x", "y"],
+ outputs=["z"],
+ name="where_node",
+ )
+
+ graph = helper.make_graph(
+ [node],
+ "quant_where_op_test",
+ [input_condition, input_x, input_y],
+ [out_put],
+ initializer=initializers,
+ )
+ model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 16)])
+ save(model, model_path)
+
+ def quantize_where_test(self, activation_type, weight_type, extra_options={}):
+
+ model_fp32_path = "where_fp32.onnx"
+ input_shape = [2, 2]
+ self.construct_model(model_fp32_path, input_shape)
+ data_reader = self.input_feeds_for_where(
+ 1,
+ {
+ "condition": input_shape,
+ "x": input_shape,
+ "y": input_shape,
+ },
+ )
+ activation_proto_qtype = TensorProto.UINT8 if activation_type == QuantType.QUInt8 else TensorProto.INT8
+ activation_type_str = "u8" if (activation_type == QuantType.QUInt8) else "s8"
+ weight_type_str = "u8" if (weight_type == QuantType.QUInt8) else "s8"
+ model_uint8_path = f"where_{activation_type_str}{weight_type_str}_{'QNoInCk' if extra_options['ForceQuantizeNoInputCheck'] else 'NoQNoInCk'}.onnx"
+ model_uint8_qdq_path = f"where_{activation_type_str}{weight_type_str}_{'QNoInCk' if extra_options['ForceQuantizeNoInputCheck'] else 'NoQNoInCk'}_qdq.onnx"
+
+ # Verify QOperator mode
+ data_reader.rewind()
+ quantize_static(
+ model_fp32_path,
+ model_uint8_path,
+ data_reader,
+ quant_format=QuantFormat.QOperator,
+ activation_type=activation_type,
+ weight_type=weight_type,
+ extra_options=extra_options,
+ )
+ qnode_counts = (
+ {
+ "QLinearWhere": 1,
+ "QuantizeLinear": 2,
+ "DequantizeLinear": 1,
+ }
+ if extra_options["ForceQuantizeNoInputCheck"]
+ else {
+ "Where": 1,
+ "QuantizeLinear": 0,
+ "DequantizeLinear": 0,
+ }
+ )
+ check_op_type_count(self, model_uint8_path, **qnode_counts)
+ qnode_io_qtypes = {
+ "QuantizeLinear": [
+ ["i", 2, activation_proto_qtype],
+ ["o", 0, activation_proto_qtype],
+ ]
+ }
+ check_qtype_by_node_type(self, model_uint8_path, qnode_io_qtypes)
+ data_reader.rewind()
+ check_model_correctness(self, model_fp32_path, model_uint8_path, data_reader.get_next())
+
+ # Verify QDQ mode
+ data_reader.rewind()
+ quantize_static(
+ model_fp32_path,
+ model_uint8_qdq_path,
+ data_reader,
+ quant_format=QuantFormat.QDQ,
+ activation_type=activation_type,
+ weight_type=weight_type,
+ extra_options=extra_options,
+ )
+ qdqnode_counts = (
+ {
+ "Where": 1,
+ "QuantizeLinear": 3,
+ "DequantizeLinear": 3,
+ }
+ if extra_options["ForceQuantizeNoInputCheck"]
+ else {
+ "Where": 1,
+ "QuantizeLinear": 0,
+ "DequantizeLinear": 0,
+ }
+ )
+ check_op_type_count(self, model_uint8_qdq_path, **qdqnode_counts)
+ qnode_io_qtypes = {
+ "QuantizeLinear": [
+ ["i", 2, activation_proto_qtype],
+ ["o", 0, activation_proto_qtype],
+ ]
+ }
+ check_qtype_by_node_type(self, model_uint8_qdq_path, qnode_io_qtypes)
+ data_reader.rewind()
+ check_model_correctness(self, model_fp32_path, model_uint8_qdq_path, data_reader.get_next())
+
+ def test_quantize_where_u8u8(self):
+ self.quantize_where_test(QuantType.QUInt8, QuantType.QUInt8, extra_options={"ForceQuantizeNoInputCheck": True})
+ print(__name__)
+
+ def test_quantize_where_u8u8_no_ForceQuantizeNoInputCheck(self):
+ self.quantize_where_test(QuantType.QUInt8, QuantType.QUInt8, extra_options={"ForceQuantizeNoInputCheck": False})
+ print(__name__)
+
+
+if __name__ == "__main__":
+ unittest.main()