feat(optim): add algebraic graph rewrites#1126
Conversation
Add opt-in static Split-to-Slice conversion and safe Conv affine/BatchNormalization folding with topology-based guards and generated-graph coverage.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| from typing import Any, ClassVar | ||
|
|
||
| import numpy as np | ||
| import onnx |
| if not self.should_process(config): | ||
| return model | ||
|
|
||
| import onnx |
| if not self.should_process(config): | ||
| return model | ||
|
|
||
| import onnx |
| from typing import TYPE_CHECKING | ||
|
|
||
| import numpy as np | ||
| import onnx |
xieofxie
left a comment
xieofxie
left a comment
There was a problem hiding this comment.
Deep-dived the three rewrites (math + topology guards), ran the new suite (82 passed) and ruff (clean), and probed edge cases empirically. The core algebra is correct and the guards (graph outputs, subgraph captures, shared consumers, overlapping/routed channel intervals, static-vs-dynamic split sizes, repeated runs) are impressively thorough. The channel-affine and Conv+Add+BN folds are numerically sound, and idempotency holds.
Main thing worth discussing is the full-static-shape requirement, which quietly disables the affine fold on any dynamic-batch model. A few smaller robustness/cleanliness notes inline. None are blocking correctness bugs — they're about applicability, precision on low-precision models, and code hygiene.
| conv = index.producers.get(conv_output) | ||
| if conv is None or conv.op_type != "Conv": | ||
| continue | ||
| conv_shape = _static_shape(index, conv_output) |
There was a problem hiding this comment.
_static_shape requires every dimension (including batch) to be static, so the whole affine fold bails on any model with a dynamic batch dim — which is extremely common. Channel-wise folding only needs the channel axis (conv_shape[1]) to be known; the direct Conv → Mul/Add(channel-const) path is mathematically valid for any N/H/W. Was requiring full-static shapes intentional, or could the direct path relax to only needing a static channel dim? As written the optimization silently no-ops for most non-fixed-batch models.
Verified empirically: Conv → Mul([1,C,1,1]) with input [None,1,3,3] is left unfolded, while the identical fully-static graph folds.
| denominator = variance + epsilon | ||
| if epsilon < 0 or np.any(denominator <= 0) or not np.all(np.isfinite(denominator)): | ||
| continue | ||
| gamma = scale / np.sqrt(denominator) |
There was a problem hiding this comment.
The BN math runs entirely in the weight dtype. Because of the guards above, scale/variance/mean/beta are all weight_dtype, so for an fp16 model denominator, np.sqrt, gamma, and the new_bias = bias*gamma + beta - gamma*mean in _copy_conv_parameters_for_bn all execute in fp16. That contrasts with the affine fold, which deliberately upcasts to result_type(weight_dtype, float32) (line 857) before folding. ONNX BatchNormalization is typically evaluated with an fp32 accumulator, so folding at fp16 can drift beyond tolerance on the low-precision models this PR targets for QNN. Consider computing gamma/new_bias in result_type(weight_dtype, float32) and casting back to match the affine path.
| return | ||
| replacements: dict[int, list[onnx.NodeProto]] = {} | ||
| for split in list(model.graph.node): | ||
| if split.op_type != "Split" or len(split.input) < 1 or not split.input[0]: |
There was a problem hiding this comment.
There's no domain guard here (same for the Conv/Add/BatchNormalization checks in the fold functions). An op is (domain, op_type) — a custom-domain node literally named Split (or Conv/BatchNormalization) would be rewritten with standard-op semantics. ONNX/QNN models can carry non-"" domain nodes. Cheap to add node.domain in ("", "ai.onnx") to each pattern match for safety.
| result = onnx.ModelProto() | ||
| result.CopyFrom(model) | ||
| allocator = _NameAllocator(result) | ||
| introduced: set[str] = set() |
There was a problem hiding this comment.
introduced is populated (every _new_initializer does introduced.add(name)) but never read — only introduced_nodes is consumed, by _prune_introduced_nodes, and initializer cleanup goes through _prune_unused_initializers (reference counting) instead. So this set is dead state threaded through all three folds. Either drop it or wire it into the pruning if it was meant to bound initializer removal.
| if input_shape is None or len(node.output) == 0: | ||
| return None | ||
|
|
||
| axis_input_index = 2 if len(node.input) > 2 else None |
There was a problem hiding this comment.
Minor: Split has no axis input in any opset — axis is always an attribute, and the only optional input is split (index 1). Treating input index 2 as a possible axis source is dead/defensive code that misreads the schema and could mislead a future reader. Reading axis purely from the attribute would be clearer.
2 participants
Summary
Splitnodes withSlicenodesConvweights and biasConv + static Add + BatchNormalizationpatternsQNN result
On the evaluated fixed-batch U-Net, the combined rewrites plus
htp_graph_finalization_optimization_mode=3reduced p50 latency from 31.125 ms to 22.729 ms and reduced QNN DDR traffic by roughly 30%.