Skip to content
9 changes: 9 additions & 0 deletions onnxruntime/core/graph/graph_utils.cc
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#include "core/common/logging/logging.h"

#include <algorithm>
#include <limits>

Check warning on line 11 in onnxruntime/core/graph/graph_utils.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++

[cpplint] reported by reviewdog 🐶 Found C++ system header after other header. Should be: graph_utils.h, c system, c++ system, other. [build/include_order] [4] Raw Output: onnxruntime/core/graph/graph_utils.cc:11: Found C++ system header after other header. Should be: graph_utils.h, c system, c++ system, other. [build/include_order] [4]
#include <queue>
#include <string>
#include <vector>
Expand Down Expand Up @@ -411,6 +412,14 @@
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();
Expand Down
5 changes: 5 additions & 0 deletions onnxruntime/core/graph/graph_utils.h
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#include "core/graph/onnx_protobuf.h"
#include "core/graph/graph.h"

#include <limits>

Check warning on line 11 in onnxruntime/core/graph/graph_utils.h

View workflow job for this annotation

GitHub Actions / Optional Lint C++

[cpplint] reported by reviewdog 🐶 Found C++ system header after other header. Should be: graph_utils.h, c system, c++ system, other. [build/include_order] [4] Raw Output: onnxruntime/core/graph/graph_utils.h:11: Found C++ system header after other header. Should be: graph_utils.h, c system, c++ system, other. [build/include_order] [4]
#include <string>
#include <vector>

Expand All @@ -31,6 +32,10 @@
/** 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.
Expand Down
58 changes: 29 additions & 29 deletions onnxruntime/core/optimizer/attention_fusion.cc
Original file line number Diff line number Diff line change
Expand Up @@ -349,22 +349,22 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul,

const Node* sequence_transpose = graph_utils::GetInputNode(qkv_matmul, 0);
if (sequence_transpose == nullptr ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*sequence_transpose, "Transpose", {1, 13}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*sequence_transpose, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain) ||
!HasExpectedPerm(*sequence_transpose, {0, 2, 1}) ||
!optimizer_utils::CheckOutputEdges(graph, *sequence_transpose, 1)) {
return false;
}

const Node* input_reshape = graph_utils::GetInputNode(*sequence_transpose, 0);
if (input_reshape == nullptr ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*input_reshape, "Reshape", {5, 13, 14}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*input_reshape, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain) ||
!optimizer_utils::CheckOutputEdges(graph, *input_reshape, 1)) {
return fail("missing input Reshape before sequence transpose");
}

Node* qkv_reshape = GetOnlyChildByOutputIndex(graph, qkv_matmul, 0, "Reshape");
if (qkv_reshape == nullptr ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*qkv_reshape, "Reshape", {5, 13, 14}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*qkv_reshape, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain) ||
!optimizer_utils::CheckOutputEdges(graph, *qkv_reshape, 1)) {
return fail("qkv Reshape after MatMul not matched");
}
Expand All @@ -379,9 +379,9 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul,
Node* k_squeeze = GetOnlyChildByOutputIndex(graph, *split, 1, "Squeeze");
Node* v_transpose = GetOnlyChildByOutputIndex(graph, *split, 2, "Transpose");
if (q_transpose == nullptr || k_squeeze == nullptr || v_transpose == nullptr ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*q_transpose, "Transpose", {1, 13}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*k_squeeze, "Squeeze", {13}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*v_transpose, "Transpose", {1, 13}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*q_transpose, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*k_squeeze, "Squeeze", {13, 21, 23, 24, 25}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*v_transpose, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain) ||
!HasExpectedPerm(*q_transpose, {2, 0, 3, 1, 4}) ||
!HasExpectedPerm(*v_transpose, {2, 0, 3, 1, 4}) ||
!HasExpectedAxesInput(graph, *k_squeeze, {2})) {
Expand All @@ -391,8 +391,8 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul,
Node* q_squeeze = GetOnlyChildByOutputIndex(graph, *q_transpose, 0, "Squeeze");
Node* v_squeeze = GetOnlyChildByOutputIndex(graph, *v_transpose, 0, "Squeeze");
if (q_squeeze == nullptr || v_squeeze == nullptr ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*q_squeeze, "Squeeze", {13}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*v_squeeze, "Squeeze", {13}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*q_squeeze, "Squeeze", {13, 21, 23, 24, 25}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*v_squeeze, "Squeeze", {13, 21, 23, 24, 25}, kOnnxDomain) ||
!HasExpectedAxesInput(graph, *q_squeeze, {0}) ||
!HasExpectedAxesInput(graph, *v_squeeze, {0})) {
return fail("q/v squeeze pattern not matched");
Expand All @@ -402,7 +402,7 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul,
Node* k_transpose = GetOnlyChildByOutputIndex(graph, *k_squeeze, 0, "Transpose");
if (q_scale_mul == nullptr || k_transpose == nullptr ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*q_scale_mul, "Mul", {7, 13, 14}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*k_transpose, "Transpose", {1, 13}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*k_transpose, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain) ||
!HasExpectedPerm(*k_transpose, {0, 2, 3, 1})) {
return fail("q scale Mul or k Transpose(0,2,3,1) not matched");
}
Expand Down Expand Up @@ -460,15 +460,15 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul,

Node* transpose_3 = GetOnlyChildByOutputIndex(graph, *qkv_matmul_1, 0, "Transpose");
if (transpose_3 == nullptr ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*transpose_3, "Transpose", {1, 13}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*transpose_3, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain) ||
!HasExpectedPerm(*transpose_3, {0, 2, 1, 3}) ||
!optimizer_utils::CheckOutputEdges(graph, *transpose_3, 1)) {
return fail("output Transpose(0,2,1,3) not matched");
}

Node* reshape_2 = GetOnlyChildByOutputIndex(graph, *transpose_3, 0, "Reshape");
if (reshape_2 == nullptr ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*reshape_2, "Reshape", {5, 13, 14}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*reshape_2, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain) ||
!optimizer_utils::CheckOutputEdges(graph, *reshape_2, 1)) {
return fail("output Reshape not matched");
}
Expand Down Expand Up @@ -497,7 +497,7 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul,
if (proj_gemm == nullptr) {
proj_gemm_input_reshape = GetOnlyChildByOutputIndex(graph, *reshape_2, 0, "Reshape");
if (proj_gemm_input_reshape == nullptr ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*proj_gemm_input_reshape, "Reshape", {5, 13, 14}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*proj_gemm_input_reshape, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain) ||
!optimizer_utils::CheckOutputEdges(graph, *proj_gemm_input_reshape, 1)) {
Comment thread
yuslepukhin marked this conversation as resolved.
return fail("projection MatMul/Gemm not matched");
}
Expand All @@ -511,7 +511,7 @@ static bool TryFuseMobileClipMHA(Node& qkv_matmul,

proj_gemm_output_reshape = GetOnlyChildByOutputIndex(graph, *proj_gemm, 0, "Reshape");
if (proj_gemm_output_reshape == nullptr ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*proj_gemm_output_reshape, "Reshape", {5, 13, 14}, kOnnxDomain) ||
!graph_utils::IsSupportedOptypeVersionAndDomain(*proj_gemm_output_reshape, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain) ||
!optimizer_utils::CheckOutputEdges(graph, *proj_gemm_output_reshape, 1)) {
return fail("normalized projection Gemm output Reshape not matched");
}
Expand Down Expand Up @@ -920,11 +920,11 @@ static bool FuseSubGraphQKImpl(Node& layer_norm,
}

std::vector<graph_utils::EdgeEndToMatch> q_path{
{0, 0, "Transpose", {1, 13}, kOnnxDomain},
{0, 0, "Reshape", {5, 13}, kOnnxDomain},
{0, 0, "Add", {7, 13}, kOnnxDomain},
{0, 0, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain},
{0, 0, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain},
{0, 0, "Add", {7, 13, 14}, kOnnxDomain},
{0, 0, "MatMul", {1, 9, 13}, kOnnxDomain},
{0, 0, "LayerNormalization", {1}, kOnnxDomain}};
{0, 0, "LayerNormalization", {1, 17}, kOnnxDomain}};
if (!graph_utils::FindPath(edges[edges.size() - 1]->GetNode(), true, q_path, edges, logger)) {
DEBUG_LOG("Failed to find path for q");
return false;
Expand Down Expand Up @@ -953,9 +953,9 @@ static bool FuseSubGraphQKImpl(Node& layer_norm,
}

std::vector<graph_utils::EdgeEndToMatch> k_path{
{0, 1, "Transpose", {1, 13}, kOnnxDomain},
{0, 0, "Reshape", {5, 13}, kOnnxDomain},
{0, 0, "Add", {7, 13}, kOnnxDomain},
{0, 1, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain},
{0, 0, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain},
{0, 0, "Add", {7, 13, 14}, kOnnxDomain},
{0, 0, "MatMul", {1, 9, 13}, kOnnxDomain},
{0, 0, "LayerNormalization", {1, 17}, kOnnxDomain}};

Expand Down Expand Up @@ -1070,8 +1070,8 @@ static bool FuseSubGraphQK(Node& layer_norm,
const logging::Logger& logger) {
// path to q
std::vector<graph_utils::EdgeEndToMatch> q_varience_path{
{0, 0, "Div", {7, 13}, kOnnxDomain},
{0, 0, "MatMul", {1, 9}, kOnnxDomain}};
{0, 0, "Div", {7, 13, 14}, kOnnxDomain},
{0, 0, "MatMul", {1, 9, 13}, kOnnxDomain}};
std::vector<const Node::EdgeEnd*> edges;
if (!graph_utils::FindPath(*(mask_nodes.add), true, q_varience_path, edges, logger)) {
DEBUG_LOG("Failed to find path for q");
Expand Down Expand Up @@ -1163,7 +1163,7 @@ static bool FuseSubGraphQKDistilBert(Node& layer_norm,
// path to q
std::vector<graph_utils::EdgeEndToMatch> q_varience_path{
{0, 2, "MatMul", {1, 9, 13}, kOnnxDomain},
{0, 0, "Div", {7, 13}, kOnnxDomain}};
{0, 0, "Div", {7, 13, 14}, kOnnxDomain}};
std::vector<const Node::EdgeEnd*> edges;
if (!graph_utils::FindPath(*(mask_nodes.where), true, q_varience_path, edges, logger)) {
DEBUG_LOG("Failed to find path for q");
Expand Down Expand Up @@ -1265,14 +1265,14 @@ bool AttentionFusion::FuseSubGraph(Node& layer_norm,
std::map<std::string, NodeArg*>& mask_int32_map,
const logging::Logger& logger) {
std::vector<graph_utils::EdgeEndToMatch> parent_path{
{0, 0, "Add", {7, 13}, kOnnxDomain},
{0, 0, "Add", {7, 13, 14}, kOnnxDomain},
{0, 0, "MatMul", {1, 9, 13}, kOnnxDomain},
{0, 0, "Reshape", {5, 13}, kOnnxDomain},
{0, 0, "Transpose", {1, 13}, kOnnxDomain},
{0, 0, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain},
{0, 0, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain},
{0, 0, "MatMul", {1, 9, 13}, kOnnxDomain},
{0, 1, "Transpose", {1, 13}, kOnnxDomain},
{0, 0, "Reshape", {5, 13}, kOnnxDomain},
{0, 0, "Add", {7, 13}, kOnnxDomain},
{0, 1, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain},
{0, 0, "Reshape", {5, 13, 14, 19, 21, 23, 24, 25}, kOnnxDomain},
{0, 0, "Add", {7, 13, 14}, kOnnxDomain},
{0, 0, "MatMul", {1, 9, 13}, kOnnxDomain},
{0, 0, "LayerNormalization", {1, 17}, kOnnxDomain}};

Expand Down
Loading
Loading