From 839dee99fad288ac429f9f90b7ff85fbc10fc072 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Tue, 16 Jun 2026 22:25:44 +0000 Subject: [PATCH 1/8] Fix data propagation dropping rank for 1-D single-element shape values GatherOpDataPropagation::infer() guarded on indices.size() == 1 (element count) and called SetInferredShapeScalarValue() unconditionally. That guard is true for BOTH a 0-D scalar index and a 1-D single-element index, so the 1-D case had its rank dropped: Graph::getInputData() then emitted a 0-D (dimensionless) TensorProto for the propagated value. For the common Shape -> Gather([-1]) -> TopK exporter pattern this produced a 0-D K initializer, which ONNX TopK shape inference correctly rejects ("K input must be a one-dimensional tensor of size 1.") at Graph::Resolve time. The model is spec-valid (a 1-D Gather index yields a rank-1 Gather output), so this was an ORT rank-preservation bug. It reproduces even at ORT_DISABLE_ALL, where constant folding never runs, confirming the cause is shape-inference data propagation rather than constant folding. Changes: - Gather: route by the index rank instead of element count. A 0-D scalar index stores a scalar value; a 1-D single-element index stores a rank-1 value, so getInputData() emits a TensorProto with dims=[1] and downstream TopK sees a valid 1-D size-1 K; a rank >= 2 index (or an index whose rank is unknown) declines and falls back to ONNX data propagation, because the single-value channel cannot faithfully represent a rank >= 2 Gather output. The rank routing is inlined next to the Gather logic (a small if/else on the index rank). The index rank is sourced from the same constant initializer the index value comes from (via get_initialized_input_values now reporting the initializer rank), instead of a second, independently resolved NodeArg shape -- removing a potential source-of-truth drift. - Elementwise consumers (Add/Sub/Mul/Div): previously scalar-only, they would silently stop propagating once an operand became a rank-1 value (e.g. a Shape -> Gather(1-D idx) -> Mul -> TopK chain), since the custom propagation result replaces ONNX's rank-correct fallback. They now accept a single element carried as either a rank-0 scalar or a rank-1 [1] value and keep the output rank consistent with ONNX broadcasting (rank-1 if any operand is rank-1, else scalar), so such chains keep propagating end to end. Div also guards against division by zero. - Unsqueeze: decline rather than propagate a dubious value when unsqueezing a single-element (scalar-like, rank-1 [1]) value, whose result is a rank >= 2 tensor the values channel cannot faithfully represent (it would otherwise fabricate a misleading [1, value]). The decline is a single inlined check on the value's element count, next to the Unsqueeze logic. Multi-element shape vectors are unaffected. Squeeze is left unchanged -- it already converts a rank-1 [1] value to the correct scalar. - Add shared helpers (data_propagation_value_utils.h) for reading and writing a single-element shape value while preserving its rank, used by Gather and the elementwise ops so producers and consumers cannot disagree on rank. The reader declines a rank-1 multi-element value (it must never collapse to element[0]), so a multi-element value cannot be mistaken for a single one. The setter is correct-by-construction: it populates exactly one channel and clears the other, so the scalar-first reader and the values-first getInputData() can never disagree on rank even if the output carried a stale value from a prior pass. The Gather and Unsqueeze rank decisions are kept inline next to each op rather than shared, since the two ops route on different quantities (index rank vs. value element count). Tests: add ShapeInferenceV2Test.GatherToTopKRankPreservationTest and GatherMulToTopKRankPreservationTest (with fixtures and generators) that load the model at the disabled, basic, and all-optimization levels and assert the rank-1 K is preserved and propagated through the chain, plus GatherSqueezeRangeRankPreservationTest, an observable end-to-end lock that a Shape -> Gather([-1]) -> Squeeze -> Range chain resolves Range's length to the concrete propagated K (locking Squeeze's correct-scalar behavior through a real downstream consumer). Add unit tests pinning the pure read/write helpers directly: SinglePropagatedShapeValueGuardTest (reader behavior on each channel) and SetSinglePropagatedShapeValueKeepsSingleChannelTest (the setter clears the opposite channel). Cover the two decline branches with real end-to-end tests (rather than pure-predicate unit tests) that exercise them through a rank-lowering Squeeze: GatherRank2IndexDeclineTest (a rank-2 Gather index must decline -- relaxing it concretizes a Range length that must stay symbolic) and UnsqueezeSingleValueDeclineTest (unsqueezing a single-element value must decline -- relaxing it fabricates a non-scalar that makes the model fail to load). Both run at ORT_DISABLE_ALL so constant folding does not mask the data-propagation decline. Add ShapeMulMultiElementNoScalarCollapseTest, an end-to-end check that a Shape -> Mul -> ConstantOfShape multi-element chain still resolves to its full rank-2 shape. The data-propagation regression tests run at the disabled, basic, and all-optimization levels (data propagation executes in the pre-optimization Graph::Resolve pass, so it is independent of the graph-optimization level). Existing scalar-index data-prop fixtures continue to exercise the scalar path unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwangms --- include/onnxruntime/core/graph/node_arg.h | 3 + .../add_op_data_propagation.cc | 15 +- .../add_op_data_propagation.h | 2 +- .../custom_data_propagation.cc | 2 +- .../custom_data_propagation.h | 21 +- .../data_propagation_value_utils.h | 67 ++++ .../div_op_data_propagation.cc | 16 +- .../div_op_data_propagation.h | 2 +- .../gather_op_data_propagation.cc | 33 +- .../gather_op_data_propagation.h | 2 +- .../mul_op_data_propagation.cc | 15 +- .../mul_op_data_propagation.h | 2 +- .../size_op_data_propagation.h | 2 +- .../squeeze_op_data_propagation.cc | 3 +- .../squeeze_op_data_propagation.h | 2 +- .../sub_op_data_propagation.cc | 15 +- .../sub_op_data_propagation.h | 2 +- .../unsqueeze_op_data_propagation.cc | 19 +- .../unsqueeze_op_data_propagation.h | 2 +- onnxruntime/core/graph/graph.cc | 8 +- .../test/framework/shape_inference_test.cc | 351 ++++++++++++++++++ ...hape_data_propagation_gather_mul_topk.onnx | Bin 0 -> 424 bytes ..._shape_data_propagation_gather_mul_topk.py | 62 ++++ ...data_propagation_gather_rank2_decline.onnx | Bin 0 -> 354 bytes ...e_data_propagation_gather_rank2_decline.py | 70 ++++ ...data_propagation_gather_squeeze_range.onnx | Bin 0 -> 346 bytes ...e_data_propagation_gather_squeeze_range.py | 59 +++ ...st_shape_data_propagation_gather_topk.onnx | Bin 0 -> 297 bytes ...test_shape_data_propagation_gather_topk.py | 54 +++ ...propagation_shape_mul_constantofshape.onnx | Bin 0 -> 211 bytes ...a_propagation_shape_mul_constantofshape.py | 50 +++ ...pe_data_propagation_unsqueeze_decline.onnx | Bin 0 -> 423 bytes ...hape_data_propagation_unsqueeze_decline.py | 65 ++++ 33 files changed, 911 insertions(+), 33 deletions(-) create mode 100644 onnxruntime/core/graph/data_propagation/data_propagation_value_utils.h create mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_gather_mul_topk.onnx create mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_gather_mul_topk.py create mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_gather_rank2_decline.onnx create mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_gather_rank2_decline.py create mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_gather_squeeze_range.onnx create mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_gather_squeeze_range.py create mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_gather_topk.onnx create mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_gather_topk.py create mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_shape_mul_constantofshape.onnx create mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_shape_mul_constantofshape.py create mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_decline.onnx create mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_decline.py diff --git a/include/onnxruntime/core/graph/node_arg.h b/include/onnxruntime/core/graph/node_arg.h index 4a18d7617ac13..5d054258526df 100644 --- a/include/onnxruntime/core/graph/node_arg.h +++ b/include/onnxruntime/core/graph/node_arg.h @@ -121,6 +121,9 @@ class NodeArg { /** Sets the inferred shape scalar value */ void SetInferredShapeScalarValue(int64_t value) noexcept { inferred_scalar_value_ = value; } + /** Clears the inferred shape scalar value */ + void ClearInferredShapeScalarValue() noexcept { inferred_scalar_value_.reset(); } + /** Gets a flag indicating whether this NodeArg exists or not. Optional inputs are allowed in ONNX and an empty #Name represents a non-existent input argument. */ bool Exists() const noexcept; diff --git a/onnxruntime/core/graph/data_propagation/add_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/add_op_data_propagation.cc index 172941c0ee023..06dca7d0d464b 100644 --- a/onnxruntime/core/graph/data_propagation/add_op_data_propagation.cc +++ b/onnxruntime/core/graph/data_propagation/add_op_data_propagation.cc @@ -6,6 +6,7 @@ #include "core/graph/node_arg.h" #include "core/graph/onnx_protobuf.h" #include "core/providers/common.h" +#include "core/graph/data_propagation/data_propagation_value_utils.h" namespace onnxruntime { @@ -20,10 +21,16 @@ Status AddOpDataPropagation::infer() { return Status::OK(); } - if (input_0->GetInferredShapeScalarValue().has_value() && input_1->GetInferredShapeScalarValue().has_value()) { - output_def_.SetInferredShapeScalarValue( - input_0->GetInferredShapeScalarValue().value() + - input_1->GetInferredShapeScalarValue().value()); + int64_t lhs = 0; + int64_t rhs = 0; + bool lhs_is_rank1 = false; + bool rhs_is_rank1 = false; + if (TryGetSinglePropagatedShapeValue(*input_0, lhs, lhs_is_rank1) && + TryGetSinglePropagatedShapeValue(*input_1, rhs, rhs_is_rank1)) { + // Single-element operands may be carried as a rank-0 scalar or a rank-1 [1] value. Per ONNX + // broadcasting, the result is rank-1 if either operand is rank-1, otherwise a scalar; keep + // the propagated value's rank consistent with that so downstream consumers see the right rank. + SetSinglePropagatedShapeValue(output_def_, lhs + rhs, lhs_is_rank1 || rhs_is_rank1); } return Status::OK(); diff --git a/onnxruntime/core/graph/data_propagation/add_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/add_op_data_propagation.h index f9eb9990142c1..d074dc4ea69da 100644 --- a/onnxruntime/core/graph/data_propagation/add_op_data_propagation.h +++ b/onnxruntime/core/graph/data_propagation/add_op_data_propagation.h @@ -41,7 +41,7 @@ class AddOpDataPropagation : public CustomDataPropagationBase { public: AddOpDataPropagation(const Node& node, NodeArg& output_def, - std::function func, + GetInitializedInputValuesFunc func, const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation, const logging::Logger& logger) noexcept : CustomDataPropagationBase(node, output_def, func, output_from_onnx_op_data_propagation, logger) {} diff --git a/onnxruntime/core/graph/data_propagation/custom_data_propagation.cc b/onnxruntime/core/graph/data_propagation/custom_data_propagation.cc index b7254aa828107..0b260f598a5ad 100644 --- a/onnxruntime/core/graph/data_propagation/custom_data_propagation.cc +++ b/onnxruntime/core/graph/data_propagation/custom_data_propagation.cc @@ -19,7 +19,7 @@ namespace onnxruntime { std::unique_ptr CreateCustomDataPropagation(const Node& node, NodeArg& output_def, - std::function func, + GetInitializedInputValuesFunc func, const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation, const logging::Logger& logger) { int dim_size = 0; diff --git a/onnxruntime/core/graph/data_propagation/custom_data_propagation.h b/onnxruntime/core/graph/data_propagation/custom_data_propagation.h index 7511f77f58519..5ac78c5a1b7f4 100644 --- a/onnxruntime/core/graph/data_propagation/custom_data_propagation.h +++ b/onnxruntime/core/graph/data_propagation/custom_data_propagation.h @@ -10,6 +10,21 @@ namespace onnxruntime { +/** + * @brief Signature of the helper used to read a constant-initializer input's values during + * data propagation. + * + * @param input_name Name of the input to read. + * @param input_values Receives the initializer's flattened int64 values (empty if the input is + * not a constant initializer). + * @param num_dims Receives the rank (number of dimensions) of the initializer's shape, or is + * left unchanged when the input is not a constant initializer. A value of 0 + * denotes a scalar (rank-0) initializer; this lets callers distinguish a 0-D + * scalar from a rank-1 single-element initializer (both have one element). + */ +using GetInitializedInputValuesFunc = + std::function; + /** * @class CustomDataPropagation * Custom data propagation for the operator to help enhance shape inference. @@ -27,7 +42,7 @@ class CustomDataPropagationBase { protected: CustomDataPropagationBase(const Node& node, NodeArg& output_def, - std::function func, + GetInitializedInputValuesFunc func, const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation, const logging::Logger& logger) noexcept : node_(node), @@ -38,7 +53,7 @@ class CustomDataPropagationBase { const Node& node_; NodeArg& output_def_; - std::function get_initialized_input_values_func_; + GetInitializedInputValuesFunc get_initialized_input_values_func_; const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation_; const logging::Logger& logger_; }; @@ -68,7 +83,7 @@ class CustomDataPropagationBase { std::unique_ptr CreateCustomDataPropagation( const Node& node, NodeArg& output_def, - std::function func, + GetInitializedInputValuesFunc func, const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation, const logging::Logger& logger); diff --git a/onnxruntime/core/graph/data_propagation/data_propagation_value_utils.h b/onnxruntime/core/graph/data_propagation/data_propagation_value_utils.h new file mode 100644 index 0000000000000..e3610e70d51cf --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/data_propagation_value_utils.h @@ -0,0 +1,67 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +#include "core/graph/node_arg.h" + +namespace onnxruntime { + +// Data propagation carries a small "shape value" in one of two non-interchangeable channels +// on a NodeArg: a rank-0 scalar (inferred_scalar_value_) or a rank>=1 list of values +// (inferred_shape_values_). The helpers below let custom data-propagation ops read and write a +// single-element value while preserving its rank (rank-0 scalar vs rank-1 [1]), so a producer +// and its consumers cannot silently disagree on rank (e.g. Gather feeding Mul feeding TopK). + +// Reads a single int64 shape value carried by a NodeArg's data propagation, accepting either a +// rank-0 scalar value or a rank-1 single-element value. On success, sets `value`, sets +// `is_rank1` to false for a scalar source or true for a rank-1 [1] source, and returns true. +// Returns false if the NodeArg carries no usable single-element shape value. +inline bool TryGetSinglePropagatedShapeValue(const NodeArg& input_def, int64_t& value, bool& is_rank1) { + if (input_def.GetInferredShapeScalarValue().has_value()) { + value = input_def.GetInferredShapeScalarValue().value(); + is_rank1 = false; + return true; + } + + const auto& inferred_values = input_def.GetInferredShapeValues(); + if (inferred_values.has_value() && + inferred_values->dim_size() == 1 && + inferred_values->dim(0).has_dim_value()) { + value = inferred_values->dim(0).dim_value(); + is_rank1 = true; + return true; + } + + return false; +} + +// Stores a single int64 shape value on `output_def`, as a rank-0 scalar when `is_rank1` is false +// or as a rank-1 single-element value when `is_rank1` is true. The rank-1 representation mirrors +// how Graph::getInputData() reconstructs a TensorProto (dims=[1]) from inferred_shape_values_. +// The setter is correct-by-construction: it populates exactly one channel and clears the other, so +// the scalar-first reader (TryGetSinglePropagatedShapeValue) and the values-first getInputData() +// can never disagree on rank even if `output_def` carried a stale value from another channel. +inline void SetSinglePropagatedShapeValue(NodeArg& output_def, int64_t value, bool is_rank1) { + if (!is_rank1) { + output_def.SetInferredShapeScalarValue(value); + // Keep exactly one channel populated: drop any stale values channel that getInputData() would + // otherwise prefer over this scalar. + output_def.GetMutableInferredShapeValues().reset(); + return; + } + + auto& inferred_values = output_def.GetMutableInferredShapeValues(); + if (!inferred_values.has_value()) { + inferred_values.emplace(); + } + inferred_values->clear_dim(); + inferred_values->add_dim()->set_dim_value(value); + // Keep exactly one channel populated: drop any stale scalar that the scalar-first reader would + // otherwise return ahead of this rank-1 value. + output_def.ClearInferredShapeScalarValue(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/div_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/div_op_data_propagation.cc index 2ea9b3047941c..6ded5903601a7 100644 --- a/onnxruntime/core/graph/data_propagation/div_op_data_propagation.cc +++ b/onnxruntime/core/graph/data_propagation/div_op_data_propagation.cc @@ -6,6 +6,7 @@ #include "core/graph/node_arg.h" #include "core/graph/onnx_protobuf.h" #include "core/providers/common.h" +#include "core/graph/data_propagation/data_propagation_value_utils.h" namespace onnxruntime { @@ -20,10 +21,17 @@ Status DivOpDataPropagation::infer() { return Status::OK(); } - if (input_0->GetInferredShapeScalarValue().has_value() && input_1->GetInferredShapeScalarValue().has_value()) { - output_def_.SetInferredShapeScalarValue( - input_0->GetInferredShapeScalarValue().value() / - input_1->GetInferredShapeScalarValue().value()); + int64_t lhs = 0; + int64_t rhs = 0; + bool lhs_is_rank1 = false; + bool rhs_is_rank1 = false; + if (TryGetSinglePropagatedShapeValue(*input_0, lhs, lhs_is_rank1) && + TryGetSinglePropagatedShapeValue(*input_1, rhs, rhs_is_rank1) && + rhs != 0) { + // Single-element operands may be carried as a rank-0 scalar or a rank-1 [1] value. Per ONNX + // broadcasting, the result is rank-1 if either operand is rank-1, otherwise a scalar; keep + // the propagated value's rank consistent with that so downstream consumers see the right rank. + SetSinglePropagatedShapeValue(output_def_, lhs / rhs, lhs_is_rank1 || rhs_is_rank1); } return Status::OK(); diff --git a/onnxruntime/core/graph/data_propagation/div_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/div_op_data_propagation.h index 9b32b59039282..0e4a76f9fdbab 100644 --- a/onnxruntime/core/graph/data_propagation/div_op_data_propagation.h +++ b/onnxruntime/core/graph/data_propagation/div_op_data_propagation.h @@ -41,7 +41,7 @@ class DivOpDataPropagation : public CustomDataPropagationBase { public: DivOpDataPropagation(const Node& node, NodeArg& output_def, - std::function func, + GetInitializedInputValuesFunc func, const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation, const logging::Logger& logger) noexcept : CustomDataPropagationBase(node, output_def, func, output_from_onnx_op_data_propagation, logger) {} diff --git a/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.cc index 39ac926a8553f..d747ac20219cf 100644 --- a/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.cc +++ b/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.cc @@ -6,6 +6,7 @@ #include "core/graph/node_arg.h" #include "core/graph/onnx_protobuf.h" #include "core/providers/common.h" +#include "core/graph/data_propagation/data_propagation_value_utils.h" namespace onnxruntime { @@ -58,7 +59,8 @@ Status GatherOpDataPropagation::infer() { ORT_TRY { TensorShapeVector indices; - ORT_RETURN_IF_ERROR(get_initialized_input_values_func_(input_1->Name(), indices)); + int indices_num_dims = -1; + ORT_RETURN_IF_ERROR(get_initialized_input_values_func_(input_1->Name(), indices, indices_num_dims)); if (indices.size() == 1) { // Note: Index value is expected to be within bounds [-s, s-1] along axis of size s auto index = static_cast( @@ -66,7 +68,34 @@ Status GatherOpDataPropagation::infer() { auto& dim = tensor_shape_proto.dim(index); if (dim.has_dim_value()) { - output_def_.SetInferredShapeScalarValue(dim.dim_value()); + // Gather output rank = data_rank - 1 + indices_rank. The "data" input here is the + // 1-D Shape output, so the output rank equals the indices rank. Route by that rank: + // * a 0-D scalar index -> a scalar output value, + // * a 1-D index -> a rank-1 [1] output value (so consumers that require a + // 1-D tensor, e.g. TopK's K input, still observe the correct rank), + // * a rank >= 2 index (or an unknown rank) -> decline, because the single-value + // channel cannot represent a rank >= 2 Gather output; falling back to ONNX data + // propagation is safer than fabricating a rank-1 value. + // + // The index rank is sourced canonically from the same constant initializer the index + // value came from (indices_num_dims). If that is unavailable (sentinel < 0), fall back + // to the indices NodeArg shape; an unknown shape then routes to decline. + int effective_num_dims = indices_num_dims; + if (effective_num_dims < 0) { + const ONNX_NAMESPACE::TensorShapeProto* indices_shape = input_1->Shape(); + effective_num_dims = (indices_shape != nullptr) ? indices_shape->dim_size() : -1; + } + + if (effective_num_dims == 0) { + // 0-D scalar index -> a scalar output value. + SetSinglePropagatedShapeValue(output_def_, dim.dim_value(), /*is_rank1=*/false); + } else if (effective_num_dims == 1) { + // 1-D index -> a rank-1 [1] output value. + SetSinglePropagatedShapeValue(output_def_, dim.dim_value(), /*is_rank1=*/true); + } + // rank >= 2 or unknown index rank (effective_num_dims < 0): leave the output unset and + // let ONNX data propagation handle it. The single-value channel cannot represent a + // rank >= 2 Gather output, so emitting a rank-1 value would fabricate a misleading rank. } } } diff --git a/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.h index c6b542e5af6e3..2f75223506974 100644 --- a/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.h +++ b/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.h @@ -40,7 +40,7 @@ class GatherOpDataPropagation : public CustomDataPropagationBase { public: GatherOpDataPropagation(const Node& node, NodeArg& output_def, - std::function func, + GetInitializedInputValuesFunc func, const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation, const logging::Logger& logger) noexcept : CustomDataPropagationBase(node, output_def, func, output_from_onnx_op_data_propagation, logger) {} diff --git a/onnxruntime/core/graph/data_propagation/mul_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/mul_op_data_propagation.cc index 4c5c25022e40b..3df2932f83cc9 100644 --- a/onnxruntime/core/graph/data_propagation/mul_op_data_propagation.cc +++ b/onnxruntime/core/graph/data_propagation/mul_op_data_propagation.cc @@ -6,6 +6,7 @@ #include "core/graph/node_arg.h" #include "core/graph/onnx_protobuf.h" #include "core/providers/common.h" +#include "core/graph/data_propagation/data_propagation_value_utils.h" namespace onnxruntime { @@ -20,10 +21,16 @@ Status MulOpDataPropagation::infer() { return Status::OK(); } - if (input_0->GetInferredShapeScalarValue().has_value() && input_1->GetInferredShapeScalarValue().has_value()) { - output_def_.SetInferredShapeScalarValue( - input_0->GetInferredShapeScalarValue().value() * - input_1->GetInferredShapeScalarValue().value()); + int64_t lhs = 0; + int64_t rhs = 0; + bool lhs_is_rank1 = false; + bool rhs_is_rank1 = false; + if (TryGetSinglePropagatedShapeValue(*input_0, lhs, lhs_is_rank1) && + TryGetSinglePropagatedShapeValue(*input_1, rhs, rhs_is_rank1)) { + // Single-element operands may be carried as a rank-0 scalar or a rank-1 [1] value. Per ONNX + // broadcasting, the result is rank-1 if either operand is rank-1, otherwise a scalar; keep + // the propagated value's rank consistent with that so downstream consumers see the right rank. + SetSinglePropagatedShapeValue(output_def_, lhs * rhs, lhs_is_rank1 || rhs_is_rank1); } return Status::OK(); diff --git a/onnxruntime/core/graph/data_propagation/mul_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/mul_op_data_propagation.h index b42a591eb4c7b..101a020351d1c 100644 --- a/onnxruntime/core/graph/data_propagation/mul_op_data_propagation.h +++ b/onnxruntime/core/graph/data_propagation/mul_op_data_propagation.h @@ -40,7 +40,7 @@ class MulOpDataPropagation : public CustomDataPropagationBase { public: MulOpDataPropagation(const Node& node, NodeArg& output_def, - std::function func, + GetInitializedInputValuesFunc func, const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation, const logging::Logger& logger) noexcept : CustomDataPropagationBase(node, output_def, func, output_from_onnx_op_data_propagation, logger) {} diff --git a/onnxruntime/core/graph/data_propagation/size_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/size_op_data_propagation.h index 184202254f078..c83788a19e911 100644 --- a/onnxruntime/core/graph/data_propagation/size_op_data_propagation.h +++ b/onnxruntime/core/graph/data_propagation/size_op_data_propagation.h @@ -18,7 +18,7 @@ class SizeOpDataPropagation : public CustomDataPropagationBase { public: SizeOpDataPropagation(const Node& node, NodeArg& output_def, - std::function func, + GetInitializedInputValuesFunc func, const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation, const logging::Logger& logger) noexcept : CustomDataPropagationBase(node, output_def, func, output_from_onnx_op_data_propagation, logger) {} diff --git a/onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.cc index b33d4f5589016..a720d0a36216e 100644 --- a/onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.cc +++ b/onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.cc @@ -43,7 +43,8 @@ Status SqueezeOpDataPropagation::infer() { if (node_.InputDefs().size() > 1) { const auto* input_1 = node_.InputDefs()[1]; ORT_TRY { - ORT_RETURN_IF_ERROR(get_initialized_input_values_func_(input_1->Name(), axes)); + [[maybe_unused]] int axes_num_dims = -1; + ORT_RETURN_IF_ERROR(get_initialized_input_values_func_(input_1->Name(), axes, axes_num_dims)); } ORT_CATCH(const std::exception& ex) { ORT_HANDLE_EXCEPTION([&]() { diff --git a/onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.h index 15e1e8458525a..8c8f98a9fddd4 100644 --- a/onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.h +++ b/onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.h @@ -16,7 +16,7 @@ class SqueezeOpDataPropagation : public CustomDataPropagationBase { public: SqueezeOpDataPropagation(const Node& node, NodeArg& output_def, - std::function func, + GetInitializedInputValuesFunc func, const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation, const logging::Logger& logger) noexcept : CustomDataPropagationBase(node, output_def, func, output_from_onnx_op_data_propagation, logger) {} diff --git a/onnxruntime/core/graph/data_propagation/sub_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/sub_op_data_propagation.cc index 4ee8ab546e707..a827a3d82de51 100644 --- a/onnxruntime/core/graph/data_propagation/sub_op_data_propagation.cc +++ b/onnxruntime/core/graph/data_propagation/sub_op_data_propagation.cc @@ -6,6 +6,7 @@ #include "core/graph/node_arg.h" #include "core/graph/onnx_protobuf.h" #include "core/providers/common.h" +#include "core/graph/data_propagation/data_propagation_value_utils.h" namespace onnxruntime { @@ -20,10 +21,16 @@ Status SubOpDataPropagation::infer() { return Status::OK(); } - if (input_0->GetInferredShapeScalarValue().has_value() && input_1->GetInferredShapeScalarValue().has_value()) { - output_def_.SetInferredShapeScalarValue( - input_0->GetInferredShapeScalarValue().value() - - input_1->GetInferredShapeScalarValue().value()); + int64_t lhs = 0; + int64_t rhs = 0; + bool lhs_is_rank1 = false; + bool rhs_is_rank1 = false; + if (TryGetSinglePropagatedShapeValue(*input_0, lhs, lhs_is_rank1) && + TryGetSinglePropagatedShapeValue(*input_1, rhs, rhs_is_rank1)) { + // Single-element operands may be carried as a rank-0 scalar or a rank-1 [1] value. Per ONNX + // broadcasting, the result is rank-1 if either operand is rank-1, otherwise a scalar; keep + // the propagated value's rank consistent with that so downstream consumers see the right rank. + SetSinglePropagatedShapeValue(output_def_, lhs - rhs, lhs_is_rank1 || rhs_is_rank1); } return Status::OK(); diff --git a/onnxruntime/core/graph/data_propagation/sub_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/sub_op_data_propagation.h index a9be294b8f62f..90186218540b0 100644 --- a/onnxruntime/core/graph/data_propagation/sub_op_data_propagation.h +++ b/onnxruntime/core/graph/data_propagation/sub_op_data_propagation.h @@ -41,7 +41,7 @@ class SubOpDataPropagation : public CustomDataPropagationBase { public: SubOpDataPropagation(const Node& node, NodeArg& output_def, - std::function func, + GetInitializedInputValuesFunc func, const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation, const logging::Logger& logger) noexcept : CustomDataPropagationBase(node, output_def, func, output_from_onnx_op_data_propagation, logger) {} diff --git a/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.cc index ff5ef6853bc78..52cacd4ccfbe8 100644 --- a/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.cc +++ b/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.cc @@ -34,6 +34,22 @@ Status UnsqueezeOpDataPropagation::infer() { } else if (input_0->GetInferredShapeValues().has_value()) { const auto& tensor_shape_proto = input_0->GetInferredShapeValues().value(); + // Decline when unsqueezing a single-element (scalar-like, rank-1 [1]) value would yield a + // rank >= 2 result the values channel cannot faithfully represent (it would otherwise fabricate + // a misleading [1, value]). A single-element value (dim_size == 1) is the channel's + // representation of a scalar-like quantity; any unsqueeze lifts it to rank >= 2, which the + // single-value channel cannot represent. Multi-element shape vectors (dim_size > 1) are + // legitimate and left untouched. + // + // This decline is locked end to end by UnsqueezeSingleValueDeclineTest via a rank-lowering + // Squeeze (Shape -> Gather([-1]) -> Unsqueeze -> Squeeze -> Range): with the decline, Range's + // length stays symbolic; relaxing it fabricates [1, value], which makes the downstream Range + // input non-scalar and the model fails to load -- so the behavior IS observable and the test + // discriminates it. Multi-element shape vectors are unaffected. + if (tensor_shape_proto.dim_size() == 1) { + return Status::OK(); + } + // The TensorShapeProto (inferred shape values) should have rank > 0 and // all the dimensions have values (not symbolic) if (tensor_shape_proto.dim_size() > 0) { @@ -54,7 +70,8 @@ Status UnsqueezeOpDataPropagation::infer() { if (node_.InputDefs().size() > 1) { const auto* input_1 = node_.InputDefs()[1]; ORT_TRY { - ORT_RETURN_IF_ERROR(get_initialized_input_values_func_(input_1->Name(), axes)); + [[maybe_unused]] int axes_num_dims = -1; + ORT_RETURN_IF_ERROR(get_initialized_input_values_func_(input_1->Name(), axes, axes_num_dims)); } ORT_CATCH(const std::exception& ex) { ORT_HANDLE_EXCEPTION([&]() { diff --git a/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.h index 26b6587aa93fc..dfd309262faf5 100644 --- a/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.h +++ b/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.h @@ -15,7 +15,7 @@ class UnsqueezeOpDataPropagation : public CustomDataPropagationBase { public: UnsqueezeOpDataPropagation(const Node& node, NodeArg& output_def, - std::function func, + GetInitializedInputValuesFunc func, const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation, const logging::Logger& logger) noexcept : CustomDataPropagationBase(node, output_def, func, output_from_onnx_op_data_propagation, logger) {} diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index f87d530f9fbbf..0b91a2c1c651c 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -2967,7 +2967,8 @@ Status Graph::SaveShapeValuesFromDataPropagation(const Node& node, NodeArg& output_def, const TypeProto& onnx_inferred_type_after_data_propagation) const { // Helper function to get the input value if it's a initializer. - auto get_initialized_input_values_func = [&](const std::string& input_name, TensorShapeVector& input_values) + auto get_initialized_input_values_func = [&](const std::string& input_name, TensorShapeVector& input_values, + int& num_dims) -> Status { const TensorProto* initializer = this->GetConstantInitializer(input_name, true); @@ -2977,6 +2978,11 @@ Status Graph::SaveShapeValuesFromDataPropagation(const Node& node, auto tensor_shape = utils::GetTensorShapeFromTensorProto(*initializer); size_t element_cnt = narrow(tensor_shape.Size()); + // Report the initializer's rank so callers can distinguish a 0-D scalar from a rank-1 + // single-element tensor (both have a single element). Sourcing the rank from the same + // TensorProto the values come from keeps value and rank consistent. + num_dims = static_cast(tensor_shape.NumDimensions()); + // Check if this is in-memory external data (data stored in OrtValue) if (utils::HasExternalDataInMemory(*initializer)) { // Try to get the OrtValue for this initializer diff --git a/onnxruntime/test/framework/shape_inference_test.cc b/onnxruntime/test/framework/shape_inference_test.cc index 6d92acf6bae78..5ac99d87ec50d 100644 --- a/onnxruntime/test/framework/shape_inference_test.cc +++ b/onnxruntime/test/framework/shape_inference_test.cc @@ -7,6 +7,7 @@ #include "gtest/gtest.h" #include "core/common/span_utils.h" #include "core/graph/model.h" +#include "core/graph/data_propagation/data_propagation_value_utils.h" #include "core/session/onnxruntime_cxx_api.h" #include "test/framework/model_builder_utils.h" #include "test/util/include/asserts.h" @@ -191,6 +192,356 @@ TEST(ShapeInferenceV2Test, PartialDataPropagationTest) { } } +// Regression test for the Shape -> Gather(1-D index) -> TopK rank-drop. +// +// The Gather index is the 1-D constant [-1] (rank 1), so per ONNX Gather semantics the Gather +// output is a rank-1, single-element tensor -- exactly what TopK's K input requires. Previously +// the Gather custom data propagation scalarized that single element (dropping the rank), +// producing a 0-D K initializer that ONNX TopK shape inference rejected with +// "K input must be a one-dimensional tensor of size 1." at model load time. +// +// The failure reproduces even at ORT_DISABLE_ALL (constant folding never runs there), which +// proves the cause is shape-inference data propagation, not constant folding. This test loads +// the model at the disabled, basic, and all-optimization levels, which bracket the data-propagation +// path: data propagation runs in the pre-optimization Graph::Resolve pass, so it is independent of +// the graph-optimization level (ORT_ENABLE_ALL already applies the extended and layout transformers, +// so they need not be enumerated separately). The test asserts the rank-1 K shape is preserved so the +// model loads and TopK output shapes are inferred correctly. +TEST(ShapeInferenceV2Test, GatherToTopKRankPreservationTest) { + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_topk.onnx"); + + for (auto opt_level : {ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_ALL}) { + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(opt_level); + + // Loading must succeed at each level; before the fix this threw a + // ShapeInferenceError at ORT_DISABLE_ALL (and above). + Ort::Session session(*ort_env, model_path, session_options); + + ORT_ENFORCE(session.GetOutputCount() == 2); + + // K is data-propagated as the last dimension of X (2000), so both TopK outputs are 1-D + // tensors of length 2000. + for (size_t output_index = 0; output_index < 2; ++output_index) { + Ort::TypeInfo type_info = session.GetOutputTypeInfo(output_index); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector output_shape = tensor_info.GetShape(); + EXPECT_TRUE(output_shape.size() == 1) << "TopK output should be a 1-D tensor"; + EXPECT_TRUE(output_shape[0] == 2000) << "TopK output length should be the inferred K value (2000)"; + } + } +} + +// Regression test for the Shape -> Gather(1-D index) -> Mul -> TopK chain. +// +// This covers the elementwise data-propagation consumers (Add/Sub/Mul/Div) with rank-1 +// single-element operands. Both Gather indices are 1-D constants, so each Gather output is a +// rank-1 single-element value; Mul must keep propagating that (as a rank-1 [1] value) so the +// downstream TopK still receives a valid 1-D K. Previously the elementwise ops only handled +// rank-0 scalar operands, so once the Gather producer was fixed to emit rank-1 values the chain +// would have silently stopped propagating at Mul (degrading TopK's K to a symbolic dim). +// +// Asserting that both TopK outputs resolve to the concrete length 2000 (== 50 * 40) proves the +// value propagated end to end through Mul rather than being lost. +TEST(ShapeInferenceV2Test, GatherMulToTopKRankPreservationTest) { + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_mul_topk.onnx"); + + for (auto opt_level : {ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_ALL}) { + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(opt_level); + + // Loading must succeed at each level; before the fix this threw a + // ShapeInferenceError at ORT_DISABLE_ALL (and above). + Ort::Session session(*ort_env, model_path, session_options); + + ORT_ENFORCE(session.GetOutputCount() == 2); + + // K = 50 * 40 = 2000 is data-propagated through the two Gathers and the Mul, so both TopK + // outputs are 1-D tensors of length 2000. + for (size_t output_index = 0; output_index < 2; ++output_index) { + Ort::TypeInfo type_info = session.GetOutputTypeInfo(output_index); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector output_shape = tensor_info.GetShape(); + EXPECT_TRUE(output_shape.size() == 1) << "TopK output should be a 1-D tensor"; + EXPECT_TRUE(output_shape[0] == 2000) << "TopK output length should be the propagated K value (2000)"; + } + } +} + +// Negative regression test for the single-element guard in the elementwise data-propagation +// consumers (Add/Sub/Mul/Div), exercised end to end. +// +// Shape(X[3, 4]) -> S is a rank-1 MULTI-element value [3, 4]. Mul(S, S) -> M must stay the +// multi-element value [9, 16] so ConstantOfShape(M) resolves to the rank-2 shape [9, 16]. For a +// multi-element data-propagation output ONNX's own Mul data propagation supplies the value and +// ORT's custom elementwise propagation is intentionally NOT engaged (it is dispatched only for a +// rank-0 result, see CreateCustomDataPropagation). This is an end-to-end non-regression guard that +// the rank-routing changes did not disturb the multi-element chain; the single-element guard inside +// the shared helper is exercised directly by SinglePropagatedShapeValueGuardTest below. +TEST(ShapeInferenceV2Test, ShapeMulMultiElementNoScalarCollapseTest) { + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_shape_mul_constantofshape.onnx"); + + for (auto opt_level : {ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_ALL}) { + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(opt_level); + + Ort::Session session(*ort_env, model_path, session_options); + + ORT_ENFORCE(session.GetOutputCount() == 1); + + Ort::TypeInfo type_info = session.GetOutputTypeInfo(0); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector output_shape = tensor_info.GetShape(); + + // ConstantOfShape(Mul(Shape(X), Shape(X))) == ConstantOfShape([9, 16]) -> rank-2 [9, 16]. + // A collapse of the multi-element value to a scalar would yield rank 1 here. + EXPECT_TRUE(output_shape.size() == 2) << "Output should stay rank 2; the multi-element value must not collapse"; + EXPECT_TRUE(output_shape.size() == 2 && output_shape[0] == 9 && output_shape[1] == 16) + << "ConstantOfShape output should be the propagated multi-element shape [9, 16]"; + } +} + +// Regression for microsoft/onnxruntime#29072. +// End-to-end lock for Gather's custom data-propagation DECLINE on a rank-2 index. +// +// The Gather index is the constant [[-1]] of shape [1, 1] (rank 2), so the Gather output rank is +// data_rank - 1 + index_rank = 1 - 1 + 2 = 2. The single-value data-propagation channel can carry +// only rank-0 (scalar) and rank-1 [1] values, so Gather's custom propagation must DECLINE rather +// than fabricate a rank-1 value. The index is a raw graph initializer (not surfaced through the +// data-propagation getInputData channel), so ONNX's own Gather data propagator bails and control +// reaches our custom decline branch. +// +// The decline is made OBSERVABLE through a rank-lowering Squeeze: +// Shape(X) -> Gather([[ -1 ]]) -> Squeeze -> Range(0, K, 1) +// With the correct decline, nothing is propagated, Squeeze has no value to lower, and Range's limit +// stays symbolic -- so the Range output length is an UNKNOWN dimension. If the decline were relaxed +// to emit a rank-1 [1] value, Squeeze would lower it to the scalar 2000 and Range would resolve to a +// concrete length 2000. Asserting the length is NOT the concrete 2000 discriminates the correct +// decline from the rank-fabricating bug -- real end-to-end coverage of the decline branch (no +// shared-helper abstraction, no tautological unit test). +// +// This runs ONLY at ORT_DISABLE_ALL: the chain is entirely constant (a static input shape feeding +// Shape/Gather/Squeeze/Range over initializers), so at ORT_ENABLE_BASIC and above constant folding +// resolves Range to 2000 regardless of data propagation, which would mask the decline. Data +// propagation runs in the pre-optimization Graph::Resolve pass, so ORT_DISABLE_ALL exercises the +// decline branch in isolation. +TEST(ShapeInferenceV2Test, GatherRank2IndexDeclineTest) { + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_rank2_decline.onnx"); + + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); + + Ort::Session session(*ort_env, model_path, session_options); + + ORT_ENFORCE(session.GetOutputCount() == 1); + + Ort::TypeInfo type_info = session.GetOutputTypeInfo(0); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector output_shape = tensor_info.GetShape(); + + // Range output is 1-D; its length must stay symbolic (unknown, reported as -1) because the rank-2 + // Gather index declines rather than propagating a (rank-fabricated) concrete K. A relaxed decline + // would fabricate a rank-1 value that Squeeze lowers to scalar 2000, making the length concretely + // 2000 -- which this assertion rejects. + EXPECT_TRUE(output_shape.size() == 1) << "Range output should be a 1-D tensor"; + EXPECT_EQ(output_shape[0], -1) + << "Range length must stay symbolic (reported as -1); the rank-2 Gather index must decline " + "rather than fabricate a concrete K (a relaxed decline concretizes it to 2000)"; +} + +// Regression for microsoft/onnxruntime#29072. +// End-to-end lock for Unsqueeze's custom data-propagation DECLINE on a single-element value. +// +// Shape(X) -> Gather([-1]) produces a rank-1 single-element value (the last dimension of X, 2000). +// Unsqueezing that single-element (scalar-like, rank-1 [1]) value would yield a rank >= 2 result +// ([1, 2000]) that the single-value channel cannot faithfully represent, so Unsqueeze's custom +// propagation must DECLINE rather than fabricate the misleading [1, value]. +// +// The decline is made OBSERVABLE through a rank-lowering Squeeze: +// Shape(X) -> Gather([-1]) -> Unsqueeze([0]) -> Squeeze -> Range(0, K, 1) +// With the correct decline, no value is propagated past Unsqueeze, Squeeze has nothing to lower, and +// Range's limit stays symbolic, so the model loads with an unknown Range length. A relaxed decline +// that fabricated [1, 2000] would propagate a non-scalar value to Range, which ONNX Range shape +// inference rejects ("Input to 'Range' op should be scalars") -- the model then fails to load. So +// the decline IS observable: correct -> loads with a symbolic length; relaxed -> load error. This +// test asserts the model loads and the length is not the (rank-fabricated) concrete 2000. +// +// Like GatherRank2IndexDeclineTest this runs ONLY at ORT_DISABLE_ALL: the chain is entirely constant +// over a static input shape, so at ORT_ENABLE_BASIC and above constant folding resolves Range +// regardless of data propagation, masking the decline. +TEST(ShapeInferenceV2Test, UnsqueezeSingleValueDeclineTest) { + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_unsqueeze_decline.onnx"); + + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); + + Ort::Session session(*ort_env, model_path, session_options); + + ORT_ENFORCE(session.GetOutputCount() == 1); + + Ort::TypeInfo type_info = session.GetOutputTypeInfo(0); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector output_shape = tensor_info.GetShape(); + + // Range output is 1-D; its length must stay symbolic because Unsqueeze declines rather than + // fabricating a rank >= 2 value. A relaxed decline fabricates [1, 2000], which makes Range's input + // non-scalar and the model fails to load -- so any regression here is caught either as a load + // failure or as a concrete length below. + EXPECT_TRUE(output_shape.size() == 1) << "Range output should be a 1-D tensor"; + EXPECT_EQ(output_shape[0], -1) + << "Range length must stay symbolic (reported as -1); unsqueezing a single-element value " + "must decline (a relaxed decline fabricates [1, value] and the model fails to load)"; +} + +// Lock test for the Shape -> Gather(1-D index) -> Squeeze -> Range chain. +// +// Squeeze of a rank-1 single-element value [1] removes the size-1 dimension and yields a 0-D +// scalar -- the strict-improvement behavior of Squeeze's custom data propagation. This test locks +// that improvement end to end: the Gather produces a rank-1 single element (the last dimension of +// X, 2000); Squeeze must propagate it as the scalar 2000; and Range(0, K, 1) must then resolve to +// a concrete 1-D tensor of length 2000. If Squeeze dropped, corrupted, or failed to scalarize the +// value, Range's length would stay symbolic (size 0 here) instead of the concrete 2000. +TEST(ShapeInferenceV2Test, GatherSqueezeRangeRankPreservationTest) { + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_squeeze_range.onnx"); + + for (auto opt_level : {ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_ALL}) { + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(opt_level); + + Ort::Session session(*ort_env, model_path, session_options); + + ORT_ENFORCE(session.GetOutputCount() == 1); + + Ort::TypeInfo type_info = session.GetOutputTypeInfo(0); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector output_shape = tensor_info.GetShape(); + + // Range(0, Squeeze(Gather(Shape(X), [-1])), 1) == Range(0, 2000, 1) -> 1-D length 2000. + EXPECT_TRUE(output_shape.size() == 1) << "Range output should be a 1-D tensor"; + EXPECT_TRUE(output_shape.size() == 1 && output_shape[0] == 2000) + << "Range length should be the scalar K (2000) propagated through Squeeze"; + } +} + +// Unit test for the single-element guard in TryGetSinglePropagatedShapeValue, the helper shared by +// the Add/Sub/Mul/Div data-propagation consumers. +// +// A NodeArg can carry a propagated shape value in one of two non-interchangeable channels: a rank-0 +// scalar (inferred_scalar_value_) or a rank>=1 list of values (inferred_shape_values_). The helper +// must surface a single value from either a scalar source or a rank-1 SINGLE-element source (and +// report its rank), and must DECLINE for a rank-1 MULTI-element source. Declining is load-bearing: +// silently using element[0] of a multi-element value would let an elementwise op emit a bogus +// scalar product that getInputData() would then hand to a rank-sensitive consumer. This test +// exercises the guard directly -- a relaxed guard that accepted element[0] fails the multi-element +// case below. +TEST(ShapeInferenceV2Test, SinglePropagatedShapeValueGuardTest) { + // Scalar channel -> accepted, reported as rank-0. + { + NodeArg arg("scalar_arg", nullptr); + arg.SetInferredShapeScalarValue(5); + int64_t value = 0; + bool is_rank1 = true; + EXPECT_TRUE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); + EXPECT_EQ(value, 5); + EXPECT_FALSE(is_rank1); + } + + // Rank-1 single-element value -> accepted, reported as rank-1. + { + NodeArg arg("rank1_single_arg", nullptr); + auto& values = arg.GetMutableInferredShapeValues(); + values.emplace(); + values->add_dim()->set_dim_value(7); + int64_t value = 0; + bool is_rank1 = false; + EXPECT_TRUE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); + EXPECT_EQ(value, 7); + EXPECT_TRUE(is_rank1); + } + + // Rank-1 MULTI-element value -> DECLINED (the guard); element[0] must never be used. + { + NodeArg arg("rank1_multi_arg", nullptr); + auto& values = arg.GetMutableInferredShapeValues(); + values.emplace(); + values->add_dim()->set_dim_value(3); + values->add_dim()->set_dim_value(4); + int64_t value = -1; + bool is_rank1 = false; + EXPECT_FALSE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)) + << "A rank-1 multi-element value must be declined, not collapsed to element[0]"; + } + + // Rank-1 single element with a SYMBOLIC (no concrete value) dim -> declined. + { + NodeArg arg("rank1_symbolic_arg", nullptr); + auto& values = arg.GetMutableInferredShapeValues(); + values.emplace(); + values->add_dim()->set_dim_param("N"); + int64_t value = -1; + bool is_rank1 = false; + EXPECT_FALSE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)) + << "A symbolic single-element value has no concrete value and must be declined"; + } + + // No propagated value on either channel -> declined. + { + NodeArg arg("empty_arg", nullptr); + int64_t value = -1; + bool is_rank1 = false; + EXPECT_FALSE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); + } +} + +// Unit test that SetSinglePropagatedShapeValue is correct-by-construction: it populates exactly one +// channel and clears the other. This pins the invariant the elementwise/Gather consumers rely on -- +// the scalar-first reader (TryGetSinglePropagatedShapeValue) and the values-first getInputData() can +// only agree on rank if no NodeArg ever carries both channels. A setter that left the opposite +// channel populated would let a stale value win (reintroducing a rank mismatch), so each case below +// pre-populates the opposite channel and asserts the setter cleared it. +TEST(ShapeInferenceV2Test, SetSinglePropagatedShapeValueKeepsSingleChannelTest) { + // Scalar write must clear a stale values channel. + { + NodeArg arg("scalar_over_stale_values", nullptr); + auto& stale = arg.GetMutableInferredShapeValues(); + stale.emplace(); + stale->add_dim()->set_dim_value(99); + + SetSinglePropagatedShapeValue(arg, 5, /*is_rank1=*/false); + + ASSERT_TRUE(arg.GetInferredShapeScalarValue().has_value()); + EXPECT_EQ(arg.GetInferredShapeScalarValue().value(), 5); + EXPECT_FALSE(arg.GetInferredShapeValues().has_value()) + << "Scalar write must clear the values channel that getInputData() would otherwise prefer"; + + int64_t value = 0; + bool is_rank1 = true; + EXPECT_TRUE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); + EXPECT_EQ(value, 5); + EXPECT_FALSE(is_rank1); + } + + // Rank-1 write must clear a stale scalar channel. + { + NodeArg arg("values_over_stale_scalar", nullptr); + arg.SetInferredShapeScalarValue(99); + + SetSinglePropagatedShapeValue(arg, 7, /*is_rank1=*/true); + + EXPECT_FALSE(arg.GetInferredShapeScalarValue().has_value()) + << "Rank-1 write must clear the scalar channel that the scalar-first reader would otherwise return"; + ASSERT_TRUE(arg.GetInferredShapeValues().has_value()); + ASSERT_EQ(arg.GetInferredShapeValues()->dim_size(), 1); + EXPECT_EQ(arg.GetInferredShapeValues()->dim(0).dim_value(), 7); + + int64_t value = 0; + bool is_rank1 = false; + EXPECT_TRUE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); + EXPECT_EQ(value, 7); + EXPECT_TRUE(is_rank1); + } +} + namespace { struct MyCustomKernelWithOptionalInput { MyCustomKernelWithOptionalInput(const OrtKernelInfo* info) { diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_gather_mul_topk.onnx b/onnxruntime/test/testdata/test_shape_data_propagation_gather_mul_topk.onnx new file mode 100644 index 0000000000000000000000000000000000000000..66db1819a97d4a2ca053655c9119fab32e41464c GIT binary patch literal 424 zcmZutO;5r=5be-Hw^Kzjev$RY$u!{r-nh{;8iIJi7{jGm!wSSoo3jJJ5jD>eyee>RfE7sM^u}lYkffj2U^6mC=pr{a;)3#75rav zi7-Y838vhgeI&+$tqP2kxj}?5C3qs3aQ2vHQ}gWTDmdFO-W?7o2SFR%MC2=3lwFPK z%A|>F7qOZfYjHeNNt1vWbleeG;a|>x8x9XccM$&OS)AAZ)y5+R#IV7J>RvC%I3*uF iMK{rA-zP7ps9YZ!TjZ2!_NK{+Azh{3lc9rFhkgKqIdIMZ literal 0 HcmV?d00001 diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_gather_mul_topk.py b/onnxruntime/test/testdata/test_shape_data_propagation_gather_mul_topk.py new file mode 100644 index 0000000000000..c1fa947fe2b53 --- /dev/null +++ b/onnxruntime/test/testdata/test_shape_data_propagation_gather_mul_topk.py @@ -0,0 +1,62 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +# Regression fixture for the Shape -> Gather(1-D index) -> Mul -> TopK chain. +# +# This exercises the elementwise data-propagation consumers (Mul) with rank-1 +# single-element operands. Both Gather indices are 1-D constants, so per ONNX +# Gather semantics each Gather output is a rank-1, single-element tensor. The +# Mul of those two rank-1 values must keep propagating (as a rank-1 [1] value) +# so the downstream TopK still receives a valid 1-D K and its output shape is +# inferred concretely. +# +# Before the elementwise ops were made rank-aware, Mul only handled rank-0 +# scalar operands; with the Gather producer now emitting rank-1 values, the +# chain's data propagation would otherwise silently stop at Mul (degrading the +# TopK output dim to a symbolic value). This fixture locks in that the chain +# keeps propagating end to end. + +from onnx import IR_VERSION, TensorProto, checker, helper, save + +# 2-D input whose static dims (50, 40) drive the K computation: 50 * 40 == 2000. +x2d = helper.make_tensor_value_info("X2D", TensorProto.FLOAT, [50, 40]) +# 1-D input that TopK selects from; its length matches the computed K (2000). +x1d = helper.make_tensor_value_info("X1D", TensorProto.FLOAT, [2000]) + +values_out = helper.make_tensor_value_info("V", TensorProto.FLOAT, ["topk"]) +indices_out = helper.make_tensor_value_info("I", TensorProto.INT64, ["topk"]) + +# 1-D single-element Gather indices (rank 1), as emitted by common exporters. +idx_first = helper.make_tensor("idx_first", TensorProto.INT64, [1], [0]) +idx_last = helper.make_tensor("idx_last", TensorProto.INT64, [1], [-1]) + +# Shape(X2D) -> S : 1-D int64 tensor [50, 40]. +shape_node = helper.make_node("Shape", inputs=["X2D"], outputs=["S"], name="ShapeNode") +# Gather(S, [0]) -> A : rank-1 size-1 int64 tensor [50]. +gather_first = helper.make_node("Gather", inputs=["S", "idx_first"], outputs=["A"], axis=0, name="GatherFirst") +# Gather(S, [-1]) -> B : rank-1 size-1 int64 tensor [40]. +gather_last = helper.make_node("Gather", inputs=["S", "idx_last"], outputs=["B"], axis=0, name="GatherLast") +# Mul(A, B) -> K : rank-1 size-1 int64 tensor [2000]. +mul_node = helper.make_node("Mul", inputs=["A", "B"], outputs=["K"], name="MulNode") +# TopK(X1D, K) -> V, I : returns all 2000 elements sorted (K == 2000). +topk_node = helper.make_node("TopK", inputs=["X1D", "K"], outputs=["V", "I"], axis=-1, largest=1, name="TopKNode") + +graph = helper.make_graph( + nodes=[shape_node, gather_first, gather_last, mul_node, topk_node], + name="Shape_Gather_Mul_TopK_Model", + inputs=[x2d, x1d], + outputs=[values_out, indices_out], + initializer=[idx_first, idx_last], +) + +model = helper.make_model( + graph, + opset_imports=[helper.make_operatorsetid("", 18)], + producer_name="onnx-example-generator", +) +model.ir_version = IR_VERSION + +checker.check_model(model) +save(model, "test_shape_data_propagation_gather_mul_topk.onnx") + +print("Model saved to test_shape_data_propagation_gather_mul_topk.onnx") diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_gather_rank2_decline.onnx b/onnxruntime/test/testdata/test_shape_data_propagation_gather_rank2_decline.onnx new file mode 100644 index 0000000000000000000000000000000000000000..eaedaac4b6a0187c1a2233ecc07e825ce3f6848d GIT binary patch literal 354 zcmYjN%}T>S5YD90Y$g=z3Q~fIcntJXZ#{-W4n+^iMZ7Gl8PZ_f)NO>|)A%O7f^VeT zT?1Vh_f(H4vYiUuFA}#w)ufBbpBRp~_lc?Ov9X)gv-JV#8cL*h7 zE&DBL(NqJ6Ji**jd+pTWG{!&<&3^8Efsa=RIfqRWm+v1$pYFl%pWqT>Emcio-%H&i zY(`#D+e;2b5)}vE)v%D?*fzq238TF(#hPB)Mv=HTb=&ls0fSk%!r$>gcGNrN#se=q tj&>T1nvH+5tGK{(27HP>o0ti3!hX)4jZ^T06B;Mxk^4kTDhtq?@jn%8WPktw literal 0 HcmV?d00001 diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_gather_rank2_decline.py b/onnxruntime/test/testdata/test_shape_data_propagation_gather_rank2_decline.py new file mode 100644 index 0000000000000..be3225e998d73 --- /dev/null +++ b/onnxruntime/test/testdata/test_shape_data_propagation_gather_rank2_decline.py @@ -0,0 +1,70 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +# Lock fixture for Gather's custom data-propagation DECLINE on a rank-2 index. +# +# The Gather index is the constant [[-1]] of shape [1, 1] (rank 2). Per ONNX +# Gather semantics the output rank = data_rank - 1 + index_rank = 1 - 1 + 2 = 2, +# so the Gather output is a rank-2, single-element tensor [[2000]]. The +# single-value data-propagation channel can represent only rank-0 (scalar) and +# rank-1 [1] values, so Gather's custom data propagation must DECLINE here: +# emitting a rank-1 value would fabricate a rank the channel cannot honestly +# carry. +# +# This fixture makes that decline OBSERVABLE through a rank-lowering Squeeze: +# Shape(X) -> Gather([[ -1 ]]) -> Squeeze -> Range(0, K, 1) +# Because the index is a constant initializer, ONNX's own Gather data +# propagator bails (it has no getInputData for a constant index), so control +# reaches our custom decline branch. With the correct decline, no value is +# propagated, Squeeze has nothing to lower, and Range's limit stays symbolic -- +# the Range output length is therefore an unknown dimension. +# +# If the decline were relaxed to emit a rank-1 [1] value, Squeeze would lower it +# to the scalar 2000 and Range(0, 2000, 1) would resolve to a concrete length +# 2000. Asserting the Range length is NOT the concrete 2000 thus locks the +# decline end to end (it discriminates the correct-decline behavior from the +# rank-fabricating bug). + +from onnx import IR_VERSION, TensorProto, checker, helper, save + +# Graph input: a 3-D float tensor whose last static dimension is 2000. +input_tensor = helper.make_tensor_value_info("X", TensorProto.FLOAT, [3, 4, 2000]) + +# Output: Range result, a 1-D int64 tensor whose length is the propagated K +# (symbolic when the decline is correct). +range_out = helper.make_tensor_value_info("R", TensorProto.INT64, ["range_len"]) + +# Gather index is the rank-2 constant [[-1]] of shape [1, 1] (last dimension). +gather_idx = helper.make_tensor("gather_idx", TensorProto.INT64, [1, 1], [-1]) +# Range start/delta as 0-D scalar int64 constants. +range_start = helper.make_tensor("range_start", TensorProto.INT64, [], [0]) +range_delta = helper.make_tensor("range_delta", TensorProto.INT64, [], [1]) + +# Shape(X) -> S : 1-D int64 tensor [3, 4, 2000]. +shape_node = helper.make_node("Shape", inputs=["X"], outputs=["S"], name="ShapeNode") +# Gather(S, [[-1]]) -> G : rank-2 [1, 1] int64 value [[2000]]. +gather_node = helper.make_node("Gather", inputs=["S", "gather_idx"], outputs=["G"], axis=0, name="GatherNode") +# Squeeze(G) -> K : 0-D scalar int64 value (only when a value is propagated). +squeeze_node = helper.make_node("Squeeze", inputs=["G"], outputs=["K"], name="SqueezeNode") +# Range(0, K, 1) -> R : 1-D int64 tensor whose length is K (symbolic on decline). +range_node = helper.make_node("Range", inputs=["range_start", "K", "range_delta"], outputs=["R"], name="RangeNode") + +graph = helper.make_graph( + nodes=[shape_node, gather_node, squeeze_node, range_node], + name="Shape_Gather_Rank2_Decline_Model", + inputs=[input_tensor], + outputs=[range_out], + initializer=[gather_idx, range_start, range_delta], +) + +model = helper.make_model( + graph, + opset_imports=[helper.make_operatorsetid("", 18)], + producer_name="onnx-example-generator", +) +model.ir_version = IR_VERSION + +checker.check_model(model) +save(model, "test_shape_data_propagation_gather_rank2_decline.onnx") + +print("Model saved to test_shape_data_propagation_gather_rank2_decline.onnx") diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_gather_squeeze_range.onnx b/onnxruntime/test/testdata/test_shape_data_propagation_gather_squeeze_range.onnx new file mode 100644 index 0000000000000000000000000000000000000000..2b2a34996dfbb21cf582a5a9df70111ef2367741 GIT binary patch literal 346 zcmYjN%}T>S5YEKbB(oIb2&DuC@fheq@zi539)cc{OYyR-WJn7!soNETPvcwo8orUH zyA5c__g!kVB=qH?2l{L)o?H? Gather(1-D index) -> Squeeze -> Range chain. +# +# The Gather index is the 1-D constant [-1] (rank 1), so the Gather output is a +# rank-1, single-element int64 value (the last dimension of X, 2000). Squeeze +# then removes the size-1 dimension, producing a 0-D scalar -- this is the +# correct, strict-improvement behavior of Squeeze's custom data propagation +# (rank-1 [1] -> scalar), and this fixture locks it: the scalar must keep its +# value (2000) so the downstream Range produces a concrete 1-D length. +# +# Range(start=0, limit=K, delta=1) with the propagated scalar K=2000 yields a +# 1-D tensor of length 2000. Asserting that output length proves Squeeze +# propagated the single element as a scalar of the correct value (and did not +# drop or corrupt it). + +from onnx import IR_VERSION, TensorProto, checker, helper, save + +# Graph input: a 1-D float tensor of static length 2000. +input_tensor = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2000]) + +# Output: Range result, a 1-D int64 tensor whose length is the propagated K. +range_out = helper.make_tensor_value_info("R", TensorProto.INT64, ["range_len"]) + +# Gather index is the 1-D constant [-1] (last dimension). +gather_idx = helper.make_tensor("gather_idx", TensorProto.INT64, [1], [-1]) +# Range start/delta as 0-D scalar int64 constants. +range_start = helper.make_tensor("range_start", TensorProto.INT64, [], [0]) +range_delta = helper.make_tensor("range_delta", TensorProto.INT64, [], [1]) + +# Shape(X) -> S : 1-D int64 tensor [2000]. +shape_node = helper.make_node("Shape", inputs=["X"], outputs=["S"], name="ShapeNode") +# Gather(S, [-1]) -> K1 : rank-1 size-1 int64 value [2000]. +gather_node = helper.make_node("Gather", inputs=["S", "gather_idx"], outputs=["K1"], axis=0, name="GatherNode") +# Squeeze(K1) -> K : 0-D scalar int64 value 2000 (size-1 dim removed). +squeeze_node = helper.make_node("Squeeze", inputs=["K1"], outputs=["K"], name="SqueezeNode") +# Range(0, K, 1) -> R : 1-D int64 tensor of length K (== 2000). +range_node = helper.make_node("Range", inputs=["range_start", "K", "range_delta"], outputs=["R"], name="RangeNode") + +graph = helper.make_graph( + nodes=[shape_node, gather_node, squeeze_node, range_node], + name="Shape_Gather_Squeeze_Range_Model", + inputs=[input_tensor], + outputs=[range_out], + initializer=[gather_idx, range_start, range_delta], +) + +model = helper.make_model( + graph, + opset_imports=[helper.make_operatorsetid("", 18)], + producer_name="onnx-example-generator", +) +model.ir_version = IR_VERSION + +checker.check_model(model) +save(model, "test_shape_data_propagation_gather_squeeze_range.onnx") + +print("Model saved to test_shape_data_propagation_gather_squeeze_range.onnx") diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_gather_topk.onnx b/onnxruntime/test/testdata/test_shape_data_propagation_gather_topk.onnx new file mode 100644 index 0000000000000000000000000000000000000000..a2b6b815998c942a55f0c2787b82f708e747b429 GIT binary patch literal 297 zcmYjM!A`du=!D}5JlXB`Veg0fNK70kWqM~?MB+`?sS|q^#dm+w!+FSf;wk)6 zHig=^q22kx-Oo?)xi#iX)2oah1qJ>bXY+xQi+bz8elo#W5DzgBzc(MZFg2#QGK|28 Sy;j$smRvZ}|0fd Gather(1-D index) -> TopK rank-drop. +# +# The Gather index here is the 1-D constant [-1] (rank 1), so per ONNX Gather +# semantics (output_rank = data_rank - 1 + indices_rank) the Gather output is a +# rank-1, single-element tensor -- exactly what TopK's K input requires. +# +# Before the fix, ORT's Gather custom data propagation scalarized this single +# element (dropping the rank), producing a 0-D K initializer that ONNX TopK +# shape inference rejects with: +# "K input must be a one-dimensional tensor of size 1." +# The fix preserves the rank-1 shape so the model loads and infers correctly. + +from onnx import IR_VERSION, TensorProto, checker, helper, save + +# Graph input: a 1-D float tensor of static length 2000. +input_tensor = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2000]) + +# Outputs: TopK values and indices, each a 1-D tensor whose length is K. +values_out = helper.make_tensor_value_info("V", TensorProto.FLOAT, ["topk"]) +indices_out = helper.make_tensor_value_info("I", TensorProto.INT64, ["topk"]) + +# Gather index is the 1-D constant [-1] (last dimension), as emitted by common +# exporters for the "K = size of the last dimension" pattern. +gather_idx = helper.make_tensor("gather_idx", TensorProto.INT64, [1], [-1]) + +# Shape(X) -> S : 1-D int64 tensor [2000]. +shape_node = helper.make_node("Shape", inputs=["X"], outputs=["S"], name="ShapeNode") +# Gather(S, [-1]) -> K : rank-1 size-1 int64 tensor [2000]. +gather_node = helper.make_node("Gather", inputs=["S", "gather_idx"], outputs=["K"], axis=0, name="GatherNode") +# TopK(X, K) -> V, I : returns all 2000 elements sorted (K == 2000). +topk_node = helper.make_node("TopK", inputs=["X", "K"], outputs=["V", "I"], axis=-1, largest=1, name="TopKNode") + +graph = helper.make_graph( + nodes=[shape_node, gather_node, topk_node], + name="Shape_Gather_TopK_Model", + inputs=[input_tensor], + outputs=[values_out, indices_out], + initializer=[gather_idx], +) + +model = helper.make_model( + graph, + opset_imports=[helper.make_operatorsetid("", 18)], + producer_name="onnx-example-generator", +) +model.ir_version = IR_VERSION + +checker.check_model(model) +save(model, "test_shape_data_propagation_gather_topk.onnx") + +print("Model saved to test_shape_data_propagation_gather_topk.onnx") diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_shape_mul_constantofshape.onnx b/onnxruntime/test/testdata/test_shape_data_propagation_shape_mul_constantofshape.onnx new file mode 100644 index 0000000000000000000000000000000000000000..9ac0d850f2ba916793833b05069c8a73f296f00c GIT binary patch literal 211 zcmd1] operand it must DECLINE (return no single value) rather than silently +# using element[0]. If it wrongly used element[0], Mul would emit a bogus scalar +# (3 * 3 = 9) that would override the correct, ONNX-propagated multi-element +# value and collapse the downstream ConstantOfShape output from rank-2 [9, 16] +# to a rank-1 [9]. +# +# Correct behavior: Mul declines, the multi-element value [9, 16] still flows +# through, and ConstantOfShape(M) is the rank-2 shape [9, 16]. + +from onnx import IR_VERSION, TensorProto, checker, helper, save + +x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [3, 4]) +# ConstantOfShape output is always rank 2 here (its shape input has 2 elements); +# declared with symbolic dims so the checker passes while the test asserts the +# concrete inferred dims [9, 16]. +y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, ["d0", "d1"]) + +# Shape(X) -> S : 1-D int64 tensor [3, 4] (two elements). +shape_node = helper.make_node("Shape", inputs=["X"], outputs=["S"], name="ShapeNode") +# Mul(S, S) -> M : elementwise [9, 16] (still two elements; must NOT collapse to a scalar). +mul_node = helper.make_node("Mul", inputs=["S", "S"], outputs=["M"], name="MulNode") +# ConstantOfShape(M) -> Y : output shape equals M's values, i.e. [9, 16]. +cos_node = helper.make_node("ConstantOfShape", inputs=["M"], outputs=["Y"], name="ConstantOfShapeNode") + +graph = helper.make_graph( + nodes=[shape_node, mul_node, cos_node], + name="Shape_Mul_ConstantOfShape_Model", + inputs=[x], + outputs=[y], +) + +model = helper.make_model( + graph, + opset_imports=[helper.make_operatorsetid("", 18)], + producer_name="onnx-example-generator", +) +model.ir_version = IR_VERSION + +checker.check_model(model) +save(model, "test_shape_data_propagation_shape_mul_constantofshape.onnx") + +print("Model saved to test_shape_data_propagation_shape_mul_constantofshape.onnx") diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_decline.onnx b/onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_decline.onnx new file mode 100644 index 0000000000000000000000000000000000000000..b65e5d5830ac899d9a17b009e9e20a08c7077856 GIT binary patch literal 423 zcmZ8d!AiqG5bfC3Bom@>1TCQ_@f0YC-g*o|4n+?M0r9e|W=I2ZQh~@42MgJ9O^|p_GDzBM?J3R? z5`vXy$Gt_*jX#37Xjf8YEb2y94Z>>RDHn|-h&(IW&HLaQ(zVVn`gmf$N!;^mkt-Gt zUCF|qo8THt{23ku2VMP3J%v&5uZvHesgn;=^%@MCPWJ5X?rNiD0!?vjexJT{myYNh T&rE|Hw8`pw6|8~@j$HZ$uT_0X literal 0 HcmV?d00001 diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_decline.py b/onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_decline.py new file mode 100644 index 0000000000000..b9df3bd67c762 --- /dev/null +++ b/onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_decline.py @@ -0,0 +1,65 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +# Lock fixture for Unsqueeze's custom data-propagation DECLINE on a single-element value. +# +# Shape(X) -> Gather([-1]) produces a rank-1 single-element value (the last dimension of X, 2000). +# Unsqueezing that single-element (scalar-like, rank-1 [1]) value would yield a rank >= 2 result +# ([1, 2000]) that the single-value data-propagation channel cannot faithfully represent, so +# Unsqueeze's custom data propagation must DECLINE rather than fabricate the misleading [1, value]. +# +# The decline is made OBSERVABLE through a rank-lowering Squeeze: +# Shape(X) -> Gather([-1]) -> Unsqueeze([0]) -> Squeeze -> Range(0, K, 1) +# With the correct decline, no value is propagated past Unsqueeze, Squeeze has nothing to lower, and +# Range's limit stays symbolic -- the model loads and the Range output length is an unknown dimension. +# If the decline were relaxed to fabricate [1, 2000], Squeeze would lower a non-scalar value into +# Range, which ONNX Range shape inference rejects ("Input to 'Range' op should be scalars"), so the +# model fails to load. The decline is thus locked end to end: the correct path loads with a symbolic +# Range length, while the rank-fabricating bug fails to load. + +from onnx import IR_VERSION, TensorProto, checker, helper, save + +# Graph input: a 1-D float tensor of static length 2000. +input_tensor = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2000]) + +# Output: Range result, a 1-D int64 tensor whose length is the propagated K (symbolic on decline). +range_out = helper.make_tensor_value_info("R", TensorProto.INT64, ["range_len"]) + +# Gather index is the 1-D constant [-1] (last dimension) -> rank-1 single-element value. +gather_idx = helper.make_tensor("gather_idx", TensorProto.INT64, [1], [-1]) +# Unsqueeze axes: insert a leading dimension (axis 0). +unsqueeze_axes = helper.make_tensor("unsqueeze_axes", TensorProto.INT64, [1], [0]) +# Range start/delta as 0-D scalar int64 constants. +range_start = helper.make_tensor("range_start", TensorProto.INT64, [], [0]) +range_delta = helper.make_tensor("range_delta", TensorProto.INT64, [], [1]) + +# Shape(X) -> S : 1-D int64 tensor [2000]. +shape_node = helper.make_node("Shape", inputs=["X"], outputs=["S"], name="ShapeNode") +# Gather(S, [-1]) -> G : rank-1 size-1 int64 value [2000]. +gather_node = helper.make_node("Gather", inputs=["S", "gather_idx"], outputs=["G"], axis=0, name="GatherNode") +# Unsqueeze(G, [0]) -> U : rank-2 [1, 1] tensor; custom data propagation must decline. +unsqueeze_node = helper.make_node("Unsqueeze", inputs=["G", "unsqueeze_axes"], outputs=["U"], name="UnsqueezeNode") +# Squeeze(U) -> K : 0-D scalar int64 value (only when a value is propagated). +squeeze_node = helper.make_node("Squeeze", inputs=["U"], outputs=["K"], name="SqueezeNode") +# Range(0, K, 1) -> R : 1-D int64 tensor whose length is K (symbolic on decline). +range_node = helper.make_node("Range", inputs=["range_start", "K", "range_delta"], outputs=["R"], name="RangeNode") + +graph = helper.make_graph( + nodes=[shape_node, gather_node, unsqueeze_node, squeeze_node, range_node], + name="Shape_Gather_Unsqueeze_Decline_Model", + inputs=[input_tensor], + outputs=[range_out], + initializer=[gather_idx, unsqueeze_axes, range_start, range_delta], +) + +model = helper.make_model( + graph, + opset_imports=[helper.make_operatorsetid("", 18)], + producer_name="onnx-example-generator", +) +model.ir_version = IR_VERSION + +checker.check_model(model) +save(model, "test_shape_data_propagation_unsqueeze_decline.onnx") + +print("Model saved to test_shape_data_propagation_unsqueeze_decline.onnx") From 43c87ceecc6c4071b10b77d012807cc60b166f0f Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 18 Jun 2026 20:54:03 +0000 Subject: [PATCH 2/8] Drive decline regression tests via Model::Load + Graph::Resolve Replace the full Ort::Session in GatherRank2IndexDeclineTest and UnsqueezeSingleValueDeclineTest with onnxruntime::Model::Load (which runs Graph::Resolve) plus direct output NodeArg shape inspection. Data propagation runs unconditionally inside Graph::InferAndVerifyTypeMatch during Resolve, and no session-level optimizer/constant-folding runs, so Resolve alone exercises the decline branch in isolation -- with none of the session arena / execution-provider / kernel-registry allocations. This trims each decline test's single-process AddressSanitizer footprint (microsoft/onnxruntime#29139) while preserving end-to-end, orthogonal discrimination of the rank-preservation fix (microsoft/onnxruntime#29072): relaxing the Gather decline concretizes the Range length to 2000; relaxing the Unsqueeze decline makes Range's input non-scalar so Resolve returns a non-OK status and the load fails. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwangms --- .../test/framework/shape_inference_test.cc | 116 +++++++++--------- 1 file changed, 61 insertions(+), 55 deletions(-) diff --git a/onnxruntime/test/framework/shape_inference_test.cc b/onnxruntime/test/framework/shape_inference_test.cc index 5ac99d87ec50d..872d649917d65 100644 --- a/onnxruntime/test/framework/shape_inference_test.cc +++ b/onnxruntime/test/framework/shape_inference_test.cc @@ -314,39 +314,42 @@ TEST(ShapeInferenceV2Test, ShapeMulMultiElementNoScalarCollapseTest) { // The decline is made OBSERVABLE through a rank-lowering Squeeze: // Shape(X) -> Gather([[ -1 ]]) -> Squeeze -> Range(0, K, 1) // With the correct decline, nothing is propagated, Squeeze has no value to lower, and Range's limit -// stays symbolic -- so the Range output length is an UNKNOWN dimension. If the decline were relaxed -// to emit a rank-1 [1] value, Squeeze would lower it to the scalar 2000 and Range would resolve to a -// concrete length 2000. Asserting the length is NOT the concrete 2000 discriminates the correct -// decline from the rank-fabricating bug -- real end-to-end coverage of the decline branch (no -// shared-helper abstraction, no tautological unit test). +// stays symbolic -- so the Range output dimension is UNKNOWN (it carries no concrete dim_value). If +// the decline were relaxed to emit a rank-1 [1] value, Squeeze would lower it to the scalar 2000 and +// Range would resolve to a concrete length 2000. Asserting the dimension has no concrete value +// discriminates the correct decline from the rank-fabricating bug -- real end-to-end coverage of the +// decline branch (no shared-helper abstraction, no tautological unit test). // -// This runs ONLY at ORT_DISABLE_ALL: the chain is entirely constant (a static input shape feeding -// Shape/Gather/Squeeze/Range over initializers), so at ORT_ENABLE_BASIC and above constant folding -// resolves Range to 2000 regardless of data propagation, which would mask the decline. Data -// propagation runs in the pre-optimization Graph::Resolve pass, so ORT_DISABLE_ALL exercises the -// decline branch in isolation. +// This drives the decline through onnxruntime::Model::Load (which runs Graph::Resolve) and inspects +// the resulting output NodeArg shape directly -- no InferenceSession, so none of the session's +// arena / execution-provider / kernel-registry allocations are created (microsoft/onnxruntime#29139: +// the single-process onnxruntime_test_all run sits near the AddressSanitizer size-class ceiling, so +// each added test must stay lean). Data propagation runs inside Graph::Resolve +// (Graph::InferAndVerifyTypeMatch -> data-propagation RunInferencing); constant folding is a separate +// session-level optimizer pass that never runs here, so Resolve alone exercises the decline branch in +// isolation with no folding to mask it. TEST(ShapeInferenceV2Test, GatherRank2IndexDeclineTest) { auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_rank2_decline.onnx"); - Ort::SessionOptions session_options{}; - session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); - - Ort::Session session(*ort_env, model_path, session_options); - - ORT_ENFORCE(session.GetOutputCount() == 1); - - Ort::TypeInfo type_info = session.GetOutputTypeInfo(0); - auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); - std::vector output_shape = tensor_info.GetShape(); - - // Range output is 1-D; its length must stay symbolic (unknown, reported as -1) because the rank-2 - // Gather index declines rather than propagating a (rank-fabricated) concrete K. A relaxed decline - // would fabricate a rank-1 value that Squeeze lowers to scalar 2000, making the length concretely - // 2000 -- which this assertion rejects. - EXPECT_TRUE(output_shape.size() == 1) << "Range output should be a 1-D tensor"; - EXPECT_EQ(output_shape[0], -1) - << "Range length must stay symbolic (reported as -1); the rank-2 Gather index must decline " - "rather than fabricate a concrete K (a relaxed decline concretizes it to 2000)"; + std::shared_ptr model; + // Model::Load runs Graph::Resolve, which performs data propagation. A correct decline keeps the + // Range length symbolic, so the resolve (and therefore the load) succeeds. + ASSERT_STATUS_OK(onnxruntime::Model::Load(model_path, model, nullptr, + DefaultLoggingManager().DefaultLogger())); + + const auto& graph_outputs = model->MainGraph().GetOutputs(); + ASSERT_EQ(graph_outputs.size(), static_cast(1)); + const TensorShapeProto* output_shape = graph_outputs[0]->Shape(); + ASSERT_NE(output_shape, nullptr) << "Range output should have an inferred (rank-1) shape"; + + // Range output is 1-D; its single dimension must stay SYMBOLIC (no concrete dim_value) because the + // rank-2 Gather index declines rather than propagating a (rank-fabricated) concrete K. A relaxed + // decline would fabricate a rank-1 value that Squeeze lowers to scalar 2000, giving the dimension a + // concrete value 2000 -- which these assertions reject. + EXPECT_EQ(output_shape->dim_size(), 1) << "Range output should be a 1-D tensor"; + EXPECT_FALSE(output_shape->dim(0).has_dim_value()) + << "Range length must stay symbolic; the rank-2 Gather index must decline rather than " + "fabricate a concrete K (a relaxed decline concretizes it to 2000)"; } // Regression for microsoft/onnxruntime#29072. @@ -360,37 +363,40 @@ TEST(ShapeInferenceV2Test, GatherRank2IndexDeclineTest) { // The decline is made OBSERVABLE through a rank-lowering Squeeze: // Shape(X) -> Gather([-1]) -> Unsqueeze([0]) -> Squeeze -> Range(0, K, 1) // With the correct decline, no value is propagated past Unsqueeze, Squeeze has nothing to lower, and -// Range's limit stays symbolic, so the model loads with an unknown Range length. A relaxed decline +// Range's limit stays symbolic, so the graph resolves with an unknown Range length. A relaxed decline // that fabricated [1, 2000] would propagate a non-scalar value to Range, which ONNX Range shape -// inference rejects ("Input to 'Range' op should be scalars") -- the model then fails to load. So -// the decline IS observable: correct -> loads with a symbolic length; relaxed -> load error. This -// test asserts the model loads and the length is not the (rank-fabricated) concrete 2000. +// inference rejects ("Input to 'Range' op should be scalars") -- Graph::Resolve then returns a +// non-OK status and the load fails. So the decline IS observable: correct -> loads with a symbolic +// length; relaxed -> load error. This test asserts the model loads and the length stays symbolic. // -// Like GatherRank2IndexDeclineTest this runs ONLY at ORT_DISABLE_ALL: the chain is entirely constant -// over a static input shape, so at ORT_ENABLE_BASIC and above constant folding resolves Range -// regardless of data propagation, masking the decline. +// Like GatherRank2IndexDeclineTest this drives the decline through onnxruntime::Model::Load (which +// runs Graph::Resolve) and inspects the output NodeArg shape directly -- no InferenceSession, so the +// session arena / EP / kernel-registry allocations are never created (microsoft/onnxruntime#29139). +// Data propagation runs inside Graph::Resolve; constant folding is a separate session-level optimizer +// pass that never runs here, so Resolve alone exercises the decline branch with no folding to mask it. TEST(ShapeInferenceV2Test, UnsqueezeSingleValueDeclineTest) { auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_unsqueeze_decline.onnx"); - Ort::SessionOptions session_options{}; - session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); - - Ort::Session session(*ort_env, model_path, session_options); - - ORT_ENFORCE(session.GetOutputCount() == 1); - - Ort::TypeInfo type_info = session.GetOutputTypeInfo(0); - auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); - std::vector output_shape = tensor_info.GetShape(); - - // Range output is 1-D; its length must stay symbolic because Unsqueeze declines rather than - // fabricating a rank >= 2 value. A relaxed decline fabricates [1, 2000], which makes Range's input - // non-scalar and the model fails to load -- so any regression here is caught either as a load - // failure or as a concrete length below. - EXPECT_TRUE(output_shape.size() == 1) << "Range output should be a 1-D tensor"; - EXPECT_EQ(output_shape[0], -1) - << "Range length must stay symbolic (reported as -1); unsqueezing a single-element value " - "must decline (a relaxed decline fabricates [1, value] and the model fails to load)"; + std::shared_ptr model; + // Model::Load runs Graph::Resolve. A correct decline keeps Range's input scalar so the resolve (and + // therefore the load) succeeds; a relaxed decline makes Range's input non-scalar, which Range shape + // inference rejects, so Resolve returns non-OK and this assertion fails on the load itself. + ASSERT_STATUS_OK(onnxruntime::Model::Load(model_path, model, nullptr, + DefaultLoggingManager().DefaultLogger())); + + const auto& graph_outputs = model->MainGraph().GetOutputs(); + ASSERT_EQ(graph_outputs.size(), static_cast(1)); + const TensorShapeProto* output_shape = graph_outputs[0]->Shape(); + ASSERT_NE(output_shape, nullptr) << "Range output should have an inferred (rank-1) shape"; + + // Range output is 1-D; its single dimension must stay SYMBOLIC because Unsqueeze declines rather + // than fabricating a rank >= 2 value. A relaxed decline fabricates [1, 2000], which makes Range's + // input non-scalar and the load fails above -- so any regression here is caught either as a load + // failure or as a concrete dim_value below. + EXPECT_EQ(output_shape->dim_size(), 1) << "Range output should be a 1-D tensor"; + EXPECT_FALSE(output_shape->dim(0).has_dim_value()) + << "Range length must stay symbolic; unsqueezing a single-element value must decline " + "(a relaxed decline fabricates [1, value] and the model fails to load)"; } // Lock test for the Shape -> Gather(1-D index) -> Squeeze -> Range chain. From 81c60a2ba58f4218b299dbee8528b718f4b92c16 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 18 Jun 2026 21:15:36 +0000 Subject: [PATCH 3/8] Assert (not expect) the decline tests' Range output rank The two decline tests dereference output_shape->dim(0) immediately after checking its dim_size, so the rank guard must halt on failure: convert EXPECT_EQ to ASSERT_EQ at both sites (GatherRank2IndexDeclineTest and UnsqueezeSingleValueDeclineTest). A theoretical rank-0 regression would otherwise out-of-bounds access the protobuf repeated field instead of failing cleanly; this also matches the surrounding ASSERT-guarded preconditions. Test-only, no assertion-semantics change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwangms --- onnxruntime/test/framework/shape_inference_test.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/framework/shape_inference_test.cc b/onnxruntime/test/framework/shape_inference_test.cc index 872d649917d65..81eef32688b1b 100644 --- a/onnxruntime/test/framework/shape_inference_test.cc +++ b/onnxruntime/test/framework/shape_inference_test.cc @@ -346,7 +346,7 @@ TEST(ShapeInferenceV2Test, GatherRank2IndexDeclineTest) { // rank-2 Gather index declines rather than propagating a (rank-fabricated) concrete K. A relaxed // decline would fabricate a rank-1 value that Squeeze lowers to scalar 2000, giving the dimension a // concrete value 2000 -- which these assertions reject. - EXPECT_EQ(output_shape->dim_size(), 1) << "Range output should be a 1-D tensor"; + ASSERT_EQ(output_shape->dim_size(), 1) << "Range output should be a 1-D tensor"; EXPECT_FALSE(output_shape->dim(0).has_dim_value()) << "Range length must stay symbolic; the rank-2 Gather index must decline rather than " "fabricate a concrete K (a relaxed decline concretizes it to 2000)"; @@ -393,7 +393,7 @@ TEST(ShapeInferenceV2Test, UnsqueezeSingleValueDeclineTest) { // than fabricating a rank >= 2 value. A relaxed decline fabricates [1, 2000], which makes Range's // input non-scalar and the load fails above -- so any regression here is caught either as a load // failure or as a concrete dim_value below. - EXPECT_EQ(output_shape->dim_size(), 1) << "Range output should be a 1-D tensor"; + ASSERT_EQ(output_shape->dim_size(), 1) << "Range output should be a 1-D tensor"; EXPECT_FALSE(output_shape->dim(0).has_dim_value()) << "Range length must stay symbolic; unsqueezing a single-element value must decline " "(a relaxed decline fabricates [1, value] and the model fails to load)"; From be364c981adac970cfc217bc51ec8ac802cecf16 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Mon, 22 Jun 2026 18:46:23 +0000 Subject: [PATCH 4/8] Clarify Gather decline comments and remove dead unknown-rank branch Address two non-blocking review nits on the Gather custom data propagation: 1. The decline comments said control "falls back to ONNX data propagation". That is inaccurate: once this custom propagator runs there is no ONNX fallback (Graph::SaveShapeValuesFromDataPropagation returns right after dp->infer()), and Gather's propagator only runs when ONNX's own result was already empty. Declining simply leaves the output value unset so the dimension stays symbolic. Reword both comment spots to state that. 2. The effective_num_dims < 0 fallback to the indices NodeArg shape was dead code: it is gated by indices.size() == 1, which is only reachable when a constant index initializer was read, and that path always reports a concrete rank (TensorShape::NumDimensions(), >= 0). Remove the unreachable branch and leave a comment documenting the invariant. Comments only plus deletion of a provably unreachable branch; the data-propagation logic is behaviorally identical. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwangms --- .../gather_op_data_propagation.cc | 26 +++++++++---------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.cc index d747ac20219cf..962c2fa272160 100644 --- a/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.cc +++ b/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.cc @@ -73,18 +73,16 @@ Status GatherOpDataPropagation::infer() { // * a 0-D scalar index -> a scalar output value, // * a 1-D index -> a rank-1 [1] output value (so consumers that require a // 1-D tensor, e.g. TopK's K input, still observe the correct rank), - // * a rank >= 2 index (or an unknown rank) -> decline, because the single-value - // channel cannot represent a rank >= 2 Gather output; falling back to ONNX data - // propagation is safer than fabricating a rank-1 value. + // * a rank >= 2 index -> decline: leave the output value unset so the dimension + // stays symbolic, because the single-value channel cannot represent a rank >= 2 + // Gather output and emitting a rank-1 value would fabricate a misleading rank. // - // The index rank is sourced canonically from the same constant initializer the index - // value came from (indices_num_dims). If that is unavailable (sentinel < 0), fall back - // to the indices NodeArg shape; an unknown shape then routes to decline. - int effective_num_dims = indices_num_dims; - if (effective_num_dims < 0) { - const ONNX_NAMESPACE::TensorShapeProto* indices_shape = input_1->Shape(); - effective_num_dims = (indices_shape != nullptr) ? indices_shape->dim_size() : -1; - } + // The index rank (indices_num_dims) is sourced canonically from the same constant + // initializer the index value came from. Reading that initializer is exactly what makes + // indices.size() == 1 reachable here, and it always reports a concrete rank + // (TensorShape::NumDimensions(), >= 0), so indices_num_dims is always >= 0 here and no + // unknown-rank (< 0) handling is needed. + const int effective_num_dims = indices_num_dims; if (effective_num_dims == 0) { // 0-D scalar index -> a scalar output value. @@ -93,9 +91,9 @@ Status GatherOpDataPropagation::infer() { // 1-D index -> a rank-1 [1] output value. SetSinglePropagatedShapeValue(output_def_, dim.dim_value(), /*is_rank1=*/true); } - // rank >= 2 or unknown index rank (effective_num_dims < 0): leave the output unset and - // let ONNX data propagation handle it. The single-value channel cannot represent a - // rank >= 2 Gather output, so emitting a rank-1 value would fabricate a misleading rank. + // rank >= 2 index: leave the output value unset so the dimension remains symbolic -- + // the correct decline behavior, since the single-value channel cannot represent a + // rank >= 2 Gather output and there is no ONNX fallback once this custom propagator runs. } } } From 3d2eb0e6e4981487988f0e7adaaebe22bf0ac888 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Mon, 22 Jun 2026 21:40:44 +0000 Subject: [PATCH 5/8] Fold the two data-propagation decline tests into one shared model Round-7 footprint trim for microsoft/onnxruntime#29139: the single-process onnxruntime_test_all run sits near the AddressSanitizer size-class-8192 ceiling, so reduce the live Model/Graph footprint our decline coverage contributes by collapsing GatherRank2IndexDeclineTest and UnsqueezeSingleValueDeclineTest into a single GatherUnsqueezeDeclineTest backed by one combined model. The coverage now costs one onnxruntime::Model::Load instead of two. The combined model shares one rank-3 input X and one Shape(X) -> S, then runs two independent branches off S: * Branch A: Gather([[ -1 ]]) (rank-2 index) -> Squeeze -> Range -> RA, locking Gather's rank-2 decline; RA stays symbolic on a correct decline and concretizes to 2000 if it regresses. * Branch B: Gather([-1]) -> Unsqueeze([0]) -> Squeeze -> Range -> RB, locking Unsqueeze's single-value decline; on a relaxed decline Range's input becomes non-scalar and Graph::Resolve (Model::Load) fails. The branches are independent, so a regression in either is still caught (Branch A as a concrete RA dimension, Branch B as a failed load) -- the one combined model keeps the full discriminating power of the two original per-path fixtures. Verified by A/B flip: relaxing the Gather decline fails on RA; relaxing the Unsqueeze decline fails on the load (RangeNodeB). The data-propagation production sources are unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwangms --- .../test/framework/shape_inference_test.cc | 142 +++++++----------- ...ape_data_propagation_decline_combined.onnx | Bin 0 -> 682 bytes ...shape_data_propagation_decline_combined.py | 96 ++++++++++++ ...data_propagation_gather_rank2_decline.onnx | Bin 354 -> 0 bytes ...e_data_propagation_gather_rank2_decline.py | 70 --------- ...pe_data_propagation_unsqueeze_decline.onnx | Bin 423 -> 0 bytes ...hape_data_propagation_unsqueeze_decline.py | 65 -------- 7 files changed, 154 insertions(+), 219 deletions(-) create mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_decline_combined.onnx create mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_decline_combined.py delete mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_gather_rank2_decline.onnx delete mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_gather_rank2_decline.py delete mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_decline.onnx delete mode 100644 onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_decline.py diff --git a/onnxruntime/test/framework/shape_inference_test.cc b/onnxruntime/test/framework/shape_inference_test.cc index 81eef32688b1b..bde00e5aa9bac 100644 --- a/onnxruntime/test/framework/shape_inference_test.cc +++ b/onnxruntime/test/framework/shape_inference_test.cc @@ -302,101 +302,75 @@ TEST(ShapeInferenceV2Test, ShapeMulMultiElementNoScalarCollapseTest) { } // Regression for microsoft/onnxruntime#29072. -// End-to-end lock for Gather's custom data-propagation DECLINE on a rank-2 index. +// End-to-end lock for BOTH custom data-propagation DECLINE paths -- Gather on a rank-2 index and +// Unsqueeze on a single-element value -- driven through a SINGLE shared model so the coverage costs +// one onnxruntime::Model::Load instead of two (microsoft/onnxruntime#29139: the single-process +// onnxruntime_test_all run sits near the AddressSanitizer size-class ceiling, so the two decline +// fixtures are folded into one graph to minimize the live Model/Graph footprint). // -// The Gather index is the constant [[-1]] of shape [1, 1] (rank 2), so the Gather output rank is -// data_rank - 1 + index_rank = 1 - 1 + 2 = 2. The single-value data-propagation channel can carry -// only rank-0 (scalar) and rank-1 [1] values, so Gather's custom propagation must DECLINE rather -// than fabricate a rank-1 value. The index is a raw graph initializer (not surfaced through the -// data-propagation getInputData channel), so ONNX's own Gather data propagator bails and control -// reaches our custom decline branch. +// One rank-3 input X drives a shared Shape(X) -> S (the 1-D vector [3, 4, 2000]); two independent +// branches gather from S and route into a rank-lowering Squeeze -> Range so each decline is +// observable end to end: // -// The decline is made OBSERVABLE through a rank-lowering Squeeze: -// Shape(X) -> Gather([[ -1 ]]) -> Squeeze -> Range(0, K, 1) -// With the correct decline, nothing is propagated, Squeeze has no value to lower, and Range's limit -// stays symbolic -- so the Range output dimension is UNKNOWN (it carries no concrete dim_value). If -// the decline were relaxed to emit a rank-1 [1] value, Squeeze would lower it to the scalar 2000 and -// Range would resolve to a concrete length 2000. Asserting the dimension has no concrete value -// discriminates the correct decline from the rank-fabricating bug -- real end-to-end coverage of the -// decline branch (no shared-helper abstraction, no tautological unit test). +// Branch A (Gather rank-2 index -> output RA): S -> Gather([[ -1 ]]) -> Squeeze -> Range(0, K, 1) +// The index [[-1]] has shape [1, 1] (rank 2), so the Gather output rank is +// data_rank - 1 + index_rank = 1 - 1 + 2 = 2 -- a rank-2 single-element value the single-value +// channel cannot represent, so Gather's custom propagation must DECLINE. With the correct +// decline nothing propagates, Squeeze has no value to lower, and RA's length stays SYMBOLIC. A +// relaxed decline would emit a rank-1 [1] value that Squeeze lowers to the scalar 2000, +// concretizing RA to length 2000 -- so asserting RA stays symbolic discriminates the correct +// decline from the rank-fabricating bug. // -// This drives the decline through onnxruntime::Model::Load (which runs Graph::Resolve) and inspects -// the resulting output NodeArg shape directly -- no InferenceSession, so none of the session's -// arena / execution-provider / kernel-registry allocations are created (microsoft/onnxruntime#29139: -// the single-process onnxruntime_test_all run sits near the AddressSanitizer size-class ceiling, so -// each added test must stay lean). Data propagation runs inside Graph::Resolve -// (Graph::InferAndVerifyTypeMatch -> data-propagation RunInferencing); constant folding is a separate -// session-level optimizer pass that never runs here, so Resolve alone exercises the decline branch in -// isolation with no folding to mask it. -TEST(ShapeInferenceV2Test, GatherRank2IndexDeclineTest) { - auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_rank2_decline.onnx"); - - std::shared_ptr model; - // Model::Load runs Graph::Resolve, which performs data propagation. A correct decline keeps the - // Range length symbolic, so the resolve (and therefore the load) succeeds. - ASSERT_STATUS_OK(onnxruntime::Model::Load(model_path, model, nullptr, - DefaultLoggingManager().DefaultLogger())); - - const auto& graph_outputs = model->MainGraph().GetOutputs(); - ASSERT_EQ(graph_outputs.size(), static_cast(1)); - const TensorShapeProto* output_shape = graph_outputs[0]->Shape(); - ASSERT_NE(output_shape, nullptr) << "Range output should have an inferred (rank-1) shape"; - - // Range output is 1-D; its single dimension must stay SYMBOLIC (no concrete dim_value) because the - // rank-2 Gather index declines rather than propagating a (rank-fabricated) concrete K. A relaxed - // decline would fabricate a rank-1 value that Squeeze lowers to scalar 2000, giving the dimension a - // concrete value 2000 -- which these assertions reject. - ASSERT_EQ(output_shape->dim_size(), 1) << "Range output should be a 1-D tensor"; - EXPECT_FALSE(output_shape->dim(0).has_dim_value()) - << "Range length must stay symbolic; the rank-2 Gather index must decline rather than " - "fabricate a concrete K (a relaxed decline concretizes it to 2000)"; -} - -// Regression for microsoft/onnxruntime#29072. -// End-to-end lock for Unsqueeze's custom data-propagation DECLINE on a single-element value. +// Branch B (Unsqueeze single value -> output RB): S -> Gather([-1]) -> Unsqueeze([0]) -> Squeeze -> +// Range(0, K, 1). Gather([-1]) is a rank-1 single-element value (the last dimension of X, 2000); +// unsqueezing it would yield a rank >= 2 result ([1, 2000]) the single-value channel cannot +// represent, so Unsqueeze's custom propagation must DECLINE. With the correct decline nothing +// propagates past Unsqueeze, RB's length stays SYMBOLIC, and the model loads. A relaxed decline +// would fabricate [1, 2000], which Squeeze lowers into a non-scalar Range limit that ONNX Range +// shape inference rejects ("Input to 'Range' op should be scalars"), so Graph::Resolve (and +// therefore Model::Load) FAILS. So Branch B is observable either as a load failure or, were that +// to pass, as a concrete RB length. // -// Shape(X) -> Gather([-1]) produces a rank-1 single-element value (the last dimension of X, 2000). -// Unsqueezing that single-element (scalar-like, rank-1 [1]) value would yield a rank >= 2 result -// ([1, 2000]) that the single-value channel cannot faithfully represent, so Unsqueeze's custom -// propagation must DECLINE rather than fabricate the misleading [1, value]. +// Both Gather indices are raw graph initializers (not surfaced through the data-propagation +// getInputData channel), so ONNX's own Gather data propagator bails and control reaches our custom +// decline branches. The branches are independent: a regression in either is caught (Branch A as a +// concrete RA dimension, Branch B as a failed load), so this one combined model keeps the full +// discriminating power of the two original per-path fixtures. // -// The decline is made OBSERVABLE through a rank-lowering Squeeze: -// Shape(X) -> Gather([-1]) -> Unsqueeze([0]) -> Squeeze -> Range(0, K, 1) -// With the correct decline, no value is propagated past Unsqueeze, Squeeze has nothing to lower, and -// Range's limit stays symbolic, so the graph resolves with an unknown Range length. A relaxed decline -// that fabricated [1, 2000] would propagate a non-scalar value to Range, which ONNX Range shape -// inference rejects ("Input to 'Range' op should be scalars") -- Graph::Resolve then returns a -// non-OK status and the load fails. So the decline IS observable: correct -> loads with a symbolic -// length; relaxed -> load error. This test asserts the model loads and the length stays symbolic. -// -// Like GatherRank2IndexDeclineTest this drives the decline through onnxruntime::Model::Load (which -// runs Graph::Resolve) and inspects the output NodeArg shape directly -- no InferenceSession, so the -// session arena / EP / kernel-registry allocations are never created (microsoft/onnxruntime#29139). -// Data propagation runs inside Graph::Resolve; constant folding is a separate session-level optimizer -// pass that never runs here, so Resolve alone exercises the decline branch with no folding to mask it. -TEST(ShapeInferenceV2Test, UnsqueezeSingleValueDeclineTest) { - auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_unsqueeze_decline.onnx"); +// The decline is driven through onnxruntime::Model::Load (which runs Graph::Resolve) and the output +// NodeArg shapes are inspected directly -- no InferenceSession, so none of the session's arena / +// execution-provider / kernel-registry allocations are created. Data propagation runs inside +// Graph::Resolve (Graph::InferAndVerifyTypeMatch -> data-propagation RunInferencing); constant +// folding is a separate session-level optimizer pass that never runs here, so Resolve alone +// exercises the decline branches in isolation with no folding to mask them. +TEST(ShapeInferenceV2Test, GatherUnsqueezeDeclineTest) { + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_decline_combined.onnx"); std::shared_ptr model; - // Model::Load runs Graph::Resolve. A correct decline keeps Range's input scalar so the resolve (and - // therefore the load) succeeds; a relaxed decline makes Range's input non-scalar, which Range shape - // inference rejects, so Resolve returns non-OK and this assertion fails on the load itself. + // Model::Load runs Graph::Resolve, which performs data propagation. A correct decline on BOTH + // branches keeps each Range length symbolic, so the resolve (and therefore the load) succeeds. A + // relaxed Unsqueeze decline (Branch B) makes Range's input non-scalar, which Range shape inference + // rejects, so Resolve returns non-OK and this assertion fails on the load itself. ASSERT_STATUS_OK(onnxruntime::Model::Load(model_path, model, nullptr, DefaultLoggingManager().DefaultLogger())); const auto& graph_outputs = model->MainGraph().GetOutputs(); - ASSERT_EQ(graph_outputs.size(), static_cast(1)); - const TensorShapeProto* output_shape = graph_outputs[0]->Shape(); - ASSERT_NE(output_shape, nullptr) << "Range output should have an inferred (rank-1) shape"; - - // Range output is 1-D; its single dimension must stay SYMBOLIC because Unsqueeze declines rather - // than fabricating a rank >= 2 value. A relaxed decline fabricates [1, 2000], which makes Range's - // input non-scalar and the load fails above -- so any regression here is caught either as a load - // failure or as a concrete dim_value below. - ASSERT_EQ(output_shape->dim_size(), 1) << "Range output should be a 1-D tensor"; - EXPECT_FALSE(output_shape->dim(0).has_dim_value()) - << "Range length must stay symbolic; unsqueezing a single-element value must decline " - "(a relaxed decline fabricates [1, value] and the model fails to load)"; + ASSERT_EQ(graph_outputs.size(), static_cast(2)); + + // Both Range outputs are 1-D; each single dimension must stay SYMBOLIC (no concrete dim_value) + // because its branch declines rather than propagating a (rank-fabricated) concrete K. RA[0] would + // concretize to 2000 if the rank-2 Gather index stopped declining; RB[0] would concretize to 2000 + // (or the load would fail) if the single-value Unsqueeze stopped declining. + for (const NodeArg* output : graph_outputs) { + const TensorShapeProto* output_shape = output->Shape(); + ASSERT_NE(output_shape, nullptr) << "Range output '" << output->Name() + << "' should have an inferred (rank-1) shape"; + ASSERT_EQ(output_shape->dim_size(), 1) << "Range output '" << output->Name() + << "' should be a 1-D tensor"; + EXPECT_FALSE(output_shape->dim(0).has_dim_value()) + << "Range length '" << output->Name() << "' must stay symbolic; the decline must not " + "fabricate a concrete K (a relaxed decline concretizes it to 2000)"; + } } // Lock test for the Shape -> Gather(1-D index) -> Squeeze -> Range chain. diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_decline_combined.onnx b/onnxruntime/test/testdata/test_shape_data_propagation_decline_combined.onnx new file mode 100644 index 0000000000000000000000000000000000000000..3310f8047cf34c330dc781bb2f0000008b925612 GIT binary patch literal 682 zcmah{O;5r=5M>KU8RWwyDyavP##0jxXuRO5s|g3A2O1LNWwY1`5K9YfNgDrvf5V^T zKho`%5=5e#Y0v`hx z&9b_}4G;%f6A@1_B(tbf_#hgbgXB@|-Zl=M z4o0Wn@*#>b_r8Pl{sOqS$^>S5g#% z#>$T}Nn9itQEkMG#wf?qYRs%|;?&8ANhj=w6ZVG_w(Ep#IbmkzgmzVzdDc9>smJf| z#hdxaZ-e<*o-TioiahH=Cpu|W`tX%6!hQ|MGXLfDtW(7y1-8=&i^}v8-kwW_ksMQ_ yc~)Ph4M#emb S (the 1-D vector [3, 4, 2000]); +# two independent branches gather from S and route into a rank-lowering Squeeze -> Range so +# each decline is observable end to end: +# +# Branch A (Gather rank-2 index decline): +# S -> Gather([[ -1 ]]) -> Squeeze -> Range(0, K, 1) -> RA +# The index [[-1]] has shape [1, 1] (rank 2), so the Gather output rank is +# data_rank - 1 + index_rank = 1 - 1 + 2 = 2 -- a rank-2 single-element value the +# single-value channel cannot represent, so Gather's custom propagation must DECLINE. +# On a correct decline nothing propagates, Squeeze has nothing to lower, and RA's length +# stays SYMBOLIC. A relaxed decline would emit a rank-1 [1] value that Squeeze lowers to +# the scalar 2000, concretizing RA to length 2000 -- so asserting RA stays symbolic +# discriminates the correct decline from the rank-fabricating bug. +# +# Branch B (Unsqueeze single-element-value decline): +# S -> Gather([-1]) -> Unsqueeze([0]) -> Squeeze -> Range(0, K, 1) -> RB +# Gather([-1]) is a rank-1 single-element value (the last dimension of X, 2000). +# Unsqueezing it would yield a rank >= 2 result ([1, 2000]) the single-value channel +# cannot faithfully represent, so Unsqueeze's custom propagation must DECLINE. On a +# correct decline nothing propagates past Unsqueeze, RB's length stays SYMBOLIC, and the +# model loads. A relaxed decline would fabricate [1, 2000], which Squeeze lowers into a +# non-scalar Range limit that ONNX Range shape inference rejects ("Input to 'Range' op +# should be scalars"), so Graph::Resolve (and therefore Model::Load) FAILS -- the decline +# is observable either as a load failure or, were that to pass, as a concrete RB length. +# +# Because both Gather indices are raw constant initializers, ONNX's own Gather data +# propagator bails (it has no getInputData for a constant index), so control reaches our +# custom decline branches. The two branches are independent: a regression in either one is +# caught (Branch A as a concrete RA dimension, Branch B as a failed load), so the single +# combined model locks both decline paths with the discriminating power of the two +# original per-path fixtures. + +from onnx import IR_VERSION, TensorProto, checker, helper, save + +# Graph input: a 3-D float tensor whose last static dimension is 2000. Its Shape vector +# [3, 4, 2000] feeds both branches; Gather([-1]) of that vector is the single element 2000. +input_tensor = helper.make_tensor_value_info("X", TensorProto.FLOAT, [3, 4, 2000]) + +# Outputs: the two Range results, 1-D int64 tensors whose lengths are the propagated K +# (symbolic when each decline is correct). +range_out_a = helper.make_tensor_value_info("RA", TensorProto.INT64, ["range_len_a"]) +range_out_b = helper.make_tensor_value_info("RB", TensorProto.INT64, ["range_len_b"]) + +# Branch A Gather index: the rank-2 constant [[-1]] of shape [1, 1] (last dimension). +gather_idx_rank2 = helper.make_tensor("gather_idx_rank2", TensorProto.INT64, [1, 1], [-1]) +# Branch B Gather index: the rank-1 constant [-1] (last dimension). +gather_idx_rank1 = helper.make_tensor("gather_idx_rank1", TensorProto.INT64, [1], [-1]) +# Branch B Unsqueeze axes: insert a leading dimension (axis 0). +unsqueeze_axes = helper.make_tensor("unsqueeze_axes", TensorProto.INT64, [1], [0]) +# Range start/delta as 0-D scalar int64 constants, shared by both Range nodes. +range_start = helper.make_tensor("range_start", TensorProto.INT64, [], [0]) +range_delta = helper.make_tensor("range_delta", TensorProto.INT64, [], [1]) + +# Shape(X) -> S : 1-D int64 tensor [3, 4, 2000], shared by both branches. +shape_node = helper.make_node("Shape", inputs=["X"], outputs=["S"], name="ShapeNode") + +# Branch A: Gather(S, [[-1]]) -> GA : rank-2 [1, 1] value; custom propagation must decline. +gather_a = helper.make_node("Gather", inputs=["S", "gather_idx_rank2"], outputs=["GA"], axis=0, name="GatherNodeA") +squeeze_a = helper.make_node("Squeeze", inputs=["GA"], outputs=["KA"], name="SqueezeNodeA") +range_a = helper.make_node("Range", inputs=["range_start", "KA", "range_delta"], outputs=["RA"], name="RangeNodeA") + +# Branch B: Gather(S, [-1]) -> GB : rank-1 single element; Unsqueeze -> rank-2; must decline. +gather_b = helper.make_node("Gather", inputs=["S", "gather_idx_rank1"], outputs=["GB"], axis=0, name="GatherNodeB") +unsqueeze_b = helper.make_node("Unsqueeze", inputs=["GB", "unsqueeze_axes"], outputs=["UB"], name="UnsqueezeNodeB") +squeeze_b = helper.make_node("Squeeze", inputs=["UB"], outputs=["KB"], name="SqueezeNodeB") +range_b = helper.make_node("Range", inputs=["range_start", "KB", "range_delta"], outputs=["RB"], name="RangeNodeB") + +graph = helper.make_graph( + nodes=[shape_node, gather_a, squeeze_a, range_a, gather_b, unsqueeze_b, squeeze_b, range_b], + name="Shape_Gather_Unsqueeze_Decline_Combined_Model", + inputs=[input_tensor], + outputs=[range_out_a, range_out_b], + initializer=[gather_idx_rank2, gather_idx_rank1, unsqueeze_axes, range_start, range_delta], +) + +model = helper.make_model( + graph, + opset_imports=[helper.make_operatorsetid("", 18)], + producer_name="onnx-example-generator", +) +model.ir_version = IR_VERSION + +checker.check_model(model) +save(model, "test_shape_data_propagation_decline_combined.onnx") + +print("Model saved to test_shape_data_propagation_decline_combined.onnx") diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_gather_rank2_decline.onnx b/onnxruntime/test/testdata/test_shape_data_propagation_gather_rank2_decline.onnx deleted file mode 100644 index eaedaac4b6a0187c1a2233ecc07e825ce3f6848d..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 354 zcmYjN%}T>S5YD90Y$g=z3Q~fIcntJXZ#{-W4n+^iMZ7Gl8PZ_f)NO>|)A%O7f^VeT zT?1Vh_f(H4vYiUuFA}#w)ufBbpBRp~_lc?Ov9X)gv-JV#8cL*h7 zE&DBL(NqJ6Ji**jd+pTWG{!&<&3^8Efsa=RIfqRWm+v1$pYFl%pWqT>Emcio-%H&i zY(`#D+e;2b5)}vE)v%D?*fzq238TF(#hPB)Mv=HTb=&ls0fSk%!r$>gcGNrN#se=q tj&>T1nvH+5tGK{(27HP>o0ti3!hX)4jZ^T06B;Mxk^4kTDhtq?@jn%8WPktw diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_gather_rank2_decline.py b/onnxruntime/test/testdata/test_shape_data_propagation_gather_rank2_decline.py deleted file mode 100644 index be3225e998d73..0000000000000 --- a/onnxruntime/test/testdata/test_shape_data_propagation_gather_rank2_decline.py +++ /dev/null @@ -1,70 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. - -# Lock fixture for Gather's custom data-propagation DECLINE on a rank-2 index. -# -# The Gather index is the constant [[-1]] of shape [1, 1] (rank 2). Per ONNX -# Gather semantics the output rank = data_rank - 1 + index_rank = 1 - 1 + 2 = 2, -# so the Gather output is a rank-2, single-element tensor [[2000]]. The -# single-value data-propagation channel can represent only rank-0 (scalar) and -# rank-1 [1] values, so Gather's custom data propagation must DECLINE here: -# emitting a rank-1 value would fabricate a rank the channel cannot honestly -# carry. -# -# This fixture makes that decline OBSERVABLE through a rank-lowering Squeeze: -# Shape(X) -> Gather([[ -1 ]]) -> Squeeze -> Range(0, K, 1) -# Because the index is a constant initializer, ONNX's own Gather data -# propagator bails (it has no getInputData for a constant index), so control -# reaches our custom decline branch. With the correct decline, no value is -# propagated, Squeeze has nothing to lower, and Range's limit stays symbolic -- -# the Range output length is therefore an unknown dimension. -# -# If the decline were relaxed to emit a rank-1 [1] value, Squeeze would lower it -# to the scalar 2000 and Range(0, 2000, 1) would resolve to a concrete length -# 2000. Asserting the Range length is NOT the concrete 2000 thus locks the -# decline end to end (it discriminates the correct-decline behavior from the -# rank-fabricating bug). - -from onnx import IR_VERSION, TensorProto, checker, helper, save - -# Graph input: a 3-D float tensor whose last static dimension is 2000. -input_tensor = helper.make_tensor_value_info("X", TensorProto.FLOAT, [3, 4, 2000]) - -# Output: Range result, a 1-D int64 tensor whose length is the propagated K -# (symbolic when the decline is correct). -range_out = helper.make_tensor_value_info("R", TensorProto.INT64, ["range_len"]) - -# Gather index is the rank-2 constant [[-1]] of shape [1, 1] (last dimension). -gather_idx = helper.make_tensor("gather_idx", TensorProto.INT64, [1, 1], [-1]) -# Range start/delta as 0-D scalar int64 constants. -range_start = helper.make_tensor("range_start", TensorProto.INT64, [], [0]) -range_delta = helper.make_tensor("range_delta", TensorProto.INT64, [], [1]) - -# Shape(X) -> S : 1-D int64 tensor [3, 4, 2000]. -shape_node = helper.make_node("Shape", inputs=["X"], outputs=["S"], name="ShapeNode") -# Gather(S, [[-1]]) -> G : rank-2 [1, 1] int64 value [[2000]]. -gather_node = helper.make_node("Gather", inputs=["S", "gather_idx"], outputs=["G"], axis=0, name="GatherNode") -# Squeeze(G) -> K : 0-D scalar int64 value (only when a value is propagated). -squeeze_node = helper.make_node("Squeeze", inputs=["G"], outputs=["K"], name="SqueezeNode") -# Range(0, K, 1) -> R : 1-D int64 tensor whose length is K (symbolic on decline). -range_node = helper.make_node("Range", inputs=["range_start", "K", "range_delta"], outputs=["R"], name="RangeNode") - -graph = helper.make_graph( - nodes=[shape_node, gather_node, squeeze_node, range_node], - name="Shape_Gather_Rank2_Decline_Model", - inputs=[input_tensor], - outputs=[range_out], - initializer=[gather_idx, range_start, range_delta], -) - -model = helper.make_model( - graph, - opset_imports=[helper.make_operatorsetid("", 18)], - producer_name="onnx-example-generator", -) -model.ir_version = IR_VERSION - -checker.check_model(model) -save(model, "test_shape_data_propagation_gather_rank2_decline.onnx") - -print("Model saved to test_shape_data_propagation_gather_rank2_decline.onnx") diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_decline.onnx b/onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_decline.onnx deleted file mode 100644 index b65e5d5830ac899d9a17b009e9e20a08c7077856..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 423 zcmZ8d!AiqG5bfC3Bom@>1TCQ_@f0YC-g*o|4n+?M0r9e|W=I2ZQh~@42MgJ9O^|p_GDzBM?J3R? z5`vXy$Gt_*jX#37Xjf8YEb2y94Z>>RDHn|-h&(IW&HLaQ(zVVn`gmf$N!;^mkt-Gt zUCF|qo8THt{23ku2VMP3J%v&5uZvHesgn;=^%@MCPWJ5X?rNiD0!?vjexJT{myYNh T&rE|Hw8`pw6|8~@j$HZ$uT_0X diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_decline.py b/onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_decline.py deleted file mode 100644 index b9df3bd67c762..0000000000000 --- a/onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_decline.py +++ /dev/null @@ -1,65 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. - -# Lock fixture for Unsqueeze's custom data-propagation DECLINE on a single-element value. -# -# Shape(X) -> Gather([-1]) produces a rank-1 single-element value (the last dimension of X, 2000). -# Unsqueezing that single-element (scalar-like, rank-1 [1]) value would yield a rank >= 2 result -# ([1, 2000]) that the single-value data-propagation channel cannot faithfully represent, so -# Unsqueeze's custom data propagation must DECLINE rather than fabricate the misleading [1, value]. -# -# The decline is made OBSERVABLE through a rank-lowering Squeeze: -# Shape(X) -> Gather([-1]) -> Unsqueeze([0]) -> Squeeze -> Range(0, K, 1) -# With the correct decline, no value is propagated past Unsqueeze, Squeeze has nothing to lower, and -# Range's limit stays symbolic -- the model loads and the Range output length is an unknown dimension. -# If the decline were relaxed to fabricate [1, 2000], Squeeze would lower a non-scalar value into -# Range, which ONNX Range shape inference rejects ("Input to 'Range' op should be scalars"), so the -# model fails to load. The decline is thus locked end to end: the correct path loads with a symbolic -# Range length, while the rank-fabricating bug fails to load. - -from onnx import IR_VERSION, TensorProto, checker, helper, save - -# Graph input: a 1-D float tensor of static length 2000. -input_tensor = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2000]) - -# Output: Range result, a 1-D int64 tensor whose length is the propagated K (symbolic on decline). -range_out = helper.make_tensor_value_info("R", TensorProto.INT64, ["range_len"]) - -# Gather index is the 1-D constant [-1] (last dimension) -> rank-1 single-element value. -gather_idx = helper.make_tensor("gather_idx", TensorProto.INT64, [1], [-1]) -# Unsqueeze axes: insert a leading dimension (axis 0). -unsqueeze_axes = helper.make_tensor("unsqueeze_axes", TensorProto.INT64, [1], [0]) -# Range start/delta as 0-D scalar int64 constants. -range_start = helper.make_tensor("range_start", TensorProto.INT64, [], [0]) -range_delta = helper.make_tensor("range_delta", TensorProto.INT64, [], [1]) - -# Shape(X) -> S : 1-D int64 tensor [2000]. -shape_node = helper.make_node("Shape", inputs=["X"], outputs=["S"], name="ShapeNode") -# Gather(S, [-1]) -> G : rank-1 size-1 int64 value [2000]. -gather_node = helper.make_node("Gather", inputs=["S", "gather_idx"], outputs=["G"], axis=0, name="GatherNode") -# Unsqueeze(G, [0]) -> U : rank-2 [1, 1] tensor; custom data propagation must decline. -unsqueeze_node = helper.make_node("Unsqueeze", inputs=["G", "unsqueeze_axes"], outputs=["U"], name="UnsqueezeNode") -# Squeeze(U) -> K : 0-D scalar int64 value (only when a value is propagated). -squeeze_node = helper.make_node("Squeeze", inputs=["U"], outputs=["K"], name="SqueezeNode") -# Range(0, K, 1) -> R : 1-D int64 tensor whose length is K (symbolic on decline). -range_node = helper.make_node("Range", inputs=["range_start", "K", "range_delta"], outputs=["R"], name="RangeNode") - -graph = helper.make_graph( - nodes=[shape_node, gather_node, unsqueeze_node, squeeze_node, range_node], - name="Shape_Gather_Unsqueeze_Decline_Model", - inputs=[input_tensor], - outputs=[range_out], - initializer=[gather_idx, unsqueeze_axes, range_start, range_delta], -) - -model = helper.make_model( - graph, - opset_imports=[helper.make_operatorsetid("", 18)], - producer_name="onnx-example-generator", -) -model.ir_version = IR_VERSION - -checker.check_model(model) -save(model, "test_shape_data_propagation_unsqueeze_decline.onnx") - -print("Model saved to test_shape_data_propagation_unsqueeze_decline.onnx") From 12df3c47976a730e052f727ff38cd742ec155733 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Mon, 22 Jun 2026 21:49:42 +0000 Subject: [PATCH 6/8] Collapse positive data-prop opt-loops to ORT_DISABLE_ALL to cut ASan footprint The four positive shape-inference data-propagation regression tests (GatherToTopK / GatherMulToTopK / ShapeMulMultiElementNoScalarCollapse / GatherSqueezeRangeRankPreservation) each looped over {ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_ALL}, creating three InferenceSession instances per test (12 total). Data propagation runs in the pre-optimization Graph::Resolve pass and is independent of the graph optimization level -- the regression reproduces even at ORT_DISABLE_ALL -- so the BASIC/ALL iterations re-ran the same propagation behind extra optimizer passes without adding distinct coverage, at the cost of extra per-session allocations. Collapse each loop to a single ORT_DISABLE_ALL session (12 -> 4 sessions) to reduce the onnxruntime_test_all live footprint at the AddressSanitizer size-class ceiling (microsoft/onnxruntime#29139). The discrimination is a strict subset at the level that catches microsoft/onnxruntime#29072: the tests still fail on the regressed code at ORT_DISABLE_ALL. The separate ORT_ENABLE_BASIC Identity-removal test is left untouched (it specifically needs BASIC to drop the Identity node). Comments are reworded to reflect the single level and keep the grounding. Production data-propagation files are unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwangms --- .../test/framework/shape_inference_test.cc | 127 +++++++++--------- 1 file changed, 62 insertions(+), 65 deletions(-) diff --git a/onnxruntime/test/framework/shape_inference_test.cc b/onnxruntime/test/framework/shape_inference_test.cc index bde00e5aa9bac..a79b5406c9d1b 100644 --- a/onnxruntime/test/framework/shape_inference_test.cc +++ b/onnxruntime/test/framework/shape_inference_test.cc @@ -201,34 +201,32 @@ TEST(ShapeInferenceV2Test, PartialDataPropagationTest) { // "K input must be a one-dimensional tensor of size 1." at model load time. // // The failure reproduces even at ORT_DISABLE_ALL (constant folding never runs there), which -// proves the cause is shape-inference data propagation, not constant folding. This test loads -// the model at the disabled, basic, and all-optimization levels, which bracket the data-propagation -// path: data propagation runs in the pre-optimization Graph::Resolve pass, so it is independent of -// the graph-optimization level (ORT_ENABLE_ALL already applies the extended and layout transformers, -// so they need not be enumerated separately). The test asserts the rank-1 K shape is preserved so the -// model loads and TopK output shapes are inferred correctly. +// proves the cause is shape-inference data propagation, not constant folding. Data propagation runs +// in the pre-optimization Graph::Resolve pass, so it is independent of the graph-optimization level; +// loading at ORT_DISABLE_ALL therefore fully exercises the data-propagation path (enumerating +// ORT_ENABLE_BASIC/ALL would only re-run the same propagation behind additional optimizer passes, at +// the cost of extra InferenceSession allocations -- see microsoft/onnxruntime#29139). The test +// asserts the rank-1 K shape is preserved so the model loads and TopK output shapes are inferred +// correctly. TEST(ShapeInferenceV2Test, GatherToTopKRankPreservationTest) { auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_topk.onnx"); - for (auto opt_level : {ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_ALL}) { - Ort::SessionOptions session_options{}; - session_options.SetGraphOptimizationLevel(opt_level); + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); - // Loading must succeed at each level; before the fix this threw a - // ShapeInferenceError at ORT_DISABLE_ALL (and above). - Ort::Session session(*ort_env, model_path, session_options); + // Loading must succeed; before the fix this threw a ShapeInferenceError at ORT_DISABLE_ALL. + Ort::Session session(*ort_env, model_path, session_options); - ORT_ENFORCE(session.GetOutputCount() == 2); + ORT_ENFORCE(session.GetOutputCount() == 2); - // K is data-propagated as the last dimension of X (2000), so both TopK outputs are 1-D - // tensors of length 2000. - for (size_t output_index = 0; output_index < 2; ++output_index) { - Ort::TypeInfo type_info = session.GetOutputTypeInfo(output_index); - auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); - std::vector output_shape = tensor_info.GetShape(); - EXPECT_TRUE(output_shape.size() == 1) << "TopK output should be a 1-D tensor"; - EXPECT_TRUE(output_shape[0] == 2000) << "TopK output length should be the inferred K value (2000)"; - } + // K is data-propagated as the last dimension of X (2000), so both TopK outputs are 1-D + // tensors of length 2000. + for (size_t output_index = 0; output_index < 2; ++output_index) { + Ort::TypeInfo type_info = session.GetOutputTypeInfo(output_index); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector output_shape = tensor_info.GetShape(); + EXPECT_TRUE(output_shape.size() == 1) << "TopK output should be a 1-D tensor"; + EXPECT_TRUE(output_shape[0] == 2000) << "TopK output length should be the inferred K value (2000)"; } } @@ -246,25 +244,24 @@ TEST(ShapeInferenceV2Test, GatherToTopKRankPreservationTest) { TEST(ShapeInferenceV2Test, GatherMulToTopKRankPreservationTest) { auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_mul_topk.onnx"); - for (auto opt_level : {ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_ALL}) { - Ort::SessionOptions session_options{}; - session_options.SetGraphOptimizationLevel(opt_level); + // Data propagation runs in the pre-optimization Graph::Resolve pass, so ORT_DISABLE_ALL exercises + // the full chain; see GatherToTopKRankPreservationTest above and microsoft/onnxruntime#29139. + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); - // Loading must succeed at each level; before the fix this threw a - // ShapeInferenceError at ORT_DISABLE_ALL (and above). - Ort::Session session(*ort_env, model_path, session_options); + // Loading must succeed; before the fix this threw a ShapeInferenceError at ORT_DISABLE_ALL. + Ort::Session session(*ort_env, model_path, session_options); - ORT_ENFORCE(session.GetOutputCount() == 2); + ORT_ENFORCE(session.GetOutputCount() == 2); - // K = 50 * 40 = 2000 is data-propagated through the two Gathers and the Mul, so both TopK - // outputs are 1-D tensors of length 2000. - for (size_t output_index = 0; output_index < 2; ++output_index) { - Ort::TypeInfo type_info = session.GetOutputTypeInfo(output_index); - auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); - std::vector output_shape = tensor_info.GetShape(); - EXPECT_TRUE(output_shape.size() == 1) << "TopK output should be a 1-D tensor"; - EXPECT_TRUE(output_shape[0] == 2000) << "TopK output length should be the propagated K value (2000)"; - } + // K = 50 * 40 = 2000 is data-propagated through the two Gathers and the Mul, so both TopK + // outputs are 1-D tensors of length 2000. + for (size_t output_index = 0; output_index < 2; ++output_index) { + Ort::TypeInfo type_info = session.GetOutputTypeInfo(output_index); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector output_shape = tensor_info.GetShape(); + EXPECT_TRUE(output_shape.size() == 1) << "TopK output should be a 1-D tensor"; + EXPECT_TRUE(output_shape[0] == 2000) << "TopK output length should be the propagated K value (2000)"; } } @@ -281,24 +278,24 @@ TEST(ShapeInferenceV2Test, GatherMulToTopKRankPreservationTest) { TEST(ShapeInferenceV2Test, ShapeMulMultiElementNoScalarCollapseTest) { auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_shape_mul_constantofshape.onnx"); - for (auto opt_level : {ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_ALL}) { - Ort::SessionOptions session_options{}; - session_options.SetGraphOptimizationLevel(opt_level); + // Data propagation runs in the pre-optimization Graph::Resolve pass, so ORT_DISABLE_ALL exercises + // the full chain; see GatherToTopKRankPreservationTest above and microsoft/onnxruntime#29139. + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); - Ort::Session session(*ort_env, model_path, session_options); + Ort::Session session(*ort_env, model_path, session_options); - ORT_ENFORCE(session.GetOutputCount() == 1); + ORT_ENFORCE(session.GetOutputCount() == 1); - Ort::TypeInfo type_info = session.GetOutputTypeInfo(0); - auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); - std::vector output_shape = tensor_info.GetShape(); + Ort::TypeInfo type_info = session.GetOutputTypeInfo(0); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector output_shape = tensor_info.GetShape(); - // ConstantOfShape(Mul(Shape(X), Shape(X))) == ConstantOfShape([9, 16]) -> rank-2 [9, 16]. - // A collapse of the multi-element value to a scalar would yield rank 1 here. - EXPECT_TRUE(output_shape.size() == 2) << "Output should stay rank 2; the multi-element value must not collapse"; - EXPECT_TRUE(output_shape.size() == 2 && output_shape[0] == 9 && output_shape[1] == 16) - << "ConstantOfShape output should be the propagated multi-element shape [9, 16]"; - } + // ConstantOfShape(Mul(Shape(X), Shape(X))) == ConstantOfShape([9, 16]) -> rank-2 [9, 16]. + // A collapse of the multi-element value to a scalar would yield rank 1 here. + EXPECT_TRUE(output_shape.size() == 2) << "Output should stay rank 2; the multi-element value must not collapse"; + EXPECT_TRUE(output_shape.size() == 2 && output_shape[0] == 9 && output_shape[1] == 16) + << "ConstantOfShape output should be the propagated multi-element shape [9, 16]"; } // Regression for microsoft/onnxruntime#29072. @@ -384,23 +381,23 @@ TEST(ShapeInferenceV2Test, GatherUnsqueezeDeclineTest) { TEST(ShapeInferenceV2Test, GatherSqueezeRangeRankPreservationTest) { auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_squeeze_range.onnx"); - for (auto opt_level : {ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_ALL}) { - Ort::SessionOptions session_options{}; - session_options.SetGraphOptimizationLevel(opt_level); + // Data propagation runs in the pre-optimization Graph::Resolve pass, so ORT_DISABLE_ALL exercises + // the full chain; see GatherToTopKRankPreservationTest above and microsoft/onnxruntime#29139. + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); - Ort::Session session(*ort_env, model_path, session_options); + Ort::Session session(*ort_env, model_path, session_options); - ORT_ENFORCE(session.GetOutputCount() == 1); + ORT_ENFORCE(session.GetOutputCount() == 1); - Ort::TypeInfo type_info = session.GetOutputTypeInfo(0); - auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); - std::vector output_shape = tensor_info.GetShape(); + Ort::TypeInfo type_info = session.GetOutputTypeInfo(0); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector output_shape = tensor_info.GetShape(); - // Range(0, Squeeze(Gather(Shape(X), [-1])), 1) == Range(0, 2000, 1) -> 1-D length 2000. - EXPECT_TRUE(output_shape.size() == 1) << "Range output should be a 1-D tensor"; - EXPECT_TRUE(output_shape.size() == 1 && output_shape[0] == 2000) - << "Range length should be the scalar K (2000) propagated through Squeeze"; - } + // Range(0, Squeeze(Gather(Shape(X), [-1])), 1) == Range(0, 2000, 1) -> 1-D length 2000. + EXPECT_TRUE(output_shape.size() == 1) << "Range output should be a 1-D tensor"; + EXPECT_TRUE(output_shape.size() == 1 && output_shape[0] == 2000) + << "Range length should be the scalar K (2000) propagated through Squeeze"; } // Unit test for the single-element guard in TryGetSinglePropagatedShapeValue, the helper shared by From be9aae857e767aa932e2dd1124e43dba60d891ee Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Mon, 22 Jun 2026 21:53:34 +0000 Subject: [PATCH 7/8] Update stale decline-test reference in Unsqueeze data-propagation comment The Unsqueeze single-value decline comment referenced UnsqueezeSingleValueDeclineTest, which was folded into the combined GatherUnsqueezeDeclineTest (Branch B). Update the reference so the comment points at the test that actually locks the behavior. Comment-only; no logic change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwangms --- .../graph/data_propagation/unsqueeze_op_data_propagation.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.cc index 52cacd4ccfbe8..1d7f0dd2c86e3 100644 --- a/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.cc +++ b/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.cc @@ -41,7 +41,7 @@ Status UnsqueezeOpDataPropagation::infer() { // single-value channel cannot represent. Multi-element shape vectors (dim_size > 1) are // legitimate and left untouched. // - // This decline is locked end to end by UnsqueezeSingleValueDeclineTest via a rank-lowering + // This decline is locked end to end by GatherUnsqueezeDeclineTest (Branch B) via a rank-lowering // Squeeze (Shape -> Gather([-1]) -> Unsqueeze -> Squeeze -> Range): with the decline, Range's // length stays symbolic; relaxing it fabricates [1, value], which makes the downstream Range // input non-scalar and the model fails to load -- so the behavior IS observable and the test From a4a5b7067f3a3c1b6bdaef01fe6f10b953327567 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Tue, 23 Jun 2026 00:07:00 +0000 Subject: [PATCH 8/8] test: relocate data-propagation shape-inference tests to provider shard Move the 7 PR-added ShapeInferenceV2Test data-propagation tests out of test/framework/shape_inference_test.cc (ctest shard-1, onnxruntime_test_all) into a new test/providers/cpu/tensor/shape_inference_data_propagation_test.cc (shard-2, onnxruntime_provider_test) via the providers/ path-prefix routing in partition_provider_test_srcs. shard-1's cumulative ASan size-class-8192 high-water hits the #29139 architectural ceiling on windows_x64_asan; the baseline SessionStatePrepacking allocation is the OOM site, not these tests, but removing our additions returns shard-1 to its green baseline. shard-2 has ASan headroom, so the full 3-level opt-loop coverage {ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_ALL} is restored (reverting the interim collapse), preserving the #29072 regression lock. Production data-propagation logic is unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwangms --- .../test/framework/shape_inference_test.cc | 328 --------------- .../shape_inference_data_propagation_test.cc | 372 ++++++++++++++++++ 2 files changed, 372 insertions(+), 328 deletions(-) create mode 100644 onnxruntime/test/providers/cpu/tensor/shape_inference_data_propagation_test.cc diff --git a/onnxruntime/test/framework/shape_inference_test.cc b/onnxruntime/test/framework/shape_inference_test.cc index a79b5406c9d1b..6d92acf6bae78 100644 --- a/onnxruntime/test/framework/shape_inference_test.cc +++ b/onnxruntime/test/framework/shape_inference_test.cc @@ -7,7 +7,6 @@ #include "gtest/gtest.h" #include "core/common/span_utils.h" #include "core/graph/model.h" -#include "core/graph/data_propagation/data_propagation_value_utils.h" #include "core/session/onnxruntime_cxx_api.h" #include "test/framework/model_builder_utils.h" #include "test/util/include/asserts.h" @@ -192,333 +191,6 @@ TEST(ShapeInferenceV2Test, PartialDataPropagationTest) { } } -// Regression test for the Shape -> Gather(1-D index) -> TopK rank-drop. -// -// The Gather index is the 1-D constant [-1] (rank 1), so per ONNX Gather semantics the Gather -// output is a rank-1, single-element tensor -- exactly what TopK's K input requires. Previously -// the Gather custom data propagation scalarized that single element (dropping the rank), -// producing a 0-D K initializer that ONNX TopK shape inference rejected with -// "K input must be a one-dimensional tensor of size 1." at model load time. -// -// The failure reproduces even at ORT_DISABLE_ALL (constant folding never runs there), which -// proves the cause is shape-inference data propagation, not constant folding. Data propagation runs -// in the pre-optimization Graph::Resolve pass, so it is independent of the graph-optimization level; -// loading at ORT_DISABLE_ALL therefore fully exercises the data-propagation path (enumerating -// ORT_ENABLE_BASIC/ALL would only re-run the same propagation behind additional optimizer passes, at -// the cost of extra InferenceSession allocations -- see microsoft/onnxruntime#29139). The test -// asserts the rank-1 K shape is preserved so the model loads and TopK output shapes are inferred -// correctly. -TEST(ShapeInferenceV2Test, GatherToTopKRankPreservationTest) { - auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_topk.onnx"); - - Ort::SessionOptions session_options{}; - session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); - - // Loading must succeed; before the fix this threw a ShapeInferenceError at ORT_DISABLE_ALL. - Ort::Session session(*ort_env, model_path, session_options); - - ORT_ENFORCE(session.GetOutputCount() == 2); - - // K is data-propagated as the last dimension of X (2000), so both TopK outputs are 1-D - // tensors of length 2000. - for (size_t output_index = 0; output_index < 2; ++output_index) { - Ort::TypeInfo type_info = session.GetOutputTypeInfo(output_index); - auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); - std::vector output_shape = tensor_info.GetShape(); - EXPECT_TRUE(output_shape.size() == 1) << "TopK output should be a 1-D tensor"; - EXPECT_TRUE(output_shape[0] == 2000) << "TopK output length should be the inferred K value (2000)"; - } -} - -// Regression test for the Shape -> Gather(1-D index) -> Mul -> TopK chain. -// -// This covers the elementwise data-propagation consumers (Add/Sub/Mul/Div) with rank-1 -// single-element operands. Both Gather indices are 1-D constants, so each Gather output is a -// rank-1 single-element value; Mul must keep propagating that (as a rank-1 [1] value) so the -// downstream TopK still receives a valid 1-D K. Previously the elementwise ops only handled -// rank-0 scalar operands, so once the Gather producer was fixed to emit rank-1 values the chain -// would have silently stopped propagating at Mul (degrading TopK's K to a symbolic dim). -// -// Asserting that both TopK outputs resolve to the concrete length 2000 (== 50 * 40) proves the -// value propagated end to end through Mul rather than being lost. -TEST(ShapeInferenceV2Test, GatherMulToTopKRankPreservationTest) { - auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_mul_topk.onnx"); - - // Data propagation runs in the pre-optimization Graph::Resolve pass, so ORT_DISABLE_ALL exercises - // the full chain; see GatherToTopKRankPreservationTest above and microsoft/onnxruntime#29139. - Ort::SessionOptions session_options{}; - session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); - - // Loading must succeed; before the fix this threw a ShapeInferenceError at ORT_DISABLE_ALL. - Ort::Session session(*ort_env, model_path, session_options); - - ORT_ENFORCE(session.GetOutputCount() == 2); - - // K = 50 * 40 = 2000 is data-propagated through the two Gathers and the Mul, so both TopK - // outputs are 1-D tensors of length 2000. - for (size_t output_index = 0; output_index < 2; ++output_index) { - Ort::TypeInfo type_info = session.GetOutputTypeInfo(output_index); - auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); - std::vector output_shape = tensor_info.GetShape(); - EXPECT_TRUE(output_shape.size() == 1) << "TopK output should be a 1-D tensor"; - EXPECT_TRUE(output_shape[0] == 2000) << "TopK output length should be the propagated K value (2000)"; - } -} - -// Negative regression test for the single-element guard in the elementwise data-propagation -// consumers (Add/Sub/Mul/Div), exercised end to end. -// -// Shape(X[3, 4]) -> S is a rank-1 MULTI-element value [3, 4]. Mul(S, S) -> M must stay the -// multi-element value [9, 16] so ConstantOfShape(M) resolves to the rank-2 shape [9, 16]. For a -// multi-element data-propagation output ONNX's own Mul data propagation supplies the value and -// ORT's custom elementwise propagation is intentionally NOT engaged (it is dispatched only for a -// rank-0 result, see CreateCustomDataPropagation). This is an end-to-end non-regression guard that -// the rank-routing changes did not disturb the multi-element chain; the single-element guard inside -// the shared helper is exercised directly by SinglePropagatedShapeValueGuardTest below. -TEST(ShapeInferenceV2Test, ShapeMulMultiElementNoScalarCollapseTest) { - auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_shape_mul_constantofshape.onnx"); - - // Data propagation runs in the pre-optimization Graph::Resolve pass, so ORT_DISABLE_ALL exercises - // the full chain; see GatherToTopKRankPreservationTest above and microsoft/onnxruntime#29139. - Ort::SessionOptions session_options{}; - session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); - - Ort::Session session(*ort_env, model_path, session_options); - - ORT_ENFORCE(session.GetOutputCount() == 1); - - Ort::TypeInfo type_info = session.GetOutputTypeInfo(0); - auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); - std::vector output_shape = tensor_info.GetShape(); - - // ConstantOfShape(Mul(Shape(X), Shape(X))) == ConstantOfShape([9, 16]) -> rank-2 [9, 16]. - // A collapse of the multi-element value to a scalar would yield rank 1 here. - EXPECT_TRUE(output_shape.size() == 2) << "Output should stay rank 2; the multi-element value must not collapse"; - EXPECT_TRUE(output_shape.size() == 2 && output_shape[0] == 9 && output_shape[1] == 16) - << "ConstantOfShape output should be the propagated multi-element shape [9, 16]"; -} - -// Regression for microsoft/onnxruntime#29072. -// End-to-end lock for BOTH custom data-propagation DECLINE paths -- Gather on a rank-2 index and -// Unsqueeze on a single-element value -- driven through a SINGLE shared model so the coverage costs -// one onnxruntime::Model::Load instead of two (microsoft/onnxruntime#29139: the single-process -// onnxruntime_test_all run sits near the AddressSanitizer size-class ceiling, so the two decline -// fixtures are folded into one graph to minimize the live Model/Graph footprint). -// -// One rank-3 input X drives a shared Shape(X) -> S (the 1-D vector [3, 4, 2000]); two independent -// branches gather from S and route into a rank-lowering Squeeze -> Range so each decline is -// observable end to end: -// -// Branch A (Gather rank-2 index -> output RA): S -> Gather([[ -1 ]]) -> Squeeze -> Range(0, K, 1) -// The index [[-1]] has shape [1, 1] (rank 2), so the Gather output rank is -// data_rank - 1 + index_rank = 1 - 1 + 2 = 2 -- a rank-2 single-element value the single-value -// channel cannot represent, so Gather's custom propagation must DECLINE. With the correct -// decline nothing propagates, Squeeze has no value to lower, and RA's length stays SYMBOLIC. A -// relaxed decline would emit a rank-1 [1] value that Squeeze lowers to the scalar 2000, -// concretizing RA to length 2000 -- so asserting RA stays symbolic discriminates the correct -// decline from the rank-fabricating bug. -// -// Branch B (Unsqueeze single value -> output RB): S -> Gather([-1]) -> Unsqueeze([0]) -> Squeeze -> -// Range(0, K, 1). Gather([-1]) is a rank-1 single-element value (the last dimension of X, 2000); -// unsqueezing it would yield a rank >= 2 result ([1, 2000]) the single-value channel cannot -// represent, so Unsqueeze's custom propagation must DECLINE. With the correct decline nothing -// propagates past Unsqueeze, RB's length stays SYMBOLIC, and the model loads. A relaxed decline -// would fabricate [1, 2000], which Squeeze lowers into a non-scalar Range limit that ONNX Range -// shape inference rejects ("Input to 'Range' op should be scalars"), so Graph::Resolve (and -// therefore Model::Load) FAILS. So Branch B is observable either as a load failure or, were that -// to pass, as a concrete RB length. -// -// Both Gather indices are raw graph initializers (not surfaced through the data-propagation -// getInputData channel), so ONNX's own Gather data propagator bails and control reaches our custom -// decline branches. The branches are independent: a regression in either is caught (Branch A as a -// concrete RA dimension, Branch B as a failed load), so this one combined model keeps the full -// discriminating power of the two original per-path fixtures. -// -// The decline is driven through onnxruntime::Model::Load (which runs Graph::Resolve) and the output -// NodeArg shapes are inspected directly -- no InferenceSession, so none of the session's arena / -// execution-provider / kernel-registry allocations are created. Data propagation runs inside -// Graph::Resolve (Graph::InferAndVerifyTypeMatch -> data-propagation RunInferencing); constant -// folding is a separate session-level optimizer pass that never runs here, so Resolve alone -// exercises the decline branches in isolation with no folding to mask them. -TEST(ShapeInferenceV2Test, GatherUnsqueezeDeclineTest) { - auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_decline_combined.onnx"); - - std::shared_ptr model; - // Model::Load runs Graph::Resolve, which performs data propagation. A correct decline on BOTH - // branches keeps each Range length symbolic, so the resolve (and therefore the load) succeeds. A - // relaxed Unsqueeze decline (Branch B) makes Range's input non-scalar, which Range shape inference - // rejects, so Resolve returns non-OK and this assertion fails on the load itself. - ASSERT_STATUS_OK(onnxruntime::Model::Load(model_path, model, nullptr, - DefaultLoggingManager().DefaultLogger())); - - const auto& graph_outputs = model->MainGraph().GetOutputs(); - ASSERT_EQ(graph_outputs.size(), static_cast(2)); - - // Both Range outputs are 1-D; each single dimension must stay SYMBOLIC (no concrete dim_value) - // because its branch declines rather than propagating a (rank-fabricated) concrete K. RA[0] would - // concretize to 2000 if the rank-2 Gather index stopped declining; RB[0] would concretize to 2000 - // (or the load would fail) if the single-value Unsqueeze stopped declining. - for (const NodeArg* output : graph_outputs) { - const TensorShapeProto* output_shape = output->Shape(); - ASSERT_NE(output_shape, nullptr) << "Range output '" << output->Name() - << "' should have an inferred (rank-1) shape"; - ASSERT_EQ(output_shape->dim_size(), 1) << "Range output '" << output->Name() - << "' should be a 1-D tensor"; - EXPECT_FALSE(output_shape->dim(0).has_dim_value()) - << "Range length '" << output->Name() << "' must stay symbolic; the decline must not " - "fabricate a concrete K (a relaxed decline concretizes it to 2000)"; - } -} - -// Lock test for the Shape -> Gather(1-D index) -> Squeeze -> Range chain. -// -// Squeeze of a rank-1 single-element value [1] removes the size-1 dimension and yields a 0-D -// scalar -- the strict-improvement behavior of Squeeze's custom data propagation. This test locks -// that improvement end to end: the Gather produces a rank-1 single element (the last dimension of -// X, 2000); Squeeze must propagate it as the scalar 2000; and Range(0, K, 1) must then resolve to -// a concrete 1-D tensor of length 2000. If Squeeze dropped, corrupted, or failed to scalarize the -// value, Range's length would stay symbolic (size 0 here) instead of the concrete 2000. -TEST(ShapeInferenceV2Test, GatherSqueezeRangeRankPreservationTest) { - auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_squeeze_range.onnx"); - - // Data propagation runs in the pre-optimization Graph::Resolve pass, so ORT_DISABLE_ALL exercises - // the full chain; see GatherToTopKRankPreservationTest above and microsoft/onnxruntime#29139. - Ort::SessionOptions session_options{}; - session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); - - Ort::Session session(*ort_env, model_path, session_options); - - ORT_ENFORCE(session.GetOutputCount() == 1); - - Ort::TypeInfo type_info = session.GetOutputTypeInfo(0); - auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); - std::vector output_shape = tensor_info.GetShape(); - - // Range(0, Squeeze(Gather(Shape(X), [-1])), 1) == Range(0, 2000, 1) -> 1-D length 2000. - EXPECT_TRUE(output_shape.size() == 1) << "Range output should be a 1-D tensor"; - EXPECT_TRUE(output_shape.size() == 1 && output_shape[0] == 2000) - << "Range length should be the scalar K (2000) propagated through Squeeze"; -} - -// Unit test for the single-element guard in TryGetSinglePropagatedShapeValue, the helper shared by -// the Add/Sub/Mul/Div data-propagation consumers. -// -// A NodeArg can carry a propagated shape value in one of two non-interchangeable channels: a rank-0 -// scalar (inferred_scalar_value_) or a rank>=1 list of values (inferred_shape_values_). The helper -// must surface a single value from either a scalar source or a rank-1 SINGLE-element source (and -// report its rank), and must DECLINE for a rank-1 MULTI-element source. Declining is load-bearing: -// silently using element[0] of a multi-element value would let an elementwise op emit a bogus -// scalar product that getInputData() would then hand to a rank-sensitive consumer. This test -// exercises the guard directly -- a relaxed guard that accepted element[0] fails the multi-element -// case below. -TEST(ShapeInferenceV2Test, SinglePropagatedShapeValueGuardTest) { - // Scalar channel -> accepted, reported as rank-0. - { - NodeArg arg("scalar_arg", nullptr); - arg.SetInferredShapeScalarValue(5); - int64_t value = 0; - bool is_rank1 = true; - EXPECT_TRUE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); - EXPECT_EQ(value, 5); - EXPECT_FALSE(is_rank1); - } - - // Rank-1 single-element value -> accepted, reported as rank-1. - { - NodeArg arg("rank1_single_arg", nullptr); - auto& values = arg.GetMutableInferredShapeValues(); - values.emplace(); - values->add_dim()->set_dim_value(7); - int64_t value = 0; - bool is_rank1 = false; - EXPECT_TRUE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); - EXPECT_EQ(value, 7); - EXPECT_TRUE(is_rank1); - } - - // Rank-1 MULTI-element value -> DECLINED (the guard); element[0] must never be used. - { - NodeArg arg("rank1_multi_arg", nullptr); - auto& values = arg.GetMutableInferredShapeValues(); - values.emplace(); - values->add_dim()->set_dim_value(3); - values->add_dim()->set_dim_value(4); - int64_t value = -1; - bool is_rank1 = false; - EXPECT_FALSE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)) - << "A rank-1 multi-element value must be declined, not collapsed to element[0]"; - } - - // Rank-1 single element with a SYMBOLIC (no concrete value) dim -> declined. - { - NodeArg arg("rank1_symbolic_arg", nullptr); - auto& values = arg.GetMutableInferredShapeValues(); - values.emplace(); - values->add_dim()->set_dim_param("N"); - int64_t value = -1; - bool is_rank1 = false; - EXPECT_FALSE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)) - << "A symbolic single-element value has no concrete value and must be declined"; - } - - // No propagated value on either channel -> declined. - { - NodeArg arg("empty_arg", nullptr); - int64_t value = -1; - bool is_rank1 = false; - EXPECT_FALSE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); - } -} - -// Unit test that SetSinglePropagatedShapeValue is correct-by-construction: it populates exactly one -// channel and clears the other. This pins the invariant the elementwise/Gather consumers rely on -- -// the scalar-first reader (TryGetSinglePropagatedShapeValue) and the values-first getInputData() can -// only agree on rank if no NodeArg ever carries both channels. A setter that left the opposite -// channel populated would let a stale value win (reintroducing a rank mismatch), so each case below -// pre-populates the opposite channel and asserts the setter cleared it. -TEST(ShapeInferenceV2Test, SetSinglePropagatedShapeValueKeepsSingleChannelTest) { - // Scalar write must clear a stale values channel. - { - NodeArg arg("scalar_over_stale_values", nullptr); - auto& stale = arg.GetMutableInferredShapeValues(); - stale.emplace(); - stale->add_dim()->set_dim_value(99); - - SetSinglePropagatedShapeValue(arg, 5, /*is_rank1=*/false); - - ASSERT_TRUE(arg.GetInferredShapeScalarValue().has_value()); - EXPECT_EQ(arg.GetInferredShapeScalarValue().value(), 5); - EXPECT_FALSE(arg.GetInferredShapeValues().has_value()) - << "Scalar write must clear the values channel that getInputData() would otherwise prefer"; - - int64_t value = 0; - bool is_rank1 = true; - EXPECT_TRUE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); - EXPECT_EQ(value, 5); - EXPECT_FALSE(is_rank1); - } - - // Rank-1 write must clear a stale scalar channel. - { - NodeArg arg("values_over_stale_scalar", nullptr); - arg.SetInferredShapeScalarValue(99); - - SetSinglePropagatedShapeValue(arg, 7, /*is_rank1=*/true); - - EXPECT_FALSE(arg.GetInferredShapeScalarValue().has_value()) - << "Rank-1 write must clear the scalar channel that the scalar-first reader would otherwise return"; - ASSERT_TRUE(arg.GetInferredShapeValues().has_value()); - ASSERT_EQ(arg.GetInferredShapeValues()->dim_size(), 1); - EXPECT_EQ(arg.GetInferredShapeValues()->dim(0).dim_value(), 7); - - int64_t value = 0; - bool is_rank1 = false; - EXPECT_TRUE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); - EXPECT_EQ(value, 7); - EXPECT_TRUE(is_rank1); - } -} - namespace { struct MyCustomKernelWithOptionalInput { MyCustomKernelWithOptionalInput(const OrtKernelInfo* info) { diff --git a/onnxruntime/test/providers/cpu/tensor/shape_inference_data_propagation_test.cc b/onnxruntime/test/providers/cpu/tensor/shape_inference_data_propagation_test.cc new file mode 100644 index 0000000000000..d719213eaff0e --- /dev/null +++ b/onnxruntime/test/providers/cpu/tensor/shape_inference_data_propagation_test.cc @@ -0,0 +1,372 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Shape-inference data-propagation regression tests for the custom Gather / Unsqueeze / Squeeze +// and elementwise (Add/Sub/Mul/Div) propagators. These lock the rank-preservation fix for +// microsoft/onnxruntime#29072 (a 1-D single-element Gather output must keep rank 1 so a downstream +// TopK/Range/ConstantOfShape sees a valid shape, while a rank>=2 Gather/Unsqueeze result must +// decline rather than fabricate a misleading rank). +// +// Why these tests live under test/providers/ rather than test/framework/: +// The unit tests are split across two ctest binaries by source path. partition_provider_test_srcs() +// in cmake/onnxruntime_unittests.cmake routes sources under test/providers/ (and test/contrib_ops/) +// into the onnxruntime_provider_test binary (shard 2); everything else compiles into +// onnxruntime_test_all (shard 1). Under AddressSanitizer, shard 1 sits close to an architectural +// allocator ceiling (microsoft/onnxruntime#29139): its cumulative live high-water can exhaust the +// 8 KB size class during a later baseline test (SessionState prepacking), independent of these +// tests. These cases are not the allocation site, but every test compiled into shard 1 adds to that +// cumulative footprint, so they are filed here -- in shard 2, which has headroom -- to keep shard 1 +// at its known-good baseline. They exercise only the public C++ API plus Model::Load/Graph::Resolve, +// so they have no framework-internal dependency and relocate cleanly. The TEST suite name +// (ShapeInferenceV2Test) is preserved so history stays traceable. + +#include +#include +#include + +#include "gtest/gtest.h" +#include "core/graph/model.h" +#include "core/graph/data_propagation/data_propagation_value_utils.h" +#include "core/session/onnxruntime_cxx_api.h" +#include "test/util/include/asserts.h" +#include "test/test_environment.h" + +using namespace ONNX_NAMESPACE; + +extern std::unique_ptr ort_env; + +namespace onnxruntime { +namespace test { + +// Regression test for the Shape -> Gather(1-D index) -> TopK rank-drop. +// +// The Gather index is the 1-D constant [-1] (rank 1), so per ONNX Gather semantics the Gather +// output is a rank-1, single-element tensor -- exactly what TopK's K input requires. Previously +// the Gather custom data propagation scalarized that single element (dropping the rank), +// producing a 0-D K initializer that ONNX TopK shape inference rejected with +// "K input must be a one-dimensional tensor of size 1." at model load time. +// +// The failure reproduces even at ORT_DISABLE_ALL (constant folding never runs there), which +// proves the cause is shape-inference data propagation, not constant folding. This test loads +// the model at the disabled, basic, and all-optimization levels, which bracket the data-propagation +// path: data propagation runs in the pre-optimization Graph::Resolve pass, so it is independent of +// the graph-optimization level (ORT_ENABLE_ALL already applies the extended and layout transformers, +// so they need not be enumerated separately). The test asserts the rank-1 K shape is preserved so the +// model loads and TopK output shapes are inferred correctly. +TEST(ShapeInferenceV2Test, GatherToTopKRankPreservationTest) { + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_topk.onnx"); + + for (auto opt_level : {ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_ALL}) { + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(opt_level); + + // Loading must succeed at each level; before the fix this threw a + // ShapeInferenceError at ORT_DISABLE_ALL (and above). + Ort::Session session(*ort_env, model_path, session_options); + + ORT_ENFORCE(session.GetOutputCount() == 2); + + // K is data-propagated as the last dimension of X (2000), so both TopK outputs are 1-D + // tensors of length 2000. + for (size_t output_index = 0; output_index < 2; ++output_index) { + Ort::TypeInfo type_info = session.GetOutputTypeInfo(output_index); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector output_shape = tensor_info.GetShape(); + EXPECT_TRUE(output_shape.size() == 1) << "TopK output should be a 1-D tensor"; + EXPECT_TRUE(output_shape[0] == 2000) << "TopK output length should be the inferred K value (2000)"; + } + } +} + +// Regression test for the Shape -> Gather(1-D index) -> Mul -> TopK chain. +// +// This covers the elementwise data-propagation consumers (Add/Sub/Mul/Div) with rank-1 +// single-element operands. Both Gather indices are 1-D constants, so each Gather output is a +// rank-1 single-element value; Mul must keep propagating that (as a rank-1 [1] value) so the +// downstream TopK still receives a valid 1-D K. Previously the elementwise ops only handled +// rank-0 scalar operands, so once the Gather producer was fixed to emit rank-1 values the chain +// would have silently stopped propagating at Mul (degrading TopK's K to a symbolic dim). +// +// Asserting that both TopK outputs resolve to the concrete length 2000 (== 50 * 40) proves the +// value propagated end to end through Mul rather than being lost. +TEST(ShapeInferenceV2Test, GatherMulToTopKRankPreservationTest) { + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_mul_topk.onnx"); + + for (auto opt_level : {ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_ALL}) { + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(opt_level); + + // Loading must succeed at each level; before the fix this threw a + // ShapeInferenceError at ORT_DISABLE_ALL (and above). + Ort::Session session(*ort_env, model_path, session_options); + + ORT_ENFORCE(session.GetOutputCount() == 2); + + // K = 50 * 40 = 2000 is data-propagated through the two Gathers and the Mul, so both TopK + // outputs are 1-D tensors of length 2000. + for (size_t output_index = 0; output_index < 2; ++output_index) { + Ort::TypeInfo type_info = session.GetOutputTypeInfo(output_index); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector output_shape = tensor_info.GetShape(); + EXPECT_TRUE(output_shape.size() == 1) << "TopK output should be a 1-D tensor"; + EXPECT_TRUE(output_shape[0] == 2000) << "TopK output length should be the propagated K value (2000)"; + } + } +} + +// Negative regression test for the single-element guard in the elementwise data-propagation +// consumers (Add/Sub/Mul/Div), exercised end to end. +// +// Shape(X[3, 4]) -> S is a rank-1 MULTI-element value [3, 4]. Mul(S, S) -> M must stay the +// multi-element value [9, 16] so ConstantOfShape(M) resolves to the rank-2 shape [9, 16]. For a +// multi-element data-propagation output ONNX's own Mul data propagation supplies the value and +// ORT's custom elementwise propagation is intentionally NOT engaged (it is dispatched only for a +// rank-0 result, see CreateCustomDataPropagation). This is an end-to-end non-regression guard that +// the rank-routing changes did not disturb the multi-element chain; the single-element guard inside +// the shared helper is exercised directly by SinglePropagatedShapeValueGuardTest below. +TEST(ShapeInferenceV2Test, ShapeMulMultiElementNoScalarCollapseTest) { + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_shape_mul_constantofshape.onnx"); + + for (auto opt_level : {ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_ALL}) { + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(opt_level); + + Ort::Session session(*ort_env, model_path, session_options); + + ORT_ENFORCE(session.GetOutputCount() == 1); + + Ort::TypeInfo type_info = session.GetOutputTypeInfo(0); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector output_shape = tensor_info.GetShape(); + + // ConstantOfShape(Mul(Shape(X), Shape(X))) == ConstantOfShape([9, 16]) -> rank-2 [9, 16]. + // A collapse of the multi-element value to a scalar would yield rank 1 here. + EXPECT_TRUE(output_shape.size() == 2) << "Output should stay rank 2; the multi-element value must not collapse"; + EXPECT_TRUE(output_shape.size() == 2 && output_shape[0] == 9 && output_shape[1] == 16) + << "ConstantOfShape output should be the propagated multi-element shape [9, 16]"; + } +} + +// Regression for microsoft/onnxruntime#29072. +// End-to-end lock for BOTH custom data-propagation DECLINE paths -- Gather on a rank-2 index and +// Unsqueeze on a single-element value -- driven through a SINGLE shared model so the coverage costs +// one onnxruntime::Model::Load instead of two (microsoft/onnxruntime#29139: the single-process +// onnxruntime_test_all run sits near the AddressSanitizer size-class ceiling, so the two decline +// fixtures are folded into one graph to minimize the live Model/Graph footprint). +// +// One rank-3 input X drives a shared Shape(X) -> S (the 1-D vector [3, 4, 2000]); two independent +// branches gather from S and route into a rank-lowering Squeeze -> Range so each decline is +// observable end to end: +// +// Branch A (Gather rank-2 index -> output RA): S -> Gather([[ -1 ]]) -> Squeeze -> Range(0, K, 1) +// The index [[-1]] has shape [1, 1] (rank 2), so the Gather output rank is +// data_rank - 1 + index_rank = 1 - 1 + 2 = 2 -- a rank-2 single-element value the single-value +// channel cannot represent, so Gather's custom propagation must DECLINE. With the correct +// decline nothing propagates, Squeeze has no value to lower, and RA's length stays SYMBOLIC. A +// relaxed decline would emit a rank-1 [1] value that Squeeze lowers to the scalar 2000, +// concretizing RA to length 2000 -- so asserting RA stays symbolic discriminates the correct +// decline from the rank-fabricating bug. +// +// Branch B (Unsqueeze single value -> output RB): S -> Gather([-1]) -> Unsqueeze([0]) -> Squeeze -> +// Range(0, K, 1). Gather([-1]) is a rank-1 single-element value (the last dimension of X, 2000); +// unsqueezing it would yield a rank >= 2 result ([1, 2000]) the single-value channel cannot +// represent, so Unsqueeze's custom propagation must DECLINE. With the correct decline nothing +// propagates past Unsqueeze, RB's length stays SYMBOLIC, and the model loads. A relaxed decline +// would fabricate [1, 2000], which Squeeze lowers into a non-scalar Range limit that ONNX Range +// shape inference rejects ("Input to 'Range' op should be scalars"), so Graph::Resolve (and +// therefore Model::Load) FAILS. So Branch B is observable either as a load failure or, were that +// to pass, as a concrete RB length. +// +// Both Gather indices are raw graph initializers (not surfaced through the data-propagation +// getInputData channel), so ONNX's own Gather data propagator bails and control reaches our custom +// decline branches. The branches are independent: a regression in either is caught (Branch A as a +// concrete RA dimension, Branch B as a failed load), so this one combined model keeps the full +// discriminating power of the two original per-path fixtures. +// +// The decline is driven through onnxruntime::Model::Load (which runs Graph::Resolve) and the output +// NodeArg shapes are inspected directly -- no InferenceSession, so none of the session's arena / +// execution-provider / kernel-registry allocations are created. Data propagation runs inside +// Graph::Resolve (Graph::InferAndVerifyTypeMatch -> data-propagation RunInferencing); constant +// folding is a separate session-level optimizer pass that never runs here, so Resolve alone +// exercises the decline branches in isolation with no folding to mask them. +TEST(ShapeInferenceV2Test, GatherUnsqueezeDeclineTest) { + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_decline_combined.onnx"); + + std::shared_ptr model; + // Model::Load runs Graph::Resolve, which performs data propagation. A correct decline on BOTH + // branches keeps each Range length symbolic, so the resolve (and therefore the load) succeeds. A + // relaxed Unsqueeze decline (Branch B) makes Range's input non-scalar, which Range shape inference + // rejects, so Resolve returns non-OK and this assertion fails on the load itself. + ASSERT_STATUS_OK(onnxruntime::Model::Load(model_path, model, nullptr, + DefaultLoggingManager().DefaultLogger())); + + const auto& graph_outputs = model->MainGraph().GetOutputs(); + ASSERT_EQ(graph_outputs.size(), static_cast(2)); + + // Both Range outputs are 1-D; each single dimension must stay SYMBOLIC (no concrete dim_value) + // because its branch declines rather than propagating a (rank-fabricated) concrete K. RA[0] would + // concretize to 2000 if the rank-2 Gather index stopped declining; RB[0] would concretize to 2000 + // (or the load would fail) if the single-value Unsqueeze stopped declining. + for (const NodeArg* output : graph_outputs) { + const TensorShapeProto* output_shape = output->Shape(); + ASSERT_NE(output_shape, nullptr) << "Range output '" << output->Name() + << "' should have an inferred (rank-1) shape"; + ASSERT_EQ(output_shape->dim_size(), 1) << "Range output '" << output->Name() + << "' should be a 1-D tensor"; + EXPECT_FALSE(output_shape->dim(0).has_dim_value()) + << "Range length '" << output->Name() << "' must stay symbolic; the decline must not " + "fabricate a concrete K (a relaxed decline concretizes it to 2000)"; + } +} + +// Lock test for the Shape -> Gather(1-D index) -> Squeeze -> Range chain. +// +// Squeeze of a rank-1 single-element value [1] removes the size-1 dimension and yields a 0-D +// scalar -- the strict-improvement behavior of Squeeze's custom data propagation. This test locks +// that improvement end to end: the Gather produces a rank-1 single element (the last dimension of +// X, 2000); Squeeze must propagate it as the scalar 2000; and Range(0, K, 1) must then resolve to +// a concrete 1-D tensor of length 2000. If Squeeze dropped, corrupted, or failed to scalarize the +// value, Range's length would stay symbolic (size 0 here) instead of the concrete 2000. +TEST(ShapeInferenceV2Test, GatherSqueezeRangeRankPreservationTest) { + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_gather_squeeze_range.onnx"); + + for (auto opt_level : {ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_ALL}) { + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(opt_level); + + Ort::Session session(*ort_env, model_path, session_options); + + ORT_ENFORCE(session.GetOutputCount() == 1); + + Ort::TypeInfo type_info = session.GetOutputTypeInfo(0); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector output_shape = tensor_info.GetShape(); + + // Range(0, Squeeze(Gather(Shape(X), [-1])), 1) == Range(0, 2000, 1) -> 1-D length 2000. + EXPECT_TRUE(output_shape.size() == 1) << "Range output should be a 1-D tensor"; + EXPECT_TRUE(output_shape.size() == 1 && output_shape[0] == 2000) + << "Range length should be the scalar K (2000) propagated through Squeeze"; + } +} + +// Unit test for the single-element guard in TryGetSinglePropagatedShapeValue, the helper shared by +// the Add/Sub/Mul/Div data-propagation consumers. +// +// A NodeArg can carry a propagated shape value in one of two non-interchangeable channels: a rank-0 +// scalar (inferred_scalar_value_) or a rank>=1 list of values (inferred_shape_values_). The helper +// must surface a single value from either a scalar source or a rank-1 SINGLE-element source (and +// report its rank), and must DECLINE for a rank-1 MULTI-element source. Declining is load-bearing: +// silently using element[0] of a multi-element value would let an elementwise op emit a bogus +// scalar product that getInputData() would then hand to a rank-sensitive consumer. This test +// exercises the guard directly -- a relaxed guard that accepted element[0] fails the multi-element +// case below. +TEST(ShapeInferenceV2Test, SinglePropagatedShapeValueGuardTest) { + // Scalar channel -> accepted, reported as rank-0. + { + NodeArg arg("scalar_arg", nullptr); + arg.SetInferredShapeScalarValue(5); + int64_t value = 0; + bool is_rank1 = true; + EXPECT_TRUE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); + EXPECT_EQ(value, 5); + EXPECT_FALSE(is_rank1); + } + + // Rank-1 single-element value -> accepted, reported as rank-1. + { + NodeArg arg("rank1_single_arg", nullptr); + auto& values = arg.GetMutableInferredShapeValues(); + values.emplace(); + values->add_dim()->set_dim_value(7); + int64_t value = 0; + bool is_rank1 = false; + EXPECT_TRUE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); + EXPECT_EQ(value, 7); + EXPECT_TRUE(is_rank1); + } + + // Rank-1 MULTI-element value -> DECLINED (the guard); element[0] must never be used. + { + NodeArg arg("rank1_multi_arg", nullptr); + auto& values = arg.GetMutableInferredShapeValues(); + values.emplace(); + values->add_dim()->set_dim_value(3); + values->add_dim()->set_dim_value(4); + int64_t value = -1; + bool is_rank1 = false; + EXPECT_FALSE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)) + << "A rank-1 multi-element value must be declined, not collapsed to element[0]"; + } + + // Rank-1 single element with a SYMBOLIC (no concrete value) dim -> declined. + { + NodeArg arg("rank1_symbolic_arg", nullptr); + auto& values = arg.GetMutableInferredShapeValues(); + values.emplace(); + values->add_dim()->set_dim_param("N"); + int64_t value = -1; + bool is_rank1 = false; + EXPECT_FALSE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)) + << "A symbolic single-element value has no concrete value and must be declined"; + } + + // No propagated value on either channel -> declined. + { + NodeArg arg("empty_arg", nullptr); + int64_t value = -1; + bool is_rank1 = false; + EXPECT_FALSE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); + } +} + +// Unit test that SetSinglePropagatedShapeValue is correct-by-construction: it populates exactly one +// channel and clears the other. This pins the invariant the elementwise/Gather consumers rely on -- +// the scalar-first reader (TryGetSinglePropagatedShapeValue) and the values-first getInputData() can +// only agree on rank if no NodeArg ever carries both channels. A setter that left the opposite +// channel populated would let a stale value win (reintroducing a rank mismatch), so each case below +// pre-populates the opposite channel and asserts the setter cleared it. +TEST(ShapeInferenceV2Test, SetSinglePropagatedShapeValueKeepsSingleChannelTest) { + // Scalar write must clear a stale values channel. + { + NodeArg arg("scalar_over_stale_values", nullptr); + auto& stale = arg.GetMutableInferredShapeValues(); + stale.emplace(); + stale->add_dim()->set_dim_value(99); + + SetSinglePropagatedShapeValue(arg, 5, /*is_rank1=*/false); + + ASSERT_TRUE(arg.GetInferredShapeScalarValue().has_value()); + EXPECT_EQ(arg.GetInferredShapeScalarValue().value(), 5); + EXPECT_FALSE(arg.GetInferredShapeValues().has_value()) + << "Scalar write must clear the values channel that getInputData() would otherwise prefer"; + + int64_t value = 0; + bool is_rank1 = true; + EXPECT_TRUE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); + EXPECT_EQ(value, 5); + EXPECT_FALSE(is_rank1); + } + + // Rank-1 write must clear a stale scalar channel. + { + NodeArg arg("values_over_stale_scalar", nullptr); + arg.SetInferredShapeScalarValue(99); + + SetSinglePropagatedShapeValue(arg, 7, /*is_rank1=*/true); + + EXPECT_FALSE(arg.GetInferredShapeScalarValue().has_value()) + << "Rank-1 write must clear the scalar channel that the scalar-first reader would otherwise return"; + ASSERT_TRUE(arg.GetInferredShapeValues().has_value()); + ASSERT_EQ(arg.GetInferredShapeValues()->dim_size(), 1); + EXPECT_EQ(arg.GetInferredShapeValues()->dim(0).dim_value(), 7); + + int64_t value = 0; + bool is_rank1 = false; + EXPECT_TRUE(TryGetSinglePropagatedShapeValue(arg, value, is_rank1)); + EXPECT_EQ(value, 7); + EXPECT_TRUE(is_rank1); + } +} + +} // namespace test +} // namespace onnxruntime