Skip to content

Add fuse_consecutive_mul and batched MatMul+Add->Gemm passes - #4

Closed
take-cheeze wants to merge 1 commit into
claude/initializers-non-constants-option-fun941from
claude/onnx-optimizer-passes-tx9y3n
Closed

Add fuse_consecutive_mul and batched MatMul+Add->Gemm passes#4
take-cheeze wants to merge 1 commit into
claude/initializers-non-constants-option-fun941from
claude/onnx-optimizer-passes-tx9y3n

Conversation

@take-cheeze

Copy link
Copy Markdown
Member

Implements two passes from the docs/proposed_passes design drafts.

Base branch: this PR is stacked on claude/initializers-non-constants-option-fun941 so its diff shows only the two new passes. onnxsim's submodule (and its SetInitializersAsConstants usage) already depends on that branch's commit, so basing here keeps the submodule bump non-regressing. Once that branch merges, retarget this PR to main.

fuse_consecutive_mul (default Fuse pass)

Folds Mul(Mul(X, C1), C2) -> Mul(X, C1 * C2) when C1/C2 are constants (initializers or Constant nodes) and the inner Mul feeds only the outer one. The combined scale is materialised with numpy-style broadcasting, so the rewrite is numerically identical (X * C1 * C2 at every position). Covers e.g. PoolFormer's LayerScale, exported as a per-channel (C,1,1) scale times a scalar factor. Conservatively limited to FLOAT/DOUBLE constants; other dtypes are left to constant folding.

fuse_matmul_add_bias_into_gemm_batched (PassType::Other, opt-in)

Rewrites batched MatMul(X[>=3D], W[K,N]) + b into:

X2 = Reshape(X, [-1, K])
G  = Gemm(X2, W, b)      # alpha=beta=1, transA=transB=0
A  = Reshape(G, [d0, ..., d_{r-2}, N])

extending the 2-D-only fuse_matmul_add_bias_into_gemm to the rank-≥3 activations transformer linear layers use. Static leading dims emit a constant output shape; dynamic leading dims rebuild it via Shape/Slice/Concat (requires opset ≥ 10). W must be a 2-D constant and b a 1-D bias broadcastable over the N axis.

Registered as PassType::Other so it stays out of the default fuse set (GetFuseAndEliminationPass) — it is a graph-shape rewrite that trades a batched MatMul for Reshape + Gemm + Reshape and is not guaranteed to be faster, so it is invoked explicitly by name.

Tests

Adds unit tests for both passes — fuse cases, no-fuse guards (multiple-use, non-constant operand, rank-2, non-constant weight), per-channel and dynamic-shape cases — with onnxruntime numeric-equivalence checks. Full suite: 175 passed, 5 skipped (torchvision, torch not installed), 6 pre-existing xfails.


Generated by Claude Code

Implements two passes from the proposed_passes design drafts:

fuse_consecutive_mul (default Fuse pass):
  Folds Mul(Mul(X, C1), C2) -> Mul(X, C1*C2) when C1 and C2 are constants
  and the inner Mul feeds only the outer one. The combined scale is
  materialised with numpy-style broadcasting, so the rewrite is numerically
  equivalent (X*C1*C2 at every position). Covers e.g. PoolFormer LayerScale
  exported as a per-channel (C,1,1) scale times a scalar factor.

fuse_matmul_add_bias_into_gemm_batched (PassType::Other, opt-in only):
  Rewrites batched MatMul(X[>=3D], W[K,N]) + b into
  Reshape(->[-1,K]) -> Gemm -> Reshape(->[...,N]), extending the 2-D-only
  fuse_matmul_add_bias_into_gemm to the rank>=3 activations used by
  transformer linear layers. Static leading dims emit a constant output
  shape; dynamic dims rebuild it via Shape/Slice/Concat. Registered as
  PassType::Other so it stays out of the default fuse set, since it is a
  graph-shape rewrite that is not guaranteed to be faster.

Adds unit tests for both passes (fuse, no-fuse guards, per-channel and
dynamic-shape cases) with onnxruntime numeric-equivalence checks.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01V63N6PXYEgNEe1BWbKi6cU
@take-cheeze
take-cheeze force-pushed the claude/onnx-optimizer-passes-tx9y3n branch from 55fb61a to c9a482c Compare July 29, 2026 22:13
@take-cheeze

Copy link
Copy Markdown
Member Author

Superseded by #5 (fuse_consecutive_mul) and #6 (fuse_matmul_add_bias_into_gemm_batched) — the two passes are now split into one PR each so they can be reviewed and merged independently. Closing this combined PR.

Note: I'm intentionally not deleting the claude/onnx-optimizer-passes-tx9y3n branch. onnxsim/onnxsim's submodule was already bumped to this branch's tip (c9a482c, which contains both passes) in onnxsim/onnxsim#523, so the commit needs to stay reachable until that submodule is repointed at the merged commits.


Generated by Claude Code

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.

2 participants