onnxslim 0.1.83__tar.gz → 0.1.84__tar.gz

onnxslim-0.1.84/.github/ISSUE_TEMPLATE/bug_report.md

@@ -0,0 +1,30 @@
+---
+name: Bug report
+about: Create a report to help us improve
+title: ''
+labels: bug
+assignees: inisis
+
+---
+
+**Describe the bug**
+A clear and concise description of what the bug is.
+
+**To Reproduce**
+Steps to reproduce the bug:
+
+**Expected behavior**
+A clear and concise description of what you expected to happen.
+
+**Screenshots**
+If applicable, add screenshots to help explain your problem.
+
+**Environment:**
+run the script and paste the output.
+```
+python -c "import importlib; pkgs=['onnx','onnxruntime','onnxslim']; \
+[print(f'{p}: {importlib.import_module(p).__version__}') if importlib.util.find_spec(p) else print(f'{p}: missing') for p in pkgs]"
+```
+
+**Additional context**
+Add any other context about the problem here.

onnxslim-0.1.84/.github/ISSUE_TEMPLATE/custom.md

@@ -0,0 +1,10 @@
+---
+name: Custom issue template
+about: Describe this issue template's purpose here.
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+

onnxslim-0.1.84/.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,20 @@
+---
+name: Feature request
+about: Suggest an idea for this project
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+**Is your feature request related to a problem? Please describe.**
+A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
+
+**Describe the solution you'd like**
+A clear and concise description of what you want to happen.
+
+**Describe alternatives you've considered**
+A clear and concise description of any alternative solutions or features you've considered.
+
+**Additional context**
+Add any other context or screenshots about the feature request here.

{onnxslim-0.1.83 → onnxslim-0.1.84}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnxslim
-Version: 0.1.83
+Version: 0.1.84
 Summary: OnnxSlim: A Toolkit to Help Optimize Onnx Model
 Project-URL: homepage, https://github.com/inisis/OnnxSlim
 Project-URL: issues, https://github.com/inisis/OnnxSlim/issues

{onnxslim-0.1.83 → onnxslim-0.1.84}/onnxslim/core/pattern/elimination/slice.py

@@ -39,6 +39,16 @@ class SlicePatternMatcher(PatternMatcher):
                 first_slice_node_axes = first_slice_node_inputs[3].values.tolist()
                 first_slice_node_steps = first_slice_node_inputs[4].values.tolist()
 
+                # Check all users upfront before modifying the graph.
+                # If any user has overlapping axes, skip the optimization entirely
+                # to avoid corrupting the graph (fixes GitHub issue #277).
+                for user_node in first_slice_node_users:
+                    second_slice_node_inputs = list(user_node.inputs)
+                    second_slice_node_axes = second_slice_node_inputs[3].values.tolist()
+                    new_axes = first_slice_node_axes + second_slice_node_axes
+                    if len(new_axes) != len(set(new_axes)):
+                        return match_case
+
                 for user_node in first_slice_node_users:
                     second_slice_node = user_node
                     second_slice_node_inputs = list(second_slice_node.inputs)
@@ -52,33 +62,30 @@ class SlicePatternMatcher(PatternMatcher):
                     new_axes = first_slice_node_axes + second_slice_node_axes
                     new_steps = first_slice_node_steps + second_slice_node_steps
 
-                    if len(new_axes) != len(set(new_axes)):
-                        continue
-
                     inputs = []
+                    output_name = second_slice_node.outputs[0].name
                     inputs.extend(
                         (
                             next(iter(first_slice_node.inputs)),
                             gs.Constant(
-                                second_slice_node_inputs[1].name + "_starts",
+                                output_name + "_starts",
                                 values=np.array(new_starts, dtype=np.int64),
                             ),
                             gs.Constant(
-                                second_slice_node_inputs[2].name + "_ends",
+                                output_name + "_ends",
                                 values=np.array(new_ends, dtype=np.int64),
                             ),
                             gs.Constant(
-                                second_slice_node_inputs[3].name + "_axes",
+                                output_name + "_axes",
                                 values=np.array(new_axes, dtype=np.int64),
                             ),
                             gs.Constant(
-                                second_slice_node_inputs[4].name + "_steps",
+                                output_name + "_steps",
                                 values=np.array(new_steps, dtype=np.int64),
                             ),
                         )
                     )
                     outputs = list(second_slice_node.outputs)
-
                     first_slice_node.outputs.clear()
                     second_slice_node.inputs.clear()
                     second_slice_node.outputs.clear()

{onnxslim-0.1.83 → onnxslim-0.1.84}/onnxslim/core/pattern/fusion/concat_reshape.py

@@ -36,9 +36,11 @@ class ConcatReshapeMatcher(PatternMatcher):
     def rewrite(self, opset=11):
         match_case = {}
         concat_node = self.concat_0
+        reshape_node = self.reshape_0
         index = next(idx for idx, i in enumerate(concat_node.inputs) if isinstance(i, gs.Variable))
+        output_name = reshape_node.outputs[0].name
         constant = gs.Constant(
-            concat_node.inputs[index].name + "_fixed",
+            output_name + "_fixed",
             values=np.array([-1], dtype=np.int64),
         )
         concat_node.inputs.pop(index)

{onnxslim-0.1.83 → onnxslim-0.1.84}/onnxslim/core/pattern/fusion/convadd.py

@@ -44,12 +44,8 @@ class ConvAddMatcher(PatternMatcher):
             inputs = []
             inputs.append(next(iter(conv_node.inputs)))
             inputs.append(conv_weight)
-            weight_name = list(conv_node.inputs)[1].name
-            if weight_name.endswith("weight"):
-                bias_name = f"{weight_name[:-6]}bias"
-            else:
-                bias_name = f"{weight_name}_bias"
-            inputs.append(gs.Constant(bias_name, values=conv_bias))
+            output_name = add_node.outputs[0].name
+            inputs.append(gs.Constant(output_name + "_bias", values=conv_bias))
             outputs = list(add_node.outputs)
 
             conv_node.outputs.clear()

{onnxslim-0.1.83 → onnxslim-0.1.84}/onnxslim/core/pattern/fusion/convbn.py

@@ -52,15 +52,11 @@ class ConvBatchNormMatcher(PatternMatcher):
 
             inputs = []
             inputs.append(next(iter(conv_transpose_node.inputs)))
-            weight_name = list(conv_transpose_node.inputs)[1].name
-            if weight_name.endswith("weight"):
-                bias_name = f"{weight_name[:-6]}bias"
-            else:
-                bias_name = f"{weight_name}_bias"
+            output_name = bn_node.outputs[0].name
             inputs.extend(
                 (
-                    gs.Constant(weight_name + "_weight", values=conv_w),
-                    gs.Constant(bias_name, values=conv_b),
+                    gs.Constant(output_name + "_weight", values=conv_w),
+                    gs.Constant(output_name + "_bias", values=conv_b),
                 )
             )
             outputs = list(bn_node.outputs)

{onnxslim-0.1.83 → onnxslim-0.1.84}/onnxslim/core/pattern/fusion/convmul.py

@@ -38,14 +38,13 @@ class ConvMulMatcher(PatternMatcher):
                 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))
+                output_name = mul_node.outputs[0].name
+                inputs.append(gs.Constant(output_name + "_weight", 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))
+                    inputs.append(gs.Constant(output_name + "_bias", values=new_bias))
 
                 outputs = list(mul_node.outputs)
 

{onnxslim-0.1.83 → onnxslim-0.1.84}/onnxslim/core/pattern/fusion/gemm.py

@@ -76,7 +76,7 @@ class MatMulAddPatternMatcher(PatternMatcher):
                 output_variable.outputs.remove(add_node)
 
                 matmul_bias_transpose_constant = gs.Constant(
-                    matmul_bias_variable.name, values=matmul_bias_variable.values.T
+                    f"{matmul_node.name}_weight", values=matmul_bias_variable.values.T
                 )
 
                 inputs = []
@@ -143,7 +143,7 @@ class MatMulAddPatternMatcher(PatternMatcher):
                 output_variable.outputs.remove(add_node)
 
                 matmul_bias_transpose_constant = gs.Constant(
-                    matmul_bias_variable.name, values=matmul_bias_variable.values.T
+                    f"{matmul_node.name}_weight", values=matmul_bias_variable.values.T
                 )
 
                 inputs = []
@@ -235,14 +235,15 @@ class GemmMulPatternMatcher(PatternMatcher):
                     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)
+                output_name = reshape_node.outputs[0].name
+                gemm_weight_fused_constant = gs.Constant(output_name + "_weight_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_bias_fused_constant = gs.Constant(output_name + "_bias_fused", values=gemm_bias_fused)
                     gemm_node.inputs[2] = gemm_bias_fused_constant
 
                 mul_node.replace_all_uses_with(reshape_node)
@@ -312,7 +313,8 @@ class GemmAddPatternMatcher(PatternMatcher):
                     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)
+                    output_name = reshape_node.outputs[0].name
+                    gemm_bias_fused_constant = gs.Constant(output_name + "_bias_fused", values=gemm_bias_fused)
                     gemm_node.inputs[2] = gemm_bias_fused_constant
                 else:
                     return match_case

{onnxslim-0.1.83 → onnxslim-0.1.84}/onnxslim/third_party/onnx_graphsurgeon/ir/graph.py

@@ -794,109 +794,6 @@ class Graph:
                     tensor.to_constant(arr)
                     tensor.inputs.clear()
 
-        # Pass 2: Run shape-tensor cast elision
-        def run_cast_elision(node):
-            """Perform cast elision optimization on an ONNX node to eliminate unnecessary cast operations."""
-            import onnx
-
-            # Search for Cast(s) (from int -> float) -> intermediate operator (with float constants) -> Cast(s) (back to int)
-            # This pattern is problematic for TensorRT since these operations may be performed on Shape Tensors, which
-            # are not allowed to be floating point type. Attempt to fold the pattern here
-            VALID_CAST_ELISION_OPS = {
-                "Add",
-                "Sub",
-                "Mul",
-                "Div",
-                "Max",
-                "Min",
-                "Equal",
-                "Greater",
-                "Less",
-                "Concat",
-            }
-
-            if node.op not in VALID_CAST_ELISION_OPS:
-                return
-
-            # If the uncasted outputs of this node have any consumers other than "Cast" nodes,
-            # then we cannot elide the cast.
-            for out_tensor in node.outputs:
-                if out_tensor in self.outputs:
-                    return
-
-                if any(out_node.op != "Cast" for out_node in out_tensor.outputs):
-                    return
-
-            # Get list of input nodes that cast to float32
-            inp_casts = [
-                inp_node
-                for inp_tensor in node.inputs
-                for inp_node in inp_tensor.inputs
-                if inp_node.op == "Cast" and inp_node.attrs["to"] == onnx.TensorProto.DataType.FLOAT
-            ]
-
-            # No cast nodes found, return early
-            if not inp_casts:
-                return
-
-            # Ensure that all input cast nodes are casting from the same type
-            inp_dtypes = [dtype_to_onnx(inp_cast.inputs[0].dtype) for inp_cast in inp_casts]
-            if len(set(inp_dtypes)) != 1:
-                return
-
-            final_type = inp_dtypes[0]
-
-            # Get list of output nodes that cast to int32 or int64
-            out_casts = [
-                out_node
-                for out_tensor in node.outputs
-                for out_node in out_tensor.outputs
-                if out_node.op == "Cast"
-                and out_node.attrs["to"] in {onnx.TensorProto.DataType.INT32, onnx.TensorProto.DataType.INT64}
-            ]
-
-            # No cast node found on outputs, return early
-            if not out_casts:
-                return
-
-            # Ensure that all output cast nodes are casting to the same type and that this
-            # matches the original type before the inputs were casted.
-            out_dtypes = [out_cast.attrs["to"] for out_cast in out_casts]
-            if len(set(out_dtypes)) != 1 or out_dtypes[0] != final_type:
-                return
-
-            # If all checks passed, reconnect inputs/outputs to the consumers/producers
-            # of the Cast nodes.
-            # Note that we need to be careful in how we rebind tensors since they may
-            # be used by multiple nodes. Thus, it is not necessarily safe to assume that
-            # `cast_node.inputs[0].outputs[0] == cast_node`.
-            for index, inp in enumerate(node.inputs):
-                if isinstance(inp, Constant):
-                    inp.values = inp.values.astype(onnx.helper.tensor_dtype_to_np_dtype(final_type))
-
-                for cast in inp_casts:
-                    if cast.outputs[0] == inp:
-                        node.inputs[index] = cast.inputs[0]
-
-            for index, out in enumerate(node.outputs):
-                for cast in out_casts:
-                    if cast.inputs[0] == out:
-                        out_tensor = cast.outputs[0]
-                        out_tensor.inputs.clear()  # Disconnect from Cast
-                        node.outputs[index] = out_tensor
-
-        if fold_shapes:
-            # Perform shape tensor cast elision prior to most other folding
-            G_LOGGER.debug(f"Performing shape tensor cast elision in {self.name}")
-            try:
-                with self.node_ids():
-                    for node in self.nodes:
-                        run_cast_elision(node)
-            except Exception as err:
-                if not error_ok:
-                    raise err
-                G_LOGGER.warning("'{:}' routine failed with: {:}".format("Shape tensor cast elision", err))
-
         # Note that most of the remaining passes operate on a clone of the original graph.
         # Pass 3: Find all descendants of constant tensors
 

{onnxslim-0.1.83 → onnxslim-0.1.84}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "onnxslim"
-version = "0.1.83"
+version = "0.1.84"
 description = "OnnxSlim: A Toolkit to Help Optimize Onnx Model"
 authors = [{ name = "inisis", email = "desmond.yao@buaa.edu.cn" }]
 license = "MIT"
@@ -45,10 +45,11 @@ exclude = '''
 dev = [
     "coverage",
     "pytest-cov",
-    "onnx<=1.16.1 ; python_version<'3.10'",
-    "onnx>=1.12.0,<=1.18.0 ; python_version>='3.10'",
+    "onnx<=1.16.1 ; python_version<'3.11'",
+    "onnx>=1.12.0 ; python_version>='3.11'",
     "onnxruntime<1.20.0 ; python_version<'3.10'",
-    "onnxruntime>=1.20.0 ; python_version>='3.10'",
+    "onnxruntime>=1.20.0,<1.22 ; python_version=='3.10'",
+    "onnxruntime>=1.20.0 ; python_version>='3.11'",
     "onnxconverter-common",    
     "pytest",
     "pytest-xdist",
@@ -85,6 +86,7 @@ testpaths = [
     "tests/test_nvidia.py",
     "tests/test_shape_folding.py",
     "tests/test_coverage.py",
+    "tests/test_amd.py",
 ]
 addopts = "-n auto" # enables pytest-xdist parallel workers
 

onnxslim-0.1.84/tests/test_amd.py

@@ -0,0 +1,103 @@
+import os
+import tempfile
+
+import numpy as np
+import onnxruntime as ort
+import pytest
+
+import onnx
+from onnxslim import slim
+from onnxslim.utils import print_model_info_as_table, summarize_model
+
+MODELZOO_PATH = "/data/modelzoo/amd"
+
+def get_input_shapes(graph):
+    input_shapes = {}
+    for input in graph.input:
+        input_shapes[input.name] = [
+            s.dim_value for s in input.type.tensor_type.shape.dim
+        ]
+    return input_shapes
+
+
+def load_dummy_data(input_shape, data_size=5):
+    """Generate dummy input dicts."""
+    for _ in range(data_size):
+        data = {}
+        for name, shape in input_shape.items():
+            data[name] = np.random.rand(*shape).astype(np.float32)
+        yield data
+
+
+def run_slim_and_compare(original_path: str, slim_path: str, data_size=5):
+    """
+    Compare ORT inference results between original and slim models.
+    Slimming must already be done outside this function.
+    """
+    # ==== Prepare input shape ====
+    model = onnx.load(original_path)
+    input_shape = get_input_shapes(model.graph)
+
+    # fix dynamic dims → 1
+    for name, shape in input_shape.items():
+        input_shape[name] = [1 if s < 1 else s for s in shape]
+
+    data_gen = load_dummy_data(input_shape, data_size=data_size)
+
+    # ==== ORT sessions ====
+    opts = ort.SessionOptions()
+    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
+    opts.log_severity_level = 3
+    EP = ["CPUExecutionProvider"]
+
+    sess_orig = ort.InferenceSession(original_path, sess_options=opts, providers=EP)
+    sess_slim = ort.InferenceSession(slim_path, sess_options=opts, providers=EP)
+
+    # ==== Compare outputs ====
+    for inp in data_gen:
+        out1 = sess_orig.run([], inp)
+        out2 = sess_slim.run([], inp)
+
+        for a, b in zip(out1, out2):
+            if not np.array_equal(a, b):
+                return False
+    return True
+
+class TestModelZoo:
+    def test_EliminationSlice(self, request):
+        name = request.node.originalname[len("test_") :]
+        filename = f"{MODELZOO_PATH}/{name}.onnx"
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            slim_path = os.path.join(tempdir, f"{name}_slim.onnx")
+            slim(filename, slim_path)
+            ok = run_slim_and_compare(filename, slim_path)
+            assert ok, f"onnxslim output mismatch for model {name}!"
+
+    def test_sub_model(self, request):
+        name = request.node.originalname[len("test_") :]
+        filename = f"{MODELZOO_PATH}/{name}.onnx"
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            slim_path = os.path.join(tempdir, f"{name}_slim.onnx")
+            slim(filename, slim_path)
+
+            # Check that the slimmed model has exactly 2 Cast nodes
+            model_info = summarize_model(slim_path)
+            cast_count = model_info.op_type_counts.get("Cast", 0)
+            assert cast_count == 2, f"Expected 2 Cast nodes, but found {cast_count}"
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(
+        pytest.main(
+            [
+                "-p",
+                "no:warnings",
+                "-sv",
+                "tests/test_amd.py",
+            ]
+        )
+    )