Add fuse_mul_into_conv pass

Fold a per-output-channel (or scalar) constant Mul that follows a Conv
into the convolution weights: for `Z = Conv(X, W[, B]) * S`, rewrite to
`Conv(X, W*S[, B*S])`. The scale is applied to the constant weights via
Mul nodes that the constant folder then materialises, so the standalone
multiply disappears; the trailing bias Add is subsequently absorbed by
fuse_add_bias_into_conv, collapsing a `Conv -> Mul -> Add` affine tail
(e.g. a BatchNorm exported as Mul/Add) back into the Conv.

Only per-channel scales aligned to the Conv output channel axis (or a
scalar) are fused; other shapes are left untouched. Only Conv (not
ConvTranspose, whose weight layout differs) is handled. Bias, when
present, must be constant so it can be scaled.

On FasterRCNN-10 this removes 53 `Conv -> Mul -> Add` chains (106 nodes)
that were previously left unfused. Adds tests covering per-channel and
scalar scales (with and without bias) and a non-per-channel negative case.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Piyb2V1gQuMBKfYdZeqqZy
Signed-off-by: take-cheeze <takechi101010@gmail.com>

Author: claude

onnxoptimizer/pass_registry.h

@@ -48,6 +48,7 @@
 #include "onnxoptimizer/passes/fuse_consecutive_squeezes.h"
 #include "onnxoptimizer/passes/fuse_consecutive_transposes.h"
 #include "onnxoptimizer/passes/fuse_matmul_add_bias_into_gemm.h"
+#include "onnxoptimizer/passes/fuse_mul_into_conv.h"
 #include "onnxoptimizer/passes/fuse_pad_into_conv.h"
 #include "onnxoptimizer/passes/fuse_pad_into_pool.h"
 #include "onnxoptimizer/passes/fuse_transpose_into_gemm.h"
@@ -102,6 +103,7 @@ struct GlobalPassRegistry {
     registerPass<FuseConsecutiveSqueezeUnsqueeze>();
     registerPass<FuseConsecutiveTransposes>();
     registerPass<FuseMatMulAddBiasIntoGemm>();
+    registerPass<FuseMulIntoConv>();
     registerPass<FusePadIntoConv>();
     registerPass<FusePadIntoPool>();
     registerPass<FuseTransposeIntoGemm>();

onnxoptimizer/passes/fuse_mul_into_conv.h

@@ -0,0 +1,224 @@
+// 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 = Conv(X, W[, B])
+//   Z = Y * S
+// After:
+//   Z = Conv(X, W', [B'])   with  W' = W * S,  B' = B * S
+//
+// S must be a constant that scales the Conv output per output channel: either a
+// single element (scalar broadcast) or a tensor whose only non-unit dimension
+// is the channel axis of the Conv output (axis 1), e.g. [C], [C, 1, 1] or [1,
+// C, 1, 1]. The scale is folded into the convolution weights (and bias, if
+// present) by creating Mul nodes on the constant weights, which the constant
+// folder then materialises; the leftover multiplicative op disappears.
+//
+// This recovers the fusion of an affine scale that a framework exported as a
+// standalone Mul -- most commonly the multiplicative half of a BatchNorm
+// written as ``Conv -> Mul -> Add``. Once the Mul is folded in, the following
+// bias Add is picked up by fuse_add_bias_into_conv, collapsing the whole affine
+// tail into the Conv. Only Conv (not ConvTranspose) is handled, because its
+// weight layout puts the output channel on axis 0.
+
+#include <numeric>
+#include <vector>
+
+#include "onnx/common/assertions.h"
+#include "onnxoptimizer/pass.h"
+#include "onnxoptimizer/passes/pass_util.h"
+
+namespace ONNX_NAMESPACE {
+namespace optimization {
+
+struct FuseMulIntoConv final : public PredicateBasedPass {
+  explicit FuseMulIntoConv()
+      : PredicateBasedPass(PassType::Fuse, PassEfficiency::Complete,
+                           PassOptimizationType::Compute) {}
+  std::string getPassName() const override {
+    return "fuse_mul_into_conv";
+  }
+
+  // Index (0 or 1) of the Mul operand that is a foldable Conv output -- a Conv
+  // used only by this Mul and whose weight is constant -- or -1 if neither is.
+  static int ConvOperandIndex(Node* mul) {
+    for (int i = 0; i < 2; ++i) {
+      Value* v = mul->inputs()[i];
+      if (v->node()->kind() == kConv && v->uses().size() == 1 &&
+          IsConstantTensor(v->node(), 1)) {
+        return i;
+      }
+    }
+    return -1;
+  }
+
+  bool patternMatchPredicate(Node* n) override {
+    if (n->kind() != kMul || n->inputs().size() != 2) {
+      return false;
+    }
+    const int conv_idx = ConvOperandIndex(n);
+    if (conv_idx < 0) {
+      return false;
+    }
+    return IsConstantTensor(n->inputs()[1 - conv_idx]);
+  }
+
+  bool runTransform(Node* n, Graph& graph,
+                    NodeDestroyType& destroy_current) override {
+    destroy_current = NodeDestroyType::DestroyZero;
+    const int conv_idx = ConvOperandIndex(n);
+    if (conv_idx < 0) {
+      return false;
+    }
+    Value* conv_out = n->inputs()[conv_idx];
+    Value* scale = n->inputs()[1 - conv_idx];
+    Node* conv = conv_out->node();
+    Value* weight = conv->inputs()[1];
+
+    const Tensor* weight_t = FetchConstantTensor(weight);
+    const Tensor* scale_t = FetchConstantTensor(scale);
+    if (weight_t == nullptr || scale_t == nullptr) {
+      return false;
+    }
+    // The scale is multiplied against the constant weight, so the element types
+    // must agree.
+    if (weight_t->elem_type() != scale_t->elem_type()) {
+      return false;
+    }
+
+    // Conv weight is [C_out, C_in / groups, k...]; the output channel is axis
+    // 0.
+    const auto& w_sizes = weight_t->sizes();
+    if (w_sizes.size() < 2) {
+      return false;
+    }
+    const int64_t C = w_sizes[0];
+    const int64_t w_rank = static_cast<int64_t>(w_sizes.size());
+
+    // Validate that `scale` scales the Conv output per output channel: a scalar
+    // (single element) or a tensor whose only non-unit axis equals C and, when
+    // right-aligned against the Conv output rank (== w_rank), lands on the
+    // channel axis (index 1).
+    const auto& s_sizes = scale_t->sizes();
+    int64_t s_numel = 1;
+    for (const int64_t d : s_sizes) {
+      s_numel *= d;
+    }
+    bool per_channel;
+    if (s_numel == 1) {
+      per_channel = false;  // scalar: broadcasts to every channel
+    } else if (s_numel == C) {
+      const int64_t s_rank = static_cast<int64_t>(s_sizes.size());
+      int64_t non_unit_axis = -1;
+      for (int64_t i = 0; i < s_rank; ++i) {
+        if (s_sizes[i] != 1) {
+          if (non_unit_axis != -1) {
+            return false;  // more than one non-unit axis
+          }
+          non_unit_axis = i;
+        }
+      }
+      if (non_unit_axis < 0 || w_rank - s_rank + non_unit_axis != 1) {
+        return false;  // the C-sized axis is not the channel axis
+      }
+      per_channel = true;
+    } else {
+      return false;
+    }
+
+    // If the Conv has a bias it must be constant too, so it can be scaled.
+    const bool has_bias = conv->inputs().size() == 3;
+    if (has_bias) {
+      Value* bias = conv->inputs()[2];
+      const Tensor* bias_t = FetchConstantTensor(bias);
+      if (bias_t == nullptr || bias_t->elem_type() != scale_t->elem_type()) {
+        return false;
+      }
+    }
+
+    // --- All checks passed; build the scaled weight (and bias). ---
+
+    // `bias_scale` is the [C] (or scalar) scale used for the bias;
+    // `weight_scale` is that scale broadcast to the weight rank [C, 1, ...] for
+    // the weight.
+    Value* bias_scale = scale;
+    Value* weight_scale = scale;
+    if (per_channel) {
+      // Flatten the scale to [C].
+      Value* scale_1d = scale;
+      if (static_cast<int64_t>(s_sizes.size()) != 1) {
+        Node* reshape = graph.create(kReshape, 1);
+        reshape->addInput(scale);
+        Tensor shape_t;
+        shape_t.elem_type() = TensorProto_DataType_INT64;
+        shape_t.sizes().push_back(1);
+        shape_t.int64s().push_back(C);
+        reshape->addInput(graph.addInitializerAndCreateValue(shape_t));
+        reshape->insertBefore(conv);
+        scale_1d = reshape->output();
+      }
+      bias_scale = scale_1d;
+      // Broadcast [C] -> [C, 1, ..., 1] so it multiplies the weight per
+      // channel.
+      weight_scale = scale_1d;
+      if (w_rank > 1) {
+        std::vector<int64_t> axes(w_rank - 1);
+        std::iota(axes.begin(), axes.end(), 1);
+        Node* unsqueeze = graph.create(kUnsqueeze, 1);
+        unsqueeze->addInput(scale_1d);
+        const int opset = getOpsetVersion(graph);
+        if (opset >= 13 || opset == 0) {
+          Tensor axes_t;
+          axes_t.elem_type() = TensorProto_DataType_INT64;
+          axes_t.sizes().push_back(static_cast<int64_t>(axes.size()));
+          axes_t.int64s() = axes;
+          unsqueeze->addInput(graph.addInitializerAndCreateValue(axes_t));
+        } else {
+          unsqueeze->is_(kaxes, std::move(axes));
+        }
+        unsqueeze->insertBefore(conv);
+        weight_scale = unsqueeze->output();
+      }
+    }
+
+    Node* mul_w = graph.create(kMul, 1);
+    mul_w->addInput(weight);
+    mul_w->addInput(weight_scale);
+    mul_w->insertBefore(conv);
+    conv->replaceInput(1, mul_w->output());
+
+    if (has_bias) {
+      Value* bias = conv->inputs()[2];
+      Node* mul_b = graph.create(kMul, 1);
+      mul_b->addInput(bias);
+      mul_b->addInput(bias_scale);
+      mul_b->insertBefore(conv);
+      conv->replaceInput(2, mul_b->output());
+    }
+
+    // The scale does not change the Conv output shape/type; carry the Mul's
+    // inferred metadata onto the Conv output.
+    if (conv_out->sizes().size() == 0 && n->output()->sizes().size() > 0) {
+      conv_out->setSizes(n->output()->sizes());
+    }
+    if (n->output()->elemType() != TensorProto_DataType_UNDEFINED) {
+      conv_out->setElemType(n->output()->elemType());
+    }
+
+    const bool replacing_success = tryReplacingAllUsesWith(n, conv);
+    if (!replacing_success) {
+      return false;
+    }
+    destroy_current = NodeDestroyType::DestroyOne;
+    return true;
+  }
+};
+
+}  // namespace optimization
+}  // namespace ONNX_NAMESPACE

onnxoptimizer/test/optimizer_test.py

@@ -1269,6 +1269,65 @@ def test_fuse_transpose_into_gemm(self):  # type: () -> None
         assert len(optimized_model.graph.node[3].attribute[0].g.node) == 1
         assert optimized_model.graph.node[3].attribute[0].g.node[0].op_type == "Gemm"
 
+    def _conv_mul_graph(self, scale_shape, with_bias):
+        # Conv(X, W[, B]) -> Mul(., S), with constant W/B/S.
+        C_out, C_in, k = 4, 3, 3
+        W = np.random.rand(C_out, C_in, k, k).astype(np.float32)
+        S = np.random.rand(*scale_shape).astype(np.float32)
+        inits = [
+            numpy_helper.from_array(W, "W"),
+            numpy_helper.from_array(S, "S"),
+        ]
+        conv_inputs = ["X", "W"]
+        if with_bias:
+            B = np.random.rand(C_out).astype(np.float32)
+            inits.append(numpy_helper.from_array(B, "B"))
+            conv_inputs.append("B")
+        nodes = [
+            helper.make_node("Conv", conv_inputs, ["Z"], pads=[1, 1, 1, 1]),
+            helper.make_node("Mul", ["Z", "S"], ["Y"]),
+        ]
+        return helper.make_graph(
+            nodes,
+            "test",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, C_in, 8, 8))],
+            [helper.make_tensor_value_info("Y", TensorProto.FLOAT, (1, C_out, 8, 8))],
+            initializer=inits,
+        )
+
+    def _conv_output_feeds_mul(self, model):
+        conv_outputs = {
+            o for n in model.graph.node if n.op_type == "Conv" for o in n.output
+        }
+        return any(
+            n.op_type == "Mul" and any(i in conv_outputs for i in n.input)
+            for n in model.graph.node
+        )
+
+    def test_fuse_mul_into_conv_per_channel(self):  # type: () -> None
+        # A per-output-channel scale (shape [1, C, 1, 1]) is folded into the Conv
+        # weights, so no Mul consumes the Conv output any more. _optimized also
+        # checks numerical equivalence via onnxruntime.
+        for with_bias in (True, False):
+            graph = self._conv_mul_graph((1, 4, 1, 1), with_bias)
+            optimized_model = self._optimized(graph, ["fuse_mul_into_conv"])
+            assert (
+                sum(n.op_type == "Conv" for n in optimized_model.graph.node) == 1
+            )
+            assert not self._conv_output_feeds_mul(optimized_model)
+
+    def test_fuse_mul_into_conv_scalar(self):  # type: () -> None
+        graph = self._conv_mul_graph((1,), True)
+        optimized_model = self._optimized(graph, ["fuse_mul_into_conv"])
+        assert not self._conv_output_feeds_mul(optimized_model)
+
+    def test_fuse_mul_into_conv_no_fuse_non_channel_scale(self):  # type: () -> None
+        # A per-pixel scale ([1, 1, H, W]) is not a per-channel scale and must not
+        # be folded into the weights.
+        graph = self._conv_mul_graph((1, 1, 8, 8), False)
+        optimized_model = self._optimized(graph, ["fuse_mul_into_conv"])
+        assert self._conv_output_feeds_mul(optimized_model)
+
     def test_fuse_add_bias_into_conv_with_scalar_bias(self):  # type: () -> None
         nodes = [
             helper.make_node("Conv", ["X", "Y"], ["Z"]),