graph_indexing#

Graph indexing and pattern matching for ONNX quantization.

Functions

`expand_node_names_from_patterns`	Expand the node names from the given patterns.
`find_mha_partitions`	Match MHA: BMM1 -> .
`get_fusible_backbone`	Returns the linear backbone node for a given node if it matches the pattern.
`get_tensor_consumer_node_indices`	Build a mapping from tensor names to the indices of nodes that use them.
`get_tensor_consumer_nodes`	Returns a dictionary of tensor name and their consumer node object mapping.
`get_tensor_from_name`	Returns a ValueInfoProto given a tensor name.
`get_tensor_producer_nodes`	Returns a dictionary of tensor name and their producer node object mapping.
`has_const_input`	Returns whether the given node has any constant input.
`has_path_type`	Checks if the given node is start/end of a given forward/backward path type.
`is_const_input`	Returns whether the given tensor is an initializer or produced by const-foldable nodes.
`match_fp8_mha_pattern`	Match FP8 fMHA v2 with the given softmax_op.

expand_node_names_from_patterns(graph, name_patterns=None)#

Expand the node names from the given patterns.

Parameters:

graph (GraphProto | Graph)
name_patterns (list[str] | None)

Return type:

list[str]

find_mha_partitions(graph)#: Match MHA: BMM1 -> … -> Softmax -> … -> BMM2.

get_fusible_backbone(node, graph)#

Returns the linear backbone node for a given node if it matches the pattern.

TensorRT fuses convolution with BN, Relu, MaxPool etc. when in some specific pattern. This rule tries to match some of those patterns. Note. BiasAdd and ConstMul are optional in path types.

Parameters:

node (Node) – Start node of the pattern.
graph (Graph) – ONNX model graph.

Returns:

Backbone node of the given node, None if not found.

Return type:

Node | None

get_tensor_consumer_node_indices(graph)#

Build a mapping from tensor names to the indices of nodes that use them.

Parameters:: graph (GraphProto | Graph) – ONNX GraphSurgeon graph to analyze
Returns:: Dictionary mapping tensor names to lists of node indices that consume them
Return type:: dict[str, list[int]]

get_tensor_consumer_nodes(graph)#

Returns a dictionary of tensor name and their consumer node object mapping.

Parameters:: graph (GraphProto) – ONNX model graph.
Returns:: Dictionary, key is tensor name and value is their consumer node object
Return type:: dict[str, list[NodeProto]]

get_tensor_from_name(graph, tensor_name)#

Returns a ValueInfoProto given a tensor name.

Parameters:

graph (GraphProto) – ONNX model graph
tensor_name (str) – String with tensor name.

Returns:

actual graph tensor.

Return type:

onnx.ValueInfoProto

get_tensor_producer_nodes(graph, get_initializer_producers=False)#

Returns a dictionary of tensor name and their producer node object mapping.

Note. we create a special Root type node as external inputs producer for ease of implementation.

Parameters:

graph (GraphProto) – ONNX model graph.
get_initializer_producers (bool)

Returns:

Dictionary, key is tensor name and value is their producer node object

Return type:

dict[str, NodeProto]

has_const_input(node)#

Returns whether the given node has any constant input.

Parameters:: node (Node)
Return type:: bool

has_path_type(node, graph, path_type, is_forward, wild_card_types=[], path_nodes=[])#

Checks if the given node is start/end of a given forward/backward path type.

Note, Path can be forward or backward wrt a node depending on the next level nodes. Additionally, this method can work with optional nodes and collect the traversed path.

Parameters:

node (Node) – Start node of the path.
graph (Graph) – ONNX model graph.
path_type (list[str]) – Path types to match from the given node.
is_forward (bool) – Whether to match forward or backward path.
wild_card_types (list[str]) – Wild card types, these type of nodes are skipped and not matched with the path_type.
path_nodes (list[Node]) – Accumulated nodes in the matched path.

Returns:

Bool, whether the given node is start/end of the given forward/backward path type.

Return type:

bool

is_const_input(tensor)#

Returns whether the given tensor is an initializer or produced by const-foldable nodes.

Parameters:: tensor (Tensor)
Return type:: bool

match_fp8_mha_pattern(graph, softmax_op, has_fp8_qdq)#

Match FP8 fMHA v2 with the given softmax_op.

If has_fp8_qdq == True, we match this FP8 fMHA v2 pattern: Q -> DQ -> BMM1 -> (Mul/Div) -> (Add) -> Softmax -> (Cast) -> Q -> DQ -> BMM2 -> Q -> DQ. If has_fp8_qdq == False, we match this FP8 fMHA v2 pattern: BMM1 -> (Mul/Div) -> (Add) -> Softmax -> (Cast) -> BMM2.

Parameters:

graph (Graph) – The graph to match FP8 MHA pattern.
softmax_op (Node) – The softmax op of FP8 MHA we want to match.
nodes_to_exclude – List of Nodes to exclude from quantization.
has_fp8_qdq (bool) – If True, match the FP8 MHA with Q/DQs. Else, match the FP8 MHA without Q/DQs.

Returns:

List of BMM1 node, Softmax node and BMM2 node.

Return type:

list[Node]