Skip to content

feat(optim): add static Split-to-Slice rewrite#1142

Merged
xieofxie merged 4 commits into
mainfrom
hualxie/add_splice_split
Jul 22, 2026
Merged

feat(optim): add static Split-to-Slice rewrite#1142
xieofxie merged 4 commits into
mainfrom
hualxie/add_splice_split

Conversation

@xieofxie

Copy link
Copy Markdown
Contributor

Summary

Add an opt-in static-split-to-slice algebraic rewrite that replaces eligible ONNX Split nodes with equivalent input-form Slice nodes while preserving output tensor names.

The rewrite only applies when the input shape, axis, and split boundaries are statically known. It supports explicit split sizes and equal splits, avoids malformed or dynamic cases, handles name collisions, preserves nested-subgraph captures, removes dead generated nodes and constants, and remains idempotent through the public optimization path.

Why this is not a pattern rewrite

The existing pattern framework is designed for fixed-topology replacements with compatible source and target schemas. Its match result exposes a single pattern output, and PatternRewriter constructs the replacement using that fixed output mapping.

A Split node is variadic: one matched node may have any number of outputs, and the replacement must generate one Slice node plus four boundary initializers for each output. The replacement topology and constants therefore depend on the specific matched node rather than on a fixed target pattern. Representing this correctly as a pattern would first require broader framework support for variadic output mappings and match-parameterized replacement topology.

For that reason, this transformation is implemented as a focused algebraic rewrite rather than adding model-specific behavior or forcing it into the current fixed-shape pattern abstraction.

@xieofxie
xieofxie requested a review from a team as a code owner July 20, 2026 08:21
Comment thread src/winml/modelkit/optim/pipes/algebraic.py Fixed
Comment thread tests/unit/optim/pipes/test_pipe_algebraic.py Fixed

@DingmaomaoBJTU DingmaomaoBJTU left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed the static Split-to-Slice rewrite. Overall design looks solid: opt-in flag, conservative bail-out for dynamic/malformed cases, dead-code pruning, and nested-subgraph capture handling are all done carefully. Left two non-blocking inline notes.

Comment thread src/winml/modelkit/optim/pipes/algebraic.py Outdated
Comment thread src/winml/modelkit/optim/pipes/algebraic.py Outdated
@xieofxie
xieofxie merged commit 76c74a7 into main Jul 22, 2026
9 checks passed
@xieofxie
xieofxie deleted the hualxie/add_splice_split branch July 22, 2026 02:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants