Add fuse_consecutive_mul pass

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<FuseConsecutiveSqueezes>();
     registerPass<FuseConsecutiveSqueezeUnsqueeze>();
     registerPass<FuseConsecutiveTransposes>();
+    registerPass<FuseConsecutiveMul>();
     registerPass<FuseMatMulAddBiasIntoGemm>();
     registerPass<FusePadIntoConv>();
     registerPass<FusePadIntoPool>();

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 <functional>
+#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 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<int64_t>& a,
+                              const std::vector<int64_t>& b,
+                              std::vector<int64_t>& 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 <typename T>
+  static std::vector<T> BroadcastMul(const std::vector<T>& va,
+                                     const std::vector<int64_t>& sa,
+                                     const std::vector<T>& vb,
+                                     const std::vector<int64_t>& sb,
+                                     const std::vector<int64_t>& so) {
+    const size_t r = so.size();
+    const int64_t total = std::accumulate(so.begin(), so.end(), (int64_t)1,
+                                          std::multiplies<int64_t>{});
+    // 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<int64_t>& s) {
+      std::vector<int64_t> full(r, 1);
+      for (size_t i = 0; i < s.size(); ++i) {
+        full[r - s.size() + i] = s[i];
+      }
+      std::vector<int64_t> 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<int64_t> stra = build_strides(sa);
+    const std::vector<int64_t> strb = build_strides(sb);
+
+    std::vector<T> out(total);
+    std::vector<int64_t> 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<int64_t> 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<float>(&c1), c1.sizes(),
+                         ParseTensorData<float>(&c2), c2.sizes(), so);
+        return true;
+      case TensorProto_DataType_DOUBLE:
+        out.doubles() =
+            BroadcastMul(ParseTensorData<double>(&c1), c1.sizes(),
+                         ParseTensorData<double>(&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<int>(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

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(