Add fuse_matmul_add_bias_into_gemm_batched pass - #6
Open
take-cheeze wants to merge 1 commit into
Open
Conversation
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Implements the
fuse_matmul_add_bias_into_gemm_batchedpass from thedocs/proposed_passesdesign draft. Split out of #4 so each pass is reviewed on its own.What it does
Rewrites batched
MatMul(X[>=3D], W[K,N]) + binto:extending the 2-D-only
fuse_matmul_add_bias_into_gemmto the rank-≥3 activations transformer linear layers use. Static leading dims emit a constant output shape; dynamic leading dims rebuild it viaShape/Slice/Concat(requires opset ≥ 10).Wmust be a 2-D constant andba 1-D bias broadcastable over theNaxis.Registered as
PassType::Otherso it stays out of the default fuse set (GetFuseAndEliminationPass) — it is a graph-shape rewrite that trades a batched MatMul forReshape + Gemm + Reshapeand 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