address revewer's comments

Author: chilo-ms

include/onnxruntime/core/graph/graph.h

@@ -1760,10 +1760,10 @@ 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(),
+  // 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(Node& node, NodeArg& output_def,
+  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.

include/onnxruntime/core/graph/node_arg.h

@@ -151,7 +151,8 @@ class NodeArg {
   //
   // 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.
+  // of shape information throughout the graph. If the concrete output shape value is not
+  // computed, then no shape value is saved here that's why this is optional.
   std::optional<ONNX_NAMESPACE::TensorShapeProto> inferred_shape_values_;
 
   // This variable stores the actual scalar value.

onnxruntime/core/graph/data_propagation/add_op_data_propagation.cc

@@ -15,8 +15,15 @@ Status AddOpDataPropagation::infer() {
   // 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());
+    output_def_.SetInferredShapeScalarValue(
+        input_0->GetInferredShapeScalarValue().value() +
+        input_1->GetInferredShapeScalarValue().value());
   }
 
   return Status::OK();

onnxruntime/core/graph/data_propagation/custom_data_propagation.h

@@ -65,10 +65,11 @@ class CustomDataPropagationBase {
  * @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);
+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

@@ -15,8 +15,15 @@ Status DivOpDataPropagation::infer() {
   // 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());
+    output_def_.SetInferredShapeScalarValue(
+        input_0->GetInferredShapeScalarValue().value() /
+        input_1->GetInferredShapeScalarValue().value());
   }
 
   return Status::OK();

onnxruntime/core/graph/data_propagation/gather_op_data_propagation.cc

@@ -10,36 +10,47 @@
 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();
+    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.");
+    }
   }
-  // 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().
-  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.
+  // 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 know.
+  //       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) {
+    return Status::OK();
+  }
+
   TensorShapeVector indices;
   ORT_RETURN_IF_ERROR(get_initialized_input_values_func_(input_1->Name(), indices));
 
@@ -51,7 +62,10 @@ Status GatherOpDataPropagation::infer() {
     if (indices.size() == 1) {
       ORT_TRY {
         // Note: Index value is expected to be within bounds [-s, s-1] along axis of size s
-        auto& dim = tensor_shape_proto.dim(static_cast<int32_t>(HandleNegativeAxis(indices[0], tensor_shape_proto.dim_size())));
+        auto index = static_cast<int32_t>(
+            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());
         }

onnxruntime/core/graph/data_propagation/mul_op_data_propagation.cc

@@ -15,8 +15,15 @@ Status MulOpDataPropagation::infer() {
   // 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());
+    output_def_.SetInferredShapeScalarValue(
+        input_0->GetInferredShapeScalarValue().value() *
+        input_1->GetInferredShapeScalarValue().value());
   }
 
   return Status::OK();

onnxruntime/core/graph/data_propagation/size_op_data_propagation.cc

@@ -12,6 +12,11 @@ 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) {
+    return Status::OK();
+  }
+
   if (input_0->GetInferredShapeValues().has_value()) {
     const auto& tensor_shape_proto = input_0->GetInferredShapeValues().value();
 

onnxruntime/core/graph/data_propagation/size_op_data_propagation.h

@@ -10,7 +10,9 @@ 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:

onnxruntime/core/graph/data_propagation/squeeze_op_data_propagation.cc

@@ -6,16 +6,18 @@
 #include "core/graph/node_arg.h"
 #include "core/graph/onnx_protobuf.h"
 #include "core/providers/common.h"
-
-#if defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
-#include <unordered_set>
-#endif
+#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) {
+    return Status::OK();
+  }
+
   if (input_0->GetInferredShapeValues().has_value()) {
     const auto& tensor_shape_proto = input_0->GetInferredShapeValues().value();
 
@@ -34,11 +36,7 @@ Status SqueezeOpDataPropagation::infer() {
     } else if (tensor_shape_proto.dim_size() > 1) {
       // Get axes value
       TensorShapeVector axes;
-#if !defined(ORT_MINIMAL_BUILD) && !defined(ORT_EXTENDED_MINIMAL_BUILD)
       InlinedHashSet<int64_t> axes_set;
-#else
-      std::unordered_set<int64_t> axes_set;
-#endif
 
       // 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.