diff --git a/onnxruntime/core/optimizer/qdq_transformer/where_dummy_dq.cc b/onnxruntime/core/optimizer/qdq_transformer/where_dummy_dq.cc index e0c080ebcac4f..9bc9fca89272e 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/where_dummy_dq.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/where_dummy_dq.cc @@ -14,112 +14,180 @@ namespace onnxruntime { bool WhereDummyDq::SatisfyCondition(const Graph& graph, const Node& node) const { + // This transformer targets a very specific pattern around `Where` when used in a QDQ graph: + // cond, DQ(xq), const_scalar -> Where -> Q(yq) + // or + // cond, const_scalar, DQ(xq) -> Where -> Q(yq) + // + // When one `Where` branch is a scalar initializer (no producer node), WhereNodeGroupSelector + // requires both data branches to be produced by DQ nodes so the `Where` can be grouped into a + // single node-unit. We insert a "dummy" DQ for the scalar branch to satisfy that requirement. if (!(node.OpType() == "Where")) { return false; } + + // ONNX Where inputs: [0]=condition, [1]=X, [2]=Y const auto& where_inputs = node.InputDefs(); + const auto& where_outputs = node.OutputDefs(); + const Node* parent_node_1 = graph.GetProducerNode(where_inputs[1]->Name()); const Node* parent_node_2 = graph.GetProducerNode(where_inputs[2]->Name()); - bool is_p1_dq = (parent_node_1 && parent_node_1->OpType() == QDQ::DQOpName); - bool is_p2_dq = (parent_node_2 && parent_node_2->OpType() == QDQ::DQOpName); + // Only apply when the `Where` output is immediately consumed by a single QuantizeLinear. + // If there are multiple consumers (or not a Q), inserting an extra DQ would not help form a + // clean QDQ node-unit and may create additional overhead. + std::vector child_nodes = graph.GetConsumerNodes(where_outputs[0]->Name()); + if (child_nodes.size() != 1 || child_nodes[0]->OpType() != QDQ::QOpName) { + return false; + } - // WhereDummyDq focus on WhereOp with one DQ input and one scalar initializer input - if (is_p1_dq && !parent_node_2) { - return (where_inputs[2]->Shape()->dim_size() == 0); + const bool is_p1_dq = (parent_node_1 && parent_node_1->OpType() == QDQ::DQOpName); + const bool is_p2_dq = (parent_node_2 && parent_node_2->OpType() == QDQ::DQOpName); + + // We require exactly one branch to be fed by a DQ and the other branch to be a scalar initializer + // (represented as a NodeArg with rank 0 shape and no producer node). + if (is_p1_dq && graph_utils::IsConstantInitializer(graph, where_inputs[2]->Name(), true)) { + return where_inputs[2]->HasTensorOrScalarShape() ? (where_inputs[2]->Shape()->dim_size() == 0) : false; } - if (!parent_node_1 && is_p2_dq) { - return (where_inputs[1]->Shape()->dim_size() == 0); + if (graph_utils::IsConstantInitializer(graph, where_inputs[1]->Name(), true) && is_p2_dq) { + return where_inputs[1]->HasTensorOrScalarShape() ? (where_inputs[1]->Shape()->dim_size() == 0) : false; } + return false; } Status WhereDummyDq::InsertDummyDQ(Node& node, Graph& graph, bool& modified, const logging::Logger& logger) const { + // Inserts a DeQuantizeLinear node on the scalar initializer branch of `Where` so that both + // data branches (X and Y) are produced by DQ nodes, enabling downstream QDQ grouping. const auto& where_inputs = node.InputDefs(); + const auto& where_outputs = node.OutputDefs(); const Node* parent_node_1 = graph.GetProducerNode(where_inputs[1]->Name()); const Node* parent_node_2 = graph.GetProducerNode(where_inputs[2]->Name()); + const Node* child_node = graph.GetConsumerNodes(where_outputs[0]->Name())[0]; - // With SatisfyCondition, we must have one DQ and one initializer + // From SatisfyCondition(): + // - exactly one of parent_node_1/parent_node_2 is a DQ node + // - the other input is a scalar initializer (rank-0 tensor) with no producer node const Node* dq_node = parent_node_1 ? parent_node_1 : parent_node_2; - int const_idx = parent_node_1 ? 2 : 1; + const int const_idx = parent_node_1 ? 2 : 1; + + // Guardrail: only insert dummy DQ when the quantized dtype matches the output Q's dtype. + // If they differ, we cannot safely synthesize quantization parameters. + const int32_t dt_input = dq_node->InputDefs()[0]->TypeAsProto()->tensor_type().elem_type(); + const int32_t dt_output = child_node->OutputDefs()[0]->TypeAsProto()->tensor_type().elem_type(); + if (dt_input != dt_output) { + LOGS(logger, WARNING) << "WhereDummyDq: skip inserting dummy DQ due to mismatched quantized dtype between input DQ " + "and output Q. DQ input dtype=" + << dt_input << ", Q output dtype=" << dt_output; + return Status::OK(); + } const ONNX_NAMESPACE::TensorProto* dq_node_scale_proto = nullptr; - graph.GetInitializedTensor(dq_node->InputDefs()[1]->Name(), dq_node_scale_proto); + if (!graph.GetInitializedTensor(dq_node->InputDefs()[1]->Name(), dq_node_scale_proto) || + dq_node_scale_proto == nullptr) { + LOGS(logger, WARNING) << "WhereDummyDq expects dq branch to have an initializer scale. " + << "DQ: " << dq_node->Name(); + return Status::OK(); + }; + const ONNX_NAMESPACE::TensorProto* dq_node_zp_proto = nullptr; - graph.GetInitializedTensor(dq_node->InputDefs()[2]->Name(), dq_node_zp_proto); + if (!graph.GetInitializedTensor(dq_node->InputDefs()[2]->Name(), dq_node_zp_proto) || + dq_node_zp_proto == nullptr) { + LOGS(logger, WARNING) << "WhereDummyDq expects dq branch to have an initializer zero point. " + << "DQ: " << dq_node->Name(); + return Status::OK(); + }; + // Create initializers for the dummy DQ input triplet: (xq, scale, zero_point). + // We choose values so that DeQuantizeLinear(dummy_xq, dummy_scale, dummy_zp) reconstructs + // the original scalar float value as closely as possible. + // + // Note: We only support float scalar constants currently. // Dummy data initializer. - ONNX_NAMESPACE::TensorProto dummy_data_proto; - dummy_data_proto.set_name(graph.GenerateNodeArgName(node.Name() + "_dummy_data")); + ONNX_NAMESPACE::TensorProto dummy_xq_proto; + dummy_xq_proto.set_name(graph.GenerateNodeArgName(node.Name() + "_dummy_xq")); // Set data type to dq node's zp dtype - dummy_data_proto.set_data_type(dq_node_zp_proto->data_type()); + dummy_xq_proto.set_data_type(dq_node_zp_proto->data_type()); // Dummy zero point initializer. ONNX_NAMESPACE::TensorProto dummy_zp_proto; dummy_zp_proto.set_name(graph.GenerateNodeArgName(node.Name() + "_dummy_zp")); dummy_zp_proto.set_data_type(dq_node_zp_proto->data_type()); + // Dummy scale initializer. + ONNX_NAMESPACE::TensorProto dummy_scale_proto; + dummy_scale_proto.set_name(graph.GenerateNodeArgName(node.Name() + "_dummy_scale")); + dummy_scale_proto.set_data_type(dq_node_scale_proto->data_type()); + + // Get original float input + const ONNX_NAMESPACE::TensorProto* const_node_data_proto = nullptr; + graph.GetInitializedTensor(where_inputs[const_idx]->Name(), const_node_data_proto); + Initializer initializer(graph, *const_node_data_proto, graph.ModelPath()); + if (dq_node_scale_proto->data_type() != const_node_data_proto->data_type()) { + // WhereDummyDq fills the const value to the dummy DQ's scale + LOGS(logger, WARNING) << "Currently only support existing DQ's scale with same datatype as scalar. " + << "DQ: " << dq_node->Name() << ", scalar(const): " << where_inputs[const_idx]->Name(); + return Status::OK(); + } + float dummy_xf = 0; + switch (initializer.data_type()) { + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: { + dummy_xf = *initializer.data(); + break; + } + default: + LOGS(logger, WARNING) << "Unsupported dtype of constant input. " + << "DQ: " << dq_node->Name() << ", scalar(const): " << where_inputs[const_idx]->Name(); + return Status::OK(); + } + + // TensorProto stores INT8/UINT8/INT16/UINT16 values via `int32_data`. + // Keep values in-range for unsigned cases (0..255 / 0..65535) before writing. + int32_t dummy_zp_i32 = 0; + int32_t dummy_xq_i32 = 0; + float dummy_scale = 1.0f; + switch (dummy_zp_proto.data_type()) { case ONNX_NAMESPACE::TensorProto_DataType_INT8: { - int8_t zp = 0; - int8_t dummy_data = 1; - utils::SetRawDataInTensorProto(dummy_zp_proto, &zp, 1); - utils::SetRawDataInTensorProto(dummy_data_proto, &dummy_data, 1); + dummy_zp_i32 = 0; + dummy_xq_i32 = (dummy_xf > 0) ? 127 : ((dummy_xf == 0) ? dummy_zp_i32 : -128); + dummy_scale = (dummy_xf == 0) ? 1 : (float)dummy_xf / (dummy_xq_i32 - dummy_zp_i32); break; } case ONNX_NAMESPACE::TensorProto_DataType_UINT8: { - uint8_t zp = 0; - uint8_t dummy_data = 1; - utils::SetRawDataInTensorProto(dummy_zp_proto, &zp, 1); - utils::SetRawDataInTensorProto(dummy_data_proto, &dummy_data, 1); + dummy_zp_i32 = 127; + dummy_xq_i32 = (dummy_xf > 0) ? 255 : ((dummy_xf == 0) ? dummy_zp_i32 : 0); + dummy_scale = (dummy_xf == 0) ? 1 : (float)dummy_xf / (dummy_xq_i32 - dummy_zp_i32); break; } case ONNX_NAMESPACE::TensorProto_DataType_INT16: { - int16_t zp = 0; - int16_t dummy_data = 1; - utils::SetRawDataInTensorProto(dummy_zp_proto, &zp, 2); - utils::SetRawDataInTensorProto(dummy_data_proto, &dummy_data, 2); + dummy_zp_i32 = 0; + dummy_xq_i32 = (dummy_xf > 0) ? 32767 : ((dummy_xf == 0) ? dummy_zp_i32 : -32768); + dummy_scale = (dummy_xf == 0) ? 1 : (float)dummy_xf / (dummy_xq_i32 - dummy_zp_i32); break; } case ONNX_NAMESPACE::TensorProto_DataType_UINT16: { - uint16_t zp = 0; - uint16_t dummy_data = 1; - utils::SetRawDataInTensorProto(dummy_zp_proto, &zp, 2); - utils::SetRawDataInTensorProto(dummy_data_proto, &dummy_data, 2); + dummy_zp_i32 = 32767; + dummy_xq_i32 = (dummy_xf > 0) ? 65535 : ((dummy_xf == 0) ? dummy_zp_i32 : 0); + dummy_scale = (dummy_xf == 0) ? 1 : (float)dummy_xf / (dummy_xq_i32 - dummy_zp_i32); break; } default: - LOGS(logger, WARNING) << "Currently support existing DQ's zero point with INT8, UINT8, INT16, UINT16"; + LOGS(logger, WARNING) << "Currently support existing DQ's zero point with INT8, UINT8, INT16, UINT16. " + << "DQ: " << dq_node->Name() << ", scalar(const): " << where_inputs[const_idx]->Name(); return Status::OK(); } - // Set dummy scale to the original value - const ONNX_NAMESPACE::TensorProto* const_node_data_proto = nullptr; - graph.GetInitializedTensor(where_inputs[const_idx]->Name(), const_node_data_proto); - Initializer initializer(graph, *const_node_data_proto, graph.ModelPath()); - if (dq_node_scale_proto->data_type() != const_node_data_proto->data_type()) { - // WhereDummyDq fills the const value to the dummy DQ's scale - LOGS(logger, WARNING) << "Currently only support existing DQ's scale with same datatype as scalar"; - return Status::OK(); - } - - // Dummy scale initializer. - ONNX_NAMESPACE::TensorProto dummy_scale_proto; - dummy_scale_proto.set_name(graph.GenerateNodeArgName(node.Name() + "_dummy_scale")); - dummy_scale_proto.set_data_type(dq_node_scale_proto->data_type()); - switch (initializer.data_type()) { - case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: { - float* where_const_scalar = initializer.data(); - utils::SetRawDataInTensorProto(dummy_scale_proto, where_const_scalar, sizeof(float)); - break; - } - default: - LOGS(logger, WARNING) << "Currently support scalar with FLOAT"; - return Status::OK(); - } + dummy_zp_proto.add_int32_data(dummy_zp_i32); + dummy_xq_proto.add_int32_data(dummy_xq_i32); + dummy_scale_proto.add_float_data(dummy_scale); - // Start editing the graph - NodeArg& dummy_data_arg = graph_utils::AddInitializerWithOrtValue(graph, dummy_data_proto); + // Start editing the graph: + // - add the initializers + // - add a DeQuantizeLinear node consuming them + // - rewire the scalar branch of `Where` to use the DQ output + // - drop the original scalar initializer if it becomes unused + NodeArg& dummy_xq_arg = graph_utils::AddInitializerWithOrtValue(graph, dummy_xq_proto); NodeArg& dummy_scale_arg = graph_utils::AddInitializerWithOrtValue(graph, dummy_scale_proto); NodeArg& dummy_zp_arg = graph_utils::AddInitializerWithOrtValue(graph, dummy_zp_proto); @@ -132,7 +200,7 @@ Status WhereDummyDq::InsertDummyDQ(Node& node, Graph& graph, bool& modified, con graph.GenerateNodeArgName(node.Name() + "_dummy_dq"), QDQ::DQOpName, "DeQuantizeLinear from WhereDummyDq GraphTransformer", - {&dummy_data_arg, &dummy_scale_arg, &dummy_zp_arg}, + {&dummy_xq_arg, &dummy_scale_arg, &dummy_zp_arg}, {&dummy_dq_arg}, node, nullptr, diff --git a/onnxruntime/core/optimizer/qdq_transformer/where_dummy_dq.h b/onnxruntime/core/optimizer/qdq_transformer/where_dummy_dq.h index 3260a865f8c4b..1c4a8cbf6ba12 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/where_dummy_dq.h +++ b/onnxruntime/core/optimizer/qdq_transformer/where_dummy_dq.h @@ -11,7 +11,29 @@ namespace onnxruntime { @Class WhereDummyDq Graph transformer that inserts a dummy DQ on Where node's initializer input - to form Node Unit when Where node has one DQ and one scalar initializer input + to form Node Unit when Where node has one DQ and one scalar initializer input. + + If `Where` gets a float scalar `xf` and a `DequantizeLinear` as its two data inputs, + `WhereDummyDq` inserts a dummy DQ so that `xf ≈ DQ(xq, scale, zp)`. + + The `xq`, `zp` are chosen per the table below (by the dtype of the existing DQ's zero-point), + and `scale` is computed from them. + + We select these values in order to keep the `scale` non-negative: + + | | uint8 | uint16 | int8 | int16 | + |-----------------|--------|--------|-------|--------| + | xf > 0 | | | | | + | xq | 255 | 65535 | 127 | 32767 | + | zp | 127 | 32767 | 0 | 0 | + | xf < 0 | | | | | + | xq | 0 | 0 | -128 | -32768 | + | zp | 127 | 32767 | 0 | 0 | + | xf = 0 | | | | | + | xq | 127 | 32767 | 0 | 0 | + | zp | 127 | 32767 | 0 | 0 | + + scale = xf / (xq - zp) if (xq != zp) else 1 */ class WhereDummyDq : public GraphTransformer { public: @@ -23,4 +45,4 @@ class WhereDummyDq : public GraphTransformer { bool SatisfyCondition(const Graph& graph, const Node& node) const; Status InsertDummyDQ(Node& node, Graph& graph, bool& modified, const logging::Logger& logger) const; }; -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/qdq_transformer_test.cc b/onnxruntime/test/optimizer/qdq_transformer_test.cc index 37da2d4247e34..85d4c51b9faae 100644 --- a/onnxruntime/test/optimizer/qdq_transformer_test.cc +++ b/onnxruntime/test/optimizer/qdq_transformer_test.cc @@ -3252,7 +3252,7 @@ TEST(QDQTransformerTests, ClipQuantFusion_MultipleInputEdges) { 18); // opset } -template +template void TestWhereWithDqInput(bool is_dq_1, bool is_dq_2, int expected_num_where, @@ -3268,9 +3268,9 @@ void TestWhereWithDqInput(bool is_dq_1, NodeArg* where_in2 = nullptr; if (is_dq_1) { // DQ - auto* dq_Input = builder.MakeInput({4, 3, 32}, 0.0, 1.0); + auto* dq_Input = builder.MakeInput({4, 3, 32}, 0.0, 1.0); auto* dq_scale = builder.MakeInitializer({}, 0.0, 1.0); - auto* dq_zp = builder.MakeInitializer({}, 0.0, 1.0); + auto* dq_zp = builder.MakeInitializer({}, 0.0, 1.0); where_in1 = builder.MakeIntermediate(); builder.AddNode("DequantizeLinear", {dq_Input, dq_scale, dq_zp}, {where_in1}); } else { @@ -3278,9 +3278,9 @@ void TestWhereWithDqInput(bool is_dq_1, } if (is_dq_2) { // DQ - auto* dq_Input = builder.MakeInput({4, 3, 32}, 0.0, 1.0); + auto* dq_Input = builder.MakeInput({4, 3, 32}, 0.0, 1.0); auto* dq_scale = builder.MakeInitializer({}, 0.0, 1.0); - auto* dq_zp = builder.MakeInitializer({}, 0.0, 1.0); + auto* dq_zp = builder.MakeInitializer({}, 0.0, 1.0); where_in2 = builder.MakeIntermediate(); builder.AddNode("DequantizeLinear", {dq_Input, dq_scale, dq_zp}, {where_in2}); } else { @@ -3294,7 +3294,7 @@ void TestWhereWithDqInput(bool is_dq_1, // Q auto* q_scale = builder.MakeInitializer({}, 0.0, 1.0); - auto* q_zp = builder.MakeInitializer({}, 0.0, 1.0); + auto* q_zp = builder.MakeInitializer({}, 0.0, 1.0); auto* q_out = builder.MakeOutput(); builder.AddNode("QuantizeLinear", {where_out, q_scale, q_zp}, {q_out}); @@ -3315,14 +3315,200 @@ void TestWhereWithDqInput(bool is_dq_1, }; TEST(QDQTransformerTests, WhereDummyDqTest) { - TestWhereWithDqInput(true, true, 1, 2, 1, false); - TestWhereWithDqInput(true, false, 1, 2, 1, true); - TestWhereWithDqInput(false, true, 1, 2, 1, true); - TestWhereWithDqInput(false, false, 1, 0, 1, false); - TestWhereWithDqInput(true, true, 1, 2, 1, false); - TestWhereWithDqInput(true, false, 1, 2, 1, true); - TestWhereWithDqInput(false, true, 1, 2, 1, true); - TestWhereWithDqInput(false, false, 1, 0, 1, false); + // is_dq_1, is_dq_2, expected_num_where, expected_num_dq, expected_num_q, expected_modified + TestWhereWithDqInput(true, true, 1, 2, 1, false); + TestWhereWithDqInput(true, false, 1, 2, 1, true); + TestWhereWithDqInput(false, true, 1, 2, 1, true); + TestWhereWithDqInput(false, false, 1, 0, 1, false); + TestWhereWithDqInput(true, true, 1, 2, 1, false); + TestWhereWithDqInput(true, false, 1, 2, 1, true); + TestWhereWithDqInput(false, true, 1, 2, 1, true); + TestWhereWithDqInput(false, false, 1, 0, 1, false); + // DQ uses uint8 but Q uses uint16 + TestWhereWithDqInput(true, false, 1, 1, 1, false); + TestWhereWithDqInput(false, true, 1, 1, 1, false); +} + +// Tests WhereDummyDq with non-QuantizeLinear consumers. +// The optimizer should NOT add dummy DQ nodes when the Where output is not consumed by a QuantizeLinear. +template +void TestWhereWithNonQLinearConsumer(bool is_dq_1, + bool is_dq_2, + int expected_num_where, + int expected_num_dq, + bool expected_modified, + bool use_op_consumer = false) { + auto& logger = DefaultLoggingManager().DefaultLogger(); + Model model("WhereDummyDqNonQLinearConsumerTester", false, logger); + Graph& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + + NodeArg* where_in1 = nullptr; + NodeArg* where_in2 = nullptr; + if (is_dq_1) { + // DQ + auto* dq_Input = builder.MakeInput({4, 3, 32}, 0.0, 1.0); + auto* dq_scale = builder.MakeInitializer({}, 0.0, 1.0); + auto* dq_zp = builder.MakeInitializer({}, 0.0, 1.0); + where_in1 = builder.MakeIntermediate(); + builder.AddNode("DequantizeLinear", {dq_Input, dq_scale, dq_zp}, {where_in1}); + } else { + where_in1 = builder.MakeInitializer({}, 0.0, 1.0); + } + if (is_dq_2) { + // DQ + auto* dq_Input = builder.MakeInput({4, 3, 32}, 0.0, 1.0); + auto* dq_scale = builder.MakeInitializer({}, 0.0, 1.0); + auto* dq_zp = builder.MakeInitializer({}, 0.0, 1.0); + where_in2 = builder.MakeIntermediate(); + builder.AddNode("DequantizeLinear", {dq_Input, dq_scale, dq_zp}, {where_in2}); + } else { + where_in2 = builder.MakeInitializer({}, 0.0, 1.0); + } + + // Where + auto* where_cond = builder.MakeInputBool({4, 3, 32}); + + if (use_op_consumer) { + // Where output consumed by another op (e.g., Add) instead of QuantizeLinear + auto* where_out = builder.MakeIntermediate(); + builder.AddNode("Where", {where_cond, where_in1, where_in2}, {where_out}); + auto* add_input = builder.MakeInput({4, 3, 32}, 0.0, 1.0); + auto* add_out = builder.MakeOutput(); + builder.AddNode("Add", {where_out, add_input}, {add_out}); + } else { + // Where output is a direct graph output (no consumer node) + auto* where_out = builder.MakeOutput(); + builder.AddNode("Where", {where_cond, where_in1, where_in2}, {where_out}); + } + + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(graph.Resolve()); + + auto where_optimizer = std::make_unique(); + bool modified = false; + ASSERT_STATUS_OK(where_optimizer->Apply(graph, modified, logger)); + + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Where"], expected_num_where); + ASSERT_EQ(op_to_count["DequantizeLinear"], expected_num_dq); + ASSERT_EQ(op_to_count["QuantizeLinear"], 0); + ASSERT_EQ(modified, expected_modified); +} + +TEST(QDQTransformerTests, WhereDummyDqTest_NonQLinearConsumer) { + // When Where output is a direct graph output (no QuantizeLinear consumer), + // the optimizer should NOT add dummy DQ nodes. + // is_dq_1, is_dq_2, expected_num_where, expected_num_dq, expected_modified + TestWhereWithNonQLinearConsumer(true, true, 1, 2, false); + TestWhereWithNonQLinearConsumer(true, false, 1, 1, false); + TestWhereWithNonQLinearConsumer(false, true, 1, 1, false); + TestWhereWithNonQLinearConsumer(false, false, 1, 0, false); + TestWhereWithNonQLinearConsumer(true, true, 1, 2, false); + TestWhereWithNonQLinearConsumer(true, false, 1, 1, false); + TestWhereWithNonQLinearConsumer(false, true, 1, 1, false); + TestWhereWithNonQLinearConsumer(false, false, 1, 0, false); + + // When Where output is consumed by another op (e.g., Add) instead of QuantizeLinear, + // the optimizer should NOT add dummy DQ nodes. + TestWhereWithNonQLinearConsumer(true, true, 1, 2, false, true /*use_op_consumer*/); + TestWhereWithNonQLinearConsumer(true, false, 1, 1, false, true /*use_op_consumer*/); + TestWhereWithNonQLinearConsumer(false, true, 1, 1, false, true /*use_op_consumer*/); + TestWhereWithNonQLinearConsumer(false, false, 1, 0, false, true /*use_op_consumer*/); +} + +// Tests WhereDummyDq with multiple consumers of the Where output. +// The optimizer should NOT add dummy DQ nodes when the Where output has multiple consumers, +// even if one of the consumers is a QuantizeLinear. +template +void TestWhereWithMultipleConsumers(bool is_dq_1, + bool is_dq_2, + int expected_num_where, + int expected_num_dq, + bool expected_modified, + bool use_two_q_consumers = true) { + auto& logger = DefaultLoggingManager().DefaultLogger(); + Model model("WhereDummyDqMultipleConsumersTester", false, logger); + Graph& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + + NodeArg* where_in1 = nullptr; + NodeArg* where_in2 = nullptr; + if (is_dq_1) { + auto* dq_input = builder.MakeInput({4, 3, 32}, 0, std::numeric_limits::max()); + auto* dq_scale = builder.MakeInitializer({}, {1.0f}); + auto* dq_zp = builder.MakeInitializer({}, {static_cast(0)}); + where_in1 = builder.MakeIntermediate(); + builder.AddNode("DequantizeLinear", {dq_input, dq_scale, dq_zp}, {where_in1}); + } else { + where_in1 = builder.MakeInput({4, 3, 32}, -1.0f, 1.0f); + } + if (is_dq_2) { + auto* dq_input = builder.MakeInput({4, 3, 32}, 0, std::numeric_limits::max()); + auto* dq_scale = builder.MakeInitializer({}, {1.0f}); + auto* dq_zp = builder.MakeInitializer({}, {static_cast(0)}); + where_in2 = builder.MakeIntermediate(); + builder.AddNode("DequantizeLinear", {dq_input, dq_scale, dq_zp}, {where_in2}); + } else { + where_in2 = builder.MakeInput({4, 3, 32}, -1.0f, 1.0f); + } + + // Where + auto* where_cond = builder.MakeInputBool({4, 3, 32}); + auto* where_out = builder.MakeIntermediate(); + builder.AddNode("Where", {where_cond, where_in1, where_in2}, {where_out}); + + // First consumer: QuantizeLinear + auto* q_scale = builder.MakeInitializer({}, {1.0f}); + auto* q_zp = builder.MakeInitializer({}, {static_cast(0)}); + auto* q_out1 = builder.MakeOutput(); + builder.AddNode("QuantizeLinear", {where_out, q_scale, q_zp}, {q_out1}); + + if (use_two_q_consumers) { + // Second consumer: another QuantizeLinear + auto* q_scale2 = builder.MakeInitializer({}, {1.0f}); + auto* q_zp2 = builder.MakeInitializer({}, {static_cast(0)}); + auto* q_out2 = builder.MakeOutput(); + builder.AddNode("QuantizeLinear", {where_out, q_scale2, q_zp2}, {q_out2}); + } else { + // Second consumer: Add op + auto* add_input = builder.MakeInput({4, 3, 32}, -1.0f, 1.0f); + auto* add_out = builder.MakeOutput(); + builder.AddNode("Add", {where_out, add_input}, {add_out}); + } + + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(graph.Resolve()); + + auto where_optimizer = std::make_unique(); + bool modified = false; + ASSERT_STATUS_OK(where_optimizer->Apply(graph, modified, logger)); + + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Where"], expected_num_where); + ASSERT_EQ(op_to_count["DequantizeLinear"], expected_num_dq); + ASSERT_EQ(modified, expected_modified); +} + +TEST(QDQTransformerTests, WhereDummyDqTest_MultipleConsumers) { + // When Where output has two QuantizeLinear consumers, + // the optimizer should NOT add dummy DQ nodes (child_nodes.size() != 1). + // is_dq_1, is_dq_2, expected_num_where, expected_num_dq, expected_modified + TestWhereWithMultipleConsumers(true, true, 1, 2, false); + TestWhereWithMultipleConsumers(true, false, 1, 1, false); + TestWhereWithMultipleConsumers(false, true, 1, 1, false); + TestWhereWithMultipleConsumers(false, false, 1, 0, false); + TestWhereWithMultipleConsumers(true, true, 1, 2, false); + TestWhereWithMultipleConsumers(true, false, 1, 1, false); + TestWhereWithMultipleConsumers(false, true, 1, 1, false); + TestWhereWithMultipleConsumers(false, false, 1, 0, false); + + // When Where output has one QuantizeLinear consumer and one non-QuantizeLinear consumer (Add), + // the optimizer should NOT add dummy DQ nodes (child_nodes.size() != 1). + TestWhereWithMultipleConsumers(true, true, 1, 2, false, false /*use_two_q_consumers*/); + TestWhereWithMultipleConsumers(true, false, 1, 1, false, false /*use_two_q_consumers*/); + TestWhereWithMultipleConsumers(false, true, 1, 1, false, false /*use_two_q_consumers*/); + TestWhereWithMultipleConsumers(false, false, 1, 0, false, false /*use_two_q_consumers*/); } TEST(QDQTransformerTests, Concat) {