From 8112db072b3cc881aa608e9566a90f934a806d2f Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 22:44:36 +0000 Subject: [PATCH] Add fuse_consecutive_mul 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, a per-channel (C,1,1) scale times a scalar factor. Conservatively limited to FLOAT/DOUBLE constants; other dtypes are left to constant folding. Registered as a default Fuse pass. Adds unit tests for the fuse cases (scalar, per-channel) and the no-fuse guards (inner multiple-use, non-constant operand) with onnxruntime numeric-equivalence checks. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01V63N6PXYEgNEe1BWbKi6cU --- onnxoptimizer/pass_registry.h | 2 + onnxoptimizer/passes/fuse_consecutive_mul.h | 223 ++++++++++++++++++++ onnxoptimizer/test/optimizer_test.py | 95 +++++++++ 3 files changed, 320 insertions(+) create mode 100644 onnxoptimizer/passes/fuse_consecutive_mul.h diff --git a/onnxoptimizer/pass_registry.h b/onnxoptimizer/pass_registry.h index 3fdfa4978..6059c7cb9 100644 --- a/onnxoptimizer/pass_registry.h +++ b/onnxoptimizer/pass_registry.h @@ -46,6 +46,7 @@ #include "onnxoptimizer/passes/fuse_consecutive_reduce_unsqueeze.h" #include "onnxoptimizer/passes/fuse_consecutive_squeeze_unsqueeze.h" #include "onnxoptimizer/passes/fuse_consecutive_squeezes.h" +#include "onnxoptimizer/passes/fuse_consecutive_mul.h" #include "onnxoptimizer/passes/fuse_consecutive_transposes.h" #include "onnxoptimizer/passes/fuse_matmul_add_bias_into_gemm.h" #include "onnxoptimizer/passes/fuse_pad_into_conv.h" @@ -101,6 +102,7 @@ struct GlobalPassRegistry { registerPass(); registerPass(); registerPass(); + registerPass(); registerPass(); registerPass(); registerPass(); diff --git a/onnxoptimizer/passes/fuse_consecutive_mul.h b/onnxoptimizer/passes/fuse_consecutive_mul.h new file mode 100644 index 000000000..3cb6d667d --- /dev/null +++ b/onnxoptimizer/passes/fuse_consecutive_mul.h @@ -0,0 +1,223 @@ +// Copyright (c) ONNX Project Contributors +// +// SPDX-License-Identifier: Apache-2.0 + +// ATTENTION: The code in this file is highly EXPERIMENTAL. +// Adventurous users should note that the APIs will probably change. + +#pragma once + +// Before: +// Y = Mul(X, C1) +// Z = Mul(Y, C2) +// After: +// Z = Mul(X, C1 * C2) +// +// where C1 and C2 are constants (initializers or Constant nodes) and the inner +// Mul (producing Y) is used only by the outer Mul. C1 * C2 is materialised as a +// new constant using numpy-style broadcasting, so the rewrite is numerically +// equivalent: at every output position the value is X * C1 * C2 regardless of +// how the two constant scales are grouped. +// +// This shows up in models that export a composed constant scale as two +// back-to-back Mul nodes, e.g. PoolFormer's LayerScale (a per-channel (C,1,1) +// scale multiplied by a scalar factor). + +#include +#include +#include + +#include "onnx/common/assertions.h" +#include "onnxoptimizer/pass.h" +#include "onnxoptimizer/passes/pass_util.h" + +namespace ONNX_NAMESPACE { +namespace optimization { + +struct FuseConsecutiveMul final : public PredicateBasedPass { + explicit FuseConsecutiveMul() + : PredicateBasedPass(PassType::Fuse, PassEfficiency::Complete, + PassOptimizationType::Compute) {} + std::string getPassName() const override { + return "fuse_consecutive_mul"; + } + + // Returns the index (0 or 1) of the single constant operand of a 2-input + // node, or -1 when the node does not have exactly one constant operand. + static int SoleConstantInput(const Node* n) { + if (n->inputs().size() != 2) { + return -1; + } + const bool c0 = IsConstantTensor(n->input(0)); + const bool c1 = IsConstantTensor(n->input(1)); + if (c0 && !c1) { + return 0; + } + if (c1 && !c0) { + return 1; + } + return -1; + } + + // numpy-style broadcast of two shapes. Returns false when incompatible. + static bool BroadcastShapes(const std::vector& a, + const std::vector& b, + std::vector& out) { + const size_t ra = a.size(); + const size_t rb = b.size(); + const size_t r = std::max(ra, rb); + out.assign(r, 1); + for (size_t i = 0; i < r; ++i) { + const int64_t da = i < r - ra ? 1 : a[i - (r - ra)]; + const int64_t db = i < r - rb ? 1 : b[i - (r - rb)]; + if (da == db || db == 1) { + out[i] = da == 1 ? db : da; + } else if (da == 1) { + out[i] = db; + } else { + return false; + } + } + return true; + } + + // Elementwise multiply of two flat buffers with numpy broadcasting, given the + // (right-aligned) input shapes and the pre-computed output shape. + template + static std::vector BroadcastMul(const std::vector& va, + const std::vector& sa, + const std::vector& vb, + const std::vector& sb, + const std::vector& so) { + const size_t r = so.size(); + const int64_t total = std::accumulate(so.begin(), so.end(), (int64_t)1, + std::multiplies{}); + // Strides into each input buffer, right-aligned to the output rank; a + // broadcast (size-1) dimension gets stride 0. + auto build_strides = [r](const std::vector& s) { + std::vector full(r, 1); + for (size_t i = 0; i < s.size(); ++i) { + full[r - s.size() + i] = s[i]; + } + std::vector stride(r, 0); + int64_t acc = 1; + for (size_t i = r; i-- > 0;) { + stride[i] = full[i] == 1 ? 0 : acc; + acc *= full[i]; + } + return stride; + }; + const std::vector stra = build_strides(sa); + const std::vector strb = build_strides(sb); + + std::vector out(total); + std::vector idx(r, 0); + for (int64_t lin = 0; lin < total; ++lin) { + int64_t ia = 0, ib = 0; + for (size_t d = 0; d < r; ++d) { + ia += idx[d] * stra[d]; + ib += idx[d] * strb[d]; + } + out[lin] = va[ia] * vb[ib]; + for (size_t d = r; d-- > 0;) { + if (++idx[d] < so[d]) { + break; + } + idx[d] = 0; + } + } + return out; + } + + // Builds C1 * C2 into `out`. Only FLOAT/DOUBLE are handled; other dtypes are + // left to constant folding and this pass conservatively declines them. + static bool MultiplyConstants(const Tensor& c1, const Tensor& c2, + Tensor& out) { + if (c1.elem_type() != c2.elem_type()) { + return false; + } + std::vector so; + if (!BroadcastShapes(c1.sizes(), c2.sizes(), so)) { + return false; + } + out.elem_type() = c1.elem_type(); + out.sizes() = so; + switch (c1.elem_type()) { + case TensorProto_DataType_FLOAT: + out.floats() = + BroadcastMul(ParseTensorData(&c1), c1.sizes(), + ParseTensorData(&c2), c2.sizes(), so); + return true; + case TensorProto_DataType_DOUBLE: + out.doubles() = + BroadcastMul(ParseTensorData(&c1), c1.sizes(), + ParseTensorData(&c2), c2.sizes(), so); + return true; + default: + return false; + } + } + + bool patternMatchPredicate(Node* node) override { + if (!CheckKind(node, kMul)) { + return false; + } + const int outer_c = SoleConstantInput(node); + if (outer_c < 0) { + return false; + } + const Value* other = node->input(1 - outer_c); + if (!CheckKind(other, kMul)) { + return false; + } + // inner Mul must feed only this outer Mul + if (other->uses().size() != 1) { + return false; + } + return SoleConstantInput(other->node()) >= 0; + } + + bool runTransform(Node* n, Graph& graph, + NodeDestroyType& destroy_current) override { + destroy_current = NodeDestroyType::DestroyZero; + + const int outer_c = SoleConstantInput(n); + Value* c2v = n->input(outer_c); + Node* inner = n->input(1 - outer_c)->node(); + const int inner_c = SoleConstantInput(inner); + if (inner_c < 0) { + return false; + } + Value* c1v = inner->input(inner_c); + Value* xv = inner->input(1 - inner_c); + + const Tensor* c1 = FetchConstantTensor(c1v); + const Tensor* c2 = FetchConstantTensor(c2v); + if (c1 == nullptr || c2 == nullptr) { + return false; + } + + Tensor fused_const; + if (!MultiplyConstants(*c1, *c2, fused_const)) { + return false; + } + Value* cv = graph.addInitializerAndCreateValue(fused_const); + + Node* fused = graph.create(kMul, n->outputs().size()); + fused->addInput(xv); + fused->addInput(cv); + for (int i = 0; i < static_cast(n->outputs().size()); ++i) { + fused->outputs()[i]->copyMetadata(n->outputs()[i]); + } + fused->insertBefore(n); + if (!tryReplacingAllUsesWith(n, fused)) { + return false; + } + // Destroy the outer Mul; DCE removes the now-dead inner Mul and constants. + destroy_current = NodeDestroyType::DestroyOne; + return true; + } +}; + +} // namespace optimization +} // namespace ONNX_NAMESPACE diff --git a/onnxoptimizer/test/optimizer_test.py b/onnxoptimizer/test/optimizer_test.py index 45ff7ca6a..f6d9cba06 100644 --- a/onnxoptimizer/test/optimizer_test.py +++ b/onnxoptimizer/test/optimizer_test.py @@ -1649,6 +1649,101 @@ def test_fuse_matmul_add_bias_into_gemm_multiple_use_no_fuse(self): assert optimized_model.graph == graph + def test_fuse_consecutive_mul_scalar(self): # type: () -> None + c1 = numpy_helper.from_array(np.array(2.0, dtype=np.float32), name="C1") + c2 = numpy_helper.from_array(np.array(3.0, dtype=np.float32), name="C2") + mul1 = helper.make_node("Mul", ["X", "C1"], ["Y"]) + mul2 = helper.make_node("Mul", ["Y", "C2"], ["Z"]) + graph = helper.make_graph( + [mul1, mul2], + "test", + [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3))], + [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (2, 3))], + initializer=[c1, c2], + ) + optimized_model = self._optimized( + graph, ["fuse_consecutive_mul", "eliminate_deadend"] + ) + + assert len(optimized_model.graph.node) == 1 + assert optimized_model.graph.node[0].op_type == "Mul" + consts = { + init.name: to_array(init) + for init in optimized_model.graph.initializer + } + fused_name = optimized_model.graph.node[0].input[1] + np.testing.assert_allclose(consts[fused_name], 6.0) + + def test_fuse_consecutive_mul_per_channel(self): # type: () -> None + # per-channel (C,1,1) scale composed with a scalar factor (LayerScale) + c1 = numpy_helper.from_array( + np.arange(4, dtype=np.float32).reshape(4, 1, 1) + 1.0, name="C1" + ) + c2 = numpy_helper.from_array(np.array(0.5, dtype=np.float32), name="C2") + mul1 = helper.make_node("Mul", ["X", "C1"], ["Y"]) + mul2 = helper.make_node("Mul", ["Y", "C2"], ["Z"]) + graph = helper.make_graph( + [mul1, mul2], + "test", + [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 4, 2, 2))], + [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (1, 4, 2, 2))], + initializer=[c1, c2], + ) + optimized_model = self._optimized( + graph, ["fuse_consecutive_mul", "eliminate_deadend"] + ) + + assert len(optimized_model.graph.node) == 1 + assert optimized_model.graph.node[0].op_type == "Mul" + consts = { + init.name: to_array(init) + for init in optimized_model.graph.initializer + } + fused_name = optimized_model.graph.node[0].input[1] + np.testing.assert_allclose( + consts[fused_name], + (np.arange(4, dtype=np.float32).reshape(4, 1, 1) + 1.0) * 0.5, + ) + + def test_fuse_consecutive_mul_inner_multiple_use_no_fuse(self): + c1 = numpy_helper.from_array(np.array(2.0, dtype=np.float32), name="C1") + c2 = numpy_helper.from_array(np.array(3.0, dtype=np.float32), name="C2") + mul1 = helper.make_node("Mul", ["X", "C1"], ["Y"]) + mul2 = helper.make_node("Mul", ["Y", "C2"], ["Z"]) + ident = helper.make_node("Identity", ["Y"], ["Y2"]) + graph = helper.make_graph( + [mul1, mul2, ident], + "test", + [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3))], + [ + helper.make_tensor_value_info("Z", TensorProto.FLOAT, (2, 3)), + helper.make_tensor_value_info("Y2", TensorProto.FLOAT, (2, 3)), + ], + initializer=[c1, c2], + ) + optimized_model = self._optimized(graph, ["fuse_consecutive_mul"]) + + assert [n.op_type for n in optimized_model.graph.node].count("Mul") == 2 + + def test_fuse_consecutive_mul_nonconst_no_fuse(self): + # inner Mul has no constant operand -> nothing to fold + c2 = numpy_helper.from_array(np.array(3.0, dtype=np.float32), name="C2") + mul1 = helper.make_node("Mul", ["X", "W"], ["Y"]) + mul2 = helper.make_node("Mul", ["Y", "C2"], ["Z"]) + graph = helper.make_graph( + [mul1, mul2], + "test", + [ + helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3)), + helper.make_tensor_value_info("W", TensorProto.FLOAT, (2, 3)), + ], + [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (2, 3))], + initializer=[c2], + ) + optimized_model = self._optimized(graph, ["fuse_consecutive_mul"]) + + assert [n.op_type for n in optimized_model.graph.node].count("Mul") == 2 + # type: () -> None def test_fuse_pad_into_conv_no_optional_value_opset10(self): pad = helper.make_node(