Skip to content

Commit ccf1073

Browse files
claudeclaude
committed
Add fuse_matmul_add_bias_into_gemm_batched pass
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
1 parent 725a177 commit ccf1073

3 files changed

Lines changed: 307 additions & 0 deletions

File tree

‎onnxoptimizer/pass_registry.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
#include "onnxoptimizer/passes/fuse_consecutive_squeezes.h"
4949
#include "onnxoptimizer/passes/fuse_consecutive_transposes.h"
5050
#include "onnxoptimizer/passes/fuse_matmul_add_bias_into_gemm.h"
51+
#include "onnxoptimizer/passes/fuse_matmul_add_bias_into_gemm_batched.h"
5152
#include "onnxoptimizer/passes/fuse_pad_into_conv.h"
5253
#include "onnxoptimizer/passes/fuse_pad_into_pool.h"
5354
#include "onnxoptimizer/passes/fuse_transpose_into_gemm.h"
@@ -102,6 +103,7 @@ struct GlobalPassRegistry {
102103
registerPass<FuseConsecutiveSqueezeUnsqueeze>();
103104
registerPass<FuseConsecutiveTransposes>();
104105
registerPass<FuseMatMulAddBiasIntoGemm>();
106+
registerPass<FuseMatMulAddBiasIntoGemmBatched>();
105107
registerPass<FusePadIntoConv>();
106108
registerPass<FusePadIntoPool>();
107109
registerPass<FuseTransposeIntoGemm>();

‎onnxoptimizer/passes/fuse_matmul_add_bias_into_gemm_batched.h‎

Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
1+
// Copyright (c) ONNX Project Contributors
2+
//
3+
// SPDX-License-Identifier: Apache-2.0
4+
5+
// ATTENTION: The code in this file is highly EXPERIMENTAL.
6+
// Adventurous users should note that the APIs will probably change.
7+
8+
#pragma once
9+
10+
// Before:
11+
// Z = MatMul(X, W) // X rank >= 3, W a 2-D constant [K, N]
12+
// A = Z + b // b broadcastable over the last (N) axis
13+
// After:
14+
// X2 = Reshape(X, [-1, K])
15+
// G = Gemm(X2, W, b) // alpha=beta=1, transA=transB=0 -> [-1, N]
16+
// A = Reshape(G, [d0, ..., d_{r-2}, N])
17+
//
18+
// `fuse_matmul_add_bias_into_gemm` only fuses the 2-D case. Transformer models
19+
// apply linear layers to rank-3 activations `[B, S, K] . [K, N]`, so the MatMul
20+
// is batched and that pass bails. This pass collapses the leading dims to a
21+
// single 2-D Gemm and reshapes back.
22+
//
23+
// NOTE: This is a node-count / graph-shape rewrite (it trades a batched MatMul
24+
// for Reshape + Gemm + Reshape). Batched MatMul and 2-D Gemm are equivalent
25+
// work and modern runtimes execute batched MatMul natively, so this is not
26+
// guaranteed to be faster. It is registered as PassType::Other so it is NOT in
27+
// the default fuse set; invoke it explicitly by name when a Gemm-centric graph
28+
// is wanted.
29+
30+
#include <vector>
31+
32+
#include "onnx/common/assertions.h"
33+
#include "onnxoptimizer/pass.h"
34+
#include "onnxoptimizer/passes/pass_util.h"
35+
36+
namespace ONNX_NAMESPACE {
37+
namespace optimization {
38+
39+
struct FuseMatMulAddBiasIntoGemmBatched final : public PredicateBasedPass {
40+
explicit FuseMatMulAddBiasIntoGemmBatched()
41+
: PredicateBasedPass(PassType::Other, PassEfficiency::Complete,
42+
PassOptimizationType::Compute) {}
43+
std::string getPassName() const override {
44+
return "fuse_matmul_add_bias_into_gemm_batched";
45+
}
46+
47+
static Value* MakeInt64Constant(Graph& graph, std::vector<int64_t> data) {
48+
Tensor t;
49+
t.sizes().push_back(static_cast<int64_t>(data.size()));
50+
t.elem_type() = TensorProto_DataType_INT64;
51+
t.int64s() = std::move(data);
52+
return graph.addInitializerAndCreateValue(t);
53+
}
54+
55+
bool patternMatchPredicate(Node* node) override {
56+
if (!CheckKind(node, kAdd, 0, kMatMul)) {
57+
return false;
58+
}
59+
Value* matmul_out = node->input(0);
60+
// MatMul result must feed only this Add.
61+
if (matmul_out->uses().size() > 1) {
62+
return false;
63+
}
64+
Node* matmul = matmul_out->node();
65+
Value* x = matmul->input(0);
66+
Value* w = matmul->input(1);
67+
68+
// X: rank >= 3 with a static trailing dim K.
69+
if (!x->has_sizes()) {
70+
return false;
71+
}
72+
const auto& x_shape = x->sizes();
73+
if (x_shape.size() < 3 || !x_shape.back().is_int) {
74+
return false;
75+
}
76+
const int64_t k = x_shape.back().dim;
77+
78+
// W: a 2-D constant [K, N] with static dims, matching K.
79+
if (!IsConstantTensor(w) || !w->has_sizes()) {
80+
return false;
81+
}
82+
const auto& w_shape = w->sizes();
83+
if (w_shape.size() != 2 || !w_shape[0].is_int || !w_shape[1].is_int) {
84+
return false;
85+
}
86+
if (w_shape[0].dim != k) {
87+
return false;
88+
}
89+
const int64_t n = w_shape[1].dim;
90+
91+
// bias: 1-D, broadcastable over the N axis only ([N] or [1]).
92+
Value* bias = node->input(1);
93+
if (!bias->has_sizes()) {
94+
return false;
95+
}
96+
const auto& bias_shape = bias->sizes();
97+
if (bias_shape.size() != 1 || !bias_shape[0].is_int) {
98+
return false;
99+
}
100+
if (bias_shape[0].dim != n && bias_shape[0].dim != 1) {
101+
return false;
102+
}
103+
104+
// Dynamic leading dims need Shape/Slice, i.e. opset >= 10.
105+
bool leading_static = true;
106+
for (size_t i = 0; i + 1 < x_shape.size(); ++i) {
107+
leading_static &= x_shape[i].is_int;
108+
}
109+
if (!leading_static) {
110+
const int opset = getOpsetVersion(*node->owningGraph());
111+
if (opset != 0 && opset < 10) {
112+
return false;
113+
}
114+
}
115+
return true;
116+
}
117+
118+
bool runTransform(Node* n, Graph& graph,
119+
NodeDestroyType& destroy_current) override {
120+
destroy_current = NodeDestroyType::DestroyZero;
121+
122+
Node* matmul = n->input(0)->node();
123+
Value* x = matmul->input(0);
124+
Value* w = matmul->input(1);
125+
Value* bias = n->input(1);
126+
127+
const auto& x_shape = x->sizes();
128+
const int64_t rank = static_cast<int64_t>(x_shape.size());
129+
const int64_t k = x_shape.back().dim;
130+
const int64_t out_n = w->sizes()[1].dim;
131+
132+
// X2 = Reshape(X, [-1, K])
133+
Node* pre = graph.create(kReshape, 1);
134+
pre->addInput(x);
135+
pre->addInput(MakeInt64Constant(graph, {-1, k}));
136+
137+
// G = Gemm(X2, W, bias)
138+
Node* gemm = graph.create(kGemm, 1);
139+
gemm->addInput(pre->output());
140+
gemm->addInput(w);
141+
gemm->addInput(bias);
142+
gemm->f_(kalpha, 1.0);
143+
gemm->f_(kbeta, 1.0);
144+
gemm->i_(ktransA, 0);
145+
gemm->i_(ktransB, 0);
146+
147+
// Reconstruct the output shape [d0, ..., d_{r-2}, N].
148+
bool leading_static = true;
149+
std::vector<int64_t> leading_dims;
150+
for (int64_t i = 0; i + 1 < rank; ++i) {
151+
leading_static &= x_shape[i].is_int;
152+
if (x_shape[i].is_int) {
153+
leading_dims.push_back(x_shape[i].dim);
154+
}
155+
}
156+
157+
// A = Reshape(G, out_shape)
158+
Node* post = graph.create(kReshape, n->outputs().size());
159+
post->addInput(gemm->output());
160+
161+
// Order nodes: pre -> gemm -> [shape ops] -> post -> n.
162+
pre->insertBefore(n);
163+
gemm->insertBefore(n);
164+
165+
if (leading_static) {
166+
std::vector<int64_t> out_shape = leading_dims;
167+
out_shape.push_back(out_n);
168+
post->addInput(MakeInt64Constant(graph, std::move(out_shape)));
169+
} else {
170+
// shape(X) -> slice off the last dim -> concat with [N]
171+
Node* shape = graph.create(Symbol("Shape"), 1);
172+
shape->addInput(x);
173+
shape->insertBefore(n);
174+
175+
Node* slice = graph.create(kSlice, 1);
176+
slice->addInput(shape->output());
177+
slice->addInput(MakeInt64Constant(graph, {0})); // starts
178+
slice->addInput(MakeInt64Constant(graph, {rank - 1})); // ends
179+
slice->addInput(MakeInt64Constant(graph, {0})); // axes
180+
slice->insertBefore(n);
181+
182+
Node* concat = graph.create(kConcat, 1);
183+
concat->addInput(slice->output());
184+
concat->addInput(MakeInt64Constant(graph, {out_n}));
185+
concat->i_(kaxis, 0);
186+
concat->insertBefore(n);
187+
188+
post->addInput(concat->output());
189+
}
190+
191+
for (int i = 0; i < static_cast<int>(n->outputs().size()); ++i) {
192+
post->outputs()[i]->copyMetadata(n->outputs()[i]);
193+
}
194+
post->insertBefore(n);
195+
196+
if (!tryReplacingAllUsesWith(n, post)) {
197+
return false;
198+
}
199+
// Destroy the Add; the now-dead MatMul is cleaned up by DCE.
200+
destroy_current = NodeDestroyType::DestroyOne;
201+
return true;
202+
}
203+
};
204+
205+
} // namespace optimization
206+
} // namespace ONNX_NAMESPACE

‎onnxoptimizer/test/optimizer_test.py‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1649,6 +1649,105 @@ def test_fuse_matmul_add_bias_into_gemm_multiple_use_no_fuse(self):
16491649

16501650
assert optimized_model.graph == graph
16511651

1652+
def test_fuse_matmul_add_bias_into_gemm_batched(self): # type: () -> None
1653+
matmul = helper.make_node("MatMul", ["X", "W"], ["Z"])
1654+
add = helper.make_node("Add", ["Z", "B"], ["A"])
1655+
w = numpy_helper.from_array(
1656+
np.random.randn(4, 5).astype(np.float32), name="W"
1657+
)
1658+
b = numpy_helper.from_array(np.random.randn(5).astype(np.float32), name="B")
1659+
graph = helper.make_graph(
1660+
[matmul, add],
1661+
"test",
1662+
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3, 4))],
1663+
[helper.make_tensor_value_info("A", TensorProto.FLOAT, (2, 3, 5))],
1664+
initializer=[w, b],
1665+
)
1666+
optimized_model = self._optimized(
1667+
graph,
1668+
["fuse_matmul_add_bias_into_gemm_batched", "eliminate_deadend"],
1669+
)
1670+
1671+
op_types = [n.op_type for n in optimized_model.graph.node]
1672+
assert op_types == ["Reshape", "Gemm", "Reshape"]
1673+
assert "MatMul" not in op_types
1674+
1675+
def test_fuse_matmul_add_bias_into_gemm_batched_dynamic(self):
1676+
# dynamic leading dims -> Shape/Slice/Concat rebuild the output shape
1677+
matmul = helper.make_node("MatMul", ["X", "W"], ["Z"])
1678+
add = helper.make_node("Add", ["Z", "B"], ["A"])
1679+
w = numpy_helper.from_array(
1680+
np.random.randn(4, 5).astype(np.float32), name="W"
1681+
)
1682+
b = numpy_helper.from_array(np.random.randn(5).astype(np.float32), name="B")
1683+
graph = helper.make_graph(
1684+
[matmul, add],
1685+
"test",
1686+
[
1687+
helper.make_tensor_value_info(
1688+
"X", TensorProto.FLOAT, ("batch", "seq", 4)
1689+
)
1690+
],
1691+
[
1692+
helper.make_tensor_value_info(
1693+
"A", TensorProto.FLOAT, ("batch", "seq", 5)
1694+
)
1695+
],
1696+
initializer=[w, b],
1697+
)
1698+
optimized_model = self._optimized(
1699+
graph,
1700+
["fuse_matmul_add_bias_into_gemm_batched", "eliminate_deadend"],
1701+
compare_result=False,
1702+
)
1703+
1704+
op_types = [n.op_type for n in optimized_model.graph.node]
1705+
assert "Gemm" in op_types
1706+
assert "MatMul" not in op_types
1707+
assert op_types[-1] == "Reshape"
1708+
1709+
def test_fuse_matmul_add_bias_into_gemm_batched_2d_no_fuse(self):
1710+
# rank-2 matmul is handled by the non-batched pass; this one must skip it
1711+
matmul = helper.make_node("MatMul", ["X", "W"], ["Z"])
1712+
add = helper.make_node("Add", ["Z", "B"], ["A"])
1713+
w = numpy_helper.from_array(
1714+
np.random.randn(4, 5).astype(np.float32), name="W"
1715+
)
1716+
b = numpy_helper.from_array(np.random.randn(5).astype(np.float32), name="B")
1717+
graph = helper.make_graph(
1718+
[matmul, add],
1719+
"test",
1720+
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (3, 4))],
1721+
[helper.make_tensor_value_info("A", TensorProto.FLOAT, (3, 5))],
1722+
initializer=[w, b],
1723+
)
1724+
optimized_model = self._optimized(
1725+
graph, ["fuse_matmul_add_bias_into_gemm_batched"]
1726+
)
1727+
1728+
assert [n.op_type for n in optimized_model.graph.node] == ["MatMul", "Add"]
1729+
1730+
def test_fuse_matmul_add_bias_into_gemm_batched_nonconst_weight_no_fuse(self):
1731+
# W is a runtime input, not a constant -> cannot form a Gemm weight
1732+
matmul = helper.make_node("MatMul", ["X", "W"], ["Z"])
1733+
add = helper.make_node("Add", ["Z", "B"], ["A"])
1734+
b = numpy_helper.from_array(np.random.randn(5).astype(np.float32), name="B")
1735+
graph = helper.make_graph(
1736+
[matmul, add],
1737+
"test",
1738+
[
1739+
helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3, 4)),
1740+
helper.make_tensor_value_info("W", TensorProto.FLOAT, (4, 5)),
1741+
],
1742+
[helper.make_tensor_value_info("A", TensorProto.FLOAT, (2, 3, 5))],
1743+
initializer=[b],
1744+
)
1745+
optimized_model = self._optimized(
1746+
graph, ["fuse_matmul_add_bias_into_gemm_batched"]
1747+
)
1748+
1749+
assert [n.op_type for n in optimized_model.graph.node] == ["MatMul", "Add"]
1750+
16521751
# type: () -> None
16531752
def test_fuse_pad_into_conv_no_optional_value_opset10(self):
16541753
pad = helper.make_node(

0 commit comments

Comments
 (0)