Skip to content

Add fuse_matmul_add_bias_into_gemm_batched pass - #6

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

Add fuse_matmul_add_bias_into_gemm_batched pass#6
take-cheeze wants to merge 1 commit into
claude/initializers-non-constants-option-fun941from
claude/onnx-optimizer-passes-tx9y3n-gemm-batched

Conversation

@take-cheeze

Copy link
Copy Markdown
Member

Implements the fuse_matmul_add_bias_into_gemm_batched pass from the docs/proposed_passes design draft. Split out of #4 so each pass is reviewed on its own.

Base branch: like #4, this is stacked on claude/initializers-non-constants-option-fun941 so its diff shows only the new pass. Retarget to main once that branch merges.

What it does

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. This means the pass does not change the default optimization pipeline's behaviour.

Tests

Adds unit tests — the fuse case, the dynamic-leading-dims case, and no-fuse guards (rank-2 MatMul, non-constant weight) — with onnxruntime numeric-equivalence checks where applicable.


Generated by Claude Code

Rewrites batched MatMul(X[>=3D], W[K,N]) + b into
Reshape(X,[-1,K]) -> Gemm(X2, W, b) -> 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 faster,
so it is invoked explicitly by name. Adds unit tests for the fuse case,
the dynamic-shape case, and the no-fuse guards (rank-2, non-constant
weight).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01V63N6PXYEgNEe1BWbKi6cU
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