yobx.xoptim.patterns.onnx_split#
- class yobx.xoptim.patterns.onnx_split.GathersSplitPattern(verbose: int = 0, priority: int = 1, min_opset: int = 1)[source]#
Merges multiple parallel gather into a split followed by unsqueeze.
Model with nodes to be fused:
graph TD classDef ioNode fill:#dfd,stroke:#333,color:#333 classDef initNode fill:#cccc00,stroke:#333,color:#333 classDef constNode fill:#f9f,stroke:#333,stroke-width:2px,color:#333 classDef opNode fill:#bbf,stroke:#333,stroke-width:2px,color:#333 I_X(["X FLOAT(a, 2)"]) Gather_0[["Gather(., 0, axis=1)"]] Gather_1[["Gather(., 1, axis=1)"]] I_X -->|"FLOAT(a, 2)"| Gather_0 I_X -->|"FLOAT(a, 2)"| Gather_1 O_x2(["x2 FLOAT(a)"]) Gather_1 --> O_x2 O_x1(["x1 FLOAT(a)"]) Gather_0 --> O_x1 class I_X,O_x2,O_x1 ioNode class Gather_0,Gather_1 opNodeOutcome of the fusion:
graph TD classDef ioNode fill:#dfd,stroke:#333,color:#333 classDef initNode fill:#cccc00,stroke:#333,color:#333 classDef constNode fill:#f9f,stroke:#333,stroke-width:2px,color:#333 classDef opNode fill:#bbf,stroke:#333,stroke-width:2px,color:#333 I_X(["X FLOAT(a, 2)"]) Split_0[["Split(., axis=1)"]] Squeeze_1[["Squeeze(., [1])"]] Squeeze_2[["Squeeze(., [1])"]] I_X -->|"FLOAT(a, 2)"| Split_0 Split_0 -->|"FLOAT(a, 1)"| Squeeze_1 Split_0 -->|"FLOAT(a, 1)"| Squeeze_2 O_x2(["x2 FLOAT(a)"]) Squeeze_2 --> O_x2 O_x1(["x1 FLOAT(a)"]) Squeeze_1 --> O_x1 class I_X,O_x2,O_x1 ioNode class Split_0,Squeeze_1,Squeeze_2 opNode- apply(g: GraphBuilder, *gather_nodes: NodeProto) List[NodeProto][source]#
The method does the rewriting. It assumes it can happen. It takes a list of nodes impacted by the rewriting assumes no other pattern optimizer will be modify them. It receives the list of nodes returned by method apply. Since it is a list of argument, method match can include None values. The method returns the new nodes. The optimizer considers that any node given to this function is removed from the graph, and any node returned by it are added. If a received node must be kept, it must be added to the list of returned node.
- Parameters:
nodes – nodes returned by method match, there are then removed
- Returns:
nodes to add to graph.
- match(g: GraphBuilderPatternOptimization, node: NodeProto, matched: List[MatchResult]) MatchResult | None[source]#
Determines nodes around node which can be rewritten.
- Parameters:
g – is a
GraphBuilderPatternOptimization, it holds all the existing nodes, is able to return any information about type, shape, the node before, the node after another one.node – the matching must determine if some nodes around this one are part of set of nodes this pattern optimizer can rewrite. From there, the function explores wherever it needs, checking any condition it needs.
matched – usually unused, it returns of nodes already matching a pattern
The method must not modify the graph. The method returns None if no match is found or an instance of class
MatchResult. It must contain:a list of nodes involved in the rewriting. It does not mean all of them will be removed but all of them are needed to do the rewriting and must not be impacted by other pattern optimizer.
A function doing the rewriting (usually method apply of the pattern class).
An existing node where the rewritten nodes can be inserted. Knowing it makes it faster to rewriter. If not specified, the optimizer will automatically determine the position of the new nodes.
- class yobx.xoptim.patterns.onnx_split.SlicesSplitPattern(verbose: int = 0, priority: int = 1, min_opset: int = 1)[source]#
Merges multiple parallel slices into a split.
Model with nodes to be fused:
graph TD classDef ioNode fill:#dfd,stroke:#333,color:#333 classDef initNode fill:#cccc00,stroke:#333,color:#333 classDef constNode fill:#f9f,stroke:#333,stroke-width:2px,color:#333 classDef opNode fill:#bbf,stroke:#333,stroke-width:2px,color:#333 I_transpose_1(["transpose_1 FLOAT16(2, 2, 1024, 512)"]) Slice_0[["Slice(., [0], [256], [3])"]] Slice_1[["Slice(., [256], [9223372036854775807], [3])"]] I_transpose_1 -->|"FLOAT16(2, 2, 1024, 512)"| Slice_0 I_transpose_1 -->|"FLOAT16(2, 2, 1024, 512)"| Slice_1 O_slice_11(["slice_11 FLOAT16(2, 2, 1024, 256)"]) Slice_0 --> O_slice_11 O_slice_12(["slice_12 FLOAT16(2, 2, 1024, 256)"]) Slice_1 --> O_slice_12 class I_transpose_1,O_slice_11,O_slice_12 ioNode class Slice_0,Slice_1 opNodeOutcome of the fusion:
graph TD classDef ioNode fill:#dfd,stroke:#333,color:#333 classDef initNode fill:#cccc00,stroke:#333,color:#333 classDef constNode fill:#f9f,stroke:#333,stroke-width:2px,color:#333 classDef opNode fill:#bbf,stroke:#333,stroke-width:2px,color:#333 I_transpose_1(["transpose_1 FLOAT16(2, 2, 1024, 512)"]) Split_0[["Split(., [256, 256], axis=3)"]] I_transpose_1 -->|"FLOAT16(2, 2, 1024, 512)"| Split_0 O_slice_11(["slice_11 FLOAT16(2, 2, 1024, 256)"]) Split_0 --> O_slice_11 O_slice_12(["slice_12 FLOAT16(2, 2, 1024, 256)"]) Split_0 --> O_slice_12 class I_transpose_1,O_slice_11,O_slice_12 ioNode class Split_0 opNode- apply(g: GraphBuilder, *nodes: NodeProto) List[NodeProto][source]#
The method does the rewriting. It assumes it can happen. It takes a list of nodes impacted by the rewriting assumes no other pattern optimizer will be modify them. It receives the list of nodes returned by method apply. Since it is a list of argument, method match can include None values. The method returns the new nodes. The optimizer considers that any node given to this function is removed from the graph, and any node returned by it are added. If a received node must be kept, it must be added to the list of returned node.
- Parameters:
nodes – nodes returned by method match, there are then removed
- Returns:
nodes to add to graph.
- match(g: GraphBuilderPatternOptimization, node: NodeProto, matched: List[MatchResult]) MatchResult | None[source]#
Determines nodes around node which can be rewritten.
- Parameters:
g – is a
GraphBuilderPatternOptimization, it holds all the existing nodes, is able to return any information about type, shape, the node before, the node after another one.node – the matching must determine if some nodes around this one are part of set of nodes this pattern optimizer can rewrite. From there, the function explores wherever it needs, checking any condition it needs.
matched – usually unused, it returns of nodes already matching a pattern
The method must not modify the graph. The method returns None if no match is found or an instance of class
MatchResult. It must contain:a list of nodes involved in the rewriting. It does not mean all of them will be removed but all of them are needed to do the rewriting and must not be impacted by other pattern optimizer.
A function doing the rewriting (usually method apply of the pattern class).
An existing node where the rewritten nodes can be inserted. Knowing it makes it faster to rewriter. If not specified, the optimizer will automatically determine the position of the new nodes.
- class yobx.xoptim.patterns.onnx_split.SplitConcatPattern(verbose: int = 0, priority: int = 1, min_opset: int = 1)[source]#
Replaces Split + Concat into identity if this is equivalent.
Model with nodes to be fused:
graph TD classDef ioNode fill:#dfd,stroke:#333,color:#333 classDef initNode fill:#cccc00,stroke:#333,color:#333 classDef constNode fill:#f9f,stroke:#333,stroke-width:2px,color:#333 classDef opNode fill:#bbf,stroke:#333,stroke-width:2px,color:#333 I_X(["X FLOAT(a, b)"]) Split_0[["Split(., axis=-1)"]] Concat_1[["Concat(., ., axis=-1)"]] I_X -->|"FLOAT(a, b)"| Split_0 Split_0 -->|"FLOAT(a, CeilToInt(b,2))"| Concat_1 Split_0 -->|"FLOAT(a, b-CeilToInt(b,2))"| Concat_1 O_Y(["Y FLOAT(a, b)"]) Concat_1 --> O_Y class I_X,O_Y ioNode class Split_0,Concat_1 opNodeOutcome of the fusion:
graph TD classDef ioNode fill:#dfd,stroke:#333,color:#333 classDef initNode fill:#cccc00,stroke:#333,color:#333 classDef constNode fill:#f9f,stroke:#333,stroke-width:2px,color:#333 classDef opNode fill:#bbf,stroke:#333,stroke-width:2px,color:#333 I_X(["X FLOAT(a, b)"]) Identity_0[["Identity(.)"]] I_X -->|"FLOAT(a, b)"| Identity_0 O_Y(["Y FLOAT(a, b)"]) Identity_0 --> O_Y class I_X,O_Y ioNode class Identity_0 opNode- apply(g: GraphBuilder, split_node: NodeProto, concat_node: NodeProto) List[NodeProto][source]#
The method does the rewriting. It assumes it can happen. It takes a list of nodes impacted by the rewriting assumes no other pattern optimizer will be modify them. It receives the list of nodes returned by method apply. Since it is a list of argument, method match can include None values. The method returns the new nodes. The optimizer considers that any node given to this function is removed from the graph, and any node returned by it are added. If a received node must be kept, it must be added to the list of returned node.
- Parameters:
nodes – nodes returned by method match, there are then removed
- Returns:
nodes to add to graph.
- match(g: GraphBuilderPatternOptimization, node: NodeProto, matched: List[MatchResult]) MatchResult | None[source]#
Determines nodes around node which can be rewritten.
- Parameters:
g – is a
GraphBuilderPatternOptimization, it holds all the existing nodes, is able to return any information about type, shape, the node before, the node after another one.node – the matching must determine if some nodes around this one are part of set of nodes this pattern optimizer can rewrite. From there, the function explores wherever it needs, checking any condition it needs.
matched – usually unused, it returns of nodes already matching a pattern
The method must not modify the graph. The method returns None if no match is found or an instance of class
MatchResult. It must contain:a list of nodes involved in the rewriting. It does not mean all of them will be removed but all of them are needed to do the rewriting and must not be impacted by other pattern optimizer.
A function doing the rewriting (usually method apply of the pattern class).
An existing node where the rewritten nodes can be inserted. Knowing it makes it faster to rewriter. If not specified, the optimizer will automatically determine the position of the new nodes.