Add routed affine and positive Exp scale folding

[Qiong Wu (qiowu) (DingmaomaoBJTU)]

Summary

• extend conv-channel-affine-folding across nested static channel Split routes and parallel non-overlapping channel Slice routes
• add opt-in exp-positive-scale-folding for finite, strictly positive constant scales after Exp
• make the rewrites fail closed for dynamic or ambiguous routes, observed intermediates, custom domains, unloaded external data, definition collisions, and cyclic graphs

The implementation is generic ONNX graph rewriting. It contains no model-name, architecture, or execution-provider special cases.

CLI

The complete optimization now runs in one invocation:

winml optimize -m model.onnx -o model_opt.onnx `
  --enable-gather-slice-to-split-fusion `
  --enable-conv-channel-affine-folding `
  --enable-exp-positive-scale-folding

ORTGraphPipe restores eligible Split nodes first. AlgebraicRewritePipe then folds routed Conv affine operations and positive post-Exp scales.

For Exp(x) * C, the new capability requires floating-point C to be immutable, finite, strictly positive, and statically broadcastable. It either combines log(C) with an existing safe pre-Exp constant bias or replaces the post-Exp Mul with a pre-Exp Add(log(C)) while preserving element order through static shape-only views.

Validation

• 93 passed in tests/unit/optim/pipes/test_pipe_algebraic.py
• 640 passed, 16 skipped, 1 xfailed in tests/unit/optim
• Ruff check and format check pass
• independent review completed after graph-safety hardening

Generated-graph coverage includes CLI composition, nested Split, parallel Slice, broadcast and dtype behavior, float32/float64 constants, idempotence, custom domains, overridable initializers, external data, shared/captured/output tensors, malformed definitions, cycles, and route-depth limits.

Real-model evidence

A single CLI invocation on the motivating source model produced:

• nodes: 195 -> 181
• both static channel Split nodes restored
• seven routed Conv affine nodes and the positive post-Exp scale Mul removed
• exact eight-output public signature
• compact spatial log-bias initializer retained
• all eight CPU outputs bit-exact to the previously confirmed C1 graph on the fixed input

The graph-equivalent C1 candidate was previously confirmed on QNN NPU with 20 warmups and 200 measured iterations across four alternating AB/BA pairs:

• baseline mean p50: 27.167 ms
• candidate mean p50: 21.786 ms
• mean gain: 5.381 ms / 19.808%
• paired Student-t 95% CI: [5.261, 5.502] ms
• target hits: 4/4 at or below 25 ms

Quality evidence for that candidate passed canonical tensor replay and 40/40 local-smoke renders (mean PSNR 89.340 dB, minimum 83.508 dB). The final FFHQ quality gate remains pending.

[xieofxie]

could you add a anyone can understand example of exp-positive-scale-folding to prove its numeric correctness?

[Qiong Wu (qiowu) (DingmaomaoBJTU)]

Addressed in the latest commits by adding an explicit numeric example and by documenting the algebraic assumptions in the tests.

For exp-positive-scale-folding, the optimized graph rewrites this shape of computation:

biased = x + bias
original = Exp(biased) * scale

into the log-domain equivalent:

folded_bias = bias + log(scale)
rewritten = Exp(x + folded_bias)

The identity is elementwise, so broadcasting-safe constants work the same way:

Exp(x + bias) * scale
= Exp(x + bias) * Exp(log(scale))        because scale > 0
= Exp((x + bias) + log(scale))           because Exp(a) * Exp(b) = Exp(a + b)
= Exp(x + (bias + log(scale)))

The new test TestExpPositiveScaleFolding.test_simple_numeric_example_matches_log_domain_identity uses small readable values:

x         = [0, 2]
bias      = [1, -1]
log_scale = [2, 0.5]
scale     = Exp(log_scale)

original[0] = Exp(0 + 1) * Exp(2)   = Exp(3)
rewritten[0] = Exp(0 + 1 + 2)       = Exp(3)

original[1] = Exp(2 - 1) * Exp(0.5) = Exp(1.5)
rewritten[1] = Exp(2 - 1 + 0.5)     = Exp(1.5)

The test verifies both sides with ONNX Runtime, and also checks that the optimizer actually removed the post-Exp Mul and replaced the pre-Exp bias with [3, -0.5].

The implementation is guarded so the rewrite only happens when the scale is a finite floating-point constant, strictly positive, broadcast-compatible with the Exp output, and the computed bias + log(scale) remains finite. If any of those conditions fail, the graph is left unchanged.

For the other algebraic rewrites in the same update:

• conv-channel-affine-folding folds per-channel affine branches after Conv. For each output channel c, if Conv_c(x) = W_c * x + b_c and the branch applies Conv_c(x) * s_c + o_c, the folded parameters are W'_c = s_c * W_c and b'_c = s_c * b_c + o_c, so Conv'_c(x) = s_c * (W_c * x + b_c) + o_c, which is identical.
• The fallback sibling Slice -> Split rewrite only handles contiguous, non-overlapping, step-1 static slices that cover the full input extent along one axis. A Split with those segment lengths returns the same tensor intervals as the original sibling Slice nodes, preserving output names.

Validation run on the branch:

python -m ruff check src\winml\modelkit\optim\pipes\algebraic.py tests\unit\optim\pipes\test_pipe_algebraic.py
python -m pytest tests\unit\optim\pipes\test_pipe_algebraic.py::TestStaticSplitToSlice tests\unit\optim\pipes\test_pipe_algebraic.py::TestConvChannelAffineFolding tests\unit\optim\pipes\test_pipe_algebraic.py::TestExpPositiveScaleFolding -q
python -m mypy -p winml.modelkit

CI is also green on the latest push.

[xieofxie]

Found one graph-corruption bug and two additional correctness/resource issues in the new Exp scale folding path.

[xieofxie]

The stale-node fix and overflow-semantics disclosure are resolved. Two issues remain in scale-shape handling.

[xieofxie]

The reshape mapping and orthogonal-broadcast initializer issues are resolved. Two additional edge cases remain.

[xieofxie]

The shape-product overflow issue is resolved. One opset edge case remains.

[xieofxie]

The strict default-opset resolution issue is resolved. Four additional compatibility/performance issues remain.

[xieofxie]

The four compatibility fixes are resolved and the batched rewrites are graph-safe. The quadratic-runtime fix remains incomplete in two paths.

[xieofxie]

Per-prefix name allocation is resolved. Three issues remain in serial scale handling and Conv parameter synthesis.

[xieofxie]

The three prior findings are resolved. Three new issues remain in generated provenance and external Constant handling.

[xieofxie]

The three prior findings are resolved. One analysis-reporting issue remains.

[xieofxie]

Quick incremental review found one analysis-state regression.

[xieofxie]

Reviewed the latest head. Previous findings are resolved and the focused final pass found no remaining blockers.