diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index 11ca73790ea79..8197d7517189d 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -1753,6 +1753,15 @@ class Graph { // NOLINT(clang-analyzer-optin.performance.Padding): preserve exi std::vector& output_types, const Graph::ResolveOptions& options); + // If ONNX operator's PartialDataPropagationFunction() infers concrete shape values in the output + // save them to the output NodeArg as a TensorShapeProto or a scalar value so that downstream (consumer) nodes + // can use them later for their TypeAndShapeInferenceFunction() and PartialDataPropagationFunction(). + common::Status SaveShapeValuesFromDataPropagation(const Node& node, NodeArg& output_def, + const ONNX_NAMESPACE::TypeProto& propagated_value_as_type_proto) const; + + // Remove intermediate inferred shape values stored in all NodeArgs to reduce memory usage. + common::Status CleanUpShapeValuesFromDataPropagation(); + // Apply type-inference and type-checking to all inputs and initializers: common::Status TypeCheckInputsAndInitializers(); diff --git a/include/onnxruntime/core/graph/node_arg.h b/include/onnxruntime/core/graph/node_arg.h index 0ddf1a2b9d3de..4a18d7617ac13 100644 --- a/include/onnxruntime/core/graph/node_arg.h +++ b/include/onnxruntime/core/graph/node_arg.h @@ -9,6 +9,8 @@ #include "core/common/status.h" #include "core/common/logging/logging.h" +#include + namespace onnxruntime { // Node argument definition, for both input and output, @@ -107,6 +109,18 @@ class NodeArg { /** Gets this NodeArg as a NodeArgInfo, AKA ValueInfoProto. */ const NodeArgInfo& ToProto() const noexcept { return node_arg_info_; } + /** Gets the inferred shape values as a TensorShapeProto. */ + const std::optional& GetInferredShapeValues() const noexcept { return inferred_shape_values_; } + + /** Gets mutable inferred shape values as a TensorShapeProto. */ + std::optional& GetMutableInferredShapeValues() noexcept { return inferred_shape_values_; } + + /** Gets the inferred shape scalar value */ + const std::optional GetInferredShapeScalarValue() const noexcept { return inferred_scalar_value_; } + + /** Sets the inferred shape scalar value */ + void SetInferredShapeScalarValue(int64_t value) noexcept { inferred_scalar_value_ = value; } + /** 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; @@ -128,6 +142,24 @@ class NodeArg { // Node arg name, type and shape. NodeArgInfo node_arg_info_; + // This variable stores the actual tensor data of the shape as a TensorShapeProto after executing + // the ONNX operator's PartialDataPropagationFunction(). It's used for shape inference purpose. + // + // Calling an operator's TypeAndShapeInferenceFunction() alone is sometimes insufficient + // for complete shape inference. For example, the Shape operator's TypeAndShapeInferenceFunction() + // only provides the output's rank which is 1 but not its actual shape values. + // + // The PartialDataPropagationFunction(), defined in the ONNX operator schema, must also + // be executed to obtain the concrete shape values output, allowing accurate propagation + // of shape information throughout the graph. If the concrete shape values output is not + // computed, nothing is stored here that's why this is optional. + std::optional inferred_shape_values_; + + // This variable stores the actual scalar value. + // It is also used for shape inference and data propagation to ensure consistent shape and + // value information throughout the graph. + std::optional inferred_scalar_value_; + // Flag indicates whether <*this> node arg exists or not. bool exists_; }; diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index d3a8856455c49..dce8c7eec4dce 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -1441,6 +1441,12 @@ struct SessionOptionsImpl : ConstSessionOptionsImpl { ///< Wraps OrtApi::SessionOptionsAppendExecutionProvider_VitisAI SessionOptionsImpl& AppendExecutionProvider_VitisAI(const std::unordered_map& provider_options = {}); + + ///< Wraps OrtApi::AddFreeDimensionOverride + SessionOptionsImpl& AddFreeDimensionOverride(const char* dim_denotation, int64_t dim_value); + + ///< Wraps OrtApi::AddFreeDimensionOverrideByName + SessionOptionsImpl& AddFreeDimensionOverrideByName(const char* dim_name, int64_t dim_value); }; } // namespace detail diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h index 8ee057f51eb20..a6d0ff236f2d8 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h @@ -1503,6 +1503,18 @@ inline SessionOptionsImpl& SessionOptionsImpl::RegisterCustomOpsUsingFunct return *this; } +template +inline SessionOptionsImpl& SessionOptionsImpl::AddFreeDimensionOverride(const char* dim_denotation, int64_t dim_value) { + ThrowOnError(GetApi().AddFreeDimensionOverrideByName(this->p_, dim_denotation, dim_value)); + return *this; +} + +template +inline SessionOptionsImpl& SessionOptionsImpl::AddFreeDimensionOverrideByName(const char* dim_name, int64_t dim_value) { + ThrowOnError(GetApi().AddFreeDimensionOverrideByName(this->p_, dim_name, dim_value)); + return *this; +} + /// Session template inline size_t ConstSessionImpl::GetInputCount() const { diff --git a/onnxruntime/core/graph/data_propagation/add_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/add_op_data_propagation.cc new file mode 100644 index 0000000000000..172941c0ee023 --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/add_op_data_propagation.cc @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "add_op_data_propagation.h" +#include "core/common/common.h" +#include "core/graph/node_arg.h" +#include "core/graph/onnx_protobuf.h" +#include "core/providers/common.h" + +namespace onnxruntime { + +Status AddOpDataPropagation::infer() { + // Get "A" input + const auto* input_0 = node_.InputDefs()[0]; + // Get "B" input + const auto* input_1 = node_.InputDefs()[1]; + + // Return and do nothing if input doesn't exist + if (!input_0 || !input_1 || !input_0->Exists() || !input_1->Exists()) { + 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()); + } + + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/add_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/add_op_data_propagation.h new file mode 100644 index 0000000000000..f9eb9990142c1 --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/add_op_data_propagation.h @@ -0,0 +1,52 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "custom_data_propagation.h" +#include "core/graph/graph.h" + +namespace onnxruntime { + +/** + * @brief Class to infer the output scalar for 'Add' operator given the input is a scalar related to shape. + * + * + * For example: + * + * (input with the shape as float32[1, 3, 64, 64]) + * | + * v + * Shape (It saves [1, 3, 64, 64] in inferred_shape_values_ in output's node_arg + * | during Graph::SaveShapeValuesFromDataPropagation()) + * | + * | ______ + * | | + * v v + * Gather Gather (First 'Gather' saves 3 in inferred_scalar_value_ in output node_arg, and + * | | second 'Gather' saves 64 in inferred_scalar_value_ in output node_arg + * | | during GatherOpDataPropagation(), if the 'index' attributes + * | | are 1 and 2 respectively) + * \ / + * \ / + * | | + * v v + * Add (It gets 3 from inferred_scalar_value_ in input A's node_arg and 64 from inferred_scalar_value_ + * | in input B's node_arg, then performs add operation to get 67 and saves in inferred_scalar_value_ + * | in output's node_arg) + * v + * ... + */ +class AddOpDataPropagation : public CustomDataPropagationBase { + public: + AddOpDataPropagation(const Node& node, + NodeArg& output_def, + std::function 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) {} + + Status infer() override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/custom_data_propagation.cc b/onnxruntime/core/graph/data_propagation/custom_data_propagation.cc new file mode 100644 index 0000000000000..b7254aa828107 --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/custom_data_propagation.cc @@ -0,0 +1,53 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "custom_data_propagation.h" +#include "core/common/common.h" +#include "core/graph/graph.h" +#include "core/common/logging/logging.h" +#include "size_op_data_propagation.h" +#include "squeeze_op_data_propagation.h" +#include "unsqueeze_op_data_propagation.h" +#include "gather_op_data_propagation.h" +#include "add_op_data_propagation.h" +#include "sub_op_data_propagation.h" +#include "mul_op_data_propagation.h" +#include "div_op_data_propagation.h" +#include + +namespace onnxruntime { + +std::unique_ptr CreateCustomDataPropagation(const Node& node, + NodeArg& output_def, + std::function func, + const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation, + const logging::Logger& logger) { + int dim_size = 0; + if (output_from_onnx_op_data_propagation.has_tensor_type() && + output_from_onnx_op_data_propagation.tensor_type().has_shape()) { + dim_size = output_from_onnx_op_data_propagation.tensor_type().shape().dim_size(); + } + + if (node.OpType() == "Size") { + return std::make_unique(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger); + } else if (node.OpType() == "Squeeze") { + return std::make_unique(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger); + } else if (node.OpType() == "Unsqueeze") { + return std::make_unique(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger); + } else if (dim_size == 0) { + if (node.OpType() == "Gather") { + return std::make_unique(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger); + } else if (node.OpType() == "Add") { + return std::make_unique(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger); + } else if (node.OpType() == "Sub") { + return std::make_unique(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger); + } else if (node.OpType() == "Mul") { + return std::make_unique(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger); + } else if (node.OpType() == "Div") { + return std::make_unique(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger); + } + } + return nullptr; +} + +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/graph/data_propagation/custom_data_propagation.h b/onnxruntime/core/graph/data_propagation/custom_data_propagation.h new file mode 100644 index 0000000000000..7511f77f58519 --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/custom_data_propagation.h @@ -0,0 +1,75 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/common/common.h" +#include "core/graph/graph.h" +#include "core/common/logging/logging.h" +#include + +namespace onnxruntime { + +/** + * @class CustomDataPropagation + * Custom data propagation for the operator to help enhance shape inference. + * + * Calling infer() can infer the output values for the specific operator given the input is shape values + * and saves the output values in output node_arg for other operators to use later. + * The purpose of this class is to make shape values being correctly inferred and propogated through the graph. + */ +class CustomDataPropagationBase { + public: + ORT_DISALLOW_COPY(CustomDataPropagationBase); + virtual ~CustomDataPropagationBase() = default; + virtual Status infer() = 0; + + protected: + CustomDataPropagationBase(const Node& node, + NodeArg& output_def, + std::function func, + const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation, + const logging::Logger& logger) noexcept + : node_(node), + output_def_(output_def), + get_initialized_input_values_func_(std::move(func)), + output_from_onnx_op_data_propagation_(output_from_onnx_op_data_propagation), + logger_(logger) {} + + const Node& node_; + NodeArg& output_def_; + std::function get_initialized_input_values_func_; + const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation_; + const logging::Logger& logger_; +}; + +/** + * @brief Create custom data propagation for the operator. + * + * For certain operators (e.g., Size, Squeeze, Unsqueeze), ONNX's + * PartialDataPropagationFunction() does not always produce complete or accurate + * inferred shape values. + * + * In particular: + * - Scalar inputs and outputs are not handled correctly. + * - Some operators require additional logic that is not covered by the default function, + e.g. PartialDataPropagationFunction. + * + * Therefore, for these cases, we perform custom data propagation to ensure + * correct and complete inference. + * + * @param node The ORT's node + * @param output_def The node's output NodeArg to save the inferred shape values if needed + * @param func Helper function to get the input value if it's a initializer + * @param output_from_onnx_op_data_propagation The result from executing ONNX operator's data propagation + * @param logger The reference to a logger + * @return std::unique_ptr Returns a CustomDataPropagation object if available + */ +std::unique_ptr CreateCustomDataPropagation( + const Node& node, + NodeArg& output_def, + std::function func, + const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation, + const logging::Logger& logger); + +} // 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 new file mode 100644 index 0000000000000..2ea9b3047941c --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/div_op_data_propagation.cc @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "div_op_data_propagation.h" +#include "core/common/common.h" +#include "core/graph/node_arg.h" +#include "core/graph/onnx_protobuf.h" +#include "core/providers/common.h" + +namespace onnxruntime { + +Status DivOpDataPropagation::infer() { + // Get "A" input + const auto* input_0 = node_.InputDefs()[0]; + // Get "B" input + const auto* input_1 = node_.InputDefs()[1]; + + // Return and do nothing if input doesn't exist + if (!input_0 || !input_1 || !input_0->Exists() || !input_1->Exists()) { + 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()); + } + + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/div_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/div_op_data_propagation.h new file mode 100644 index 0000000000000..9b32b59039282 --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/div_op_data_propagation.h @@ -0,0 +1,52 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "custom_data_propagation.h" +#include "core/graph/graph.h" + +namespace onnxruntime { + +/** + * @brief Class to infer the output scalar for 'Div' operator given the input is a scalar related to shape. + * + * + * For example: + * + * (input with the shape as float32[1, 3, 64, 64]) + * | + * v + * Shape (It saves [1, 3, 64, 64] in inferred_shape_values_ in output's node_arg + * | during graph::SaveShapeValuesFromDataPropagation()) + * | + * | ______ + * | | + * v v + * Gather Gather (First 'Gather' saves 64 in inferred_scalar_value_ in output node_arg, and + * | | second 'Gather' saves 1 in inferred_scalar_value_ in output node_arg + * | | during GatherOpDataPropagation(), if the 'index' attributes + * | | are 2 and 0 respectively) + * \ / + * \ / + * | | + * v v + * Div (It gets 64 from inferred_scalar_value_ in input A's node_arg and 1 from inferred_scalar_value_ + * | in input B's node_arg, then performs div operation to get 64 and saves in inferred_scalar_value_ + * | in output's node_arg) + * v + * ... + */ +class DivOpDataPropagation : public CustomDataPropagationBase { + public: + DivOpDataPropagation(const Node& node, + NodeArg& output_def, + std::function 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) {} + + Status infer() override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.cc new file mode 100644 index 0000000000000..39ac926a8553f --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.cc @@ -0,0 +1,85 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gather_op_data_propagation.h" +#include "core/common/common.h" +#include "core/graph/node_arg.h" +#include "core/graph/onnx_protobuf.h" +#include "core/providers/common.h" + +namespace onnxruntime { + +Status GatherOpDataPropagation::infer() { + if (output_from_onnx_op_data_propagation_.has_tensor_type() && + output_from_onnx_op_data_propagation_.tensor_type().has_shape()) { + int dim_size = output_from_onnx_op_data_propagation_.tensor_type().shape().dim_size(); + // Check there is no result from Gather's PartialDataPropagationFunction(), + // so that it can run custom data propagation below. + // Otherwise, this infer() function won't be called as the result from Gather's PartialDataPropagationFunction() + // will be used in Graph::SaveShapeValuesFromDataPropagation(). + if (dim_size != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "ORT shouldn't run Gather's custom data propagation here as Gather's" + "PartialDataPropagationFunction() already infers shape values in the output."); + } + } + + // Following code extracts an element from a 1D array if all conditions are met. + // e.g. + // shape data is [1, 3, 64, 64] -> gets 64 if the index is 2. + // shape data is [1, 3, 64, 64] -> gets 3 if the index is 1. + + // Get "data" input + // Note: The "data" input should be an one dimension array in this case. + const auto* input_0 = node_.InputDefs()[0]; + + // Get "indices" input + // Note: The "indices" input could be one of the three cases: + // 1. A tensor with rank > 0 and all tensor values are known. + // 2. A tensor with rank > 0 but not all tensor values are known. + // 3. A scalar. + // + // If it's case #1, ONNX operator's PartialDataPropagationFunction() + // should have inferred the output shape value. + // If it's case #2, neither ONNX operator's PartialDataPropagationFunction() + // nor Gather's custom data propagation can handle it. + // This Gather's custom data propagation handles case #3. + const auto* input_1 = node_.InputDefs()[1]; + + // Return and do nothing if input doesn't exist + if (!input_0 || !input_1 || !input_0->Exists() || !input_1->Exists()) { + return Status::OK(); + } + + // If input's inferred shape values is present, we then perfrom the gather operation on the shape values + // and saves the result in output's node_arg. + if (input_0->GetInferredShapeValues().has_value()) { + const auto& tensor_shape_proto = input_0->GetInferredShapeValues().value(); + + ORT_TRY { + TensorShapeVector indices; + ORT_RETURN_IF_ERROR(get_initialized_input_values_func_(input_1->Name(), indices)); + 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( + HandleNegativeAxis(indices[0], tensor_shape_proto.dim_size())); + + auto& dim = tensor_shape_proto.dim(index); + if (dim.has_dim_value()) { + output_def_.SetInferredShapeScalarValue(dim.dim_value()); + } + } + } + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + LOGS(logger_, ERROR) << ex.what(); + LOGS(logger_, INFO) << "Skip Gather op custom data propagation."; + }); + return Status::OK(); + } + } + + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.h new file mode 100644 index 0000000000000..c6b542e5af6e3 --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/gather_op_data_propagation.h @@ -0,0 +1,51 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "custom_data_propagation.h" +#include "core/graph/graph.h" + +namespace onnxruntime { + +/** + * @brief Class to infer the output scalar for 'Gather' operator given the input is shape values. + * + * + * For example: + * + * (input with the shape as float32[1, 3, 64, 64]) + * | + * v + * Shape (It saves [1, 3, 64, 64] in inferred_shape_values_ in output's node_arg + * | during graph::SaveShapeValuesFromDataPropagation()) + * | + * | ______ + * | | + * v v + * Gather Gather (First 'Gather' gets [1, 3, 64, 64] from input node_node's inferred_shape_values_, and + * | | then saves 3 in inferred_scalar_value_ in output node_args if 'index' attribute is 1. + * | | Same logic for second 'Gather', it saves 64 in inferred_scalar_value_ in output node_arga + * \ / if 'index" attribute is 2) + * \ / + * | | + * v v + * Mul (It gets 3 from inferred_scalar_value_ in input A's node_arg and 64 from inferred_scalar_value_ + * | in input B's node_arg, then performs mul operation to get 192 and saves in inferred_scalar_value_ + * | in output's node_arg) + * v + * ... + */ +class GatherOpDataPropagation : public CustomDataPropagationBase { + public: + GatherOpDataPropagation(const Node& node, + NodeArg& output_def, + std::function 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) {} + + Status infer() override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/mul_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/mul_op_data_propagation.cc new file mode 100644 index 0000000000000..4c5c25022e40b --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/mul_op_data_propagation.cc @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "mul_op_data_propagation.h" +#include "core/common/common.h" +#include "core/graph/node_arg.h" +#include "core/graph/onnx_protobuf.h" +#include "core/providers/common.h" + +namespace onnxruntime { + +Status MulOpDataPropagation::infer() { + // Get "A" input + const auto* input_0 = node_.InputDefs()[0]; + // Get "B" input + const auto* input_1 = node_.InputDefs()[1]; + + // Return and do nothing if input doesn't exist + if (!input_0 || !input_1 || !input_0->Exists() || !input_1->Exists()) { + 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()); + } + + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/mul_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/mul_op_data_propagation.h new file mode 100644 index 0000000000000..b42a591eb4c7b --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/mul_op_data_propagation.h @@ -0,0 +1,51 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "custom_data_propagation.h" +#include "core/graph/graph.h" +namespace onnxruntime { + +/** + * @brief Class to infer the output scalar for 'Mul' operator given the input is a scalar related to shape. + * + * + * For example: + * + * (input with the shape as float32[1, 3, 64, 64]) + * | + * v + * Shape (It saves [1, 3, 64, 64] in inferred_shape_values_ in output's node_arg + * | during graph::SaveShapeValuesFromDataPropagation()) + * | + * | ______ + * | | + * v v + * Gather Gather (First 'Gather' saves 3 in inferred_scalar_value_ in output node_arg, and + * | | second 'Gather' saves 64 in inferred_scalar_value_ in output node_arg + * | | during GatherOpDataPropagation(), if the 'index' attributes + * | | are 1 and 2 respectively) + * \ / + * \ / + * | | + * v v + * Mul (It gets 3 from inferred_scalar_value_ in input A's node_arg and 64 from inferred_scalar_value_ + * | in input B's node_arg, then performs mul operation to get 192 and saves in inferred_scalar_value_ + * | in output's node_arg) + * v + * ... + */ +class MulOpDataPropagation : public CustomDataPropagationBase { + public: + MulOpDataPropagation(const Node& node, + NodeArg& output_def, + std::function 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) {} + + Status infer() override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/size_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/size_op_data_propagation.cc new file mode 100644 index 0000000000000..fe8ff3864296d --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/size_op_data_propagation.cc @@ -0,0 +1,41 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "size_op_data_propagation.h" +#include "core/common/common.h" +#include "core/graph/node_arg.h" +#include "core/graph/onnx_protobuf.h" + +namespace onnxruntime { + +Status SizeOpDataPropagation::infer() { + // Size operator generates a scalar output + const auto* input_0 = node_.InputDefs()[0]; + + // Return and do nothing if input doesn't exist + if (!input_0 || !input_0->Exists()) { + return Status::OK(); + } + + if (input_0->GetInferredShapeValues().has_value()) { + const auto& tensor_shape_proto = input_0->GetInferredShapeValues().value(); + + int64_t num_elements = 1; + // The TensorShapeProto (inferred shape values) should have rank > 0 and + // all the dimensions have values (not symbolic) + if (tensor_shape_proto.dim_size() > 0) { + for (const auto& dim : tensor_shape_proto.dim()) { + if (!dim.has_dim_value()) { + return Status::OK(); // Or handle the error appropriately + } + num_elements *= dim.dim_value(); + } + + output_def_.SetInferredShapeScalarValue(num_elements); + } + } + + return Status::OK(); +} + +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/graph/data_propagation/size_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/size_op_data_propagation.h new file mode 100644 index 0000000000000..184202254f078 --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/size_op_data_propagation.h @@ -0,0 +1,29 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "custom_data_propagation.h" +#include "core/graph/graph.h" + +namespace onnxruntime { + +/** + * @brief Class to infer the output scalar for 'Size' operator given the input is shape values. + * + * 'Size' operator takes a tensor as input and outputs a int64 scalar that equals to the total + * number of elements of the input tensor. + */ +class SizeOpDataPropagation : public CustomDataPropagationBase { + public: + SizeOpDataPropagation(const Node& node, + NodeArg& output_def, + std::function 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) {} + + Status infer() override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.cc new file mode 100644 index 0000000000000..b33d4f5589016 --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.cc @@ -0,0 +1,108 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "squeeze_op_data_propagation.h" +#include "core/common/common.h" +#include "core/graph/node_arg.h" +#include "core/graph/onnx_protobuf.h" +#include "core/providers/common.h" +#include "core/common/inlined_containers.h" + +namespace onnxruntime { + +Status SqueezeOpDataPropagation::infer() { + const auto* input_0 = node_.InputDefs()[0]; + + // Return and do nothing if input doesn't exist + if (!input_0 || !input_0->Exists()) { + return Status::OK(); + } + + if (input_0->GetInferredShapeValues().has_value()) { + const auto& tensor_shape_proto = input_0->GetInferredShapeValues().value(); + + // The TensorShapeProto (inferred shape values) should have rank > 0 and + // all the dimensions have values (not symbolic) + if (tensor_shape_proto.dim_size() > 0) { + for (const auto& dim : tensor_shape_proto.dim()) { + if (!dim.has_dim_value()) { + return Status::OK(); + } + } + } + + if (tensor_shape_proto.dim_size() == 1) { + output_def_.SetInferredShapeScalarValue(tensor_shape_proto.dim(0).dim_value()); + } else if (tensor_shape_proto.dim_size() > 1) { + // Get axes value + TensorShapeVector axes; + InlinedHashSet axes_set; + + // Note: Starting from opset 13, "axes" is provided as a second input to the Squeeze operator. + // In opset 11 and earlier, "axes" is defined as a node attribute instead. + 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)); + } + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + LOGS(logger_, ERROR) << ex.what(); + LOGS(logger_, INFO) << "Skip Squeeze op custom data propagation."; + }); + return Status::OK(); + } + } else { + const auto& attrs = node_.GetAttributes(); + auto it = attrs.find("axes"); + if (it != attrs.end()) { + const auto& axes_attr = it->second; + for (const auto& i : axes_attr.ints()) { + axes.push_back(i); + } + } + } + + ORT_TRY { + for (size_t i = 0; i < axes.size(); ++i) { + // Negative value means counting dimensions from the back. + axes_set.insert(HandleNegativeAxis(axes[i], tensor_shape_proto.dim_size())); + } + } + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + LOGS(logger_, ERROR) << ex.what(); + LOGS(logger_, INFO) << "Skip Squeeze op custom data propagation."; + }); + return Status::OK(); + } + + auto& inferred_shape_values = output_def_.GetMutableInferredShapeValues(); + + if (!inferred_shape_values.has_value()) { + inferred_shape_values.emplace(); + } + inferred_shape_values->clear_dim(); + + int64_t dim_index = 0; + for (const auto& dim : tensor_shape_proto.dim()) { + auto value = dim.dim_value(); + if (axes_set.size() > 0) { + if (axes_set.find(dim_index) == axes_set.end()) { + inferred_shape_values->add_dim()->set_dim_value(value); + } + } else { + if (value != 1) { + inferred_shape_values->add_dim()->set_dim_value(value); + } + } + + dim_index++; + } + } + } + + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.h new file mode 100644 index 0000000000000..15e1e8458525a --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.h @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "custom_data_propagation.h" +#include "core/graph/graph.h" + +namespace onnxruntime { + +/** + * @brief Class to infer the output values/scalar for 'Squeeze' operator given the input is shape values. + * + */ +class SqueezeOpDataPropagation : public CustomDataPropagationBase { + public: + SqueezeOpDataPropagation(const Node& node, + NodeArg& output_def, + std::function 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) {} + + Status infer() override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/sub_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/sub_op_data_propagation.cc new file mode 100644 index 0000000000000..4ee8ab546e707 --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/sub_op_data_propagation.cc @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "sub_op_data_propagation.h" +#include "core/common/common.h" +#include "core/graph/node_arg.h" +#include "core/graph/onnx_protobuf.h" +#include "core/providers/common.h" + +namespace onnxruntime { + +Status SubOpDataPropagation::infer() { + // Get "A" input + const auto* input_0 = node_.InputDefs()[0]; + // Get "B" input + const auto* input_1 = node_.InputDefs()[1]; + + // Return and do nothing if input doesn't exist + if (!input_0 || !input_1 || !input_0->Exists() || !input_1->Exists()) { + 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()); + } + + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/sub_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/sub_op_data_propagation.h new file mode 100644 index 0000000000000..a9be294b8f62f --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/sub_op_data_propagation.h @@ -0,0 +1,52 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "custom_data_propagation.h" +#include "core/graph/graph.h" + +namespace onnxruntime { + +/** + * @brief Class to infer the output scalar for 'Sub' operator given the input is a scalar related to shape. + * + * + * For example: + * + * (input with the shape as float32[1, 3, 64, 64]) + * | + * v + * Shape (It saves [1, 3, 64, 64] in inferred_shape_values_ in output's node_arg + * | during graph::SaveShapeValuesFromDataPropagation()) + * | + * | ______ + * | | + * v v + * Gather Gather (First 'Gather' saves 64 in inferred_scalar_value_ in output node_arg, and + * | | second 'Gather' saves 3 in inferred_scalar_value_ in output node_arg + * | | during GatherOpDataPropagation(), if the 'index' attributes + * | | are 2 and 1 respectively) + * \ / + * \ / + * | | + * v v + * Sub (It gets 64 from inferred_scalar_value_ in input A's node_arg and 3 from inferred_scalar_value_ + * | in input B's node_arg, then performs sub operation to get 61 and saves in inferred_scalar_value_ + * | in output's node_arg) + * v + * ... + */ +class SubOpDataPropagation : public CustomDataPropagationBase { + public: + SubOpDataPropagation(const Node& node, + NodeArg& output_def, + std::function 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) {} + + Status infer() override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.cc b/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.cc new file mode 100644 index 0000000000000..ff5ef6853bc78 --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.cc @@ -0,0 +1,119 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "unsqueeze_op_data_propagation.h" +#include "core/common/common.h" +#include "core/graph/node_arg.h" +#include "core/graph/onnx_protobuf.h" +#include "core/providers/common.h" +#include "core/common/inlined_containers.h" + +namespace onnxruntime { + +Status UnsqueezeOpDataPropagation::infer() { + const auto* input_0 = node_.InputDefs()[0]; + + // Return and do nothing if input doesn't exist + if (!input_0 || !input_0->Exists()) { + return Status::OK(); + } + + auto dim_size = output_from_onnx_op_data_propagation_.tensor_type().shape().dim_size(); + + if (dim_size == 0 && input_0->GetInferredShapeScalarValue().has_value()) { + // Following code expands a scalr to one dimension array, e.g. shape data is 64 -> it becomes [64] + // In this case, the axis should be 0 + auto& inferred_shape_values = output_def_.GetMutableInferredShapeValues(); + + if (!inferred_shape_values.has_value()) { + inferred_shape_values.emplace(); + } + inferred_shape_values->clear_dim(); + + inferred_shape_values->add_dim()->set_dim_value(input_0->GetInferredShapeScalarValue().value()); + } else if (input_0->GetInferredShapeValues().has_value()) { + const auto& tensor_shape_proto = input_0->GetInferredShapeValues().value(); + + // The TensorShapeProto (inferred shape values) should have rank > 0 and + // all the dimensions have values (not symbolic) + if (tensor_shape_proto.dim_size() > 0) { + for (const auto& dim : tensor_shape_proto.dim()) { + if (!dim.has_dim_value()) { + return Status::OK(); + } + } + } + + if (tensor_shape_proto.dim_size() > 0) { + // Get axes value + TensorShapeVector axes; + InlinedHashSet axes_set; + + // Note: Starting from opset 13, "axes" is provided as a second input to the Squeeze operator. + // In opset 11 and earlier, "axes" is defined as a node attribute instead. + 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)); + } + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + LOGS(logger_, ERROR) << ex.what(); + LOGS(logger_, INFO) << "Skip Unsqueeze op custom data propagation."; + }); + return Status::OK(); + } + } else { + const auto& attrs = node_.GetAttributes(); + auto it = attrs.find("axes"); + if (it != attrs.end()) { + const auto& axes_attr = it->second; + for (const auto& i : axes_attr.ints()) { + axes.push_back(i); + } + } + } + + // axes is required, if not provided just do nothing and return. + if (axes.empty()) { + return Status::OK(); + } + ORT_TRY { + for (size_t i = 0; i < axes.size(); ++i) { + // Negative value means counting dimensions from the back. + axes_set.insert(HandleNegativeAxis(axes[i], tensor_shape_proto.dim_size())); + } + } + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + LOGS(logger_, ERROR) << ex.what(); + LOGS(logger_, INFO) << "Skip Unsqueeze op custom data propagation."; + }); + return Status::OK(); + } + + auto& inferred_shape_values = output_def_.GetMutableInferredShapeValues(); + + if (!inferred_shape_values.has_value()) { + inferred_shape_values.emplace(); + } + inferred_shape_values->clear_dim(); + + int64_t axis = 0; + for (const auto& dim : tensor_shape_proto.dim()) { + if (axes_set.find(axis) != axes_set.end()) { + inferred_shape_values->add_dim()->set_dim_value(1); + } + + auto value = dim.dim_value(); + inferred_shape_values->add_dim()->set_dim_value(value); + + axis += 1; + } + } + } + + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.h b/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.h new file mode 100644 index 0000000000000..26b6587aa93fc --- /dev/null +++ b/onnxruntime/core/graph/data_propagation/unsqueeze_op_data_propagation.h @@ -0,0 +1,26 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "custom_data_propagation.h" +#include "core/graph/graph.h" + +namespace onnxruntime { +/** + * @brief Class to infer the output values/scalar for 'Unsqueeze' operator given the input is shape values. + * + */ +class UnsqueezeOpDataPropagation : public CustomDataPropagationBase { + public: + UnsqueezeOpDataPropagation(const Node& node, + NodeArg& output_def, + std::function 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) {} + + Status infer() override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index 3d67314cf693a..af62077fd790b 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -16,13 +16,13 @@ #include "core/common/inlined_containers.h" #include "core/common/logging/logging.h" #include "core/common/narrow.h" +#include "core/providers/common.h" #include "core/flatbuffers/flatbuffers_utils.h" #include "core/framework/error_code_helper.h" #include "core/framework/tensor_type_and_shape.h" #include "core/flatbuffers/schema/ort.fbs.h" #include "core/framework/tensor_external_data_info.h" #include "core/framework/tensor_shape.h" -#include "core/framework/tensor_type_and_shape.h" #include "core/framework/tensorprotoutils.h" #include "core/framework/utils.h" #include "core/graph/function_utils.h" @@ -37,6 +37,7 @@ #include "core/graph/node_attr_utils.h" #include "core/graph/op.h" #include "core/graph/runtime_optimization_record_container.h" +#include "data_propagation/custom_data_propagation.h" #if !defined(ORT_MINIMAL_BUILD) #include "core/graph/function.h" @@ -44,6 +45,7 @@ #include "core/graph/schema_registry.h" #include "onnx/checker.h" #include "onnx/defs/parser.h" +#include "onnx/defs/tensor_proto_util.h" using namespace ONNX_NAMESPACE::checker; #endif @@ -2717,8 +2719,8 @@ class InferenceContextImpl : public ONNX_NAMESPACE::InferenceContext { if (!def) return nullptr; - // only return data if it's for a constant initializer. checks for outer scope initializers - // if this is a subgraph and the name isn't found locally. + // Returns if it's a constant initializer. + // Checks for outer scope initializers if this is a subgraph and the name isn't found locally. const TensorProto* initializer = graph_.GetConstantInitializer(def->Name(), true); if (initializer != nullptr) { // Check if this is in-memory external data (data stored in OrtValue) @@ -2740,8 +2742,60 @@ class InferenceContextImpl : public ONNX_NAMESPACE::InferenceContext { " has in-memory external data but cannot get OrtValue during shape inference"); } } + + return initializer; + } + + // The following code handles cases where a node stores the previously inferred output shape values in its NodeArg. + // + // For example, the Reshape operator, its shape input may come from a producer node such as a Shape operator, + // and the inferred output shape value is already stored as a TensorShapeProto in corresponding NodeArg. + // + // In such cases, the Reshape operator should convert this TensorShapeProto into a TensorProto. + // The resulting TensorProto will then be treated as an initializer during ONNX shape inference, + // allowing the real dimension values to be correctly used. + + const auto& inferred_shape_values = def->GetInferredShapeValues(); + + // Converts the inferred shape values if any to a TensorProto and returns the TensorProto. + if (inferred_shape_values.has_value() && inferred_shape_values->dim_size() > 0) { + TensorProto tensor_proto; + tensor_proto.set_data_type(TensorProto_DataType_INT64); + tensor_proto.add_dims(inferred_shape_values->dim_size()); + bool all_values = true; + for (const auto& dim : inferred_shape_values->dim()) { + if (dim.has_dim_value()) { + tensor_proto.add_int64_data(dim.dim_value()); + } else { + all_values = false; + break; + } + } + + if (all_values) { + temp_tensor_protos_.push_back(std::make_unique(std::move(tensor_proto))); + return temp_tensor_protos_.back().get(); + } } - return initializer; + + const std::optional inferred_shape_scalar_value = def->GetInferredShapeScalarValue(); + + // Converts the inferred shape scalar value if any to a TensorProto and returns the TensorProto. + // + // Note: ONNX's getShapeInput() internally calls getInputData() to retrieve a TensorProto (if available) + // and then extracts shape/dimension values from it. As a result, the scalar value may not be + // properly handled and propagated in ONNX's shape inference. + // However, Graph::SaveShapeValuesFromDataPropagation() properly handles data propagation for + // some operators. + if (inferred_shape_scalar_value.has_value()) { + TensorProto tensor_proto; + tensor_proto.set_data_type(TensorProto_DataType_INT64); + tensor_proto.add_int64_data(inferred_shape_scalar_value.value()); + temp_tensor_protos_.push_back(std::make_unique(std::move(tensor_proto))); + return temp_tensor_protos_.back().get(); + } + + return nullptr; } // ORT does not implement partial data propagation yet so just return nullptr. @@ -2784,9 +2838,223 @@ class InferenceContextImpl : public ONNX_NAMESPACE::InferenceContext { // These need to outlive the shape inference call, so we store them here // Inference is per node and the instance of this context is on the stack, // so this is safe. + // It can also be used to temporarily save the inferred shape values as a TensorProto. mutable InlinedVector> temp_tensor_protos_; }; +// An implementation of the DataPropagationContext interface optional by operator-specific +// shape inference for onnxruntime graphs. +// Please see the description and usage of ONNX's data propagation here: +// https://github.com/onnx/onnx/blob/main/onnx/defs/shape_inference.h#L117-L127 +class DataPropagationContextImpl : public ONNX_NAMESPACE::DataPropagationContext { + public: + DataPropagationContextImpl(Node& node) noexcept : node_(node) { + node_output_types_.resize(node.OutputDefs().size()); + } + + const AttributeProto* getAttribute(const std::string& name) const override { + auto& attribute_value_map = node_.GetAttributes(); + auto iter = attribute_value_map.find(name); + if (iter == attribute_value_map.end()) { + return nullptr; + } + return &iter->second; + } + + size_t getNumInputs() const noexcept override { + return node_.InputDefs().size(); + } + + const TypeProto* getInputType(size_t index) const override { + if (index >= getNumInputs()) { + return nullptr; + } + + const TypeProto* type = nullptr; + auto p_node_arg = node_.InputDefs().at(index); + if ((nullptr != p_node_arg) && p_node_arg->Exists()) { + type = p_node_arg->TypeAsProto(); + } + + return type; + } + + size_t getNumOutputs() const noexcept override { + return node_output_types_.size(); + } + + const TypeProto* getOutputType(size_t index) const override { + if (index >= getNumOutputs()) { + return nullptr; + } + + return &node_output_types_[index]; + } + + const TensorShapeProto* getInputData(size_t index) override { + if (index >= getNumInputs()) { + return nullptr; + } + + auto def = node_.InputDefs()[index]; + if (!def) + return nullptr; + + // Get the previously inferred shape values that stored in NodeArg's inferred_shape_values_ if any. + // Note: getInputData() only supports input data (shape values) that is a tensor not a scalar, + // becase the returning TensorShapeProto can't store scalar value. Therefore, op's data propagation + // defined in ONNX Op schema does not support scalar output. + // However, Graph::SaveShapeValuesFromDataPropagation() does support output scalar value for + // some operators. + const auto& tensor_shape_proto = def->GetInferredShapeValues(); + if (tensor_shape_proto.has_value()) { + return &*tensor_shape_proto; + } + + return nullptr; + } + + void addOutputData(size_t index, TensorShapeProto&& tsp) override { + if (index >= node_output_types_.size()) return; + + TypeProto& type_proto = node_output_types_[index]; + *type_proto.mutable_tensor_type()->mutable_shape() = std::move(tsp); + } + + void RunInferencing() { + auto* schema = node_.Op(); + if (nullptr != schema) { + schema->GetDataPropagationFunction()(*this); + } + } + + const std::vector& InferredOutputTypes() const { return node_output_types_; } + + private: + Node& node_; + std::vector node_output_types_; +}; + +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) + -> Status { + const TensorProto* initializer = this->GetConstantInitializer(input_name, true); + + if (initializer) { + // Get shape from TensorProto as well as element counts. + // If shape has dimension size equals zero, it means it's a scalar and has only one element. + auto tensor_shape = utils::GetTensorShapeFromTensorProto(*initializer); + size_t element_cnt = narrow(tensor_shape.Size()); + + // Check if this is in-memory external data (data stored in OrtValue) + if (utils::HasExternalDataInMemory(*initializer)) { + // Try to get the OrtValue for this initializer + OrtValue ort_value; + if (this->GetOrtValueInitializer(input_name, ort_value, true)) { + const Tensor& tensor = ort_value.Get(); + if (initializer->data_type() == TensorProto_DataType_INT32) { + auto data_span = tensor.DataAsSpan(); + ORT_ENFORCE(data_span.size() == element_cnt, + "The element counts from Tensor should be the same" + "from using utils::GetTensorShapeFromTensorProto()"); + + size_t index = 0; + input_values.resize(element_cnt); + for (const auto& v : data_span) { + input_values[index] = static_cast(v); + ++index; + } + } else if (initializer->data_type() == TensorProto_DataType_INT64) { + const int64_t* src = tensor.Data(); + memcpy(input_values.data(), src, element_cnt * sizeof(int64_t)); + } + } else { + // If we can't get the OrtValue, it is a bug + ORT_THROW("Initializer ", input_name, + " has in-memory external data but cannot get OrtValue during shape inference"); + } + } + // Unpack tensor from raw data, external data (not in memory) or the type specific data field + else { + if (initializer->data_type() == TensorProto_DataType_INT32) { + InlinedVector tmp_values; + tmp_values.resize(element_cnt); + ORT_RETURN_IF_ERROR(utils::UnpackTensor(*initializer, + this->ModelPath(), + tmp_values.data(), + element_cnt)); + + input_values.resize(element_cnt); + for (size_t i = 0; i < element_cnt; ++i) { + input_values[i] = static_cast(tmp_values[i]); // copy values + } + } else if (initializer->data_type() == TensorProto_DataType_INT64) { + input_values.resize(element_cnt); + ORT_RETURN_IF_ERROR(utils::UnpackTensor(*initializer, + this->ModelPath(), + input_values.data(), + element_cnt)); + } + } + } + + return Status::OK(); + }; + + // For certain operators (e.g., Size, Squeeze, Unsqueeze), ONNX's + // PartialDataPropagationFunction() does not always produce complete or accurate + // inferred shape values. + // + // In particular: + // - Scalar inputs and outputs are not handled correctly. + // - Some operators require additional logic that is not covered by the default function. + // + // Therefore, for these cases, we perform custom data propagation to ensure + // correct and complete inference. + auto dp = CreateCustomDataPropagation(node, output_def, + get_initialized_input_values_func, + onnx_inferred_type_after_data_propagation, + logger_); + if (dp) { + ORT_RETURN_IF_ERROR(dp->infer()); + return Status::OK(); + } + + // If no custom data propagation is defined for the operator, + // fall back to using the result of ONNX's PartialDataPropagationFunction(), if available. + + int dim_size = 0; + if (onnx_inferred_type_after_data_propagation.has_tensor_type() && + onnx_inferred_type_after_data_propagation.tensor_type().has_shape()) { + dim_size = onnx_inferred_type_after_data_propagation.tensor_type().shape().dim_size(); + } + + if (dim_size > 0) { + // Only handle that the inferred shape values (from ONNX operator's PartialDataPropagationFunction() ) has rank > 0 + // and all dimensions have concrete (non-symbolic) values. + for (int i = 0; i < dim_size; ++i) { + if (!onnx_inferred_type_after_data_propagation.tensor_type().shape().dim(i).has_dim_value()) { + return Status::OK(); + } + } + + if (!output_def.inferred_shape_values_.has_value()) { + output_def.inferred_shape_values_.emplace(); + } + + output_def.inferred_shape_values_->clear_dim(); + for (int i = 0; i < dim_size; ++i) { + auto value = onnx_inferred_type_after_data_propagation.tensor_type().shape().dim(i).dim_value(); + output_def.inferred_shape_values_->add_dim()->set_dim_value(value); + } + } + + return Status::OK(); +} + Status Graph::InferAndVerifySubgraphTypes(const Node& node, Graph& subgraph, const std::vector& input_types, std::vector& output_types, @@ -2978,11 +3246,20 @@ Status Graph::InferAndVerifyTypeMatch(Node& node, const OpSchema& op, const Reso // returned here. SubgraphInferencingFunc func(Graph::InferAndVerifySubgraphTypes); InferenceContextImpl context(node, func, *this, options); + DataPropagationContextImpl data_propagation_context(node); { auto status = Status::OK(); ORT_TRY { context.RunInferencing(); + + // Calling an operator's TypeAndShapeInferenceFunction() alone is sometimes insufficient + // for complete shape inference. For example, the Shape operator only provides the + // output's rank (1-dimensional) but not its actual dimension values. + // The PartialDataPropagationFunction(), defined in the ONNX operator schema, must also + // be executed to obtain the concrete output shape values, allowing accurate propagation + // of shape information throughout the graph. + data_propagation_context.RunInferencing(); } ORT_CATCH(const std::exception& ex) { ORT_HANDLE_EXCEPTION([&]() { @@ -2994,6 +3271,8 @@ Status Graph::InferAndVerifyTypeMatch(Node& node, const OpSchema& op, const Reso const auto& onnx_inferred_types(context.InferredOutputTypes()); + const auto& onnx_inferred_types_after_data_propagation(data_propagation_context.InferredOutputTypes()); + // Infer and verify node output arg type information. int i = -1; for (auto& output_def : node.MutableDefinitions().output_defs) { @@ -3010,6 +3289,12 @@ Status Graph::InferAndVerifyTypeMatch(Node& node, const OpSchema& op, const Reso auto op_formal_parameter = op.outputs().at(operand_index); const TypeProto& onnx_inferred_type = onnx_inferred_types[i]; + const TypeProto& onnx_inferred_type_after_data_propagation = onnx_inferred_types_after_data_propagation[i]; + + ORT_RETURN_IF_ERROR(SaveShapeValuesFromDataPropagation(node, + *output_def, + onnx_inferred_type_after_data_propagation)); + DataType existing_type = output_def->Type(); DataType inferred_type = nullptr; @@ -3322,6 +3607,26 @@ Status Graph::VerifyNodeAndOpMatch(const ResolveOptions& options) { } } + ORT_RETURN_IF_ERROR(CleanUpShapeValuesFromDataPropagation()); + + return Status::OK(); +} + +Status Graph::CleanUpShapeValuesFromDataPropagation() { + for (auto node_index : nodes_in_topological_order_) { + auto& node = *GetNode(node_index); + + for (auto node_arg : node.MutableInputDefs()) { + node_arg->inferred_shape_values_.reset(); + node_arg->inferred_scalar_value_.reset(); + } + + for (auto node_arg : node.MutableOutputDefs()) { + node_arg->inferred_shape_values_.reset(); + node_arg->inferred_scalar_value_.reset(); + } + } + return Status::OK(); } diff --git a/onnxruntime/test/framework/shape_inference_test.cc b/onnxruntime/test/framework/shape_inference_test.cc index 2d5c3a43ee8ed..3554b6eb5d21e 100644 --- a/onnxruntime/test/framework/shape_inference_test.cc +++ b/onnxruntime/test/framework/shape_inference_test.cc @@ -16,6 +16,8 @@ using namespace ONNX_NAMESPACE; +extern std::unique_ptr ort_env; + namespace onnxruntime { namespace test { @@ -76,6 +78,103 @@ TEST_F(ShapeInferenceTest, BasicTest) { CheckShapeEquality(InputShape(node), OutputShape(node)); } +TEST(ShapeInferenceV2Test, PartialDataPropagationTest) { + { + // Model #1 + // This model contains "Shape" and "Reshape" operators. + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_with_shape_related_nodes.onnx"); + + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); + session_options.AddFreeDimensionOverrideByName("batch", 1); + session_options.AddFreeDimensionOverrideByName("width", 64); + session_options.AddFreeDimensionOverrideByName("height", 64); + + // Even though all graph optimizations are disabled, the free dimension override is still enabled by default. + // The shape of graph's output should be correctly inferred by shape inference and data propagation. + Ort::Session session(*ort_env, model_path, session_options); + + // This graph only has one output + 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(); + EXPECT_TRUE(output_shape.size() == 4) << "The output shape should have 4 dimensions"; + EXPECT_TRUE(output_shape[0] == 1) << "The first dimension should have 1 as value"; + EXPECT_TRUE(output_shape[1] == 3) << "The second dimension should have 3 as value"; + EXPECT_TRUE(output_shape[2] == 64) << "The second dimension should have 64 as value"; + EXPECT_TRUE(output_shape[3] == 64) << "The second dimension should have 64 as value"; + } + + { + // Model #2 + // This model contains "Shape", "Reshape", "Gather" and "Unsqueeze" operators. + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_with_shape_related_nodes_v2.onnx"); + + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); + session_options.AddFreeDimensionOverrideByName("batch", 1); + session_options.AddFreeDimensionOverrideByName("width", 64); + session_options.AddFreeDimensionOverrideByName("height", 64); + + // Even though all graph optimizations are disabled, the free dimension override is still enabled by default. + // The shape of graph's output should be correctly inferred by shape inference and data propagation. + Ort::Session session(*ort_env, model_path, session_options); + + // This graph only has one output + 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(); + EXPECT_TRUE(output_shape.size() == 3) << "The output shape should have 3 dimensions"; + EXPECT_TRUE(output_shape[0] == 1) << "The first dimension should have 1 as value"; + EXPECT_TRUE(output_shape[1] == 3) << "The second dimension should have 3 as value"; + EXPECT_TRUE(output_shape[2] == 4096) << "The second dimension should have 4096 as value"; + } + + { + // Model #3 + // This model extends model #2 and appends Unsqueeze -> Unsqueeze -> Squeeze -> Squeeze -> Reshape to the end. + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_with_shape_related_nodes_v3.onnx"); + + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); + session_options.AddFreeDimensionOverrideByName("batch", 1); + session_options.AddFreeDimensionOverrideByName("width", 64); + session_options.AddFreeDimensionOverrideByName("height", 64); + + // Even though all graph optimizations are disabled, the free dimension override is still enabled by default. + // The shape of graph's output should be correctly inferred by shape inference and data propagation. + Ort::Session session(*ort_env, model_path, session_options); + + // This graph only has one output + 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(); + EXPECT_TRUE(output_shape.size() == 3) << "The output shape should have 3 dimensions"; + EXPECT_TRUE(output_shape[0] == 1) << "The first dimension should have 1 as value"; + EXPECT_TRUE(output_shape[1] == 3) << "The second dimension should have 3 as value"; + EXPECT_TRUE(output_shape[2] == 4096) << "The second dimension should have 4096 as value"; + } + + { + // Model #4 + // This model contains Shape, Reshape, Squeeze, Range, ReduceSum. + // It's from SoftmaxGrad_DefaultAxis test. + auto model_path = ORT_TSTR("testdata/test_shape_data_propagation_with_shape_related_nodes_v4.onnx"); + + Ort::SessionOptions session_options{}; + session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); + + // Make sure it can load the model and run shape inference without errors. + Ort::Session session(*ort_env, model_path, session_options); + } +} + namespace { struct MyCustomKernelWithOptionalInput { MyCustomKernelWithOptionalInput(const OrtKernelInfo* /*info*/) { diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes.onnx b/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes.onnx new file mode 100644 index 0000000000000..e18aa31e414e5 Binary files /dev/null and b/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes.onnx differ diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes.py b/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes.py new file mode 100644 index 0000000000000..6537a3cd357c3 --- /dev/null +++ b/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes.py @@ -0,0 +1,53 @@ +import onnx +from onnx import TensorProto, helper + +# 1. Define graph input with symbolic shape ['batch', 3, 'width', 'height'] +input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, ["batch", 3, "width", "height"]) + +# 2. Define intermediate and output tensors +shape_out = helper.make_tensor_value_info("shape_out", TensorProto.INT64, [4]) # Shape output +reshape_a_out = helper.make_tensor_value_info("reshape_a_out", TensorProto.FLOAT, None) +output_tensor = helper.make_tensor_value_info("output", TensorProto.FLOAT, None) + +# 3. Create the initializer for Reshape A's 'shape' input: [0, 32, -1] +shape_initializer = helper.make_tensor( + name="reshape_a_shape", + data_type=TensorProto.INT64, + dims=[3], + vals=[0, 32, -1], +) + +# 4. Create nodes: +# Shape node +shape_node = helper.make_node("Shape", inputs=["input"], outputs=["shape_out"], name="ShapeNode") + +# Reshape A node: takes input + constant shape +reshape_a_node = helper.make_node( + "Reshape", inputs=["input", "reshape_a_shape"], outputs=["reshape_a_out"], name="ReshapeA" +) + +# Reshape B node: takes Shape + ReshapeA outputs, outputs final output +reshape_b_node = helper.make_node("Reshape", inputs=["reshape_a_out", "shape_out"], outputs=["output"], name="ReshapeB") + +# 5. Assemble the graph +graph = helper.make_graph( + nodes=[shape_node, reshape_a_node, reshape_b_node], + name="Shape_Reshape_Model", + inputs=[input_tensor], + outputs=[output_tensor], + initializer=[shape_initializer], + value_info=[shape_out, reshape_a_out], +) + +# 6. Define the model (set IR and opset) +model = helper.make_model( + graph, + opset_imports=[helper.make_operatorsetid("", 18)], + producer_name="onnx-example-generator", +) +model.ir_version = onnx.IR_VERSION + +# 7. Save the model +onnx.save(model, "test_shape_data_propagation_with_shape_related_nodes.onnx") + +print("Model saved to test_shape_data_propagation_with_shape_related_nodes.onnx") diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v2.onnx b/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v2.onnx new file mode 100644 index 0000000000000..ff41075ff64cc Binary files /dev/null and b/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v2.onnx differ diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v2.py b/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v2.py new file mode 100644 index 0000000000000..7cfbcca8d4d03 --- /dev/null +++ b/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v2.py @@ -0,0 +1,59 @@ +import onnx +from onnx import TensorProto, helper + +# === Graph input/output === +input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, ["batch", 3, "width", "height"]) +output_tensor = helper.make_tensor_value_info("output", TensorProto.FLOAT, ["batch", 3, "width*height"]) + +# === Initializers === +B = helper.make_tensor("B", TensorProto.FLOAT, [], [1.0]) + +# Gather indices +g0_idx = helper.make_tensor("g0_idx", TensorProto.INT64, [], [0]) +g1_idx = helper.make_tensor("g1_idx", TensorProto.INT64, [], [1]) +g2_idx = helper.make_tensor("g2_idx", TensorProto.INT64, [], [2]) +g3_idx = helper.make_tensor("g3_idx", TensorProto.INT64, [], [3]) + +# Unsqueeze axes tensors +axes_unsq0 = helper.make_tensor("axes_unsq0", TensorProto.INT64, [1], [0]) +axes_unsq1 = helper.make_tensor("axes_unsq1", TensorProto.INT64, [1], [0]) +axes_unsq2 = helper.make_tensor("axes_unsq2", TensorProto.INT64, [1], [0]) + +# === Nodes === +div = helper.make_node("Div", ["input", "B"], ["div_out"]) + +# Two Shape nodes from Div +shape_left = helper.make_node("Shape", ["div_out"], ["shape_left_out"]) +shape_right = helper.make_node("Shape", ["div_out"], ["shape_right_out"]) + +# Left Shape path +gather0 = helper.make_node("Gather", ["shape_left_out", "g0_idx"], ["g0_out"]) +gather1 = helper.make_node("Gather", ["shape_left_out", "g1_idx"], ["g1_out"]) +unsq0 = helper.make_node("Unsqueeze", ["g0_out", "axes_unsq0"], ["u0_out"]) +unsq1 = helper.make_node("Unsqueeze", ["g1_out", "axes_unsq1"], ["u1_out"]) + +# Right Shape path +gather2 = helper.make_node("Gather", ["shape_right_out", "g2_idx"], ["g2_out"]) +gather3 = helper.make_node("Gather", ["shape_right_out", "g3_idx"], ["g3_out"]) +mul = helper.make_node("Mul", ["g2_out", "g3_out"], ["mul_out"]) +unsq2 = helper.make_node("Unsqueeze", ["mul_out", "axes_unsq2"], ["u2_out"]) + +# Combine +concat = helper.make_node("Concat", ["u0_out", "u1_out", "u2_out"], ["concat_out"], axis=0) +reshape = helper.make_node("Reshape", ["div_out", "concat_out"], ["output"]) + +# === Graph === +graph = helper.make_graph( + [div, shape_left, shape_right, gather0, gather1, gather2, gather3, mul, unsq0, unsq1, unsq2, concat, reshape], + "Div_Shape_Gather_Concat_Reshape", + [input_tensor], + [output_tensor], + initializer=[B, g0_idx, g1_idx, g2_idx, g3_idx, axes_unsq0, axes_unsq1, axes_unsq2], +) + +# === Model === +model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)], producer_name="onnx-example") +onnx.checker.check_model(model) +onnx.save(model, "test_shape_data_propagation_with_shape_related_nodes_v2.onnx") + +print("✅ Model saved as test_shape_data_propagation_with_shape_related_nodes_v2.onnx") diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v3.onnx b/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v3.onnx new file mode 100644 index 0000000000000..2889ec34afd41 Binary files /dev/null and b/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v3.onnx differ diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v3.py b/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v3.py new file mode 100644 index 0000000000000..75bbbe7b4557c --- /dev/null +++ b/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v3.py @@ -0,0 +1,109 @@ +import onnx +from onnx import TensorProto, helper + +# === Graph input/output === +input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, ["batch", 3, "width", "height"]) +output_tensor = helper.make_tensor_value_info("output", TensorProto.FLOAT, ["batch", 3, "width*height"]) + +# === Initializers === +B = helper.make_tensor("B", TensorProto.FLOAT, [], [1.0]) + +# Gather indices +g0_idx = helper.make_tensor("g0_idx", TensorProto.INT64, [], [0]) +g1_idx = helper.make_tensor("g1_idx", TensorProto.INT64, [], [1]) +g2_idx = helper.make_tensor("g2_idx", TensorProto.INT64, [], [2]) +g3_idx = helper.make_tensor("g3_idx", TensorProto.INT64, [], [3]) + +# Unsqueeze axes tensors +axes_unsq0 = helper.make_tensor("axes_unsq0", TensorProto.INT64, [1], [0]) +axes_unsq1 = helper.make_tensor("axes_unsq1", TensorProto.INT64, [1], [0]) +axes_unsq2 = helper.make_tensor("axes_unsq2", TensorProto.INT64, [1], [0]) + +# === Nodes === +div = helper.make_node("Div", ["input", "B"], ["div_out"]) + +# Two Shape nodes from Div +shape_left = helper.make_node("Shape", ["div_out"], ["shape_left_out"]) +shape_right = helper.make_node("Shape", ["div_out"], ["shape_right_out"]) + +# Left Shape path +gather0 = helper.make_node("Gather", ["shape_left_out", "g0_idx"], ["g0_out"]) +gather1 = helper.make_node("Gather", ["shape_left_out", "g1_idx"], ["g1_out"]) +unsq0 = helper.make_node("Unsqueeze", ["g0_out", "axes_unsq0"], ["u0_out"]) +unsq1 = helper.make_node("Unsqueeze", ["g1_out", "axes_unsq1"], ["u1_out"]) + +# Right Shape path +gather2 = helper.make_node("Gather", ["shape_right_out", "g2_idx"], ["g2_out"]) +gather3 = helper.make_node("Gather", ["shape_right_out", "g3_idx"], ["g3_out"]) +mul = helper.make_node("Mul", ["g2_out", "g3_out"], ["mul_out"]) +unsq2 = helper.make_node("Unsqueeze", ["mul_out", "axes_unsq2"], ["u2_out"]) + +# Combine +concat = helper.make_node("Concat", ["u0_out", "u1_out", "u2_out"], ["concat_out"], axis=0) + +# Axes initializers +axes_u1 = helper.make_tensor("axes_u1", TensorProto.INT64, [1], [1]) +axes_u2 = helper.make_tensor("axes_u2", TensorProto.INT64, [1], [1]) +axes_s1 = helper.make_tensor("axes_s1", TensorProto.INT64, [1], [1]) +axes_s2 = helper.make_tensor("axes_s2", TensorProto.INT64, [1], [1]) + +# First Unsqueeze +unsqueeze1 = helper.make_node("Unsqueeze", inputs=["concat_out", "axes_u1"], outputs=["u1"], name="Unsqueeze_1") + +# Second Unsqueeze +unsqueeze2 = helper.make_node("Unsqueeze", inputs=["u1", "axes_u2"], outputs=["u2"], name="Unsqueeze_2") + +# First Squeeze +squeeze1 = helper.make_node("Squeeze", inputs=["u2", "axes_s1"], outputs=["s1"], name="Squeeze_1") + +# Second Squeeze +squeeze2 = helper.make_node("Squeeze", inputs=["s1", "axes_s2"], outputs=["squeeze_output"], name="Squeeze_2") + +reshape = helper.make_node("Reshape", ["div_out", "squeeze_output"], ["output"]) + +# === Graph === +graph = helper.make_graph( + [ + div, + shape_left, + shape_right, + gather0, + gather1, + gather2, + gather3, + mul, + unsq0, + unsq1, + unsq2, + concat, + unsqueeze1, + unsqueeze2, + squeeze1, + squeeze2, + reshape, + ], + "Div_Shape_Gather_Concat_Reshape", + [input_tensor], + [output_tensor], + initializer=[ + B, + g0_idx, + g1_idx, + g2_idx, + g3_idx, + axes_unsq0, + axes_unsq1, + axes_unsq2, + axes_u1, + axes_u2, + axes_s1, + axes_s2, + ], +) + +# === Model === +model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)], producer_name="onnx-example") +onnx.checker.check_model(model) +onnx.save(model, "test_shape_data_propagation_with_shape_related_nodes_v3.onnx") + +print("✅ Model saved as test_shape_data_propagation_with_shape_related_nodes_v3.onnx") diff --git a/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v4.onnx b/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v4.onnx new file mode 100644 index 0000000000000..d13f317b3a1c8 Binary files /dev/null and b/onnxruntime/test/testdata/test_shape_data_propagation_with_shape_related_nodes_v4.onnx differ