Add fuse_gelu pass to fuse erf-based GELU into a Gelu op

onnxoptimizer/pass_registry.h

@@ -61,6 +61,7 @@
 #include "onnxoptimizer/passes/eliminate_common_subexpression.h"
 #include "onnxoptimizer/passes/fuse_qkv.h"
 #include "onnxoptimizer/passes/fuse_consecutive_unsqueezes.h"
+#include "onnxoptimizer/passes/fuse_gelu.h"
 #include "onnxoptimizer/passes/eliminate_nop_with_unit.h"
 #include "onnxoptimizer/passes/rewrite_input_dtype.h"
 #include "onnxoptimizer/passes/rewrite_where.h"
@@ -114,6 +115,7 @@ struct GlobalPassRegistry {
     registerPass<EliminateOpWithUnit>();
     registerPass<EliminateCommonSubexpression>();
     registerPass<FuseQKV>();
+    registerPass<FuseGelu>();
     registerPass<FuseConsecutiveUnsqueezes>();
     registerPass<EliminateDeadEnd>();
     registerPass<EliminateIdentity>();

onnxoptimizer/passes/fuse_gelu.h

@@ -0,0 +1,154 @@
+// 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
+
+// Fuse the erf-based (exact) GELU decomposition into a single Gelu operator,
+// mirroring the GELU fusion performed by onnxslim.
+//
+// It matches the subgraph
+//   div = x / sqrt(2)
+//   erf = Erf(div)
+//   add = erf + 1
+//   mul = x * add
+//   out = mul * 0.5
+// and rewrites it to
+//   out = Gelu(x)          // approximate = "none" (the default)
+//
+// The multiplication by 0.5 and the multiplication by (1 + erf(...)) are both
+// commutative, so either operand ordering is accepted. The division, however,
+// is not commutative: `x` must be the numerator.
+//
+// The standard-domain Gelu operator was introduced in opset 20, so the fusion
+// is only applied when the model targets opset >= 20; otherwise the resulting
+// node would be invalid.
+
+#include <cmath>
+
+#include "onnxoptimizer/pass.h"
+#include "onnxoptimizer/passes/pass_util.h"
+
+namespace ONNX_NAMESPACE {
+namespace optimization {
+
+struct FuseGelu final : public PredicateBasedPass {
+  explicit FuseGelu()
+      : PredicateBasedPass(PassType::Fuse, PassEfficiency::Complete,
+                           PassOptimizationType::Compute) {}
+
+  std::string getPassName() const override {
+    return "fuse_gelu";
+  }
+
+  // Returns true when `v` is a constant tensor holding a single float or double
+  // element approximately equal to `target`.
+  static bool IsScalarCloseTo(const Value* v, double target) {
+    float f;
+    if (FetchSoleValueOfTensor<float>(v, f)) {
+      return std::abs(static_cast<double>(f) - target) < 1e-4;
+    }
+    double d;
+    if (FetchSoleValueOfTensor<double>(v, d)) {
+      return std::abs(d - target) < 1e-4;
+    }
+    return false;
+  }
+
+  // The intermediate nodes of the pattern must feed only the next node in the
+  // chain, otherwise the decomposition is shared and cannot be safely fused.
+  static bool SingleUse(const Node* n) {
+    return n->output()->uses().size() == 1;
+  }
+
+  bool patternMatchPredicate(Node* node) override {
+    // Anchor on the final `Mul` (out = mul * 0.5).
+    return CheckKind(node, kMul) && node->inputs().size() == 2;
+  }
+
+  bool runTransform(Node* mul1, Graph& graph,
+                    NodeDestroyType& destroy_current) override {
+    destroy_current = NodeDestroyType::DestroyZero;
+
+    // The standard Gelu op lives in the default domain from opset 20 onwards.
+    if (getOpsetVersion(graph) < 20) {
+      return false;
+    }
+
+    // out = Mul(mul0, 0.5): locate the 0.5 scalar and the inner `Mul` node.
+    Node* mul0 = nullptr;
+    for (int i = 0; i < 2; ++i) {
+      Value* a = mul1->input(i);
+      Value* b = mul1->input(1 - i);
+      if (IsScalarCloseTo(b, 0.5) && CheckKind(a, kMul) &&
+          a->node()->inputs().size() == 2 && SingleUse(a->node())) {
+        mul0 = a->node();
+        break;
+      }
+    }
+    if (!mul0) {
+      return false;
+    }
+
+    // mul0 = Mul(x, add): locate the `Add` node and the data input `x`.
+    Node* add = nullptr;
+    Value* x = nullptr;
+    for (int i = 0; i < 2; ++i) {
+      Value* a = mul0->input(i);
+      Value* b = mul0->input(1 - i);
+      if (CheckKind(a, kAdd) && a->node()->inputs().size() == 2 &&
+          SingleUse(a->node())) {
+        add = a->node();
+        x = b;
+        break;
+      }
+    }
+    if (!add) {
+      return false;
+    }
+
+    // add = Add(erf, 1): locate the `Erf` node and check the addend.
+    Node* erf = nullptr;
+    for (int i = 0; i < 2; ++i) {
+      Value* a = add->input(i);
+      Value* b = add->input(1 - i);
+      if (CheckKind(a, "Erf") && a->node()->inputs().size() == 1 &&
+          SingleUse(a->node()) && IsScalarCloseTo(b, 1.0)) {
+        erf = a->node();
+        break;
+      }
+    }
+    if (!erf) {
+      return false;
+    }
+
+    // erf = Erf(div), div = Div(x, sqrt(2)). Div is not commutative, so `x`
+    // must be the numerator and the sqrt(2) constant the denominator.
+    Node* div = erf->input(0)->node();
+    if (!CheckKind(div, kDiv) || div->inputs().size() != 2 || !SingleUse(div)) {
+      return false;
+    }
+    if (div->input(0) != x ||
+        !IsScalarCloseTo(div->input(1), std::sqrt(2.0))) {
+      return false;
+    }
+
+    Node* gelu = graph.create(Symbol("Gelu"), 1);
+    gelu->addInput(x);
+    gelu->output()->copyMetadata(mul1->output());
+    gelu->insertBefore(mul1);
+
+    if (!tryReplacingAllUsesWith(mul1->output(), gelu->output())) {
+      gelu->destroy();
+      return false;
+    }
+    destroy_current = NodeDestroyType::DestroyOne;
+    return true;
+  }
+};
+
+}  // namespace optimization
+}  // namespace ONNX_NAMESPACE

onnxoptimizer/test/optimizer_test.py

@@ -18,6 +18,13 @@
 except ImportError:
     has_tv = False
 
+try:
+    import torch
+
+    has_torch = True
+except ImportError:
+    has_torch = False
+
 import onnx
 import pytest
 from onnx import (
@@ -4938,6 +4945,100 @@ def test_fuse_qkv(self):  # type: () -> None
             self._test_fuse_qkv_with_opset(opset_version)
         self._test_fuse_qkv_with_opset(LATEST_STABLE_OPSET_VERSION)
 
+    def _make_gelu_graph(self, half_last=True):  # type: (bool) -> onnx.GraphProto
+        # Builds the erf-based (exact) GELU decomposition:
+        #   out = 0.5 * x * (1 + erf(x / sqrt(2)))
+        X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4, 8])
+        Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4, 8])
+
+        sqrt2 = helper.make_tensor("sqrt2", TensorProto.FLOAT, [], [np.sqrt(2.0)])
+        one = helper.make_tensor("one", TensorProto.FLOAT, [], [1.0])
+        half = helper.make_tensor("half", TensorProto.FLOAT, [], [0.5])
+
+        nodes = [
+            helper.make_node("Div", ["X", "sqrt2"], ["div"]),
+            helper.make_node("Erf", ["div"], ["erf"]),
+            helper.make_node("Add", ["erf", "one"], ["add"]),
+            helper.make_node("Mul", ["X", "add"], ["mul"]),
+        ]
+        if half_last:
+            nodes.append(helper.make_node("Mul", ["mul", "half"], ["Y"]))
+        else:
+            nodes.append(helper.make_node("Mul", ["half", "mul"], ["Y"]))
+
+        return helper.make_graph(
+            nodes,
+            "test_gelu",
+            [X],
+            [Y],
+            [sqrt2, one, half],
+        )
+
+    def test_fuse_gelu(self):  # type: () -> None
+        for half_last in [True, False]:
+            graph = self._make_gelu_graph(half_last=half_last)
+            optimized_model = self._optimized(
+                graph,
+                ["fuse_gelu", "eliminate_deadend"],
+                opset_imports=[helper.make_opsetid("", 20)],
+            )
+            assert len(optimized_model.graph.node) == 1
+            assert optimized_model.graph.node[0].op_type == "Gelu"
+
+    def test_fuse_gelu_low_opset_is_noop(self):  # type: () -> None
+        # Gelu is only a standard-domain op from opset 20, so a lower opset must
+        # be left untouched.
+        graph = self._make_gelu_graph()
+        optimized_model = self._optimized(
+            graph,
+            ["fuse_gelu", "eliminate_deadend"],
+            opset_imports=[helper.make_opsetid("", 17)],
+        )
+        assert all(node.op_type != "Gelu" for node in optimized_model.graph.node)
+
+    @unittest.skipUnless(has_torch, "This test needs torch")
+    def test_fuse_gelu_torch_exported(self):  # type: () -> None
+        # Exercise the pass against a real graph produced by torch.onnx. With
+        # approximate="none", torch decomposes GELU into
+        # Div -> Erf -> Add -> Mul -> Mul (with the sqrt(2)/1/0.5 constants
+        # materialized as Constant *nodes* rather than initializers). torch only
+        # emits this decomposition below opset 20; at opset >= 20 it exports a
+        # Gelu op directly, so we export at opset 17 and version-convert to 20.
+        class Net(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(8, 8)
+                self.act = torch.nn.GELU(approximate="none")
+
+            def forward(self, x):
+                return self.act(self.linear(x))
+
+        net = Net().eval()
+        dummy = torch.randn(2, 4, 8)
+        buffer = io.BytesIO()
+        try:
+            torch.onnx.export(net, (dummy,), buffer, opset_version=17, dynamo=False)
+        except Exception as e:  # pragma: no cover - depends on torch version
+            self.skipTest(f"torch.onnx.export (TorchScript) unavailable: {e}")
+        model = onnx.load_from_string(buffer.getvalue())
+
+        # Confirm torch produced the erf decomposition we intend to fuse; if a
+        # future exporter emits a different structure, there is nothing to test.
+        op_types = {node.op_type for node in model.graph.node}
+        if not ({"Div", "Erf", "Add", "Mul"} <= op_types) or "Gelu" in op_types:
+            self.skipTest(f"torch did not emit the erf GELU decomposition: {op_types}")
+
+        model = onnx.version_converter.convert_version(model, 20)
+        optimized_model = self._optimized(
+            model, ["fuse_gelu", "eliminate_deadend"]
+        )
+
+        gelu_nodes = [n for n in optimized_model.graph.node if n.op_type == "Gelu"]
+        assert len(gelu_nodes) == 1
+        assert all(
+            n.op_type not in ("Erf", "Div") for n in optimized_model.graph.node
+        )
+
     def test_fuse_consecutive_unsqueezes_opset13(self):  # type: () -> None
         graph = parser.parse_graph("""
                agraph (float[4, 64, 160, 160] X) => (float[1, 1, 1, 4, 64, 1, 160, 160, 1, 1, 1] Z)