Skip to content

Commit e5f6c93

Browse files
andifeCopilot
andcommitted
Preserve bias dependency order and per-channel QDQ axes
Address correctness findings in #37 and add targeted regression coverage. Native builds and focused tests were verified locally on Windows. Signed-off-by: Andreas Fehlner <fehlner@arcor.de> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent e9e3aa4 commit e5f6c93

2 files changed

Lines changed: 130 additions & 72 deletions

File tree

‎onnxoptimizer/passes/fuse_add_bias_into_conv.h‎

Lines changed: 12 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -37,22 +37,11 @@ struct FuseAddBiasIntoConv final : public PredicateBasedPass {
3737
return CheckKind(node, kAdd, 0, kConv) &&
3838
GetInputsOfPreNode(node, 0).size() == 2;
3939
}
40-
static Node *makeSqueezeOrUnsqueeze(Graph &graph, std::vector<int64_t> &axes,
41-
Value *input, Node *target_node,
42-
BuiltinSymbol k, bool is_input_qdq) {
40+
static Node* makeSqueezeOrUnsqueeze(Graph& graph, std::vector<int64_t>& axes,
41+
Value* input, Node* target_node,
42+
BuiltinSymbol k) {
4343
assert(k == kSqueeze || k == kUnsqueeze);
4444
Node *squeeze = graph.create(k, 1);
45-
Node *dequant_node = nullptr;
46-
Node *quant_node = nullptr;
47-
// insert squeeze op before qdq
48-
if (is_input_qdq) {
49-
dequant_node = input->node();
50-
quant_node = dequant_node->input(0)->node();
51-
target_node = quant_node;
52-
input = target_node->input(0);
53-
dequant_node->output()->clearMetadata();
54-
quant_node->output()->clearMetadata();
55-
}
5645
squeeze->addInput(input);
5746
int opset_version = getOpsetVersion(graph);
5847
int version_threshold = 13;
@@ -66,13 +55,7 @@ struct FuseAddBiasIntoConv final : public PredicateBasedPass {
6655
Value *tv = graph.addInitializerAndCreateValue(t);
6756
squeeze->addInput(tv);
6857
}
69-
if (is_input_qdq) {
70-
quant_node->replaceInput(0, squeeze->output());
71-
}
7258
squeeze->insertBefore(target_node);
73-
if (is_input_qdq) {
74-
return dequant_node;
75-
}
7659
return squeeze;
7760
}
7861
bool runTransform(Node *n, Graph &graph,
@@ -125,21 +108,19 @@ struct FuseAddBiasIntoConv final : public PredicateBasedPass {
125108
return false;
126109
}
127110
if (num_el == 1) {
128-
if (orig_bias->node()->kind() != kParam &&
129-
orig_conv->node()->isBefore(orig_bias->node())) {
130-
orig_bias->node()->moveBefore(orig_conv->node());
131-
}
111+
// Conv feeds only this Add, so it can follow every bias dependency.
112+
orig_conv->node()->moveBefore(n);
132113
Value *conv_3rd_input = orig_bias;
133114
if (bias_shape.size() > 1) {
134115
std::vector<int64_t> axes(bias_shape.size() - 1);
135116
std::iota(axes.begin(), axes.end(), 0);
136-
Node *squeeze = makeSqueezeOrUnsqueeze(
137-
graph, axes, conv_3rd_input, orig_conv->node(), kSqueeze, false);
117+
Node* squeeze = makeSqueezeOrUnsqueeze(graph, axes, conv_3rd_input,
118+
orig_conv->node(), kSqueeze);
138119
conv_3rd_input = squeeze->output();
139120
} else if (bias_shape.size() == 0) {
140121
std::vector<int64_t> axes = {0};
141-
Node *unsqueeze = makeSqueezeOrUnsqueeze(
142-
graph, axes, conv_3rd_input, orig_conv->node(), kUnsqueeze, false);
122+
Node* unsqueeze = makeSqueezeOrUnsqueeze(graph, axes, conv_3rd_input,
123+
orig_conv->node(), kUnsqueeze);
143124
conv_3rd_input = unsqueeze->output();
144125
}
145126
if (M > 1) {
@@ -167,24 +148,14 @@ struct FuseAddBiasIntoConv final : public PredicateBasedPass {
167148
bias_shape[1 + bias_shape.size() - static_cast<unsigned>(rank)]
168149
.dim == M) {
169150
ONNX_ASSERT(bias_shape.size() > 1);
170-
const bool is_input_qdq =
171-
orig_bias->node()->kind() == Symbol("DequantizeLinear") &&
172-
orig_bias->node()->input(0)->node()->kind() ==
173-
Symbol("QuantizeLinear");
174-
if (orig_bias->node()->kind() != kParam &&
175-
orig_conv->node()->isBefore(orig_bias->node())) {
176-
if (is_input_qdq) {
177-
orig_bias->node()->input(0)->node()->moveBefore(orig_conv->node());
178-
}
179-
orig_bias->node()->moveBefore(orig_conv->node());
180-
}
151+
orig_conv->node()->moveBefore(n);
181152
std::vector<int64_t> axes(bias_shape.size());
182153
std::iota(axes.begin(), axes.end(), static_cast<int64_t>(0));
183154
axes.erase(axes.begin() +
184155
(1 + bias_shape.size() - static_cast<unsigned>(rank)));
185156

186-
Node *new_bias = makeSqueezeOrUnsqueeze(
187-
graph, axes, orig_bias, orig_conv->node(), kSqueeze, is_input_qdq);
157+
Node* new_bias = makeSqueezeOrUnsqueeze(graph, axes, orig_bias,
158+
orig_conv->node(), kSqueeze);
188159
orig_conv->node()->addInput(new_bias->output());
189160
} else {
190161
return false;

‎onnxoptimizer/test/optimizer_test.py‎

Lines changed: 118 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
#
33
# SPDX-License-Identifier: Apache-2.0
44

5+
from __future__ import annotations
6+
57
import io
68
import os
79
import unittest
@@ -52,7 +54,7 @@ def _compare(
5254
model_opt: onnx.ModelProto,
5355
model_ori: onnx.ModelProto,
5456
n_times: int = 5,
55-
input_shapes: Optional[TensorShapes] = None,
57+
input_shapes: TensorShapes | None = None,
5658
verbose=True,
5759
) -> bool:
5860
"""
@@ -65,7 +67,7 @@ def _compare(
6567
def get_shape_from_value_info_proto(v: onnx.ValueInfoProto) -> list[int]:
6668
return [dim.dim_value for dim in v.type.tensor_type.shape.dim]
6769

68-
def get_value_info_all(m: onnx.ModelProto, name: str) -> Optional[onnx.ValueInfoProto]:
70+
def get_value_info_all(m: onnx.ModelProto, name: str) -> onnx.ValueInfoProto | None:
6971
for v in m.graph.value_info:
7072
if v.name == name:
7173
return v
@@ -89,7 +91,7 @@ def get_shape(m: onnx.ModelProto, name: str) -> TensorShape:
8991
return get_shape_from_value_info_proto(v)
9092
raise RuntimeError(f'Cannot get shape of "{name}"')
9193

92-
def get_elem_type(m: onnx.ModelProto, name: str) -> Optional[int]:
94+
def get_elem_type(m: onnx.ModelProto, name: str) -> int | None:
9395
v = get_value_info_all(m, name)
9496
if v is not None:
9597
return v.type.tensor_type.elem_type
@@ -127,7 +129,7 @@ def get_input_names(model: onnx.ModelProto) -> list[str]:
127129
)
128130
return input_names
129131

130-
def generate_rand_input(model, input_shapes: Optional[TensorShapes] = None):
132+
def generate_rand_input(model, input_shapes: TensorShapes | None = None):
131133
if input_shapes is None:
132134
input_shapes = {}
133135
input_names = get_input_names(model)
@@ -151,7 +153,7 @@ def generate_rand_input(model, input_shapes: Optional[TensorShapes] = None):
151153
return inputs
152154

153155
def forward(
154-
model, inputs=None, input_shapes: Optional[TensorShapes] = None
156+
model, inputs=None, input_shapes: TensorShapes | None = None
155157
) -> dict[str, np.ndarray]:
156158
if input_shapes is None:
157159
input_shapes = {}
@@ -1571,38 +1573,123 @@ def test_fuse_add_bias_into_conv_with_non_constant_bias(self):
15711573

15721574
# type: () -> None
15731575
def test_fuse_add_bias_into_conv_with_quanted_bias(self):
1574-
nodes = [helper.make_node("Conv", ["X", "Y"], ["Z"]),
1575-
helper.make_node("QuantizeLinear", ["A", "scale", "zero_point"], ["B"], axis=0),
1576-
helper.make_node("DequantizeLinear", ["B", "scale", "zero_point"], ["C"], axis=0),
1577-
helper.make_node("Add", ["Z", "C"], ["D"])]
1576+
nodes = [
1577+
helper.make_node("Conv", ["X", "Y"], ["Z"]),
1578+
helper.make_node("QuantizeLinear", ["A", "scale", "zero_point"], ["B"], axis=0),
1579+
helper.make_node("DequantizeLinear", ["B", "scale", "zero_point"], ["C"], axis=0),
1580+
helper.make_node("Add", ["Z", "C"], ["D"]),
1581+
]
15781582
graph = helper.make_graph(
15791583
nodes,
15801584
"test",
1581-
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 5, 3, 3)),
1582-
helper.make_tensor_value_info(
1583-
"Y", TensorProto.FLOAT, (16, 5, 3, 3)),
1584-
helper.make_tensor_value_info("A", TensorProto.FLOAT, (16, 1, 1))],
1585-
[helper.make_tensor_value_info(
1586-
"D", TensorProto.FLOAT, (1, 16, 1, 1))],
1587-
[helper.make_tensor("scale", TensorProto.FLOAT,
1588-
dims=(16,),
1589-
vals=np.random.rand(16).astype(np.float32).tobytes(),
1590-
raw=True),
1591-
helper.make_tensor("zero_point", TensorProto.INT8,
1592-
dims=(16,),
1593-
vals=np.zeros([16]).astype(np.int8).tobytes(),
1594-
raw=True)],
1595-
value_info=[helper.make_tensor_value_info(
1596-
"C", TensorProto.FLOAT, (16, 1, 1))]
1585+
[
1586+
helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 5, 3, 3)),
1587+
helper.make_tensor_value_info("Y", TensorProto.FLOAT, (16, 5, 3, 3)),
1588+
helper.make_tensor_value_info("A", TensorProto.FLOAT, (16, 1, 1)),
1589+
],
1590+
[helper.make_tensor_value_info("D", TensorProto.FLOAT, (1, 16, 1, 1))],
1591+
[
1592+
helper.make_tensor(
1593+
"scale",
1594+
TensorProto.FLOAT,
1595+
dims=(16,),
1596+
vals=np.random.rand(16).astype(np.float32).tobytes(),
1597+
raw=True,
1598+
),
1599+
helper.make_tensor(
1600+
"zero_point",
1601+
TensorProto.INT8,
1602+
dims=(16,),
1603+
vals=np.zeros([16]).astype(np.int8).tobytes(),
1604+
raw=True,
1605+
),
1606+
],
1607+
value_info=[helper.make_tensor_value_info("C", TensorProto.FLOAT, (16, 1, 1))],
1608+
)
1609+
optimized_model = self._optimized(
1610+
graph, ["fuse_add_bias_into_conv"], opset_imports=[helper.make_opsetid("", 13)]
15971611
)
1598-
optimized_model = self._optimized(graph, ["fuse_add_bias_into_conv"], opset_imports=[helper.make_opsetid("", 13)])
15991612

16001613
assert len(list(optimized_model.graph.node)) == 4
1601-
assert optimized_model.graph.node[0].op_type == 'Squeeze'
1602-
assert optimized_model.graph.node[1].op_type == 'QuantizeLinear'
1603-
assert optimized_model.graph.node[2].op_type == 'DequantizeLinear'
1604-
assert optimized_model.graph.node[3].op_type == 'Conv'
1605-
assert optimized_model.graph.output[0].name == 'D'
1614+
assert optimized_model.graph.node[0].op_type == "QuantizeLinear"
1615+
assert optimized_model.graph.node[1].op_type == "DequantizeLinear"
1616+
assert optimized_model.graph.node[2].op_type == "Squeeze"
1617+
assert optimized_model.graph.node[3].op_type == "Conv"
1618+
assert optimized_model.graph.output[0].name == "D"
1619+
1620+
def test_fuse_add_bias_into_conv_dependency_chain(self) -> None:
1621+
for bias_shape in [[1], [2, 1, 1]]:
1622+
with self.subTest(bias_shape=bias_shape):
1623+
graph = helper.make_graph(
1624+
[
1625+
helper.make_node("Conv", ["X", "W"], ["conv"]),
1626+
helper.make_node("Sin", ["A"], ["sin"]),
1627+
helper.make_node("Cos", ["sin"], ["bias"]),
1628+
helper.make_node("Add", ["conv", "bias"], ["Y"]),
1629+
],
1630+
"bias_dependency_chain",
1631+
[
1632+
helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 2, 2]),
1633+
helper.make_tensor_value_info("A", TensorProto.FLOAT, bias_shape),
1634+
],
1635+
[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 2, 2, 2])],
1636+
initializer=[
1637+
numpy_helper.from_array(np.ones([2, 1, 1, 1], np.float32), "W")
1638+
],
1639+
value_info=[
1640+
helper.make_tensor_value_info("bias", TensorProto.FLOAT, bias_shape)
1641+
],
1642+
)
1643+
optimized = self._optimized(graph, ["fuse_add_bias_into_conv"])
1644+
assert all(n.op_type != "Add" for n in optimized.graph.node)
1645+
assert optimized.graph.node[-1].op_type == "Conv"
1646+
1647+
def test_fuse_add_bias_into_conv_per_channel_qdq(self) -> None:
1648+
for axis in [1, -3]:
1649+
with self.subTest(axis=axis):
1650+
graph = helper.make_graph(
1651+
[
1652+
helper.make_node("Conv", ["X", "W"], ["conv"]),
1653+
helper.make_node(
1654+
"QuantizeLinear", ["A", "scale", "zero"], ["quantized"], axis=axis
1655+
),
1656+
helper.make_node(
1657+
"DequantizeLinear",
1658+
["quantized", "scale", "zero"],
1659+
["bias"],
1660+
axis=axis,
1661+
),
1662+
helper.make_node("Add", ["conv", "bias"], ["Y"]),
1663+
],
1664+
"per_channel_bias",
1665+
[
1666+
helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 2, 2]),
1667+
helper.make_tensor_value_info("A", TensorProto.FLOAT, [1, 2, 1, 1]),
1668+
],
1669+
[
1670+
helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 2, 2, 2]),
1671+
helper.make_tensor_value_info(
1672+
"quantized", TensorProto.INT8, [1, 2, 1, 1]
1673+
),
1674+
],
1675+
initializer=[
1676+
numpy_helper.from_array(np.ones([2, 1, 1, 1], np.float32), "W"),
1677+
numpy_helper.from_array(np.array([0.1, 0.2], np.float32), "scale"),
1678+
numpy_helper.from_array(np.zeros(2, np.int8), "zero"),
1679+
],
1680+
value_info=[
1681+
helper.make_tensor_value_info(
1682+
"bias", TensorProto.FLOAT, [1, 2, 1, 1]
1683+
)
1684+
],
1685+
)
1686+
optimized = self._optimized(graph, ["fuse_add_bias_into_conv"])
1687+
assert [n.op_type for n in optimized.graph.node] == [
1688+
"QuantizeLinear", "DequantizeLinear", "Squeeze", "Conv"
1689+
]
1690+
assert next(
1691+
a.i for a in optimized.graph.node[0].attribute if a.name == "axis"
1692+
) == axis
16061693

16071694
def test_fuse_matmul_add_bias_into_gemm(self): # type: () -> None
16081695
matmul = helper.make_node("MatMul", ["X", "Y"], ["Z"])

0 commit comments

Comments
 (0)