Add partial data propagation to enhance shape inference (#26269)

### Description
Calling an operator's `TypeAndShapeInferenceFunction()` alone is
sometimes insufficient for complete shape inference.

For example, the `Shape` operator only infers the output’s rank (a
1-dimensional tensor) but not its actual dimension values.
For instance, given an input of shape [1, 3, 64, 64], the Shape
operator's `TypeAndShapeInferenceFunction()` produces an output shape
tensor with 1-dimension as int64[4], representing the rank of the input
tensor.

Therefore, as you can imagine, the below graph's output shape can't be
properly inferred (even though the input shape is known) because the
shape data is lost at the `Shape `operator.
<img width="563" height="488" alt="image"
src="https://github.com/user-attachments/assets/bfa9fd8f-5291-4c6d-a679-3ce4a8c48669"
/>

To solve the issue, 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.



This PR adds the support of executing operator's
`PartialDataPropagationFunction()` in ORT, and makes sure the shape
values is properly propagated throughout the graph.

### Motivation and Context

When using the Compile API to generate an EPContext model, all graph
optimizations are disabled by default except for free dimension
overrides. However, for certain models, such as a VAE decoder, the
output shape may still fail to be properly inferred even when free
dimension override values are provided beforehand.

However, you won't hit this issue if enabling all the graph
optimizations as some nodes, e.g. `Shape`, `Reshape `.. will be constant
folded.

Author: chilo-ms

include/onnxruntime/core/graph/graph.h

@@ -1753,6 +1753,15 @@ class Graph {  // NOLINT(clang-analyzer-optin.performance.Padding): preserve exi
                                                     std::vector<const ONNX_NAMESPACE::TypeProto*>& 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();
 

include/onnxruntime/core/graph/node_arg.h

@@ -9,6 +9,8 @@
 #include "core/common/status.h"
 #include "core/common/logging/logging.h"
 
+#include <optional>
+
 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<ONNX_NAMESPACE::TensorShapeProto>& GetInferredShapeValues() const noexcept { return inferred_shape_values_; }
+
+  /** Gets mutable inferred shape values as a TensorShapeProto. */
+  std::optional<ONNX_NAMESPACE::TensorShapeProto>& GetMutableInferredShapeValues() noexcept { return inferred_shape_values_; }
+
+  /** Gets the inferred shape scalar value */
+  const std::optional<int64_t> 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<ONNX_NAMESPACE::TensorShapeProto> 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<int64_t> inferred_scalar_value_;
+
   // Flag indicates whether <*this> node arg exists or not.
   bool exists_;
 };

include/onnxruntime/core/session/onnxruntime_cxx_api.h

@@ -1441,6 +1441,12 @@ struct SessionOptionsImpl : ConstSessionOptionsImpl<T> {
 
   ///< Wraps OrtApi::SessionOptionsAppendExecutionProvider_VitisAI
   SessionOptionsImpl& AppendExecutionProvider_VitisAI(const std::unordered_map<std::string, std::string>& 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
 

include/onnxruntime/core/session/onnxruntime_cxx_inline.h

@@ -1503,6 +1503,18 @@ inline SessionOptionsImpl<T>& SessionOptionsImpl<T>::RegisterCustomOpsUsingFunct
   return *this;
 }
 
+template <typename T>
+inline SessionOptionsImpl<T>& SessionOptionsImpl<T>::AddFreeDimensionOverride(const char* dim_denotation, int64_t dim_value) {
+  ThrowOnError(GetApi().AddFreeDimensionOverrideByName(this->p_, dim_denotation, dim_value));
+  return *this;
+}
+
+template <typename T>
+inline SessionOptionsImpl<T>& SessionOptionsImpl<T>::AddFreeDimensionOverrideByName(const char* dim_name, int64_t dim_value) {
+  ThrowOnError(GetApi().AddFreeDimensionOverrideByName(this->p_, dim_name, dim_value));
+  return *this;
+}
+
 /// Session
 template <typename T>
 inline size_t ConstSessionImpl<T>::GetInputCount() const {

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

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<Status(const std::string&, TensorShapeVector&)> 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

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 <onnx/onnx-ml.pb.h>
+
+namespace onnxruntime {
+
+std::unique_ptr<CustomDataPropagationBase> CreateCustomDataPropagation(const Node& node,
+                                                                       NodeArg& output_def,
+                                                                       std::function<Status(const std::string&, TensorShapeVector&)> 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<SizeOpDataPropagation>(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger);
+  } else if (node.OpType() == "Squeeze") {
+    return std::make_unique<SqueezeOpDataPropagation>(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger);
+  } else if (node.OpType() == "Unsqueeze") {
+    return std::make_unique<UnsqueezeOpDataPropagation>(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<GatherOpDataPropagation>(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger);
+    } else if (node.OpType() == "Add") {
+      return std::make_unique<AddOpDataPropagation>(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger);
+    } else if (node.OpType() == "Sub") {
+      return std::make_unique<SubOpDataPropagation>(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger);
+    } else if (node.OpType() == "Mul") {
+      return std::make_unique<MulOpDataPropagation>(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger);
+    } else if (node.OpType() == "Div") {
+      return std::make_unique<DivOpDataPropagation>(node, output_def, std::move(func), output_from_onnx_op_data_propagation, logger);
+    }
+  }
+  return nullptr;
+}
+
+}  // namespace onnxruntime

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 <onnx/onnx-ml.pb.h>
+
+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<Status(const std::string&, TensorShapeVector&)> 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<Status(const std::string&, TensorShapeVector&)> 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<CustomDataPropagation> Returns a CustomDataPropagation object if available
+ */
+std::unique_ptr<CustomDataPropagationBase> CreateCustomDataPropagation(
+    const Node& node,
+    NodeArg& output_def,
+    std::function<Status(const std::string&, TensorShapeVector&)> func,
+    const ONNX_NAMESPACE::TypeProto& output_from_onnx_op_data_propagation,
+    const logging::Logger& logger);
+
+}  // namespace onnxruntime

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