Extract IsFullShapeNode helper into graph_utils

Author: yuslepukhin

onnxruntime/core/graph/graph_utils.cc

@@ -8,6 +8,7 @@
 #include "core/common/logging/logging.h"
 
 #include <algorithm>
+#include <limits>
 #include <queue>
 #include <string>
 #include <vector>
@@ -411,6 +412,14 @@ const ONNX_NAMESPACE::AttributeProto* GetNodeAttribute(const Node& node, const s
   return iter == attrs.end() ? nullptr : &iter->second;
 }
 
+bool IsFullShapeNode(const Node& node) {
+  const auto* start_attr = GetNodeAttribute(node, "start");
+  const auto* end_attr = GetNodeAttribute(node, "end");
+  // end=INT64_MAX is the runtime default meaning "all dimensions" (full shape).
+  return (!start_attr || start_attr->i() == 0) &&
+         (!end_attr || end_attr->i() == std::numeric_limits<int64_t>::max());
+}
+
 static NodeArg& GetOrCreateNodeArg(Graph& graph, const ONNX_NAMESPACE::TensorProto& new_initializer) {
   ONNX_NAMESPACE::TypeProto new_type;
   auto* typeproto_tensor = new_type.mutable_tensor_type();

onnxruntime/core/graph/graph_utils.h

@@ -8,6 +8,7 @@
 #include "core/graph/onnx_protobuf.h"
 #include "core/graph/graph.h"
 
+#include <limits>
 #include <string>
 #include <vector>
 
@@ -31,6 +32,10 @@ bool IsSupportedOptypeVersionAndDomain(const Node& node,
 /** Returns the attribute of a Node with a given name. */
 const ONNX_NAMESPACE::AttributeProto* GetNodeAttribute(const Node& node, const std::string& attr_name);
 
+/** Checks whether a Shape node returns the full tensor shape (all dimensions).
+ * Returns false if start/end attributes restrict the output to a subset of dimensions. */
+bool IsFullShapeNode(const Node& node);
+
 /** Add a new initializer to 'graph'.
 Checks that new_initializer does not already exist in 'graph' before adding it.
 @returns The NodeArg for the new initializer.

onnxruntime/core/optimizer/attention_fusion_helper.h

@@ -98,19 +98,12 @@ bool MatchGemmSubgraph(Graph& graph,
   const Node& slice = edges[6]->GetNode();
   const Node& shape_before_slice = edges[7]->GetNode();
 
-  // Opset 15+ added optional start/end attributes to Shape, allowing it to return only a
-  // subset of dimensions ("partial shape"). The downstream Slice/Squeeze/Gather nodes assume
-  // Shape returns the full tensor shape. If start != 0 or end is not the default (INT64_MAX),
-  // the fusion would produce incorrect index mapping. Reject such cases.
-  if (shape_before_slice.SinceVersion() >= 15) {
-    const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_before_slice, "start");
-    const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_before_slice, "end");
-    // end=INT64_MAX is the runtime default, meaning "all dimensions" (i.e. full shape).
-    if (!((!start_attr || start_attr->i() == 0) &&
-          (!end_attr || end_attr->i() == std::numeric_limits<int64_t>::max()))) {
-      DEBUG_LOG("Shape node has non-default start/end attributes");
-      return false;
-    }
+  // The downstream Slice/Squeeze/Gather nodes assume Shape returns the full tensor shape so
+  // that indices map directly to tensor dimensions. A partial shape (opset 15+ start/end
+  // attributes) would produce incorrect index mapping.
+  if (!graph_utils::IsFullShapeNode(shape_before_slice)) {
+    DEBUG_LOG("Shape node has non-default start/end attributes");
+    return false;
   }
 
   const auto& subgraph_input = shape_before_slice.InputDefs()[0];

onnxruntime/core/optimizer/embed_layer_norm_fusion.cc

@@ -2,8 +2,6 @@
 // Licensed under the MIT License.
 #include "core/optimizer/embed_layer_norm_fusion.h"
 
-#include <limits>
-
 #include "core/common/span_utils.h"
 #include "core/optimizer/initializer.h"
 #include "core/graph/contrib_ops/contrib_defs.h"
@@ -140,20 +138,12 @@ static bool MatchInputToConcatSubgraph(
     }
   }
 
-  // Opset 15+ added optional start/end attributes to Shape, allowing it to return only a
-  // subset of dimensions ("partial shape"). The Gather(index=0) below assumes Shape returns
-  // the full tensor shape. If start != 0 or end is not the default (INT64_MAX), the Gather
-  // would pick the wrong dimension and fusion would be incorrect. Reject such cases.
+  // The Gather(index=0) below assumes Shape returns the full tensor shape. A partial shape
+  // (opset 15+ start/end attributes) would cause Gather to pick the wrong dimension.
   const Node& shape_node_path1 = edges[shape_index]->GetNode();
-  if (shape_node_path1.SinceVersion() >= 15) {
-    const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_node_path1, "start");
-    const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_node_path1, "end");
-    // end=INT64_MAX is the runtime default, meaning "all dimensions" (i.e. full shape).
-    if (!((!start_attr || start_attr->i() == 0) &&
-          (!end_attr || end_attr->i() == std::numeric_limits<int64_t>::max()))) {
-      DEBUG_LOG("Shape node in path 1 has non-default start/end attributes.");
-      return false;
-    }
+  if (!graph_utils::IsFullShapeNode(shape_node_path1)) {
+    DEBUG_LOG("Shape node in path 1 has non-default start/end attributes.");
+    return false;
   }
 
   Node& concat_node = *graph.GetNode(edges[0]->GetNode().Index());
@@ -185,19 +175,11 @@ static bool MatchInputToConcatSubgraph(
   Node& gather_node_1 = *graph.GetNode(edges[1]->GetNode().Index());
   Node& shape_node_1 = *graph.GetNode(edges[2]->GetNode().Index());
 
-  // Opset 15+ added optional start/end attributes to Shape, allowing it to return only a
-  // subset of dimensions ("partial shape"). The Gather(index=1) below assumes Shape returns
-  // the full tensor shape. If start != 0 or end is not the default (INT64_MAX), the Gather
-  // would pick the wrong dimension and fusion would be incorrect. Reject such cases.
-  if (shape_node_1.SinceVersion() >= 15) {
-    const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape_node_1, "start");
-    const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape_node_1, "end");
-    // end=INT64_MAX is the runtime default, meaning "all dimensions" (i.e. full shape).
-    if (!((!start_attr || start_attr->i() == 0) &&
-          (!end_attr || end_attr->i() == std::numeric_limits<int64_t>::max()))) {
-      DEBUG_LOG("Shape node in path 2 has non-default start/end attributes.");
-      return false;
-    }
+  // The Gather(index=1) below assumes Shape returns the full tensor shape. A partial shape
+  // (opset 15+ start/end attributes) would cause Gather to pick the wrong dimension.
+  if (!graph_utils::IsFullShapeNode(shape_node_1)) {
+    DEBUG_LOG("Shape node in path 2 has non-default start/end attributes.");
+    return false;
   }
 
   // The gather node (with second input indices==1) is also shared by other subgraph

onnxruntime/core/optimizer/reshape_fusion.cc

@@ -2,7 +2,6 @@
 // Licensed under the MIT License.
 
 #include <algorithm>
-#include <limits>
 
 #include "core/graph/graph_utils.h"
 #include "core/optimizer/initializer.h"
@@ -173,19 +172,11 @@ bool ReshapeFusion::Match_One_Element_Output_Subgraph_1(Graph& graph, const Node
     const Node& gather = edges[1]->GetNode();
     const Node& shape = edges[2]->GetNode();
 
-    // Opset 15+ added optional start/end attributes to Shape, allowing it to return only a
-    // subset of dimensions ("partial shape"). The fusion below assumes Shape returns the full
-    // tensor shape so that Gather indices correspond directly to tensor dimensions. If start != 0
-    // or end is set to something other than the default (INT64_MAX = all dims), the Gather index
-    // semantics change and the fusion would produce incorrect results. Reject such cases.
-    if (shape.SinceVersion() >= 15) {
-      const ONNX_NAMESPACE::AttributeProto* start_attr = graph_utils::GetNodeAttribute(shape, "start");
-      const ONNX_NAMESPACE::AttributeProto* end_attr = graph_utils::GetNodeAttribute(shape, "end");
-      // end=INT64_MAX is the runtime default, meaning "all dimensions" (i.e. full shape).
-      if (!((!start_attr || start_attr->i() == 0) &&
-            (!end_attr || end_attr->i() == std::numeric_limits<int64_t>::max()))) {
-        return false;
-      }
+    // The fusion assumes Shape returns the full tensor shape so that Gather indices correspond
+    // directly to tensor dimensions. A partial shape (opset 15+ start/end attributes) would shift
+    // the index mapping and produce incorrect results.
+    if (!graph_utils::IsFullShapeNode(shape)) {
+      return false;
     }
 
     InlinedVector<int64_t> axes;