Add partial data propagation to enhance shape inference

include/onnxruntime/core/graph/graph.h

@@ -1760,6 +1760,15 @@ class Graph {  // NOLINT(clang-analyzer-optin.performance.Padding): preserve exi
                                                     std::vector<const ONNX_NAMESPACE::TypeProto*>& output_types,
                                                     const Graph::ResolveOptions& options);
 
+  // If the shape values are inferred after executing ONNX operator's PartialDataPropagationFunction(),
+  // 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(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 @@
   /** 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,22 @@
   // Node arg name, type and shape.
   NodeArgInfo node_arg_info_;
 
+  // This variable stores the inferred shape data as a TensorShapeProto after executing
+  // the ONNX operator's PartialDataPropagationFunction().
+  //
+  // 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.
+  std::optional<ONNX_NAMESPACE::TensorShapeProto> inferred_shape_values_;
+
+  // This variable stores the inferred scalar output.
+  // 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_;
 };

onnxruntime/core/graph/data_propagation/add_op_data_propagation.cc

@@ -0,0 +1,25 @@
+// 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];
+
+  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,20 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "data_propagation.h"
+#include "core/graph/graph.h"
+namespace onnxruntime {
+
+class AddOpDataPropagation : public CustomDataPropagation {
+ 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
+      : CustomDataPropagation(node, output_def, func, output_from_onnx_op_data_propagation, logger) {}
+
+  Status infer() override;
+};
+
+}  // namespace onnxruntime

onnxruntime/core/graph/data_propagation/data_propagation.h

@@ -0,0 +1,40 @@
+// 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.
+ */
+class CustomDataPropagation {
+ public:
+  virtual ~CustomDataPropagation() = default;
+  virtual Status infer() = 0;
+
+ protected:
+  CustomDataPropagation(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_(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_;
+};
+
+}  // namespace onnxruntime

onnxruntime/core/graph/data_propagation/data_propagation_factory.cc

@@ -0,0 +1,71 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "data_propagation_factory.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 {
+/**
+ * @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.
+ *
+ * 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<CustomDataPropagation> 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, func, output_from_onnx_op_data_propagation, logger);
+  } else if (node.OpType() == "Squeeze") {
+    return std::make_unique<SqueezeOpDataPropagation>(node, output_def, func, output_from_onnx_op_data_propagation, logger);
+  } else if (node.OpType() == "Unsqueeze") {
+    return std::make_unique<UnsqueezeOpDataPropagation>(node, output_def, 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, func, output_from_onnx_op_data_propagation, logger);
+    } else if (node.OpType() == "Add") {
+      return std::make_unique<AddOpDataPropagation>(node, output_def, func, output_from_onnx_op_data_propagation, logger);
+    } else if (node.OpType() == "Sub") {
+      return std::make_unique<SubOpDataPropagation>(node, output_def, func, output_from_onnx_op_data_propagation, logger);
+    } else if (node.OpType() == "Mul") {
+      return std::make_unique<MulOpDataPropagation>(node, output_def, func, output_from_onnx_op_data_propagation, logger);
+    } else if (node.OpType() == "Div") {
+      return std::make_unique<DivOpDataPropagation>(node, output_def, func, output_from_onnx_op_data_propagation, logger);
+    }
+  }
+  return nullptr;
+}
+
+}  // namespace onnxruntime

onnxruntime/core/graph/data_propagation/data_propagation_factory.h

@@ -0,0 +1,37 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+#include <memory>
+#include "core/graph/graph.h"
+#include "data_propagation.h"
+
+namespace onnxruntime {
+/**
+ * @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.
+ *
+ * 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<CustomDataPropagation> CreateCustomDataPropagation(const Node& node,
+                                                                   NodeArg& output_def,
+                                                                   std::function<Status(const std::string&, TensorShapeVector&)> funcs,
+                                                                   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,25 @@
+// 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];
+
+  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/div_op_data_propagation.h

@@ -0,0 +1,20 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "data_propagation.h"
+#include "core/graph/graph.h"
+namespace onnxruntime {
+
+class DivOpDataPropagation : public CustomDataPropagation {
+ public:
+  DivOpDataPropagation(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
+      : CustomDataPropagation(node, output_def, func, output_from_onnx_op_data_propagation, logger) {}
+
+  Status infer() override;
+};
+
+}  // namespace onnxruntime

onnxruntime/core/graph/data_propagation/gather_op_data_propagation.cc

@@ -0,0 +1,68 @@
+// 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() {
+  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();
+  }
+  ORT_ENFORCE(dim_size == 0);
+
+  // 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 a 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 know.
+  //       3. A scalar.
+  //
+  //       If it's case #1, ONNX operator's PartialDataPropagationFunction()
+  //       should have inferred the output shape value.
+  //       This Gather's custom data propagation handles case #3.
+  const auto* input_1 = node_.InputDefs()[1];
+
+  TensorShapeVector indices;
+  ORT_RETURN_IF_ERROR(get_initialized_input_values_func_(input_1->Name(), indices));
+
+  // Save the dimension value in the NodeArg.
+  // Index value is expected to be within bounds [-s, s-1] along axis of size s
+  if (input_0->GetInferredShapeValues().has_value()) {
+    const auto& tensor_shape_proto = input_0->GetInferredShapeValues().value();
+
+    // If "indices" input is a scalar, then the size of indices is 1.
+    if (indices.size() == 1) {
+      ORT_TRY {
+        auto& dim = tensor_shape_proto.dim(static_cast<int32_t>(HandleNegativeAxis(indices[0], tensor_shape_proto.dim_size())));
+        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();
+}
+
+}  // namespace onnxruntime