diff --git a/onnxruntime/core/optimizer/transpose_optimization/onnx_transpose_optimization.cc b/onnxruntime/core/optimizer/transpose_optimization/onnx_transpose_optimization.cc index 467d0c090070f..8529ccf169f56 100755 --- a/onnxruntime/core/optimizer/transpose_optimization/onnx_transpose_optimization.cc +++ b/onnxruntime/core/optimizer/transpose_optimization/onnx_transpose_optimization.cc @@ -454,6 +454,23 @@ static std::optional GetDQWithConstInitializerInputAndSingleConsum return result; } +// Element types that ONNX QuantizeLinear can produce as its output. +static bool IsQuantizeLinearOutputType(api::DataType dtype) { + switch (dtype) { + case api::DataType::INT8: + case api::DataType::UINT8: + case api::DataType::INT16: + case api::DataType::UINT16: + case api::DataType::FLOAT8E4M3FN: + case api::DataType::FLOAT8E4M3FNUZ: + case api::DataType::FLOAT8E5M2: + case api::DataType::FLOAT8E5M2FNUZ: + return true; + default: + return false; + } +} + /// /// Insert a Q -> DQ pair after the node following the DQ by using scale and zp info from the preceding DQ node. /// DQ -> next node => DQ -> next node -> Q -> DQ. @@ -528,9 +545,26 @@ static bool MakeQDQNodeUnit(api::GraphRef& graph, const api::NodeRef& dq_node) { inputs.push_back(zp_input.value()); } + // A zero-point-less DQ with a non-uint8 type needs the new Q's output_dtype pinned, or Q type + // inference defaults to uint8 and clashes with the int8 value-info copied below. output_dtype is + // ONNX opset 21+ only; if it can't be expressed (older opset, non-ONNX domain, or a type that + // QuantizeLinear can't output), skip the push-through so the graph stays valid. + std::optional q_output_dtype; + if (!zp_input.has_value()) { + const api::DataType dq_input_dtype = graph.GetValueInfo(dq_inputs[0])->DType(); + if (dq_input_dtype != api::DataType::UNDEFINED && dq_input_dtype != api::DataType::UINT8) { + const std::optional domain_opset = graph.Opset(dq_domain); + if (!IsOnnxDomain(dq_domain) || !domain_opset || *domain_opset < 21 || + !IsQuantizeLinearOutputType(dq_input_dtype)) { + return false; + } + q_output_dtype = static_cast(dq_input_dtype); + } + } + // Add Q auto new_q_node = MakeQuantizeOp(graph, dq_domain, inputs, axis, dq_node.GetAttributeInt("block_size"), - dq_node.GetAttributeInt("output_dtype"), dq_node.GetAttributeInt("saturate")); + q_output_dtype, dq_node.GetAttributeInt("saturate")); new_q_node->SetLayeringAnnotation(dq_node.GetLayeringAnnotation()); auto q_node_outputs = new_q_node->Outputs(); diff --git a/onnxruntime/test/optimizer/transpose_optimizer_test.cc b/onnxruntime/test/optimizer/transpose_optimizer_test.cc index 080c382db5d93..2640408df7ab7 100644 --- a/onnxruntime/test/optimizer/transpose_optimizer_test.cc +++ b/onnxruntime/test/optimizer/transpose_optimizer_test.cc @@ -3773,6 +3773,37 @@ TEST(TransposeOptimizerTests, TestDequantizeLinearNoAxis) { #endif } +// Regression test for #28716: pushing a Transpose through a zero-point-less int8 DequantizeLinear +// inserts a QuantizeLinear that must set output_dtype, else it defaults to uint8 and Resolve() fails. +TEST(TransposeOptimizerTests, TestDequantizeLinearNoZeroPoint) { + auto build_test_case = [&](ModelTestBuilder& builder) { + auto* input0_arg = MakeInput(builder, {{2, -1, 6, 3}}, {2, 4, 6, 3}, -128, 127); + auto* scale_arg = MakeInput(builder, std::vector{}, std::vector{}, {0.05f}); + auto* transpose_1_out_0 = builder.MakeIntermediate(); + auto* dq_out_0 = builder.MakeIntermediate(); + auto* transpose_2_out_0 = builder.MakeOutput(); + + auto& transpose_1 = builder.AddNode("Transpose", {input0_arg}, {transpose_1_out_0}); + transpose_1.AddAttribute("perm", std::vector{0, 3, 1, 2}); + builder.AddNode("DequantizeLinear", {transpose_1_out_0, scale_arg}, {dq_out_0}); // no zero-point + auto& transpose_2 = builder.AddNode("Transpose", {dq_out_0}, {transpose_2_out_0}); + transpose_2.AddAttribute("perm", std::vector{0, 2, 3, 1}); + }; + + // opset 21: output_dtype pins the inserted Q's type and the transposes cancel. + auto check_cancelled = [](InferenceSessionWrapper& session) { + EXPECT_EQ(EstimateTransposeCost(session.GetGraph()), 0); + }; + TransformerTester(build_test_case, check_cancelled, TransformerLevel::Default, + TransformerLevel::Level1, /*opset_version*/ 21); + + // Pre-opset-21 has no output_dtype, so the optimizer must skip the push-through rather than emit + // an invalid QuantizeLinear; the model must still initialize (no type-inference crash). + auto check_valid = [](InferenceSessionWrapper& /*session*/) {}; + TransformerTester(build_test_case, check_valid, TransformerLevel::Default, + TransformerLevel::Level1, /*opset_version*/ 19); +} + TEST(TransposeOptimizerTests, TestCast) { auto build_test_case_1 = [&](ModelTestBuilder& builder) { auto* input0_arg = MakeInput(builder, {{-1, 4, -1, 5}}, {2, 4, 6, 5}, -1, 5);