Add partial data propagation to enhance shape inference

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 @@
   /** 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 @@
   // 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_;
 };

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) {
+    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) {
+    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/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<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