diff --git a/onnxoptimizer/pass_registry.h b/onnxoptimizer/pass_registry.h index 3fdfa4978..65355ce62 100644 --- a/onnxoptimizer/pass_registry.h +++ b/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(); registerPass(); registerPass(); + registerPass(); registerPass(); registerPass(); registerPass(); diff --git a/onnxoptimizer/passes/fuse_gelu.h b/onnxoptimizer/passes/fuse_gelu.h new file mode 100644 index 000000000..d9822202c --- /dev/null +++ b/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 + +#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(v, f)) { + return std::abs(static_cast(f) - target) < 1e-4; + } + double d; + if (FetchSoleValueOfTensor(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 diff --git a/onnxoptimizer/test/optimizer_test.py b/onnxoptimizer/test/optimizer_test.py index 2ee7d3003..b87e126ac 100644 --- a/onnxoptimizer/test/optimizer_test.py +++ b/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)