onnxslim 0.1.80__py3-none-any.whl

onnxslim/core/pattern/fusion/convmul.py

@@ -0,0 +1,69 @@
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.core.pattern import Pattern, PatternMatcher
+from onnxslim.core.pattern.registry import register_fusion_pattern
+
+
+class ConvMulMatcher(PatternMatcher):
+    def __init__(self, priority):
+        """Initializes the ConvMulMatcher for fusing Conv and Mul layers in an ONNX graph."""
+        pattern = Pattern(
+            """
+            input    input  0  1 conv_0
+            Conv     conv_0 1+ 1 input mul_0
+            Mul      mul_0  2  1 conv_0 ? output
+            output   output 1  0 mul_0
+            """
+        )
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the FusionConvMul pattern."""
+        return "FusionConvMul"
+
+    def rewrite(self, opset=11):
+        match_case = {}
+        conv_node = self.conv_0
+        mul_node = self.mul_0
+        conv_weight = list(conv_node.inputs)[1]
+        if len(conv_node.users) == 1 and conv_node.users[0] == mul_node and isinstance(mul_node.inputs[1], gs.Constant):
+            mul_constant = mul_node.inputs[1].values
+
+            if mul_constant.squeeze().ndim == 1 and mul_constant.squeeze().shape[0] == conv_weight.shape[0]:
+                weight_shape = conv_weight.values.shape
+                reshape_shape = [-1] + [1] * (len(weight_shape) - 1)
+
+                mul_scale_reshaped = mul_constant.squeeze().reshape(reshape_shape)
+                new_weight = conv_weight.values * mul_scale_reshaped
+
+                inputs = []
+                inputs.append(next(iter(conv_node.inputs)))
+
+                weight_name = list(conv_node.inputs)[1].name
+                inputs.append(gs.Constant(weight_name, values=new_weight))
+
+                if len(conv_node.inputs) == 3:
+                    conv_bias = conv_node.inputs[2].values
+                    new_bias = conv_bias * mul_constant.squeeze()
+                    bias_name = list(conv_node.inputs)[2].name
+                    inputs.append(gs.Constant(bias_name, values=new_bias))
+
+                outputs = list(mul_node.outputs)
+
+                conv_node.outputs.clear()
+                mul_node.inputs.clear()
+                mul_node.outputs.clear()
+
+                match_case[conv_node.name] = {
+                    "op": conv_node.op,
+                    "inputs": inputs,
+                    "outputs": outputs,
+                    "name": conv_node.name,
+                    "attrs": conv_node.attrs,
+                    "domain": None,
+                }
+
+        return match_case
+
+
+register_fusion_pattern(ConvMulMatcher(1))

onnxslim/core/pattern/fusion/gelu.py

@@ -0,0 +1,47 @@
+from onnxslim.core.pattern import Pattern, PatternMatcher
+
+
+class GeluPatternMatcher(PatternMatcher):
+    def __init__(self, priority):
+        """Initializes a `GeluPatternMatcher` to identify and fuse GELU patterns in a computational graph."""
+        pattern = Pattern(
+            """
+            input  input  0 2 mul_0 div_0
+            Div    div_0  2 1 input ? erf_0
+            Erf    erf_0  1 1 div_0 add_0
+            Add    add_0  2 1 erf_0 ? mul_0
+            Mul    mul_0  2 1 input add_0 mul_1
+            Mul    mul_1  2 1 mul_0 ? output
+            output output 1 0 mul_1
+            """
+        )
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the fusion pattern, 'FusionGelu'."""
+        return "FusionGelu"
+
+    def rewrite(self, opset=11):
+        """Rewrite the computation graph pattern to fuse GELU operations."""
+        input_variable = self.div_0.inputs[0]
+        mul_node = self.mul_0
+        div_node = self.div_0
+
+        input_variable.outputs.remove(mul_node)
+        input_variable.outputs.remove(div_node)
+
+        output_variable = self.mul_1.outputs[0]
+        output_variable.inputs.clear()
+
+        return {
+            self.mul_1.name: {
+                "op": "Gelu",
+                "inputs": [input_variable],
+                "outputs": [output_variable],
+                "domain": None,
+            }
+        }
+
+
+# register_fusion_pattern(GeluPatternMatcher(1))

onnxslim/core/pattern/fusion/gemm.py

@@ -0,0 +1,330 @@
+import numpy as np
+
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.core.optimization.dead_node_elimination import get_constant_variable
+from onnxslim.core.pattern import Pattern, PatternMatcher
+from onnxslim.core.pattern.registry import register_fusion_pattern
+
+
+class MatMulAddPatternMatcher(PatternMatcher):
+    def __init__(self, priority):
+        """Initializes a matcher for fusing MatMul and Add operations in ONNX graph optimization."""
+        pattern = Pattern(
+            """
+            input    input    0 1 matmul_0
+            MatMul   matmul_0 2 1 input ? add_0
+            Add      add_0    1* 1 matmul_0 output
+            output   output   1 0 add_0
+            """
+        )
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the fusion pattern as a string 'FusionGemm'."""
+        return "FusionGemm"
+
+    def rewrite(self, opset=11):
+        """Rewrites the graph for the fusion pattern 'FusionGemm' based on matching criteria and constant variables in
+        matmul nodes.
+        """
+        match_case = {}
+        node = self.add_0
+        matmul_node = self.matmul_0
+        matmul_bias_variable = get_constant_variable(matmul_node)
+        add_bias_variable = get_constant_variable(node)
+        input_variable = (
+            matmul_node.inputs[0] if isinstance(matmul_node.inputs[1], gs.Constant) else matmul_node.inputs[1]
+        )
+        users = matmul_node.users
+        if len(users) == 1 and matmul_bias_variable and add_bias_variable and len(matmul_bias_variable.shape) == 2:
+            if (
+                input_variable.shape
+                and len(input_variable.shape) > 2
+                and all([isinstance(value, int) for value in input_variable.shape])
+            ):
+                pre_reshape_const = gs.Constant(
+                    f"{matmul_node.name}_pre_reshape_in",
+                    values=np.array([-1, matmul_bias_variable.values.shape[0]], dtype=np.int64),
+                )
+                inputs = []
+                inputs.append(input_variable)
+                inputs.append(pre_reshape_const)
+
+                reshape_out_variable = gs.Variable(
+                    f"{matmul_node.name}_pre_reshape_out",
+                    dtype=input_variable.dtype,
+                )
+                outputs = [reshape_out_variable]
+
+                match_case.update(
+                    {
+                        f"{matmul_node.name}_pre_reshape": {
+                            "op": "Reshape",
+                            "inputs": inputs,
+                            "outputs": outputs,
+                            "name": f"{matmul_node.name}_pre_reshape",
+                            "domain": None,
+                        }
+                    }
+                )
+
+                add_node = node
+                add_bias_variable = get_constant_variable(add_node)
+
+                output_variable = add_node.inputs[0]
+                output_variable.outputs.remove(add_node)
+
+                matmul_bias_transpose_constant = gs.Constant(
+                    matmul_bias_variable.name, values=matmul_bias_variable.values.T
+                )
+
+                inputs = []
+                inputs.append(reshape_out_variable)
+                inputs.append(matmul_bias_transpose_constant)
+                inputs.append(add_bias_variable)
+
+                gemm_out_variable = gs.Variable(f"{matmul_node.name}_gemm_out", dtype=output_variable.dtype)
+                outputs = [gemm_out_variable]
+
+                match_case.update(
+                    {
+                        matmul_node.name: {
+                            "op": "Gemm",
+                            "inputs": inputs,
+                            "outputs": outputs,
+                            "name": matmul_node.name,
+                            "attrs": {
+                                "alpha": 1.0,
+                                "beta": 1.0,
+                                "transA": 0,
+                                "transB": 1,
+                            },
+                            "domain": None,
+                        }
+                    }
+                )
+
+                values = [*list(input_variable.shape[:-1]), matmul_bias_variable.values.shape[-1]]
+                post_reshape_const = gs.Constant(
+                    f"{matmul_node.name}_post_reshape_in",
+                    values=np.array(values, dtype=np.int64),
+                )
+
+                inputs = []
+                inputs.append(gemm_out_variable)
+                inputs.append(post_reshape_const)
+                outputs = list(add_node.outputs)
+
+                matmul_node.outputs.clear()
+                add_node.inputs.clear()
+                add_node.outputs.clear()
+
+                match_case.update(
+                    {
+                        f"{matmul_node.name}_post_reshape": {
+                            "op": "Reshape",
+                            "inputs": inputs,
+                            "outputs": outputs,
+                            "name": f"{matmul_node.name}_post_reshape",
+                            "domain": None,
+                        }
+                    }
+                )
+            elif (
+                input_variable.shape
+                and len(input_variable.shape) == 2
+                and all([isinstance(value, int) for value in input_variable.shape])
+            ):
+                add_node = node
+                add_bias_variable = get_constant_variable(add_node)
+
+                output_variable = add_node.inputs[0]
+                output_variable.outputs.remove(add_node)
+
+                matmul_bias_transpose_constant = gs.Constant(
+                    matmul_bias_variable.name, values=matmul_bias_variable.values.T
+                )
+
+                inputs = []
+                inputs.append(input_variable)
+                inputs.append(matmul_bias_transpose_constant)
+                inputs.append(add_bias_variable)
+
+                outputs = list(add_node.outputs)
+                add_node.inputs.clear()
+                add_node.outputs.clear()
+                match_case.update(
+                    {
+                        matmul_node.name: {
+                            "op": "Gemm",
+                            "inputs": inputs,
+                            "outputs": outputs,
+                            "name": matmul_node.name,
+                            "attrs": {
+                                "alpha": 1.0,
+                                "beta": 1.0,
+                                "transA": 0,
+                                "transB": 1,
+                            },
+                            "domain": None,
+                        }
+                    }
+                )
+        return match_case
+
+
+register_fusion_pattern(MatMulAddPatternMatcher(1))
+
+
+class GemmMulPatternMatcher(PatternMatcher):
+    def __init__(self, priority):
+        """Initializes a matcher for fusing MatMul and Add operations in ONNX graph optimization."""
+        pattern = Pattern(
+            """
+            input    input     0  1 gemm_0
+            Gemm     gemm_0    1+ 1 input reshape_0
+            Reshape  reshape_0 2  1 gemm_0 ? mul_0
+            Mul      mul_0     1* 1 reshape_0 output
+            output   output    1  0 mul_0
+            """
+        )
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the fusion pattern as a string 'FusionGemmMul'."""
+        return "FusionGemmMul"
+
+    def rewrite(self, opset=11):
+        """Rewrites the graph for the fusion pattern 'FusionGemmMul' based on matching criteria and constant variables
+        in gemm nodes.
+        """
+        match_case = {}
+        gemm_node = self.gemm_0
+        reshape_node = self.reshape_0
+        mul_node = self.mul_0
+        mul_bias_variable = get_constant_variable(mul_node)
+
+        if (
+            (
+                (len(gemm_node.inputs) == 2 and isinstance(gemm_node.inputs[1], gs.Constant))
+                or (
+                    len(gemm_node.inputs) == 3
+                    and isinstance(gemm_node.inputs[1], gs.Constant)
+                    and isinstance(gemm_node.inputs[2], gs.Constant)
+                )
+            )
+            and mul_bias_variable
+            and len(reshape_node.users) == 1
+        ):
+            gemm_attr = gemm_node.attrs
+            gemm_weight_constant = gemm_node.inputs[1]
+            gemm_bias_constant = gemm_node.inputs[2] if len(gemm_node.inputs) == 3 else None
+            if (
+                gemm_attr["transA"] == 0
+                and gemm_attr["transB"] == 1
+                and (
+                    (mul_bias_variable.values.ndim == 1 and gemm_weight_constant.shape[0] == mul_bias_variable.shape[0])
+                    or mul_bias_variable.values.ndim == 0
+                )
+            ):
+                gemm_weight = gemm_weight_constant.values
+                mul_weight = mul_bias_variable.values
+                if mul_bias_variable.values.ndim == 1:
+                    gemm_weight_fused = gemm_weight * mul_weight[:, None]
+                else:
+                    gemm_weight_fused = gemm_weight * mul_weight
+                gemm_weight_fused_constant = gs.Constant(gemm_weight_constant.name + "_fused", values=gemm_weight_fused)
+                gemm_node.inputs[1] = gemm_weight_fused_constant
+
+                if gemm_bias_constant:
+                    gemm_bias = gemm_bias_constant.values
+                    mul_bias = mul_bias_variable.values
+                    gemm_bias_fused = gemm_bias * mul_bias
+                    gemm_bias_fused_constant = gs.Constant(gemm_bias_constant.name + "_fused", values=gemm_bias_fused)
+                    gemm_node.inputs[2] = gemm_bias_fused_constant
+
+                mul_node.replace_all_uses_with(reshape_node)
+
+        return match_case
+
+
+register_fusion_pattern(GemmMulPatternMatcher(1))
+
+
+class GemmAddPatternMatcher(PatternMatcher):
+    def __init__(self, priority):
+        """Initializes a matcher for fusing MatMul and Add operations in ONNX graph optimization."""
+        pattern = Pattern(
+            """
+            input    input     0  1 gemm_0
+            Gemm     gemm_0    1+ 1 input reshape_0
+            Reshape  reshape_0 2  1 gemm_0 ? add_0
+            Add      add_0     1* 1 reshape_0 output
+            output   output    1  0 add_0
+            """
+        )
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the fusion pattern as a string 'FusionGemmAdd'."""
+        return "FusionGemmAdd"
+
+    def rewrite(self, opset=11):
+        """Rewrites the graph for the fusion pattern 'FusionGemmAdd' based on matching criteria and constant variables
+        in gemm nodes.
+        """
+        match_case = {}
+        gemm_node = self.gemm_0
+        reshape_node = self.reshape_0
+        add_node = self.add_0
+        add_bias_variable = get_constant_variable(add_node)
+
+        if (
+            (
+                (len(gemm_node.inputs) == 2)
+                or (len(gemm_node.inputs) == 3 and isinstance(gemm_node.inputs[2], gs.Constant))
+            )
+            and add_bias_variable
+            and len(reshape_node.users) == 1
+            and gemm_node.outputs[0].shape
+        ):
+
+            def can_broadcast_to(shape_from, shape_to):
+                """Return True if shape_from can broadcast to shape_to per NumPy rules."""
+                if shape_from is None or shape_to is None:
+                    return False
+                try:
+                    np.empty(shape_to, dtype=np.float32) + np.empty(shape_from, dtype=np.float32)
+                    return True
+                except ValueError:
+                    return False
+
+            gemm_bias_constant = gemm_node.inputs[2] if len(gemm_node.inputs) == 3 else None
+            if gemm_bias_constant:
+                gemm_bias = gemm_bias_constant.values
+                add_bias = add_bias_variable.values
+                if (
+                    can_broadcast_to(gemm_bias.shape, gemm_node.outputs[0].shape)
+                    and can_broadcast_to(add_bias.shape, gemm_node.outputs[0].shape)
+                    and add_bias.ndim <= 2
+                ):
+                    gemm_bias_fused = gemm_bias + add_bias
+                    gemm_bias_fused_constant = gs.Constant(gemm_bias_constant.name + "_fused", values=gemm_bias_fused)
+                    gemm_node.inputs[2] = gemm_bias_fused_constant
+                else:
+                    return match_case
+            else:
+                if can_broadcast_to(add_bias_variable.values.shape, gemm_node.outputs[0].shape):
+                    gemm_node.inputs[2] = add_bias_variable
+                else:
+                    return match_case
+
+            add_node.replace_all_uses_with(reshape_node)
+
+        return match_case
+
+
+register_fusion_pattern(GemmAddPatternMatcher(1))

onnxslim/core/pattern/fusion/padconv.py

@@ -0,0 +1,89 @@
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.core.pattern import Pattern, PatternMatcher
+from onnxslim.core.pattern.registry import register_fusion_pattern
+
+
+class PadConvMatcher(PatternMatcher):
+    def __init__(self, priority):
+        """Initializes the PadConvMatcher with a specified priority and defines its matching pattern."""
+        pattern = Pattern(
+            """
+            input  input  0  1 pad_0
+            Pad    pad_0  1+ 1 input conv_0
+            Conv   conv_0 1+ 1 pad_0 output
+            output output 1  0 conv_0
+            """
+        )
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the fusion pattern used."""
+        return "FusionPadConv"
+
+    def parameter_check(self) -> bool:
+        """Validates if the padding parameter for a convolutional node is a constant."""
+        pad_node = self.pad_0
+
+        return isinstance(pad_node.inputs[1], gs.Constant)
+
+    def rewrite(self, opset=11):
+        """Rewrites the padding parameter for a convolutional node to use a constant if the current parameter is not a
+        constant.
+        """
+        match_case = {}
+        conv_node = self.conv_0
+        pad_node = self.pad_0
+        pad_node_users = pad_node.users
+
+        pad_inputs = len(pad_node.inputs)
+        if pad_inputs < 3 or (
+            (pad_inputs >= 3 and (isinstance(pad_node.inputs[2], gs.Constant) and pad_node.inputs[2].values == 0))
+            or (pad_inputs >= 3 and (isinstance(pad_node.inputs[2], gs.Variable) and pad_node.inputs[2].name == ""))
+        ):
+            if (
+                isinstance(pad_node.inputs[1], gs.Constant)
+                and pad_node.attrs.get("mode", "constant") == "constant"
+                and conv_node.inputs[1].shape
+            ):
+                conv_weight_dim = len(conv_node.inputs[1].shape)
+                pad_value = pad_node.inputs[1].values.tolist()
+                if all(pad == 0 for pad in (pad_value[:2] + pad_value[conv_weight_dim : conv_weight_dim + 2])):
+                    conv_weight_dim - 2
+                    input_variable = self.pad_0.inputs[0]
+                    pad_variable = pad_node.outputs[0]  # pad output variable
+                    index = conv_node.inputs.index(pad_variable)
+                    conv_node.inputs.pop(index)
+                    conv_node.inputs.insert(index, input_variable)
+
+                    inputs = list(conv_node.inputs)
+                    outputs = list(conv_node.outputs)
+                    attrs = conv_node.attrs
+
+                    conv_node.inputs.clear()
+                    conv_node.outputs.clear()
+                    # remove pad node if it has only one user
+                    if len(pad_node_users) == 0:
+                        input_variable.outputs.remove(pad_node)
+                        pad_node.inputs.clear()
+                        pad_node.outputs.clear()
+
+                    pads = pad_value[2:conv_weight_dim] + pad_value[conv_weight_dim + 2 :]
+                    if hasattr(attrs, "pads"):
+                        conv_pads = attrs["pads"]
+                        pads = [pad + conv_pad for pad, conv_pad in zip(pads, conv_pads)]
+
+                    attrs["pads"] = pads
+                    match_case[conv_node.name] = {
+                        "op": "Conv",
+                        "inputs": inputs,
+                        "outputs": outputs,
+                        "name": conv_node.name,
+                        "attrs": conv_node.attrs,
+                        "domain": None,
+                    }
+
+        return match_case
+
+
+register_fusion_pattern(PadConvMatcher(1))

onnxslim/core/pattern/fusion/reduce.py

@@ -0,0 +1,67 @@
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.core.pattern import Pattern, PatternMatcher
+from onnxslim.core.pattern.registry import register_fusion_pattern
+
+
+class ReducePatternMatcher(PatternMatcher):
+    def __init__(self, priority):
+        """Initializes the ReducePatternMatcher with a specified pattern matching priority level."""
+        pattern = Pattern(
+            """
+            input     input       0   1 reduce_0
+            ReduceSum reduce_0    1+  1 input unsqueeze_0
+            Unsqueeze unsqueeze_0 1+  1 reduce_0 output
+            output    output      1   0 unsqueeze_0
+            """
+        )
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the fusion pattern 'FusionReduce'."""
+        return "FusionReduce"
+
+    def rewrite(self, opset=11):
+        """Rewrites the graph pattern based on opset version; reuses Reduce and Unsqueeze nodes if possible."""
+        match_case = {}
+        node = self.unsqueeze_0
+        reduce_node = self.reduce_0
+        reduce_node_node_users = reduce_node.users
+        if len(reduce_node_node_users) == 1:
+            unsqueeze_node = node
+
+            if opset < 13:
+                reduce_node_axes = reduce_node.attrs.get("axes", None)
+                reduce_node_keepdims = reduce_node.attrs.get("keepdims", 1)
+                unsqueeze_node_axes = unsqueeze_node.attrs.get("axes", None)
+            else:
+                reduce_node_axes_ = reduce_node.inputs[1]
+                reduce_node_keepdims = reduce_node.attrs.get("keepdims", 1)
+                unsqueeze_node_axes_ = unsqueeze_node.inputs[1]
+                if isinstance(reduce_node_axes_, gs.Constant) and isinstance(unsqueeze_node_axes_, gs.Constant):
+                    reduce_node_axes = reduce_node_axes_.values
+                    unsqueeze_node_axes = unsqueeze_node_axes_.values
+                else:
+                    return match_case
+
+            if reduce_node_axes == unsqueeze_node_axes and reduce_node_keepdims == 0:
+                inputs = list(reduce_node.inputs)
+                outputs = list(unsqueeze_node.outputs)
+                attrs = reduce_node.attrs
+                reduce_node.outputs.clear()
+                unsqueeze_node.inputs.clear()
+                unsqueeze_node.outputs.clear()
+                attrs["keepdims"] = 1
+                match_case[reduce_node.name] = {
+                    "op": reduce_node.op,
+                    "inputs": inputs,
+                    "outputs": outputs,
+                    "name": reduce_node.name,
+                    "attrs": attrs,
+                    "domain": None,
+                }
+
+        return match_case
+
+
+register_fusion_pattern(ReducePatternMatcher(1))

onnxslim/core/pattern/registry.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from collections import OrderedDict
+
+DEFAULT_FUSION_PATTERNS = OrderedDict()
+
+
+def register_fusion_pattern(fusion_pattern):
+    """Registers a fusion pattern function for a specified layer type in the DEFAULT_FUSION_PATTERNS dictionary."""
+    layer_type = fusion_pattern.name
+
+    if layer_type in DEFAULT_FUSION_PATTERNS.keys():
+        raise
+    DEFAULT_FUSION_PATTERNS[layer_type] = fusion_pattern
+
+
+def get_fusion_patterns(skip_fusion_patterns: str | None = None):
+    """Returns a copy of the default fusion patterns, optionally excluding specific patterns."""
+    default_fusion_patterns = DEFAULT_FUSION_PATTERNS.copy()
+    if skip_fusion_patterns:
+        for pattern in skip_fusion_patterns:
+            default_fusion_patterns.pop(pattern)
+
+    return default_fusion_patterns
+
+
+from .elimination import *
+from .fusion import *